diff --git a/client/src/components/Chat/Messages/Content/MemoryArtifacts.tsx b/client/src/components/Chat/Messages/Content/MemoryArtifacts.tsx index 7af4e9fcdd..6209414221 100644 --- a/client/src/components/Chat/Messages/Content/MemoryArtifacts.tsx +++ b/client/src/components/Chat/Messages/Content/MemoryArtifacts.tsx @@ -13,14 +13,25 @@ export default function MemoryArtifacts({ attachments }: { attachments?: TAttach const [isAnimating, setIsAnimating] = useState(false); const prevShowInfoRef = useRef(showInfo); - const memoryArtifacts = useMemo(() => { + const { hasErrors, memoryArtifacts } = useMemo(() => { + let hasErrors = false; const result: MemoryArtifact[] = []; - for (const attachment of attachments ?? []) { + + if (!attachments || attachments.length === 0) { + return { hasErrors, memoryArtifacts: result }; + } + + for (const attachment of attachments) { if (attachment?.[Tools.memory] != null) { result.push(attachment[Tools.memory]); + + if (!hasErrors && attachment[Tools.memory].type === 'error') { + hasErrors = true; + } } } - return result; + + return { hasErrors, memoryArtifacts: result }; }, [attachments]); useLayoutEffect(() => { @@ -75,7 +86,12 @@ export default function MemoryArtifacts({ attachments }: { attachments?: TAttach
diff --git a/client/src/components/Chat/Messages/Content/MemoryInfo.tsx b/client/src/components/Chat/Messages/Content/MemoryInfo.tsx index 574c2e8f5f..d00fac68c6 100644 --- a/client/src/components/Chat/Messages/Content/MemoryInfo.tsx +++ b/client/src/components/Chat/Messages/Content/MemoryInfo.tsx @@ -1,17 +1,47 @@ import type { MemoryArtifact } from 'librechat-data-provider'; +import { useMemo } from 'react'; import { useLocalize } from '~/hooks'; export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: MemoryArtifact[] }) { const localize = useLocalize(); + + const { updatedMemories, deletedMemories, errorMessages } = useMemo(() => { + const updated = memoryArtifacts.filter((art) => art.type === 'update'); + const deleted = memoryArtifacts.filter((art) => art.type === 'delete'); + const errors = memoryArtifacts.filter((art) => art.type === 'error'); + + const messages = errors.map((artifact) => { + try { + const errorData = JSON.parse(artifact.value as string); + + if (errorData.errorType === 'already_exceeded') { + return localize('com_ui_memory_already_exceeded', { + tokens: errorData.tokenCount, + }); + } else if (errorData.errorType === 'would_exceed') { + return localize('com_ui_memory_would_exceed', { + tokens: errorData.tokenCount, + }); + } else { + return localize('com_ui_memory_error'); + } + } catch { + return localize('com_ui_memory_error'); + } + }); + + return { + updatedMemories: updated, + deletedMemories: deleted, + errorMessages: messages, + }; + }, [memoryArtifacts, localize]); + if (memoryArtifacts.length === 0) { return null; } - // Group artifacts by type - const updatedMemories = memoryArtifacts.filter((artifact) => artifact.type === 'update'); - const deletedMemories = memoryArtifacts.filter((artifact) => artifact.type === 'delete'); - - if (updatedMemories.length === 0 && deletedMemories.length === 0) { + if (updatedMemories.length === 0 && deletedMemories.length === 0 && errorMessages.length === 0) { return null; } @@ -23,8 +53,8 @@ export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: Memor {localize('com_ui_memory_updated_items')}
- {updatedMemories.map((artifact, index) => ( -
+ {updatedMemories.map((artifact) => ( +
{artifact.key}
@@ -43,8 +73,8 @@ export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: Memor {localize('com_ui_memory_deleted_items')}
- {deletedMemories.map((artifact, index) => ( -
+ {deletedMemories.map((artifact) => ( +
{artifact.key}
@@ -56,6 +86,24 @@ export default function MemoryInfo({ memoryArtifacts }: { memoryArtifacts: Memor
)} + + {errorMessages.length > 0 && ( +
+

+ {localize('com_ui_memory_storage_full')} +

+
+ {errorMessages.map((errorMessage) => ( +
+ {errorMessage} +
+ ))} +
+
+ )}
); } diff --git a/client/src/components/Chat/Messages/Content/__tests__/MemoryArtifacts.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/MemoryArtifacts.test.tsx new file mode 100644 index 0000000000..875750a54f --- /dev/null +++ b/client/src/components/Chat/Messages/Content/__tests__/MemoryArtifacts.test.tsx @@ -0,0 +1,197 @@ +import React from 'react'; +import { render, screen, fireEvent } from '@testing-library/react'; +import '@testing-library/jest-dom/extend-expect'; +import MemoryArtifacts from '../MemoryArtifacts'; +import type { TAttachment, MemoryArtifact } from 'librechat-data-provider'; +import { Tools } from 'librechat-data-provider'; + +// Mock the localize hook +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string) => { + const translations: Record = { + com_ui_memory_updated: 'Updated saved memory', + com_ui_memory_error: 'Memory Error', + }; + return translations[key] || key; + }, +})); + +// Mock the MemoryInfo component +jest.mock('../MemoryInfo', () => ({ + __esModule: true, + default: ({ memoryArtifacts }: { memoryArtifacts: MemoryArtifact[] }) => ( +
+ {memoryArtifacts.map((artifact, index) => ( +
+ {artifact.type}: {artifact.key} +
+ ))} +
+ ), +})); + +describe('MemoryArtifacts', () => { + const createMemoryAttachment = (type: 'update' | 'delete' | 'error', key: string): TAttachment => + ({ + type: Tools.memory, + [Tools.memory]: { + type, + key, + value: + type === 'error' + ? JSON.stringify({ errorType: 'exceeded', tokenCount: 100 }) + : 'test value', + } as MemoryArtifact, + }) as TAttachment; + + describe('Error State Handling', () => { + test('displays error styling when memory artifacts contain errors', () => { + const attachments = [ + createMemoryAttachment('error', 'system'), + createMemoryAttachment('update', 'memory1'), + ]; + + render(); + + const button = screen.getByRole('button'); + expect(button).toHaveClass('text-red-500'); + expect(button).toHaveClass('hover:text-red-600'); + expect(button).toHaveClass('dark:text-red-400'); + expect(button).toHaveClass('dark:hover:text-red-500'); + }); + + test('displays normal styling when no errors present', () => { + const attachments = [ + createMemoryAttachment('update', 'memory1'), + createMemoryAttachment('delete', 'memory2'), + ]; + + render(); + + const button = screen.getByRole('button'); + expect(button).toHaveClass('text-text-secondary-alt'); + expect(button).toHaveClass('hover:text-text-primary'); + expect(button).not.toHaveClass('text-red-500'); + }); + + test('displays error message when errors are present', () => { + const attachments = [createMemoryAttachment('error', 'system')]; + + render(); + + expect(screen.getByText('Memory Error')).toBeInTheDocument(); + expect(screen.queryByText('Updated saved memory')).not.toBeInTheDocument(); + }); + + test('displays normal message when no errors are present', () => { + const attachments = [createMemoryAttachment('update', 'memory1')]; + + render(); + + expect(screen.getByText('Updated saved memory')).toBeInTheDocument(); + expect(screen.queryByText('Memory Error')).not.toBeInTheDocument(); + }); + }); + + describe('Memory Artifacts Filtering', () => { + test('filters and passes only memory-type attachments to MemoryInfo', () => { + const attachments = [ + createMemoryAttachment('update', 'memory1'), + { type: 'file' } as TAttachment, // Non-memory attachment + createMemoryAttachment('error', 'system'), + ]; + + render(); + + // Click to expand + fireEvent.click(screen.getByRole('button')); + + // Check that only memory artifacts are passed to MemoryInfo + expect(screen.getByTestId('memory-artifact-update')).toBeInTheDocument(); + expect(screen.getByTestId('memory-artifact-error')).toBeInTheDocument(); + }); + + test('correctly identifies multiple error artifacts', () => { + const attachments = [ + createMemoryAttachment('error', 'system1'), + createMemoryAttachment('error', 'system2'), + createMemoryAttachment('update', 'memory1'), + ]; + + render(); + + const button = screen.getByRole('button'); + expect(button).toHaveClass('text-red-500'); + expect(screen.getByText('Memory Error')).toBeInTheDocument(); + }); + }); + + describe('Collapse/Expand Functionality', () => { + test('toggles memory info visibility on button click', () => { + const attachments = [createMemoryAttachment('update', 'memory1')]; + + render(); + + // Initially collapsed + expect(screen.queryByTestId('memory-info')).not.toBeInTheDocument(); + + // Click to expand + fireEvent.click(screen.getByRole('button')); + expect(screen.getByTestId('memory-info')).toBeInTheDocument(); + + // Click to collapse + fireEvent.click(screen.getByRole('button')); + expect(screen.queryByTestId('memory-info')).not.toBeInTheDocument(); + }); + + test('updates aria-expanded attribute correctly', () => { + const attachments = [createMemoryAttachment('update', 'memory1')]; + + render(); + + const button = screen.getByRole('button'); + expect(button).toHaveAttribute('aria-expanded', 'false'); + + fireEvent.click(button); + expect(button).toHaveAttribute('aria-expanded', 'true'); + }); + }); + + describe('Edge Cases', () => { + test('handles empty attachments array', () => { + render(); + expect(screen.queryByRole('button')).not.toBeInTheDocument(); + }); + + test('handles undefined attachments', () => { + render(); + expect(screen.queryByRole('button')).not.toBeInTheDocument(); + }); + + test('handles attachments with no memory artifacts', () => { + const attachments = [{ type: 'file' } as TAttachment, { type: 'image' } as TAttachment]; + + render(); + expect(screen.queryByRole('button')).not.toBeInTheDocument(); + }); + + test('handles malformed memory artifacts gracefully', () => { + const attachments = [ + { + type: Tools.memory, + [Tools.memory]: { + type: 'error', + key: 'system', + // Missing value + }, + } as TAttachment, + ]; + + render(); + + const button = screen.getByRole('button'); + expect(button).toHaveClass('text-red-500'); + expect(screen.getByText('Memory Error')).toBeInTheDocument(); + }); + }); +}); diff --git a/client/src/components/Chat/Messages/Content/__tests__/MemoryInfo.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/MemoryInfo.test.tsx new file mode 100644 index 0000000000..ba8247484c --- /dev/null +++ b/client/src/components/Chat/Messages/Content/__tests__/MemoryInfo.test.tsx @@ -0,0 +1,267 @@ +import React from 'react'; +import { render, screen } from '@testing-library/react'; +import '@testing-library/jest-dom/extend-expect'; +import MemoryInfo from '../MemoryInfo'; +import type { MemoryArtifact } from 'librechat-data-provider'; + +// Mock the localize hook +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string, params?: Record) => { + const translations: Record = { + com_ui_memory_updated_items: 'Updated Memories', + com_ui_memory_deleted_items: 'Deleted Memories', + com_ui_memory_already_exceeded: `Memory storage already full - exceeded by ${params?.tokens || 0} tokens. Delete existing memories before adding new ones.`, + com_ui_memory_would_exceed: `Cannot save - would exceed limit by ${params?.tokens || 0} tokens. Delete existing memories to make space.`, + com_ui_memory_deleted: 'This memory has been deleted', + com_ui_memory_storage_full: 'Memory Storage Full', + com_ui_memory_error: 'Memory Error', + com_ui_updated_successfully: 'Updated successfully', + com_ui_none_selected: 'None selected', + }; + return translations[key] || key; + }, +})); + +describe('MemoryInfo', () => { + const createMemoryArtifact = ( + type: 'update' | 'delete' | 'error', + key: string, + value?: string, + ): MemoryArtifact => ({ + type, + key, + value: value || `test value for ${key}`, + }); + + describe('Error Memory Display', () => { + test('displays error section when memory is already exceeded', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system', + value: JSON.stringify({ errorType: 'already_exceeded', tokenCount: 150 }), + }, + ]; + + render(); + + expect(screen.getByText('Memory Storage Full')).toBeInTheDocument(); + expect( + screen.getByText( + 'Memory storage already full - exceeded by 150 tokens. Delete existing memories before adding new ones.', + ), + ).toBeInTheDocument(); + }); + + test('displays error when memory would exceed limit', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system', + value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 50 }), + }, + ]; + + render(); + + expect(screen.getByText('Memory Storage Full')).toBeInTheDocument(); + expect( + screen.getByText( + 'Cannot save - would exceed limit by 50 tokens. Delete existing memories to make space.', + ), + ).toBeInTheDocument(); + }); + + test('displays multiple error messages', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system1', + value: JSON.stringify({ errorType: 'already_exceeded', tokenCount: 100 }), + }, + { + type: 'error', + key: 'system2', + value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 25 }), + }, + ]; + + render(); + + expect( + screen.getByText( + 'Memory storage already full - exceeded by 100 tokens. Delete existing memories before adding new ones.', + ), + ).toBeInTheDocument(); + expect( + screen.getByText( + 'Cannot save - would exceed limit by 25 tokens. Delete existing memories to make space.', + ), + ).toBeInTheDocument(); + }); + + test('applies correct styling to error messages', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system', + value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 50 }), + }, + ]; + + render(); + + const errorMessage = screen.getByText( + 'Cannot save - would exceed limit by 50 tokens. Delete existing memories to make space.', + ); + const errorContainer = errorMessage.closest('div'); + + expect(errorContainer).toHaveClass('rounded-md'); + expect(errorContainer).toHaveClass('bg-red-50'); + expect(errorContainer).toHaveClass('p-3'); + expect(errorContainer).toHaveClass('text-sm'); + expect(errorContainer).toHaveClass('text-red-800'); + expect(errorContainer).toHaveClass('dark:bg-red-900/20'); + expect(errorContainer).toHaveClass('dark:text-red-400'); + }); + }); + + describe('Mixed Memory Types', () => { + test('displays all sections when different memory types are present', () => { + const memoryArtifacts: MemoryArtifact[] = [ + createMemoryArtifact('update', 'memory1', 'Updated content'), + createMemoryArtifact('delete', 'memory2'), + { + type: 'error', + key: 'system', + value: JSON.stringify({ errorType: 'would_exceed', tokenCount: 200 }), + }, + ]; + + render(); + + // Check all sections are present + expect(screen.getByText('Updated Memories')).toBeInTheDocument(); + expect(screen.getByText('Deleted Memories')).toBeInTheDocument(); + expect(screen.getByText('Memory Storage Full')).toBeInTheDocument(); + + // Check content + expect(screen.getByText('memory1')).toBeInTheDocument(); + expect(screen.getByText('Updated content')).toBeInTheDocument(); + expect(screen.getByText('memory2')).toBeInTheDocument(); + expect( + screen.getByText( + 'Cannot save - would exceed limit by 200 tokens. Delete existing memories to make space.', + ), + ).toBeInTheDocument(); + }); + + test('only displays sections with content', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system', + value: JSON.stringify({ errorType: 'already_exceeded', tokenCount: 10 }), + }, + ]; + + render(); + + // Only error section should be present + expect(screen.getByText('Memory Storage Full')).toBeInTheDocument(); + expect(screen.queryByText('Updated Memories')).not.toBeInTheDocument(); + expect(screen.queryByText('Deleted Memories')).not.toBeInTheDocument(); + }); + }); + + describe('Edge Cases', () => { + test('handles empty memory artifacts array', () => { + const { container } = render(); + expect(container.firstChild).toBeNull(); + }); + + test('handles malformed error data gracefully', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system', + value: 'invalid json', + }, + ]; + + render(); + + // Should render generic error message + expect(screen.getByText('Memory Storage Full')).toBeInTheDocument(); + expect(screen.getByText('Memory Error')).toBeInTheDocument(); + }); + + test('handles missing value in error artifact', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system', + // value is undefined + } as MemoryArtifact, + ]; + + render(); + + expect(screen.getByText('Memory Storage Full')).toBeInTheDocument(); + expect(screen.getByText('Memory Error')).toBeInTheDocument(); + }); + + test('handles unknown errorType gracefully', () => { + const memoryArtifacts: MemoryArtifact[] = [ + { + type: 'error', + key: 'system', + value: JSON.stringify({ errorType: 'unknown_type', tokenCount: 30 }), + }, + ]; + + render(); + + // Should show generic error message for unknown types + expect(screen.getByText('Memory Storage Full')).toBeInTheDocument(); + expect(screen.getByText('Memory Error')).toBeInTheDocument(); + }); + + test('returns null when no memories of any type exist', () => { + const memoryArtifacts: MemoryArtifact[] = [{ type: 'unknown' as any, key: 'test' }]; + + const { container } = render(); + expect(container.firstChild).toBeNull(); + }); + }); + + describe('Update and Delete Memory Display', () => { + test('displays updated memories correctly', () => { + const memoryArtifacts: MemoryArtifact[] = [ + createMemoryArtifact('update', 'preferences', 'User prefers dark mode'), + createMemoryArtifact('update', 'location', 'Lives in San Francisco'), + ]; + + render(); + + expect(screen.getByText('Updated Memories')).toBeInTheDocument(); + expect(screen.getByText('preferences')).toBeInTheDocument(); + expect(screen.getByText('User prefers dark mode')).toBeInTheDocument(); + expect(screen.getByText('location')).toBeInTheDocument(); + expect(screen.getByText('Lives in San Francisco')).toBeInTheDocument(); + }); + + test('displays deleted memories correctly', () => { + const memoryArtifacts: MemoryArtifact[] = [ + createMemoryArtifact('delete', 'old_preference'), + createMemoryArtifact('delete', 'outdated_info'), + ]; + + render(); + + expect(screen.getByText('Deleted Memories')).toBeInTheDocument(); + expect(screen.getByText('old_preference')).toBeInTheDocument(); + expect(screen.getByText('outdated_info')).toBeInTheDocument(); + }); + }); +}); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 8de6c323df..fd271b9923 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -860,6 +860,10 @@ "com_ui_memory_key_validation": "Memory key must only contain lowercase letters and underscores.", "com_ui_memory_updated": "Updated saved memory", "com_ui_memory_updated_items": "Updated Memories", + "com_ui_memory_storage_full": "Memory Storage Full", + "com_ui_memory_error": "Memory Error", + "com_ui_memory_already_exceeded": "Memory storage already full - exceeded by {{tokens}} tokens. Delete existing memories before adding new ones.", + "com_ui_memory_would_exceed": "Cannot save - would exceed limit by {{tokens}} tokens. Delete existing memories to make space.", "com_ui_mention": "Mention an endpoint, assistant, or preset to quickly switch to it", "com_ui_min_tags": "Cannot remove more values, a minimum of {{0}} are required.", "com_ui_misc": "Misc.", diff --git a/client/test/setupTests.js b/client/test/setupTests.js index d0e7555f49..043c7cdf67 100644 --- a/client/test/setupTests.js +++ b/client/test/setupTests.js @@ -17,6 +17,9 @@ import '@testing-library/jest-dom/extend-expect'; // 'react-lottie' uses canvas import 'jest-canvas-mock'; +// Mock ResizeObserver +import './resizeObserver.mock'; + beforeEach(() => { jest.clearAllMocks(); }); @@ -40,4 +43,4 @@ jest.mock('react-i18next', () => { init: jest.fn(), }, }; -}); \ No newline at end of file +}); diff --git a/packages/api/jest.config.mjs b/packages/api/jest.config.mjs index a65d8d8af9..eb6be102de 100644 --- a/packages/api/jest.config.mjs +++ b/packages/api/jest.config.mjs @@ -1,6 +1,7 @@ export default { collectCoverageFrom: ['src/**/*.{js,jsx,ts,tsx}', '!/node_modules/'], coveragePathIgnorePatterns: ['/node_modules/', '/dist/'], + testPathIgnorePatterns: ['/node_modules/', '/dist/'], coverageReporters: ['text', 'cobertura'], testResultsProcessor: 'jest-junit', moduleNameMapper: { diff --git a/packages/api/src/agents/__tests__/memory.test.ts b/packages/api/src/agents/__tests__/memory.test.ts new file mode 100644 index 0000000000..21dda8b0e4 --- /dev/null +++ b/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; + 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; + 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; + 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(); + }); + }); +}); diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index c9e4d14063..3e59c1d7a9 100644 --- a/packages/api/src/agents/memory.ts +++ b/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]: { + 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]: { + 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 = { @@ -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 = { @@ -269,7 +296,13 @@ export async function processMemory({ llmConfig?: Partial; }): 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, diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index 3ea250f4a4..a1494d68f5 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -534,7 +534,7 @@ export type MemoryArtifact = { key: string; value?: string; tokenCount?: number; - type: 'update' | 'delete'; + type: 'update' | 'delete' | 'error'; }; export type TAttachmentMetadata = {