feat(cody-chat-memory): Enhance chat memory management

This PR introduces the following improvements to the CodyChatMemory module:

- Transition to a static utility class with a Map-based storage approach
- Implement timestamp-based memory organization and automatic trimming to maintain the 10 most recent items
- Enhance the retrieve() method to return a formatted ContextItem with the memory content
- Add support for REMOVE command to clear the chat memory
- Introduce comprehensive unit tests to cover various chat session workflows

## Test plan

- Run the unit tests for the CodyChatMemory module
- Verify the chat memory functionality in the Cody extension, including adding, retrieving, and clearing memories

## Changelog

- Improved chat memory management with timestamp-based organization and automatic trimming
- Added support for REMOVE command to clear the chat memory
- Enhanced unit test coverage for chat memory workflows

vscode/src/chat/agentic/CodyChatMemory.test.ts

@@ -0,0 +1,157 @@
+import { ContextItemSource } from '@sourcegraph/cody-shared'
+import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import { URI } from 'vscode-uri'
+import { localStorage } from '../../services/LocalStorageProvider'
+import { CodyChatMemory } from './CodyChatMemory'
+
+// Mock localStorage
+vi.mock('../../services/LocalStorageProvider', () => ({
+    localStorage: {
+        getChatMemory: vi.fn(),
+        setChatMemory: vi.fn(),
+    },
+}))
+
+describe('CodyChatMemory Workflows', () => {
+    beforeEach(() => {
+        vi.clearAllMocks()
+        vi.useFakeTimers()
+    })
+
+    afterAll(() => {
+        vi.useRealTimers()
+    })
+
+    describe('Chat Session Workflows', () => {
+        interface TestScenario {
+            name: string
+            actions: Array<{
+                type: 'initialize' | 'load' | 'retrieve' | 'unload'
+                input?: string
+                expectedContent?: string | null
+                expectedStorageCall?: boolean
+            }>
+        }
+
+        const scenarios: TestScenario[] = [
+            {
+                name: 'New user first chat session',
+                actions: [
+                    {
+                        type: 'initialize',
+                        expectedContent: null,
+                        expectedStorageCall: true,
+                    },
+                    {
+                        type: 'load',
+                        input: 'User prefers TypeScript',
+                        // Update to match new timestamp + content format
+                        expectedContent:
+                            '## \\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}.\\d{3}Z\\nUser prefers TypeScript',
+                        expectedStorageCall: true,
+                    },
+                    {
+                        type: 'retrieve',
+                        expectedContent: 'User prefers TypeScript',
+                    },
+                ],
+            },
+            {
+                name: 'Multiple chat interactions in one session',
+                actions: [
+                    {
+                        type: 'load',
+                        input: 'User likes unit testing',
+                        expectedContent: 'User likes unit testing',
+                    },
+                    {
+                        type: 'load',
+                        input: 'User works on VS Code extensions',
+                        expectedContent: 'User works on VS Code extensions',
+                    },
+                    {
+                        type: 'retrieve',
+                        // Update regex to match new Map-based format with timestamps
+                        expectedContent:
+                            '## \\d{4}.*User likes unit testing.*## \\d{4}.*User works on VS Code extensions',
+                    },
+                ],
+            },
+            {
+                name: 'Memory capacity management with timestamps',
+                actions: [
+                    ...Array.from({ length: 10 }, (_, i) => ({
+                        type: 'load' as const,
+                        input: `Memory item ${i}`,
+                        // Verify only last 8 items are kept
+                        expectedContent: i >= 2 ? `Memory item ${i}` : null,
+                    })),
+                    {
+                        type: 'retrieve',
+                        // Verify chronological order is maintained
+                        expectedContent: 'Memory item 2.*Memory item 9',
+                    },
+                ],
+            },
+            // Add new test scenario for timestamp ordering
+            {
+                name: 'Timestamp ordering verification',
+                actions: [
+                    {
+                        type: 'load',
+                        input: 'First message',
+                    },
+                    {
+                        type: 'load',
+                        input: 'Second message',
+                    },
+                    {
+                        type: 'retrieve',
+                        // Verify messages appear in chronological order with timestamps
+                        expectedContent: '.*First message.*Second message',
+                    },
+                ],
+            },
+        ]
+
+        for (const scenario of scenarios) {
+            it(scenario.name, () => {
+                for (const action of scenario.actions) {
+                    switch (action.type) {
+                        case 'load':
+                            // Advance by 1 second to ensure unique timestamps
+                            vi.advanceTimersByTime(1000)
+                            CodyChatMemory.load(action.input!)
+                            if (action.expectedStorageCall) {
+                                expect(localStorage.setChatMemory).toHaveBeenCalled()
+                            }
+                            break
+
+                        case 'retrieve': {
+                            const retrieved = CodyChatMemory.retrieve()
+                            if (action.expectedContent === null) {
+                                expect(retrieved).toBeUndefined()
+                            } else {
+                                if (action.expectedContent) {
+                                    expect(retrieved?.content).toMatch(
+                                        new RegExp(action.expectedContent, 's')
+                                    )
+                                }
+                                expect(retrieved?.source).toBe(ContextItemSource.Agentic)
+                                expect(retrieved?.uri).toEqual(URI.file('Cody Memory'))
+                            }
+                            break
+                        }
+                        case 'unload': {
+                            const lastState = CodyChatMemory.reset()
+                            if (action.expectedContent) {
+                                expect(lastState?.content).toContain(action.expectedContent)
+                            }
+                            break
+                        }
+                    }
+                }
+            })
+        }
+    })
+})

vscode/src/chat/agentic/CodyChatMemory.ts

@@ -3,65 +3,89 @@ import { URI } from 'vscode-uri'
 import { localStorage } from '../../services/LocalStorageProvider'
 
 /**
- * CodyChatMemory is a singleton class that manages short-term memory storage for chat conversations.
- * It maintains a maximum of 8 most recent memory items in a static Store.
- * We store the memory items in local storage to persist them across sessions.
- * NOTE: The memory items set to a maximum of 8 to avoid overloading the local storage.
+ * CodyChatMemory is a static utility class that manages persistent chat memory storage.
+ * It maintains the 10 most recent memory items using a static Map and localStorage for persistence.
  *
- * @remarks
- * This class should never be instantiated directly. All operations should be performed
- * through static methods. The only instance creation happens internally during initialization.
+ * The class handles:
+ * - Memory persistence across sessions via localStorage
+ * - Automatic trimming to maintain the 10 most recent items
+ * - Timestamp-based memory organization
+ * - Context retrieval for chat interactions
  *
- * Key features:
- * - Maintains a static Set of up to 8 chat memory items
- * - Persists memory items to local storage
- * - Provides memory retrieval as ContextItem for chat context
+ * Key Features:
+ * - Static interface - all operations performed through static methods
+ * - Automatic initialization on first use
+ * - Memory items formatted with timestamps
+ * - Integration with chat context system via ContextItem format
  *
  * Usage:
- * - Call CodyChatMemory.initialize() once at startup
- * - Use static methods load(), retrieve(), and unload() for memory operations
+ * - CodyChatMemory.initialize() - Called at startup to load existing memories
+ * - CodyChatMemory.load(memory) - Add new memory
+ * - CodyChatMemory.retrieve() - Get memories as chat context
+ * - CodyChatMemory.reset() - Clear and return last state
  */
 export class CodyChatMemory {
-    private static readonly MAX_MEMORY_ITEMS = 8
-    private static Store = new Set<string>([])
+    private static readonly MAX_MEMORY_ITEMS = 10
+    private static Store = new Map<string, string>()
 
     public static initialize(): void {
         if (CodyChatMemory.Store.size === 0) {
             const newMemory = new CodyChatMemory()
-            CodyChatMemory.Store = new Set(newMemory.getChatMemory())
+            const memories = newMemory.getChatMemory()
+            CodyChatMemory.Store = new Map(
+                memories.map(memory => {
+                    const [timestamp, ...content] = memory.split('\n')
+                    return [timestamp.replace('## ', ''), content.join('\n')]
+                })
+            )
         }
     }
 
     public static load(memory: string): void {
-        CodyChatMemory.Store.add(memory)
-        // If store exceeds the max, remove oldest items
-        if (CodyChatMemory.Store.size > CodyChatMemory.MAX_MEMORY_ITEMS) {
-            const storeArray = Array.from(CodyChatMemory.Store)
-            CodyChatMemory.Store = new Set(storeArray.slice(-5))
-        }
-        // TODO - persist to local file system
-        localStorage?.setChatMemory(Array.from(CodyChatMemory.Store))
+        const timestamp = new Date().toISOString()
+        CodyChatMemory.Store.set(timestamp, memory)
+        // Convert existing entries to array for manipulation
+        const entries = Array.from(CodyChatMemory.Store.entries())
+        // Keep only the most recent MAX_MEMORY_ITEMS entries &
+        // update stores with trimmed entries
+        const trimmedEntries = entries.slice(-CodyChatMemory.MAX_MEMORY_ITEMS)
+        CodyChatMemory.Store = new Map(trimmedEntries)
+        localStorage?.setChatMemory(
+            Array.from(trimmedEntries.entries()).map(([ts, mem]) => `## ${ts}\n${mem}`)
+        )
     }
 
     public static retrieve(): ContextItem | undefined {
         return CodyChatMemory.Store.size > 0
             ? {
                   type: 'file',
-                  content: '# Chat Memory\n' + Array.from(CodyChatMemory.Store).reverse().join('\n- '),
+                  content: populateMemoryContent(CodyChatMemory.Store),
                   uri: URI.file('Cody Memory'),
                   source: ContextItemSource.Agentic,
                   title: 'Cody Chat Memory',
               }
             : undefined
     }
 
-    public static unload(): ContextItem | undefined {
+    public static reset(): ContextItem | undefined {
         const stored = CodyChatMemory.retrieve()
-        CodyChatMemory.Store = new Set<string>()
+        CodyChatMemory.Store.clear()
         return stored
     }
 
     private getChatMemory(): string[] {
         return localStorage?.getChatMemory() || []
     }
 }
+
+export const CHAT_MEMORY_CONTEXT_TEMPLATE = `# Chat Memory
+Here are the notes you made about the user (me) from previous chat:
+{memoryItems}`
+
+function populateMemoryContent(memoryMap: Map<string, string>): string {
+    const memories = Array.from(memoryMap.entries())
+        .map(([timestamp, content]) => `\n## ${timestamp}\n${content}`)
+        .join('')
+
+    return CHAT_MEMORY_CONTEXT_TEMPLATE.replace('{memoryItems}', memories)
+}

vscode/src/chat/agentic/CodyTool.ts

@@ -283,12 +283,16 @@ class MemoryTool extends CodyTool {
         const newMemories = this.parse()
         for (const memory of newMemories) {
             if (memory === 'FORGET') {
-                CodyChatMemory.unload()
+                CodyChatMemory.reset()
                 return
             }
             if (memory === 'GET') {
                 return
             }
+            if (memory.startsWith('REMOVE')) {
+                CodyChatMemory.reset()
+                return
+            }
             CodyChatMemory.load(memory)
             logDebug('Cody Memory', 'added', { verbose: memory })
         }