fix(amazonq): Address multiple history bugs

Author: tsmithsz

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

@@ -744,6 +744,7 @@ export class ChatController {
                 const toolResults: ToolResult[] = []
 
                 let response = ''
+                let shouldDisplayMessage = true
                 if (toolUseError) {
                     toolResults.push({
                         content: [{ text: toolUseError.message }],
@@ -753,6 +754,7 @@ export class ChatController {
                     if (toolUseError instanceof SyntaxError) {
                         response =
                             "Your toolUse input isn't valid. Please check the syntax and make sure the input is complete. If the input is large, break it down into multiple tool uses with smaller input."
+                        shouldDisplayMessage = false
                     }
                 } else {
                     const result = ToolUtils.tryFromToolUse(toolUse)
@@ -844,6 +846,7 @@ export class ChatController {
                         contextLengths: {
                             ...defaultContextLengths,
                         },
+                        shouldDisplayMessage: shouldDisplayMessage,
                     },
                     triggerID
                 )
@@ -1138,9 +1141,6 @@ export class ChatController {
 
         this.messenger.sendErrorMessage(errorMessage, tabID, requestID, statusCode)
         getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
-
-        this.sessionStorage.deleteSession(tabID)
-        this.chatHistoryDb.clearTab(tabID)
     }
 
     private async processContextMenuCommand(command: EditorContextCommand) {
@@ -1317,6 +1317,7 @@ export class ChatController {
             this.processException(e, message.tabID)
         }
     }
+
     private initialCleanUp(session: ChatSession) {
         // Create a fresh token for this new conversation
         session.createNewTokenSource()
@@ -1595,11 +1596,12 @@ export class ChatController {
 
         const currentMessage = request.conversationState.currentMessage
         if (currentMessage) {
-            this.chatHistoryDb.fixHistory(tabID, currentMessage)
+            this.chatHistoryDb.fixHistory(tabID, currentMessage, session.sessionIdentifier)
         }
-        request.conversationState.history = this.chatHistoryDb
-            .getMessages(tabID)
-            .map((chat) => messageToChatMessage(chat))
+        //  Do not include chatHistory for requests going to Mynah
+        request.conversationState.history = request.conversationState.currentMessage?.userInputMessage?.userIntent
+            ? []
+            : this.chatHistoryDb.getMessages(tabID).map((chat) => messageToChatMessage(chat))
         request.conversationState.conversationId = session.sessionIdentifier
 
         triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
@@ -1669,6 +1671,7 @@ export class ChatController {
                     userIntent: currentMessage.userInputMessage?.userIntent,
                     origin: currentMessage.userInputMessage?.origin,
                     userInputMessageContext: currentMessage.userInputMessage?.userInputMessageContext,
+                    shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true,
                 })
             }
 

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

@@ -558,6 +558,7 @@ export class Messenger {
                             toolUse && toolUse.input !== undefined && toolUse.input !== ''
                                 ? [{ ...toolUse }]
                                 : undefined,
+                        shouldDisplayMessage: triggerPayload.shouldDisplayMessage ?? true,
                     })
                 }
                 if (

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

@@ -238,6 +238,7 @@ export interface TriggerPayload {
     origin?: Origin
     pairProgrammingModeOn?: boolean
     history?: Message[]
+    shouldDisplayMessage?: boolean
 }
 
 export type ContextLengths = {

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

@@ -88,7 +88,9 @@ export class TabBarController {
                 selectedTab.historyId,
                 selectedTab.tabType,
                 selectedTab.conversations.flatMap((conv: Conversation) =>
-                    conv.messages.map((message) => messageToChatItem(message))
+                    conv.messages
+                        .filter((message) => message.shouldDisplayMessage !== false)
+                        .map((message) => messageToChatItem(message))
                 ),
                 exportTab
             )

packages/core/src/shared/db/chatDb/chatDb.ts

@@ -24,6 +24,8 @@ import { ChatMessage, ToolResultStatus } from '@amzn/codewhisperer-streaming'
 
 // Maximum number of characters to keep in history
 const MaxConversationHistoryCharacters = 600_000
+// Maximum number of messages to keep in history
+const MaxConversationHistoryMessages = 250
 
 /**
  * A singleton database class that manages chat history persistence using LokiJS.
@@ -144,18 +146,6 @@ export class Database {
         return undefined
     }
 
-    getActiveConversation(historyId: string): Conversation | undefined {
-        const tabCollection = this.db.getCollection<Tab>(TabCollection)
-        const tabData = tabCollection.findOne({ historyId })
-
-        if (!tabData?.conversations.length) {
-            this.logger.debug('No active conversations found')
-            return undefined
-        }
-
-        return tabData.conversations[0]
-    }
-
     clearTab(tabId: string) {
         if (this.initialized) {
             const tabCollection = this.db.getCollection<Tab>(TabCollection)
@@ -170,34 +160,6 @@ export class Database {
         }
     }
 
-    //  Removes the most recent message(s) from the chat history for a given tab
-    clearRecentHistory(tabId: string): void {
-        if (this.initialized) {
-            const historyId = this.historyIdMapping.get(tabId)
-            this.logger.info(`Clearing recent history: tabId=${tabId}, historyId=${historyId || 'undefined'}`)
-            if (historyId) {
-                const tabCollection = this.db.getCollection<Tab>(TabCollection)
-                const tabData = tabCollection.findOne({ historyId })
-                if (tabData) {
-                    const activeConversation = tabData.conversations[0]
-                    const allMessages = this.getMessages(tabId)
-                    const lastMessage = allMessages[allMessages.length - 1]
-                    this.logger.debug(`Last message type: ${lastMessage.type}`)
-
-                    if (lastMessage.type === ('prompt' as ChatItemType)) {
-                        allMessages.pop()
-                        this.logger.info(`Removed last user message`)
-                    } else {
-                        allMessages.splice(-2)
-                        this.logger.info(`Removed last assistant message and user message`)
-                    }
-                    activeConversation.messages = allMessages
-                    tabCollection.update(tabData)
-                }
-            }
-        }
-    }
-
     updateTabOpenState(tabId: string, isOpen: boolean) {
         if (this.initialized) {
             const tabCollection = this.db.getCollection<Tab>(TabCollection)
@@ -348,13 +310,14 @@ export class Database {
 
     /**
      * Fixes the history to maintain the following invariants:
-     * 1. The history character length is <= MAX_CONVERSATION_HISTORY_CHARACTERS. Oldest messages are dropped.
-     * 2. The first message is from the user. Oldest messages are dropped if needed.
-     * 3. The last message is from the assistant. The last message is dropped if it is from the user.
-     * 4. If the last message is from the assistant and it contains tool uses, and a next user
+     * 1. The history contains at most MaxConversationHistoryMessages messages. Oldest messages are dropped.
+     * 2. The history character length is <= MaxConversationHistoryCharacters. Oldest messages are dropped.
+     * 3. The first message is from the user. Oldest messages are dropped if needed.
+     * 4. The last message is from the assistant. The last message is dropped if it is from the user.
+     * 5. If the last message is from the assistant and it contains tool uses, and a next user
      *    message is set without tool results, then the user message will have cancelled tool results.
      */
-    fixHistory(tabId: string, newUserMessage: ChatMessage): void {
+    fixHistory(tabId: string, newUserMessage: ChatMessage, conversationId: string): void {
         if (!this.initialized) {
             return
         }
@@ -371,10 +334,17 @@ export class Database {
             return
         }
 
-        const activeConversation = tabData.conversations[0]
-        let allMessages = activeConversation.messages
+        let allMessages = tabData.conversations.flatMap((conversation: Conversation) => conversation.messages)
         this.logger.info(`Found ${allMessages.length} messages in conversation`)
 
+        //  Make sure we don't exceed MaxConversationHistoryMessages
+        if (allMessages.length > MaxConversationHistoryMessages) {
+            this.logger.debug(
+                `Trimming history from ${allMessages.length} to ${MaxConversationHistoryMessages} messages`
+            )
+            allMessages = allMessages.slice(-MaxConversationHistoryMessages)
+        }
+
         //  Drop empty assistant partial if it’s the last message
         this.handleEmptyAssistantMessage(allMessages)
 
@@ -387,7 +357,14 @@ export class Database {
         //  If the last message is from the assistant and it contains tool uses, and a next user message is set without tool results, then the user message will have cancelled tool results.
         this.handleToolUses(allMessages, newUserMessage)
 
-        activeConversation.messages = allMessages
+        tabData.conversations = [
+            {
+                conversationId: conversationId,
+                clientType: ClientType.VSCode,
+                messages: allMessages,
+            },
+        ]
+        tabData.updatedAt = new Date()
         tabCollection.update(tabData)
         this.logger.info(`Updated tab data in collection`)
     }

packages/core/src/shared/db/chatDb/util.ts

@@ -56,6 +56,7 @@ export type Message = {
     origin?: Origin
     userInputMessageContext?: UserInputMessageContext
     toolUses?: ToolUse[]
+    shouldDisplayMessage?: boolean
 }
 
 /**