🧠 fix: Address Codex re-review on memory capability (round 3)

- Serialize set_memory writes and advance a running token total inside createMemoryTool, so parallel batched calls in one event-driven turn can't each pass the limit check against a stale total and collectively exceed memory.tokenLimit (packages/api/src/agents/memory.ts) + tests.
- Inject the keyed memory context (withKeys) instead of withoutKeys when the running agent has the inline memory capability, so delete_memory has a visible key to target (api/server/controllers/agents/client.js).

api/server/controllers/agents/client.js

@@ -595,10 +595,16 @@ class AgentClient extends BaseClient {
     const userId = this.options.req.user.id + '';
     this.processMemory = undefined;
 
+    /** Inline memory tools (the `memory` capability on this agent) let the main
+     *  agent read/write memory itself, so it needs the keyed memory context —
+     *  otherwise `delete_memory` has no visible key to target. */
+    const agentHasMemoryTools =
+      this.options.agent?.tools?.includes(AgentCapabilities.memory) === true;
+
     if (!isMemoryAgentEnabled(memoryConfig)) {
       try {
-        const { withoutKeys } = await db.getFormattedMemories({ userId });
-        return withoutKeys;
+        const { withKeys, withoutKeys } = await db.getFormattedMemories({ userId });
+        return agentHasMemoryTools ? withKeys : withoutKeys;
       } catch (error) {
         logger.error(
           '[api/server/controllers/agents/client.js #useMemory] Error loading memories',
@@ -715,6 +721,17 @@ class AgentClient extends BaseClient {
     });
 
     this.processMemory = processMemory;
+    if (agentHasMemoryTools) {
+      try {
+        const { withKeys } = await db.getFormattedMemories({ userId });
+        return withKeys;
+      } catch (error) {
+        logger.error(
+          '[api/server/controllers/agents/client.js #useMemory] Error loading keyed memories',
+          error,
+        );
+      }
+    }
     return withoutKeys;
   }
 

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

@@ -2,7 +2,7 @@ import { Types } from 'mongoose';
 import { Run, Providers } from '@librechat/agents';
 import type { IUser } from '@librechat/data-schemas';
 import type { Response } from 'express';
-import { processMemory } from './memory';
+import { processMemory, createMemoryTool } from './memory';
 
 jest.mock('~/stream/GenerationJobManager');
 
@@ -560,3 +560,39 @@ describe('Memory Agent Header Resolution', () => {
     expect(runConfig.graphConfig.llmConfig.temperature).toBe(0.7);
   });
 });
+
+describe('createMemoryTool tokenLimit enforcement', () => {
+  it('serializes parallel set_memory calls so they cannot collectively exceed tokenLimit', async () => {
+    const setMemory = jest.fn().mockResolvedValue({ ok: true });
+    /** ~100 tokens; two of these (≈200) exceed the 150 limit, but each fits alone. */
+    const value = 'word '.repeat(100).trim();
+    const tool = createMemoryTool({
+      userId: 'user-1',
+      setMemory,
+      tokenLimit: 150,
+      totalTokens: 0,
+    });
+
+    await Promise.all([tool.invoke({ key: 'k1', value }), tool.invoke({ key: 'k2', value })]);
+
+    /** Only the first write is committed; the second is rejected against the
+     *  updated running total instead of the stale construction-time total. */
+    expect(setMemory).toHaveBeenCalledTimes(1);
+  });
+
+  it('allows sequential writes that each fit within the remaining capacity', async () => {
+    const setMemory = jest.fn().mockResolvedValue({ ok: true });
+    const value = 'word '.repeat(10).trim();
+    const tool = createMemoryTool({
+      userId: 'user-1',
+      setMemory,
+      tokenLimit: 1000,
+      totalTokens: 0,
+    });
+
+    await tool.invoke({ key: 'k1', value });
+    await tool.invoke({ key: 'k2', value });
+
+    expect(setMemory).toHaveBeenCalledTimes(2);
+  });
+});

packages/api/src/agents/memory.ts

@@ -93,6 +93,8 @@ ${tokenLimit ? `\nTOKEN LIMIT: Maximum ${tokenLimit} tokens per memory value.` :
 
 When in doubt, and the user hasn't asked to remember or forget anything, END THE TURN IMMEDIATELY.`;
 
+type MemoryArtifactRecord = Record<Tools.memory, MemoryArtifact>;
+
 /**
  * Creates a memory tool instance with user context
  */
@@ -109,82 +111,98 @@ export const createMemoryTool = ({
   tokenLimit?: number;
   totalTokens?: number;
 }): DynamicStructuredTool => {
-  const remainingTokens = tokenLimit ? tokenLimit - totalTokens : Infinity;
-  const isOverflowing = tokenLimit ? remainingTokens <= 0 : false;
+  /** Running token total, advanced after each successful write. Writes are
+   *  serialized through `writeChain` so multiple `set_memory` calls in one
+   *  event-driven batch (executed in parallel) can't each pass the limit
+   *  check against the same stale total and collectively exceed `tokenLimit`. */
+  let currentTotalTokens = totalTokens;
+  let writeChain: Promise<unknown> = Promise.resolve();
 
   return tool(
     async ({ key, value }) => {
-      try {
-        if (validKeys && validKeys.length > 0 && !validKeys.includes(key)) {
-          logger.warn(
-            `Memory Agent failed to set memory: Invalid key "${key}". Must be one of: ${validKeys.join(
-              ', ',
-            )}`,
-          );
-          return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
-        }
+      const run = async (): Promise<[string, MemoryArtifactRecord?]> => {
+        try {
+          if (validKeys && validKeys.length > 0 && !validKeys.includes(key)) {
+            logger.warn(
+              `Memory Agent failed to set memory: Invalid key "${key}". Must be one of: ${validKeys.join(
+                ', ',
+              )}`,
+            );
+            return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
+          }
 
-        const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base');
+          const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base');
+          const remainingTokens = tokenLimit ? tokenLimit - currentTotalTokens : Infinity;
 
-        if (isOverflowing) {
-          const errorArtifact: Record<Tools.memory, MemoryArtifact> = {
-            [Tools.memory]: {
-              key: 'system',
-              type: 'error',
-              value: JSON.stringify({
-                errorType: 'already_exceeded',
-                tokenCount: Math.abs(remainingTokens),
-                totalTokens: totalTokens,
-                tokenLimit: tokenLimit!,
-              }),
-              tokenCount: totalTokens,
-            },
-          };
-          return [`Memory storage exceeded. Cannot save new memories.`, errorArtifact];
-        }
-
-        if (tokenLimit) {
-          const newTotalTokens = totalTokens + tokenCount;
-          const newRemainingTokens = tokenLimit - newTotalTokens;
-
-          if (newRemainingTokens < 0) {
-            const errorArtifact: Record<Tools.memory, MemoryArtifact> = {
+          if (tokenLimit && remainingTokens <= 0) {
+            const errorArtifact: MemoryArtifactRecord = {
               [Tools.memory]: {
                 key: 'system',
                 type: 'error',
                 value: JSON.stringify({
-                  errorType: 'would_exceed',
-                  tokenCount: Math.abs(newRemainingTokens),
-                  totalTokens: newTotalTokens,
-                  tokenLimit,
+                  errorType: 'already_exceeded',
+                  tokenCount: Math.abs(remainingTokens),
+                  totalTokens: currentTotalTokens,
+                  tokenLimit: tokenLimit!,
                 }),
-                tokenCount: totalTokens,
+                tokenCount: currentTotalTokens,
               },
             };
-            return [`Memory storage would exceed limit. Cannot save this memory.`, errorArtifact];
+            return [`Memory storage exceeded. Cannot save new memories.`, errorArtifact];
           }
-        }
 
-        const artifact: Record<Tools.memory, MemoryArtifact> = {
-          [Tools.memory]: {
-            key,
-            value,
-            tokenCount,
-            type: 'update',
-          },
-        };
+          if (tokenLimit) {
+            const newTotalTokens = currentTotalTokens + tokenCount;
+            const newRemainingTokens = tokenLimit - newTotalTokens;
 
-        const result = await setMemory({ userId, key, value, tokenCount });
-        if (result.ok) {
-          logger.debug(`Memory set for key "${key}" (${tokenCount} tokens) for user "${userId}"`);
-          return [`Memory set for key "${key}" (${tokenCount} tokens)`, artifact];
+            if (newRemainingTokens < 0) {
+              const errorArtifact: MemoryArtifactRecord = {
+                [Tools.memory]: {
+                  key: 'system',
+                  type: 'error',
+                  value: JSON.stringify({
+                    errorType: 'would_exceed',
+                    tokenCount: Math.abs(newRemainingTokens),
+                    totalTokens: newTotalTokens,
+                    tokenLimit,
+                  }),
+                  tokenCount: currentTotalTokens,
+                },
+              };
+              return [`Memory storage would exceed limit. Cannot save this memory.`, errorArtifact];
+            }
+          }
+
+          const artifact: MemoryArtifactRecord = {
+            [Tools.memory]: {
+              key,
+              value,
+              tokenCount,
+              type: 'update',
+            },
+          };
+
+          const result = await setMemory({ userId, key, value, tokenCount });
+          if (result.ok) {
+            if (tokenLimit) {
+              currentTotalTokens += tokenCount;
+            }
+            logger.debug(`Memory set for key "${key}" (${tokenCount} tokens) for user "${userId}"`);
+            return [`Memory set for key "${key}" (${tokenCount} tokens)`, artifact];
+          }
+          logger.warn(`Failed to set memory for key "${key}" for user "${userId}"`);
+          return [`Failed to set memory for key "${key}"`, undefined];
+        } catch (error) {
+          logger.error('Memory Agent failed to set memory', error);
+          return [`Error setting memory for key "${key}"`, undefined];
         }
-        logger.warn(`Failed to set memory for key "${key}" for user "${userId}"`);
-        return [`Failed to set memory for key "${key}"`, undefined];
-      } catch (error) {
-        logger.error('Memory Agent failed to set memory', error);
-        return [`Error setting memory for key "${key}"`, undefined];
-      }
+      };
+
+      const resultPromise = writeChain.then(run, run);
+      /** Keep the chain alive (and non-rejecting) so the next queued call still
+       *  runs even if a prior one threw; `run` already resolves on every path. */
+      writeChain = resultPromise.catch(() => undefined);
+      return resultPromise;
     },
     {
       name: SET_MEMORY_TOOL_NAME,