From 95e1a5c9fa4fe85f48a53b672648a23184f5c1fa Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 25 Aug 2026 19:58:30 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=A0=20fix:=20Correct=20Memory=20Token?= =?UTF-8?q?=20Accounting=20(#15224)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/api/src/agents/memory.spec.ts | 58 +++++++++++++++++++ packages/api/src/agents/memory.ts | 48 +++++++++------ .../data-schemas/src/methods/memory.spec.ts | 2 + packages/data-schemas/src/methods/memory.ts | 28 +++++++-- packages/data-schemas/src/types/memory.ts | 1 + 5 files changed, 112 insertions(+), 25 deletions(-) diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts index 6ecd579422..e74851ccf0 100644 --- a/packages/api/src/agents/memory.spec.ts +++ b/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 }); diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index ff345a9955..b453dc83fa 100644 --- a/packages/api/src/agents/memory.ts +++ b/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; 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 = 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(); + /** 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 | 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; filters?: FiltersConfig; llmConfig?: Partial; 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, diff --git a/packages/data-schemas/src/methods/memory.spec.ts b/packages/data-schemas/src/methods/memory.spec.ts index 58eea80298..ec1f3d4362 100644 --- a/packages/data-schemas/src/methods/memory.spec.ts +++ b/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'); }); diff --git a/packages/data-schemas/src/methods/memory.ts b/packages/data-schemas/src/methods/memory.ts index 058a1e54de..ebff3a2160 100644 --- a/packages/data-schemas/src/methods/memory.ts +++ b/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(), + }; } 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() }, + ); 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(), + }; } } diff --git a/packages/data-schemas/src/types/memory.ts b/packages/data-schemas/src/types/memory.ts index 1629b16854..a8ccce8fe4 100644 --- a/packages/data-schemas/src/types/memory.ts +++ b/packages/data-schemas/src/types/memory.ts @@ -76,4 +76,5 @@ export interface FormattedMemoriesResult { withKeys: string; withoutKeys: string; totalTokens?: number; + tokenCountsByKey?: Map; }