🧠 feat: Enforce Token Limit for Memory Usage (#8401)

packages/api/src/agents/__tests__/memory.test.ts

@@ -0,0 +1,165 @@
+import { Tools, type MemoryArtifact } from 'librechat-data-provider';
+import { createMemoryTool } from '../memory';
+
+// Mock the logger
+jest.mock('winston', () => ({
+  createLogger: jest.fn(() => ({
+    debug: jest.fn(),
+    warn: jest.fn(),
+    error: jest.fn(),
+  })),
+  format: {
+    combine: jest.fn(),
+    colorize: jest.fn(),
+    simple: jest.fn(),
+  },
+  transports: {
+    Console: jest.fn(),
+  },
+}));
+
+// Mock the Tokenizer
+jest.mock('~/utils', () => ({
+  Tokenizer: {
+    getTokenCount: jest.fn((text: string) => text.length), // Simple mock: 1 char = 1 token
+  },
+}));
+
+describe('createMemoryTool', () => {
+  let mockSetMemory: jest.Mock;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockSetMemory = jest.fn().mockResolvedValue({ ok: true });
+  });
+
+  // Memory overflow tests
+  describe('overflow handling', () => {
+    it('should return error artifact when memory is already overflowing', async () => {
+      const tool = createMemoryTool({
+        userId: 'test-user',
+        setMemory: mockSetMemory,
+        tokenLimit: 100,
+        totalTokens: 150, // Already over limit
+      });
+
+      // Call the underlying function directly since invoke() doesn't handle responseFormat in tests
+      const result = await tool.func({ key: 'test', value: 'new memory' });
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBe('Memory storage exceeded. Cannot save new memories.');
+
+      const artifacts = result[1] as Record<Tools.memory, MemoryArtifact>;
+      expect(artifacts[Tools.memory]).toBeDefined();
+      expect(artifacts[Tools.memory].type).toBe('error');
+      expect(artifacts[Tools.memory].key).toBe('system');
+
+      const errorData = JSON.parse(artifacts[Tools.memory].value as string);
+      expect(errorData).toEqual({
+        errorType: 'already_exceeded',
+        tokenCount: 50,
+        totalTokens: 150,
+        tokenLimit: 100,
+      });
+
+      expect(mockSetMemory).not.toHaveBeenCalled();
+    });
+
+    it('should return error artifact when new memory would exceed limit', async () => {
+      const tool = createMemoryTool({
+        userId: 'test-user',
+        setMemory: mockSetMemory,
+        tokenLimit: 100,
+        totalTokens: 80,
+      });
+
+      // This would put us at 101 tokens total, exceeding the limit
+      const result = await tool.func({ key: 'test', value: 'This is a 20 char str' });
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBe('Memory storage would exceed limit. Cannot save this memory.');
+
+      const artifacts = result[1] as Record<Tools.memory, MemoryArtifact>;
+      expect(artifacts[Tools.memory]).toBeDefined();
+      expect(artifacts[Tools.memory].type).toBe('error');
+      expect(artifacts[Tools.memory].key).toBe('system');
+
+      const errorData = JSON.parse(artifacts[Tools.memory].value as string);
+      expect(errorData).toEqual({
+        errorType: 'would_exceed',
+        tokenCount: 1, // Math.abs(-1)
+        totalTokens: 101,
+        tokenLimit: 100,
+      });
+
+      expect(mockSetMemory).not.toHaveBeenCalled();
+    });
+
+    it('should successfully save memory when below limit', async () => {
+      const tool = createMemoryTool({
+        userId: 'test-user',
+        setMemory: mockSetMemory,
+        tokenLimit: 100,
+        totalTokens: 50,
+      });
+
+      const result = await tool.func({ key: 'test', value: 'small memory' });
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBe('Memory set for key "test" (12 tokens)');
+
+      const artifacts = result[1] as Record<Tools.memory, MemoryArtifact>;
+      expect(artifacts[Tools.memory]).toBeDefined();
+      expect(artifacts[Tools.memory].type).toBe('update');
+      expect(artifacts[Tools.memory].key).toBe('test');
+      expect(artifacts[Tools.memory].value).toBe('small memory');
+
+      expect(mockSetMemory).toHaveBeenCalledWith({
+        userId: 'test-user',
+        key: 'test',
+        value: 'small memory',
+        tokenCount: 12,
+      });
+    });
+  });
+
+  // Basic functionality tests
+  describe('basic functionality', () => {
+    it('should validate keys when validKeys is provided', async () => {
+      const tool = createMemoryTool({
+        userId: 'test-user',
+        setMemory: mockSetMemory,
+        validKeys: ['allowed', 'keys'],
+      });
+
+      const result = await tool.func({ key: 'invalid', value: 'some value' });
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBe('Invalid key "invalid". Must be one of: allowed, keys');
+      expect(result[1]).toBeUndefined();
+      expect(mockSetMemory).not.toHaveBeenCalled();
+    });
+
+    it('should handle setMemory failure', async () => {
+      mockSetMemory.mockResolvedValue({ ok: false });
+      const tool = createMemoryTool({
+        userId: 'test-user',
+        setMemory: mockSetMemory,
+      });
+
+      const result = await tool.func({ key: 'test', value: 'some value' });
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBe('Failed to set memory for key "test"');
+      expect(result[1]).toBeUndefined();
+    });
+
+    it('should handle exceptions', async () => {
+      mockSetMemory.mockRejectedValue(new Error('DB error'));
+      const tool = createMemoryTool({
+        userId: 'test-user',
+        setMemory: mockSetMemory,
+      });
+
+      const result = await tool.func({ key: 'test', value: 'some value' });
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBe('Error setting memory for key "test"');
+      expect(result[1]).toBeUndefined();
+    });
+  });
+});

packages/api/src/agents/memory.ts

@@ -71,7 +71,7 @@ const getDefaultInstructions = (
 /**
  * Creates a memory tool instance with user context
  */
-const createMemoryTool = ({
+export const createMemoryTool = ({
   userId,
   setMemory,
   validKeys,
@@ -84,6 +84,9 @@ const createMemoryTool = ({
   tokenLimit?: number;
   totalTokens?: number;
 }) => {
+  const remainingTokens = tokenLimit ? tokenLimit - totalTokens : Infinity;
+  const isOverflowing = tokenLimit ? remainingTokens <= 0 : false;
+
   return tool(
     async ({ key, value }) => {
       try {
@@ -93,24 +96,48 @@ const createMemoryTool = ({
               ', ',
             )}`,
           );
-          return `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');
 
-        if (tokenLimit && tokenCount > tokenLimit) {
-          logger.warn(
-            `Memory Agent failed to set memory: Value exceeds token limit. Value has ${tokenCount} tokens, but limit is ${tokenLimit}`,
-          );
-          return `Memory value too large: ${tokenCount} tokens exceeds limit of ${tokenLimit}`;
+        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 && totalTokens + tokenCount > tokenLimit) {
-          const remainingCapacity = tokenLimit - totalTokens;
-          logger.warn(
-            `Memory Agent failed to set memory: Would exceed total token limit. Current usage: ${totalTokens}, new memory: ${tokenCount} tokens, limit: ${tokenLimit}`,
-          );
-          return `Cannot add memory: would exceed token limit. Current usage: ${totalTokens}/${tokenLimit} tokens. This memory requires ${tokenCount} tokens, but only ${remainingCapacity} tokens available.`;
+        if (tokenLimit) {
+          const newTotalTokens = totalTokens + tokenCount;
+          const newRemainingTokens = tokenLimit - newTotalTokens;
+
+          if (newRemainingTokens < 0) {
+            const errorArtifact: Record<Tools.memory, MemoryArtifact> = {
+              [Tools.memory]: {
+                key: 'system',
+                type: 'error',
+                value: JSON.stringify({
+                  errorType: 'would_exceed',
+                  tokenCount: Math.abs(newRemainingTokens),
+                  totalTokens: newTotalTokens,
+                  tokenLimit,
+                }),
+                tokenCount: totalTokens,
+              },
+            };
+            return [`Memory storage would exceed limit. Cannot save this memory.`, errorArtifact];
+          }
         }
 
         const artifact: Record<Tools.memory, MemoryArtifact> = {
@@ -177,7 +204,7 @@ const createDeleteMemoryTool = ({
               ', ',
             )}`,
           );
-          return `Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`;
+          return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
         }
 
         const artifact: Record<Tools.memory, MemoryArtifact> = {
@@ -269,7 +296,13 @@ export async function processMemory({
   llmConfig?: Partial<LLMConfig>;
 }): Promise<(TAttachment | null)[] | undefined> {
   try {
-    const memoryTool = createMemoryTool({ userId, tokenLimit, setMemory, validKeys, totalTokens });
+    const memoryTool = createMemoryTool({
+      userId,
+      tokenLimit,
+      setMemory,
+      validKeys,
+      totalTokens,
+    });
     const deleteMemoryTool = createDeleteMemoryTool({
       userId,
       validKeys,