🧠 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();
+    });
+  });
+});