feat(amazonq): Use Chat History DB instead of in-memory store

Author: tsmithsz

packages/core/src/codewhispererChat/clients/chat/v0/chat.ts

@@ -16,16 +16,15 @@ import { createQDeveloperStreamingClient } from '../../../../shared/clients/qDev
 import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWrittenCodeTracker'
 import { PromptMessage } from '../../../controllers/chat/model'
 import { FsWriteBackup } from '../../../../codewhispererChat/tools/fsWrite'
+import { randomUUID } from '../../../../shared/crypto'
 
 export type ToolUseWithError = {
     toolUse: ToolUse
     error: Error | undefined
 }
-import { getLogger } from '../../../../shared/logger/logger'
-import { randomUUID } from '../../../../shared/crypto'
 
 export class ChatSession {
-    private sessionId?: string
+    private sessionId: string
     /**
      * _readFiles = list of files read from the project to gather context before generating response.
      * _showDiffOnFileWrite = Controls whether to show diff view (true) or file context view (false) to the user
@@ -46,7 +45,7 @@ export class ChatSession {
     // TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
     // e.g. root_a/file1 vs root_b/file1
     relativePathToWorkspaceRoot: Map<string, string> = new Map()
-    public get sessionIdentifier(): string | undefined {
+    public get sessionIdentifier(): string {
         return this.sessionId
     }
 
@@ -86,13 +85,14 @@ export class ChatSession {
 
     constructor() {
         this.createNewTokenSource()
+        this.sessionId = randomUUID()
     }
 
     createNewTokenSource() {
         this.tokenSource = new vscode.CancellationTokenSource()
     }
 
-    public setSessionID(id?: string) {
+    public setSessionID(id: string) {
         this.sessionId = id
     }
     public get readFiles(): string[] {
@@ -120,14 +120,6 @@ export class ChatSession {
             )
         }
 
-        const responseStream = response.sendMessageResponse
-        for await (const event of responseStream) {
-            if ('messageMetadataEvent' in event) {
-                this.sessionId = event.messageMetadataEvent?.conversationId
-                break
-            }
-        }
-
         UserWrittenCodeTracker.instance.onQFeatureInvoked()
         return response
     }
@@ -142,12 +134,6 @@ export class ChatSession {
             )
         }
 
-        this.sessionId = response.conversationId
-        if (this.sessionId?.length === 0) {
-            getLogger().debug(`Session ID: ${this.sessionId} is empty. Generating random UUID`)
-            this.sessionId = randomUUID()
-        }
-
         UserWrittenCodeTracker.instance.onQFeatureInvoked()
 
         return response

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -96,10 +96,10 @@ import { amazonQTabSuffix } from '../../../shared/constants'
 import { OutputKind } from '../../tools/toolShared'
 import { ToolUtils, Tool, ToolType } from '../../tools/toolUtils'
 import { ChatStream } from '../../tools/chatStream'
-import { ChatHistoryStorage } from '../../storages/chatHistoryStorage'
 import { tempDirPath } from '../../../shared/filesystemUtilities'
 import { Database } from '../../../shared/db/chatDb/chatDb'
 import { TabBarController } from './tabBarController'
+import { messageToChatMessage } from '../../../shared/db/chatDb/util'
 
 export interface ChatControllerMessagePublishers {
     readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -174,7 +174,6 @@ export class ChatController {
     private readonly userIntentRecognizer: UserIntentRecognizer
     private readonly telemetryHelper: CWCTelemetryHelper
     private userPromptsWatcher: vscode.FileSystemWatcher | undefined
-    private readonly chatHistoryStorage: ChatHistoryStorage
     private chatHistoryDb = Database.getInstance()
 
     public constructor(
@@ -193,7 +192,6 @@ export class ChatController {
         this.editorContentController = new EditorContentController()
         this.promptGenerator = new PromptsGenerator()
         this.userIntentRecognizer = new UserIntentRecognizer()
-        this.chatHistoryStorage = new ChatHistoryStorage()
         this.tabBarController = new TabBarController(this.messenger)
 
         onDidChangeAmazonQVisibility((visible) => {
@@ -408,7 +406,7 @@ export class ChatController {
         const session = this.sessionStorage.getSession(message.tabID)
         session.tokenSource.cancel()
         this.messenger.sendEmptyMessage(message.tabID, '', undefined)
-        this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
+        this.chatHistoryDb.clearRecentHistory(message.tabID)
     }
 
     private async processTriggerTabIDReceived(message: TriggerTabIDReceived) {
@@ -465,7 +463,6 @@ export class ChatController {
 
     private async processTabCloseMessage(message: TabClosedMessage) {
         this.sessionStorage.deleteSession(message.tabID)
-        this.chatHistoryStorage.deleteHistory(message.tabID)
         this.triggerEventsStorage.removeTabEvents(message.tabID)
         // this.telemetryHelper.recordCloseChat(message.tabID)
         this.chatHistoryDb.updateTabOpenState(message.tabID, false)
@@ -1000,7 +997,7 @@ export class ChatController {
         getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
 
         this.sessionStorage.deleteSession(tabID)
-        this.chatHistoryStorage.getTabHistory(tabID).clearRecentHistory()
+        this.chatHistoryDb.clearRecentHistory(tabID)
     }
 
     private async processContextMenuCommand(command: EditorContextCommand) {
@@ -1121,7 +1118,6 @@ export class ChatController {
         switch (message.command) {
             case 'clear':
                 this.sessionStorage.deleteSession(message.tabID)
-                this.chatHistoryStorage.getTabHistory(message.tabID).clear()
                 this.triggerEventsStorage.removeTabEvents(message.tabID)
                 recordTelemetryChatRunCommand('clear')
                 this.chatHistoryDb.clearTab(message.tabID)
@@ -1193,8 +1189,6 @@ export class ChatController {
                     context,
                 })
 
-                this.messenger.sendAsyncEventProgress(message.tabID, true, '')
-
                 await this.generateResponse(
                     {
                         message: message.message ?? '',
@@ -1410,10 +1404,6 @@ export class ChatController {
         }
 
         const session = this.sessionStorage.getSession(tabID)
-        if (!session.localHistoryHydrated) {
-            triggerPayload.history = this.chatHistoryDb.getMessages(triggerEvent.tabID, 10)
-            session.localHistoryHydrated = true
-        }
         await this.resolveContextCommandPayload(triggerPayload, session)
         triggerPayload.useRelevantDocuments = triggerPayload.context.some(
             (context) => typeof context !== 'string' && context.command === '@workspace'
@@ -1452,16 +1442,14 @@ export class ChatController {
 
         const request = triggerPayloadToChatRequest(triggerPayload)
 
-        const chatHistory = this.chatHistoryStorage.getTabHistory(tabID)
         const currentMessage = request.conversationState.currentMessage
         if (currentMessage) {
-            chatHistory.fixHistory(currentMessage)
+            this.chatHistoryDb.fixHistory(tabID, currentMessage)
         }
-        request.conversationState.history = chatHistory.getHistory()
-
-        const conversationId = chatHistory.getConversationId() || randomUUID()
-        chatHistory.setConversationId(conversationId)
-        request.conversationState.conversationId = conversationId
+        request.conversationState.history = this.chatHistoryDb
+            .getMessages(tabID, 100)
+            .map((chat) => messageToChatMessage(chat))
+        request.conversationState.conversationId = session.sessionIdentifier
 
         triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
 
@@ -1501,6 +1489,7 @@ export class ChatController {
                 session.setContext(triggerPayload.context)
             }
             this.messenger.sendInitalStream(tabID, triggerID)
+            this.messenger.sendAsyncEventProgress(tabID, true, '')
             this.telemetryHelper.setConversationStreamStartTime(tabID)
             if (isSsoConnection(AuthUtil.instance.conn)) {
                 const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)
@@ -1519,7 +1508,6 @@ export class ChatController {
             this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
 
             if (currentMessage && session.sessionIdentifier) {
-                chatHistory.appendUserMessage(currentMessage)
                 this.chatHistoryDb.addMessage(tabID, 'cwc', session.sessionIdentifier, {
                     body: triggerPayload.message,
                     type: 'prompt' as any,
@@ -1534,7 +1522,7 @@ export class ChatController {
                     response.$metadata.requestId
                 } metadata: ${inspect(response.$metadata, { depth: 12 })}`
             )
-            await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload, chatHistory)
+            await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload)
 
             // Turn off AgentLoop flag after sending the AI response
             this.sessionStorage.setAgentLoopInProgress(tabID, false)

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -56,7 +56,6 @@ import {
 } from '@aws/mynah-ui'
 import { Database } from '../../../../shared/db/chatDb/chatDb'
 import { TabType } from '../../../../amazonq/webview/ui/storages/tabsStorage'
-import { ChatHistoryManager } from '../../../storages/chatHistory'
 import { ToolType, ToolUtils } from '../../../tools/toolUtils'
 import { ChatStream } from '../../../tools/chatStream'
 import path from 'path'
@@ -182,8 +181,7 @@ export class Messenger {
         session: ChatSession,
         tabID: string,
         triggerID: string,
-        triggerPayload: TriggerPayload,
-        chatHistoryManager: ChatHistoryManager
+        triggerPayload: TriggerPayload
     ) {
         let message = ''
         const messageID = response.$metadata.requestId ?? ''
@@ -469,17 +467,6 @@ export class Messenger {
                         tabID
                     )
                 )
-
-                chatHistoryManager.pushAssistantMessage({
-                    assistantResponseMessage: {
-                        messageId: messageID,
-                        content: message,
-                        references: codeReference,
-                        ...(toolUse &&
-                            toolUse.input !== undefined &&
-                            toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
-                    },
-                })
                 const agenticLoopEnded = !eventCounts.has('toolUseEvent')
                 if (agenticLoopEnded) {
                     // Reset context for the next request