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 { DocumentReference, 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
@@ -50,7 +49,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
     }
     public get messageIdToUpdate(): string | undefined {
@@ -105,13 +104,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(): DocumentReference[] {
@@ -148,14 +148,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
     }
@@ -170,12 +162,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()
     private cancelTokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource()
 
@@ -194,7 +193,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) => {
@@ -416,7 +414,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) {
@@ -473,7 +471,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)
@@ -1023,7 +1020,7 @@ export class ChatController {
         getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
 
         this.sessionStorage.deleteSession(tabID)
-        this.chatHistoryStorage.getTabHistory(tabID).clear()
+        this.chatHistoryDb.clearRecentHistory(tabID)
     }
 
     private async processContextMenuCommand(command: EditorContextCommand) {
@@ -1145,7 +1142,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)
@@ -1437,10 +1433,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'
@@ -1479,16 +1471,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)
+            .map((chat) => messageToChatMessage(chat))
+        request.conversationState.conversationId = session.sessionIdentifier
 
         triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
 
@@ -1547,7 +1537,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,
@@ -1562,15 +1551,13 @@ export class ChatController {
                     response.$metadata.requestId
                 } metadata: ${inspect(response.$metadata, { depth: 12 })}`
             )
-
             this.cancelTokenSource = new vscode.CancellationTokenSource()
             await this.messenger.sendAIResponse(
                 response,
                 session,
                 tabID,
                 triggerID,
                 triggerPayload,
-                chatHistory,
                 this.cancelTokenSource.token
             )
 

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

@@ -58,7 +58,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'
@@ -188,7 +187,6 @@ export class Messenger {
         tabID: string,
         triggerID: string,
         triggerPayload: TriggerPayload,
-        chatHistoryManager: ChatHistoryManager,
         cancelToken: vscode.CancellationToken
     ) {
         let message = ''
@@ -528,17 +526,6 @@ export class Messenger {
                     )
                 )
 
-                chatHistoryManager.pushAssistantMessage({
-                    assistantResponseMessage: {
-                        messageId: messageID,
-                        content: message,
-                        references: codeReference,
-                        ...(toolUse &&
-                            toolUse.input !== undefined &&
-                            toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
-                    },
-                })
-
                 getLogger().info(
                     `All events received. requestId=%s counts=%s`,
                     response.$metadata.requestId,