fix(chat): fix for chat history for multuple tabs and user tool ignore (aws#6906)

ashishrp-aws · web-flow · commit 62214eecfba7 · 2025-04-01T16:42:59.000-07:00
## Problem - History storage not set up for each tabs. - If user ignores tool use and moves on, then request fails. ## Solution - Fixed chat history for multiple tab support - User cancellation message is added to request in this case. --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -87,11 +87,11 @@ import {
     defaultContextLengths,
 } from '../../constants'
 import { ChatSession } from '../../clients/chat/v0/chat'
-import { ChatHistoryManager } from '../../storages/chatHistory'
 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 { FsWrite, FsWriteParams } from '../../tools/fsWrite'
 import { tempDirPath } from '../../../shared/filesystemUtilities'
 
@@ -155,7 +155,7 @@ export class ChatController {
     private readonly userIntentRecognizer: UserIntentRecognizer
     private readonly telemetryHelper: CWCTelemetryHelper
     private userPromptsWatcher: vscode.FileSystemWatcher | undefined
-    private readonly chatHistoryManager: ChatHistoryManager
+    private readonly chatHistoryStorage: ChatHistoryStorage
 
     public constructor(
         private readonly chatControllerMessageListeners: ChatControllerMessageListeners,
@@ -173,7 +173,7 @@ export class ChatController {
         this.editorContentController = new EditorContentController()
         this.promptGenerator = new PromptsGenerator()
         this.userIntentRecognizer = new UserIntentRecognizer()
-        this.chatHistoryManager = new ChatHistoryManager()
+        this.chatHistoryStorage = new ChatHistoryStorage()
 
         onDidChangeAmazonQVisibility((visible) => {
             if (visible) {
@@ -424,7 +424,7 @@ export class ChatController {
 
     private async processTabCloseMessage(message: TabClosedMessage) {
         this.sessionStorage.deleteSession(message.tabID)
-        this.chatHistoryManager.clear()
+        this.chatHistoryStorage.deleteHistory(message.tabID)
         this.triggerEventsStorage.removeTabEvents(message.tabID)
         // this.telemetryHelper.recordCloseChat(message.tabID)
     }
@@ -711,7 +711,7 @@ export class ChatController {
                         customization: getSelectedCustomization(),
                         toolResults: toolResults,
                         origin: Origin.IDE,
-                        chatHistory: this.chatHistoryManager.getHistory(),
+                        chatHistory: this.chatHistoryStorage.getTabHistory(tabID).getHistory(),
                         context: session.context ?? [],
                         relevantTextDocuments: [],
                         additionalContents: [],
@@ -891,7 +891,6 @@ export class ChatController {
         getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
 
         this.sessionStorage.deleteSession(tabID)
-        this.chatHistoryManager.clear()
     }
 
     private async processContextMenuCommand(command: EditorContextCommand) {
@@ -965,7 +964,6 @@ export class ChatController {
                         codeQuery: context?.focusAreaContext?.names,
                         userIntent: this.userIntentRecognizer.getFromContextMenuCommand(command),
                         customization: getSelectedCustomization(),
-                        chatHistory: this.chatHistoryManager.getHistory(),
                         additionalContents: [],
                         relevantTextDocuments: [],
                         documentReferences: [],
@@ -1013,7 +1011,7 @@ export class ChatController {
         switch (message.command) {
             case 'clear':
                 this.sessionStorage.deleteSession(message.tabID)
-                this.chatHistoryManager.clear()
+                this.chatHistoryStorage.getTabHistory(message.tabID).clear()
                 this.triggerEventsStorage.removeTabEvents(message.tabID)
                 recordTelemetryChatRunCommand('clear')
                 return
@@ -1052,7 +1050,7 @@ export class ChatController {
                     codeQuery: lastTriggerEvent.context?.focusAreaContext?.names,
                     userIntent: message.userIntent,
                     customization: getSelectedCustomization(),
-                    chatHistory: this.chatHistoryManager.getHistory(),
+                    chatHistory: this.chatHistoryStorage.getTabHistory(message.tabID).getHistory(),
                     contextLengths: {
                         ...defaultContextLengths,
                     },
@@ -1101,7 +1099,7 @@ export class ChatController {
                         codeQuery: context?.focusAreaContext?.names,
                         userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
                         customization: getSelectedCustomization(),
-                        chatHistory: this.chatHistoryManager.getHistory(),
+                        chatHistory: this.chatHistoryStorage.getTabHistory(message.tabID).getHistory(),
                         origin: Origin.IDE,
                         context: message.context ?? [],
                         relevantTextDocuments: [],
@@ -1328,16 +1326,28 @@ export class ChatController {
 
         triggerPayload.contextLengths.userInputContextLength = triggerPayload.message.length
         triggerPayload.contextLengths.focusFileContextLength = triggerPayload.fileText.length
-        const request = triggerPayloadToChatRequest(triggerPayload)
-        if (
-            this.chatHistoryManager.getConversationId() !== undefined &&
-            this.chatHistoryManager.getConversationId() !== ''
-        ) {
-            request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
-        } else {
-            this.chatHistoryManager.setConversationId(randomUUID())
-            request.conversationState.conversationId = this.chatHistoryManager.getConversationId()
+
+        const chatHistory = this.chatHistoryStorage.getTabHistory(tabID)
+        const newUserMessage = {
+            userInputMessage: {
+                content: triggerPayload.message,
+                userIntent: triggerPayload.userIntent,
+                ...(triggerPayload.origin && { origin: triggerPayload.origin }),
+                userInputMessageContext: {
+                    tools: tools,
+                    ...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
+                },
+            },
+        }
+        const fixedHistoryMessage = chatHistory.fixHistory(newUserMessage)
+        if (fixedHistoryMessage.userInputMessage?.userInputMessageContext) {
+            triggerPayload.toolResults = fixedHistoryMessage.userInputMessage.userInputMessageContext.toolResults
         }
+        const request = triggerPayloadToChatRequest(triggerPayload)
+        const conversationId = chatHistory.getConversationId() || randomUUID()
+        chatHistory.setConversationId(conversationId)
+        request.conversationState.conversationId = conversationId
+
         triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments)
 
         // Update context transparency after it's truncated dynamically to show users only the context sent.
@@ -1387,32 +1397,14 @@ export class ChatController {
             }
             this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
             this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
-
-            this.chatHistoryManager.appendUserMessage({
-                userInputMessage: {
-                    content: triggerPayload.message,
-                    userIntent: triggerPayload.userIntent,
-                    ...(triggerPayload.origin && { origin: triggerPayload.origin }),
-                    userInputMessageContext: {
-                        tools: tools,
-                        ...(triggerPayload.toolResults && { toolResults: triggerPayload.toolResults }),
-                    },
-                },
-            })
+            chatHistory.appendUserMessage(fixedHistoryMessage)
 
             getLogger().info(
                 `response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${
                     response.$metadata.requestId
                 } metadata: ${inspect(response.$metadata, { depth: 12 })}`
             )
-            await this.messenger.sendAIResponse(
-                response,
-                session,
-                tabID,
-                triggerID,
-                triggerPayload,
-                this.chatHistoryManager
-            )
+            await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload, chatHistory)
         } catch (e: any) {
             this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)
             // clears session, record telemetry before this call
diff --git a/packages/core/src/codewhispererChat/storages/chatHistory.ts b/packages/core/src/codewhispererChat/storages/chatHistory.ts
@@ -23,14 +23,15 @@ const MaxConversationHistoryLength = 100
  */
 export class ChatHistoryManager {
     private conversationId: string
+    private tabId: string
     private history: ChatMessage[] = []
     private logger = getLogger()
     private lastUserMessage?: ChatMessage
     private tools: Tool[] = []
 
-    constructor() {
+    constructor(tabId?: string) {
         this.conversationId = randomUUID()
-        this.logger.info(`Generated new conversation id: ${this.conversationId}`)
+        this.tabId = tabId ?? randomUUID()
         this.tools = tools
     }
 
@@ -45,6 +46,20 @@ export class ChatHistoryManager {
         this.conversationId = conversationId
     }
 
+    /**
+     * Get the tab ID
+     */
+    public getTabId(): string {
+        return this.tabId
+    }
+
+    /**
+     * Set the tab ID
+     */
+    public setTabId(tabId: string) {
+        this.tabId = tabId
+    }
+
     /**
      * Get the full chat history
      */
@@ -65,7 +80,6 @@ export class ChatHistoryManager {
      */
     public appendUserMessage(newMessage: ChatMessage): void {
         this.lastUserMessage = newMessage
-        this.fixHistory()
         if (!newMessage.userInputMessage?.content || newMessage.userInputMessage?.content.trim() === '') {
             this.logger.warn('input must not be empty when adding new messages')
         }
@@ -90,7 +104,7 @@ export class ChatHistoryManager {
      * 4. 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.
      */
-    public fixHistory(): void {
+    public fixHistory(newUserMessage: ChatMessage): ChatMessage {
         // Trim the conversation history if it exceeds the maximum length
         if (this.history.length > MaxConversationHistoryLength) {
             // Find the second oldest user message without tool results
@@ -123,22 +137,22 @@ export class ChatHistoryManager {
             this.history.pop()
         }
 
-        // TODO: If the last message from the assistant contains tool uses, ensure the next user message contains tool results
+        // If the last message from the assistant contains tool uses, ensure the next user message contains tool results
 
         const lastHistoryMessage = this.history[this.history.length - 1]
 
         if (
             lastHistoryMessage &&
             (lastHistoryMessage.assistantResponseMessage ||
                 lastHistoryMessage.assistantResponseMessage !== undefined) &&
-            this.lastUserMessage
+            newUserMessage
         ) {
             const toolUses = lastHistoryMessage.assistantResponseMessage.toolUses
 
             if (toolUses && toolUses.length > 0) {
-                if (this.lastUserMessage.userInputMessage) {
-                    if (this.lastUserMessage.userInputMessage.userInputMessageContext) {
-                        const ctx = this.lastUserMessage.userInputMessage.userInputMessageContext
+                if (newUserMessage.userInputMessage) {
+                    if (newUserMessage.userInputMessage.userInputMessageContext) {
+                        const ctx = newUserMessage.userInputMessage.userInputMessageContext
 
                         if (!ctx.toolResults || ctx.toolResults.length === 0) {
                             ctx.toolResults = toolUses.map((toolUse) => ({
@@ -164,16 +178,21 @@ export class ChatHistoryManager {
                             status: ToolResultStatus.ERROR,
                         }))
 
-                        this.lastUserMessage.userInputMessage.userInputMessageContext = {
+                        newUserMessage.userInputMessage.userInputMessageContext = {
                             shellState: undefined,
                             envState: undefined,
                             toolResults: toolResults,
                             tools: this.tools.length === 0 ? undefined : [...this.tools],
                         }
+
+                        return newUserMessage
                     }
                 }
             }
         }
+
+        // Always return the message to fix the TypeScript error
+        return newUserMessage
     }
 
     /**
@@ -197,6 +216,59 @@ export class ChatHistoryManager {
         }
     }
 
+    /**
+     * Checks if the latest message in history is an Assistant Message.
+     * If it is and doesn't have toolUse, it will be removed.
+     * If it has toolUse, an assistantResponse message with cancelled tool status will be added.
+     */
+    public checkLatestAssistantMessage(): void {
+        if (this.history.length === 0) {
+            return
+        }
+
+        const lastMessage = this.history[this.history.length - 1]
+
+        if (lastMessage.assistantResponseMessage) {
+            const toolUses = lastMessage.assistantResponseMessage.toolUses
+
+            if (!toolUses || toolUses.length === 0) {
+                // If there are no tool uses, remove the assistant message
+                this.logger.debug('Removing assistant message without tool uses')
+                this.history.pop()
+            } else {
+                // If there are tool uses, add cancelled tool results
+                const toolResults = toolUses.map((toolUse) => ({
+                    toolUseId: toolUse.toolUseId,
+                    content: [
+                        {
+                            type: 'Text',
+                            text: 'Tool use was cancelled by the user',
+                        },
+                    ],
+                    status: ToolResultStatus.ERROR,
+                }))
+
+                // Create a new user message with cancelled tool results
+                const userInputMessageContext: UserInputMessageContext = {
+                    shellState: undefined,
+                    envState: undefined,
+                    toolResults: toolResults,
+                    tools: this.tools.length === 0 ? undefined : [...this.tools],
+                }
+
+                const userMessage: ChatMessage = {
+                    userInputMessage: {
+                        content: '',
+                        userInputMessageContext: userInputMessageContext,
+                    },
+                }
+
+                this.history.push(this.formatChatHistoryMessage(userMessage))
+                this.logger.debug('Added user message with cancelled tool results')
+            }
+        }
+    }
+
     private formatChatHistoryMessage(message: ChatMessage): ChatMessage {
         if (message.userInputMessage !== undefined) {
             return {
diff --git a/packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts b/packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts
@@ -0,0 +1,43 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { ChatHistoryManager } from './chatHistory'
+
+/**
+ * ChatHistoryStorage manages ChatHistoryManager instances for multiple tabs.
+ * Each tab has its own ChatHistoryManager to maintain separate chat histories.
+ */
+export class ChatHistoryStorage {
+    private histories: Map<string, ChatHistoryManager> = new Map()
+
+    /**
+     * Gets the ChatHistoryManager for a specific tab.
+     * If no history exists for the tab, creates a new one.
+     *
+     * @param tabId The ID of the tab
+     * @returns The ChatHistoryManager for the specified tab
+     */
+    public getTabHistory(tabId: string): ChatHistoryManager {
+        const historyFromStorage = this.histories.get(tabId)
+        if (historyFromStorage !== undefined) {
+            return historyFromStorage
+        }
+
+        // Create a new ChatHistoryManager with the tabId
+        const newHistory = new ChatHistoryManager(tabId)
+        this.histories.set(tabId, newHistory)
+
+        return newHistory
+    }
+
+    /**
+     * Deletes the ChatHistoryManager for a specific tab.
+     *
+     * @param tabId The ID of the tab
+     */
+    public deleteHistory(tabId: string) {
+        this.histories.delete(tabId)
+    }
+}
