🧠 fix: Correct Memory Token Accounting (#15224)

packages/api/src/agents/memory.spec.ts

@@ -712,6 +712,32 @@ describe('createMemoryTool tokenLimit enforcement', () => {
     expect(setMemory).toHaveBeenCalledTimes(2);
   });
 
+  it('treats a persisted key as a replacement in a new tool instance', async () => {
+    const setMemory = jest.fn().mockResolvedValue({ ok: true });
+    /** ~100 tokens; adding this value to its previous version would exceed the
+     *  limit, while replacing it correctly remains within the limit. */
+    const value = 'word '.repeat(100).trim();
+    const firstTool = createMemoryTool({
+      userId: 'user-1',
+      setMemory,
+      tokenLimit: 150,
+    });
+
+    await firstTool.invoke({ key: 'k1', value });
+
+    const secondTool = createMemoryTool({
+      userId: 'user-1',
+      setMemory,
+      tokenLimit: 150,
+      totalTokens: 100,
+      tokenCountsByKey: new Map([['k1', 100]]),
+    });
+
+    await secondTool.invoke({ key: 'k1', value });
+
+    expect(setMemory).toHaveBeenCalledTimes(2);
+  });
+
   it('fires onWrite after a successful set, but not when the write fails', async () => {
     const onWrite = jest.fn();
     const okTool = createMemoryTool({
@@ -746,6 +772,38 @@ describe('createMemoryTool tokenLimit enforcement', () => {
   });
 });
 
+describe('memory token limit guidance', () => {
+  it('describes the aggregate limit and never reports negative remaining capacity', async () => {
+    const [, process] = await createMemoryProcessor({
+      res: { headersSent: false, write: jest.fn() } as unknown as Response,
+      userId: 'user-1',
+      messageId: 'message-1',
+      conversationId: 'conversation-1',
+      config: { tokenLimit: 100 },
+      memoryMethods: {
+        setMemory: jest.fn().mockResolvedValue({ ok: true }),
+        deleteMemory: jest.fn().mockResolvedValue({ ok: true }),
+        getUserMemories: jest.fn().mockResolvedValue([]),
+        getFormattedMemories: jest.fn().mockResolvedValue({
+          withKeys: 'existing memory',
+          withoutKeys: 'existing memory',
+          totalTokens: 150,
+          tokenCountsByKey: new Map([['preferences', 150]]),
+        }),
+      },
+    });
+
+    await process([]);
+
+    const runCalls = (Run.create as jest.Mock).mock.calls;
+    const runConfig = runCalls[runCalls.length - 1][0];
+    expect(runConfig.graphConfig.instructions).toContain(
+      'Maximum 100 tokens across all memory values.',
+    );
+    expect(runConfig.graphConfig.additional_instructions).toContain('Remaining capacity: 0 tokens');
+  });
+});
+
 describe('buildInlineMemoryTool content filtering', () => {
   it('keeps a legacy-only message filter scoped to ingress messages', async () => {
     const setMemory = jest.fn().mockResolvedValue({ ok: true });

packages/api/src/agents/memory.ts

@@ -113,7 +113,7 @@ The \`delete_memory\` tool should only be used in two scenarios:
 
 ${validKeys && validKeys.length > 0 ? `\nVALID KEYS: ${validKeys.join(', ')}` : ''}
 
-${tokenLimit ? `\nTOKEN LIMIT: Maximum ${tokenLimit} tokens per memory value.` : ''}
+${tokenLimit ? `\nTOKEN LIMIT: Maximum ${tokenLimit} tokens across all memory values.` : ''}
 
 When in doubt, and the user hasn't asked to remember or forget anything, END THE TURN IMMEDIATELY.`;
 
@@ -130,6 +130,7 @@ export const createMemoryTool = ({
   charLimit,
   tokenLimit,
   totalTokens = 0,
+  tokenCountsByKey,
   filters,
   onWrite,
 }: {
@@ -141,6 +142,7 @@ export const createMemoryTool = ({
   charLimit?: number;
   tokenLimit?: number;
   totalTokens?: number;
+  tokenCountsByKey?: ReadonlyMap<string, number>;
   filters?: FiltersConfig;
   onWrite?: () => void;
 }): DynamicStructuredTool => {
@@ -150,10 +152,10 @@ export const createMemoryTool = ({
    *  check against the same stale total and collectively exceed `tokenLimit`. */
   let currentTotalTokens = totalTokens;
   let writeChain: Promise<unknown> = Promise.resolve();
-  /** Tokens this instance has already committed per key. `set_memory` upserts,
-   *  so a repeat write to the same key REPLACES its value — the running total
-   *  must swap the prior contribution for the new one, not add both. */
-  const writtenTokensByKey = new Map<string, number>();
+  /** Token counts are seeded from persisted memory and advanced after each
+   *  successful write. `set_memory` upserts, so every write replaces this
+   *  key's prior contribution instead of adding both values to the total. */
+  const currentTokensByKey = new Map(tokenCountsByKey);
 
   return tool(
     async ({ key, value }) => {
@@ -188,9 +190,9 @@ export const createMemoryTool = ({
           }
 
           const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base');
-          /** Total excluding this key's prior in-instance write, so a same-key
-           *  rewrite is measured as a replacement rather than an addition. */
-          const baseTotalTokens = currentTotalTokens - (writtenTokensByKey.get(key) ?? 0);
+          /** Total excluding this key's prior persisted or in-instance value,
+           *  so a rewrite is measured as a replacement rather than an addition. */
+          const baseTotalTokens = currentTotalTokens - (currentTokensByKey.get(key) ?? 0);
           const remainingTokens = tokenLimit ? tokenLimit - baseTotalTokens : Infinity;
 
           if (tokenLimit && remainingTokens <= 0) {
@@ -247,7 +249,7 @@ export const createMemoryTool = ({
           if (result.ok) {
             if (tokenLimit) {
               currentTotalTokens = newTotalTokens;
-              writtenTokensByKey.set(key, tokenCount);
+              currentTokensByKey.set(key, tokenCount);
             }
             onWrite?.();
             logger.debug(`Memory set for key "${key}" (${tokenCount} tokens) for user "${userId}"`);
@@ -680,6 +682,7 @@ export async function buildInlineMemoryTool({
   const charLimit = memoryConfig?.charLimit as number | undefined;
   const tokenLimit = memoryConfig?.tokenLimit as number | undefined;
   let totalTokens = 0;
+  let tokenCountsByKey: ReadonlyMap<string, number> | undefined;
   if (tokenLimit) {
     try {
       const formatted = await getRequestMemories({
@@ -689,6 +692,7 @@ export async function buildInlineMemoryTool({
         getFormattedMemories: memoryMethods.getFormattedMemories,
       });
       totalTokens = formatted?.totalTokens ?? 0;
+      tokenCountsByKey = formatted?.tokenCountsByKey;
     } catch (error) {
       logger.error(
         '[memory] Failed to load memory token count for set_memory',
@@ -708,6 +712,7 @@ export async function buildInlineMemoryTool({
     charLimit,
     tokenLimit,
     totalTokens,
+    tokenCountsByKey,
     filters: req.config?.filters,
     onWrite: () => invalidateRequestMemories(req, memoryAgentId),
   });
@@ -754,6 +759,7 @@ export async function processMemory({
   llmConfig,
   tokenLimit,
   totalTokens = 0,
+  tokenCountsByKey,
   filters,
   streamId = null,
   jobCreatedAt,
@@ -780,6 +786,7 @@ export async function processMemory({
   }[];
   tokenLimit?: number;
   totalTokens?: number;
+  tokenCountsByKey?: ReadonlyMap<string, number>;
   filters?: FiltersConfig;
   llmConfig?: Partial<LLMConfig>;
   streamId?: string | null;
@@ -813,6 +820,7 @@ export async function processMemory({
       setMemory,
       validKeys,
       totalTokens,
+      tokenCountsByKey,
       filters,
     });
     const deleteMemoryTool = createDeleteMemoryTool({
@@ -827,7 +835,7 @@ export async function processMemory({
     let memoryStatus = `# Existing memory:\n${memory ?? 'No existing memories'}`;
 
     if (tokenLimit) {
-      const remainingTokens = tokenLimit - currentMemoryTokens;
+      const remainingTokens = Math.max(tokenLimit - currentMemoryTokens, 0);
       memoryStatus = `# Memory Status:
 Current memory usage: ${currentMemoryTokens} tokens
 Token limit: ${tokenLimit} tokens
@@ -1044,15 +1052,16 @@ export async function createMemoryProcessor({
   const { validKeys, instructions, llmConfig, tokenLimit } = config;
   const finalInstructions = instructions || getDefaultInstructions(validKeys, tokenLimit);
 
-  const [{ withKeys, withoutKeys, totalTokens }, memoryEntries] = await Promise.all([
-    memoryMethods.getFormattedMemories({
-      userId,
-      agentId,
-    }),
-    hasActivePiiPatterns(filters?.memories?.pii)
-      ? memoryMethods.getUserMemories({ userId, agentId })
-      : Promise.resolve(undefined),
-  ]);
+  const [{ withKeys, withoutKeys, totalTokens, tokenCountsByKey }, memoryEntries] =
+    await Promise.all([
+      memoryMethods.getFormattedMemories({
+        userId,
+        agentId,
+      }),
+      hasActivePiiPatterns(filters?.memories?.pii)
+        ? memoryMethods.getUserMemories({ userId, agentId })
+        : Promise.resolve(undefined),
+    ]);
 
   return [
     withoutKeys,
@@ -1077,6 +1086,7 @@ export async function createMemoryProcessor({
           memory: withKeys,
           memoryEntries,
           totalTokens: totalTokens || 0,
+          tokenCountsByKey,
           filters,
           instructions: finalInstructions,
           setMemory: memoryMethods.setMemory,

packages/data-schemas/src/methods/memory.spec.ts

@@ -245,11 +245,13 @@ describe('memory partitions', () => {
 
     const personal = await methods.getFormattedMemories({ userId });
     expect(personal.totalTokens).toBe(5);
+    expect(personal.tokenCountsByKey).toEqual(new Map([['one', 5]]));
     expect(personal.withoutKeys).toContain('personal one');
     expect(personal.withoutKeys).not.toContain('agent two');
 
     const partitionA = await methods.getFormattedMemories({ userId, agentId });
     expect(partitionA.totalTokens).toBe(7);
+    expect(partitionA.tokenCountsByKey).toEqual(new Map([['two', 7]]));
     expect(partitionA.withKeys).toContain('"key": "two"');
     expect(partitionA.withKeys).not.toContain('personal one');
   });

packages/data-schemas/src/methods/memory.ts

@@ -271,16 +271,27 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
       const memories = await getUserMemories({ userId, agentId });
 
       if (!memories || memories.length === 0) {
-        return { withKeys: '', withoutKeys: '', totalTokens: 0 };
+        return {
+          withKeys: '',
+          withoutKeys: '',
+          totalTokens: 0,
+          tokenCountsByKey: new Map<string, number>(),
+        };
       }
 
       const sortedMemories = memories.sort(
         (a, b) => new Date(a.updated_at!).getTime() - new Date(b.updated_at!).getTime(),
       );
 
-      const totalTokens = sortedMemories.reduce((sum, memory) => {
-        return sum + (memory.tokenCount || 0);
-      }, 0);
+      const { totalTokens, tokenCountsByKey } = sortedMemories.reduce(
+        (counts, memory) => {
+          const tokenCount = memory.tokenCount || 0;
+          counts.totalTokens += tokenCount;
+          counts.tokenCountsByKey.set(memory.key, tokenCount);
+          return counts;
+        },
+        { totalTokens: 0, tokenCountsByKey: new Map<string, number>() },
+      );
 
       const withKeys = sortedMemories
         .map((memory, index) => {
@@ -297,10 +308,15 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
         })
         .join('\n\n');
 
-      return { withKeys, withoutKeys, totalTokens };
+      return { withKeys, withoutKeys, totalTokens, tokenCountsByKey };
     } catch (error) {
       logger.error('Failed to get formatted memories:', error);
-      return { withKeys: '', withoutKeys: '', totalTokens: 0 };
+      return {
+        withKeys: '',
+        withoutKeys: '',
+        totalTokens: 0,
+        tokenCountsByKey: new Map<string, number>(),
+      };
     }
   }
 

packages/data-schemas/src/types/memory.ts

@@ -76,4 +76,5 @@ export interface FormattedMemoriesResult {
   withKeys: string;
   withoutKeys: string;
   totalTokens?: number;
+  tokenCountsByKey?: Map<string, number>;
 }