From 4754c981aa049242c3891bd6364d2bb5a575c312 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 21 Jun 2026 08:37:33 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=A0=20fix:=20Invalidate=20request=20me?= =?UTF-8?q?mory=20cache=20after=20inline=20writes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inline set_memory/delete_memory now invalidate the request-scoped getFormattedMemories cache on a successful write, so a later tool round in the same response is seeded with the post-write usage total instead of the stale pre-write one (multi-round writes no longer collectively exceed tokenLimit, and a set after a delete is not over-counted). The within-round sharing across multiple memory-enabled agents is preserved. --- packages/api/src/agents/memory.spec.ts | 61 +++++++++++++++++++++++++- packages/api/src/agents/memory.ts | 24 +++++++++- 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts index 312988ff15..f1a8f33e38 100644 --- a/packages/api/src/agents/memory.spec.ts +++ b/packages/api/src/agents/memory.spec.ts @@ -2,7 +2,13 @@ import { Types } from 'mongoose'; import { Run, Providers } from '@librechat/agents'; import type { IUser } from '@librechat/data-schemas'; import type { Response } from 'express'; -import { processMemory, createMemoryTool } from './memory'; +import { + processMemory, + createMemoryTool, + createDeleteMemoryTool, + getRequestMemories, + invalidateRequestMemories, +} from './memory'; jest.mock('~/stream/GenerationJobManager'); @@ -622,4 +628,57 @@ describe('createMemoryTool tokenLimit enforcement', () => { 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({ + userId: 'user-1', + setMemory: jest.fn().mockResolvedValue({ ok: true }), + onWrite, + }); + await okTool.invoke({ key: 'k1', value: 'a fact' }); + expect(onWrite).toHaveBeenCalledTimes(1); + + onWrite.mockClear(); + const failTool = createMemoryTool({ + userId: 'user-1', + setMemory: jest.fn().mockResolvedValue({ ok: false }), + onWrite, + }); + await failTool.invoke({ key: 'k1', value: 'a fact' }); + expect(onWrite).not.toHaveBeenCalled(); + }); + + it('fires onWrite after a successful delete', async () => { + const onWrite = jest.fn(); + const tool = createDeleteMemoryTool({ + userId: 'user-1', + deleteMemory: jest.fn().mockResolvedValue({ ok: true }), + onWrite, + }); + + await tool.invoke({ key: 'k1' }); + + expect(onWrite).toHaveBeenCalledTimes(1); + }); +}); + +describe('getRequestMemories caching', () => { + it('memoizes per request, then re-fetches after invalidation', async () => { + const getFormattedMemories = jest + .fn() + .mockResolvedValue({ withKeys: '', withoutKeys: '', totalTokens: 10 }); + const req = {}; + + await getRequestMemories({ req, userId: 'user-1', getFormattedMemories }); + await getRequestMemories({ req, userId: 'user-1', getFormattedMemories }); + /** A second memory-enabled agent in the same run reuses the first fetch. */ + expect(getFormattedMemories).toHaveBeenCalledTimes(1); + + /** A successful inline write invalidates the cache so a later tool round in + * the same response re-reads the post-write usage total. */ + invalidateRequestMemories(req); + await getRequestMemories({ req, userId: 'user-1', getFormattedMemories }); + expect(getFormattedMemories).toHaveBeenCalledTimes(2); + }); }); diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index 91f53e6170..aa53907d5a 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -122,6 +122,7 @@ export const createMemoryTool = ({ charLimit, tokenLimit, totalTokens = 0, + onWrite, }: { userId: string | ObjectId; setMemory: MemoryMethods['setMemory']; @@ -129,6 +130,7 @@ export const createMemoryTool = ({ charLimit?: number; tokenLimit?: number; totalTokens?: number; + onWrite?: () => void; }): DynamicStructuredTool => { /** Running token total, advanced after each successful write. Writes are * serialized through `writeChain` so multiple `set_memory` calls in one @@ -227,6 +229,7 @@ export const createMemoryTool = ({ currentTotalTokens = newTotalTokens; writtenTokensByKey.set(key, tokenCount); } + onWrite?.(); logger.debug(`Memory set for key "${key}" (${tokenCount} tokens) for user "${userId}"`); return [`Memory set for key "${key}" (${tokenCount} tokens)`, artifact]; } @@ -273,10 +276,12 @@ export const createDeleteMemoryTool = ({ userId, deleteMemory, validKeys, + onWrite, }: { userId: string | ObjectId; deleteMemory: MemoryMethods['deleteMemory']; validKeys?: string[]; + onWrite?: () => void; }): DynamicStructuredTool => { return tool( async ({ key }) => { @@ -299,6 +304,7 @@ export const createDeleteMemoryTool = ({ const result = await deleteMemory({ userId, key }); if (result.ok) { + onWrite?.(); logger.debug(`Memory deleted for key "${key}" for user "${userId}"`); return [`Memory deleted for key "${key}"`, artifact]; } @@ -474,6 +480,16 @@ export function getRequestMemories({ return cached; } +/** + * Drops the cached memories for a request so the next {@link getRequestMemories} + * re-fetches. Inline `set_memory`/`delete_memory` writes call this on success so + * a later tool round in the same response is seeded with the post-write usage + * total instead of a stale pre-write one. + */ +export function invalidateRequestMemories(req: object): void { + requestMemoriesCache.delete(req); +} + /** * Re-checks the run-level memory gate at tool-execution time: the agents * `memory` capability is enabled, memory is configured, the user hasn't opted @@ -553,7 +569,12 @@ export async function buildInlineMemoryTool({ if (!allowed) { return null; } - return createDeleteMemoryTool({ userId, deleteMemory: memoryMethods.deleteMemory, validKeys }); + return createDeleteMemoryTool({ + userId, + deleteMemory: memoryMethods.deleteMemory, + validKeys, + onWrite: () => invalidateRequestMemories(req), + }); } const allowed = await isMemoryToolAllowed({ @@ -591,6 +612,7 @@ export async function buildInlineMemoryTool({ charLimit, tokenLimit, totalTokens, + onWrite: () => invalidateRequestMemories(req), }); }