feat(chat): Adding state Management to stop the response for new user input prompt (aws#6934)

## Problem
Responses need to stop for new user prompt and fix user history

## Solution
- featureflag to check the response state to stop the responses.
- Fixing history incase of exception or stopped responses.

---

- 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.

Author: ashishrp-aws

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

@@ -368,6 +368,7 @@ export class ChatController {
     private async processStopResponseMessage(message: StopResponseMessage) {
         const session = this.sessionStorage.getSession(message.tabID)
         session.tokenSource.cancel()
+        this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
     }
 
     private async processTriggerTabIDReceived(message: TriggerTabIDReceived) {
@@ -650,6 +651,8 @@ export class ChatController {
                 const session = this.sessionStorage.getSession(tabID)
                 const toolUse = session.toolUse
                 if (!toolUse || !toolUse.input) {
+                    // Turn off AgentLoop flag if there's no tool use
+                    this.sessionStorage.setAgentLoopInProgress(tabID, false)
                     return
                 }
                 session.setToolUse(undefined)
@@ -723,7 +726,6 @@ export class ChatController {
                         customization: getSelectedCustomization(),
                         toolResults: toolResults,
                         origin: Origin.IDE,
-                        chatHistory: this.chatHistoryStorage.getTabHistory(tabID).getHistory(),
                         context: session.context ?? [],
                         relevantTextDocuments: [],
                         additionalContents: [],
@@ -899,10 +901,16 @@ export class ChatController {
             errorMessage = e.message
         }
 
+        // Turn off AgentLoop flag in case of exception
+        if (tabID) {
+            this.sessionStorage.setAgentLoopInProgress(tabID, false)
+        }
+
         this.messenger.sendErrorMessage(errorMessage, tabID, requestID)
         getLogger().error(`error: ${errorMessage} tabID: ${tabID} requestID: ${requestID}`)
 
         this.sessionStorage.deleteSession(tabID)
+        this.chatHistoryStorage.getTabHistory(tabID).clearRecentHistory()
     }
 
     private async processContextMenuCommand(command: EditorContextCommand) {
@@ -1062,7 +1070,6 @@ export class ChatController {
                     codeQuery: lastTriggerEvent.context?.focusAreaContext?.names,
                     userIntent: message.userIntent,
                     customization: getSelectedCustomization(),
-                    chatHistory: this.chatHistoryStorage.getTabHistory(message.tabID).getHistory(),
                     contextLengths: {
                         ...defaultContextLengths,
                     },
@@ -1111,7 +1118,6 @@ export class ChatController {
                         codeQuery: context?.focusAreaContext?.names,
                         userIntent: undefined,
                         customization: getSelectedCustomization(),
-                        chatHistory: this.chatHistoryStorage.getTabHistory(message.tabID).getHistory(),
                         origin: Origin.IDE,
                         context: message.context ?? [],
                         relevantTextDocuments: [],
@@ -1293,6 +1299,16 @@ export class ChatController {
         }
 
         const tabID = triggerEvent.tabID
+        if (this.sessionStorage.isAgentLoopInProgress(tabID)) {
+            // If a response is already in progress, stop it first
+            const stopResponseMessage: StopResponseMessage = {
+                tabID: tabID,
+            }
+            await this.processStopResponseMessage(stopResponseMessage)
+        }
+
+        // Ensure AgentLoop flag is set to true during response generation
+        this.sessionStorage.setAgentLoopInProgress(tabID, true)
 
         const credentialsState = await AuthUtil.instance.getChatAuthState()
 
@@ -1355,6 +1371,7 @@ export class ChatController {
         if (fixedHistoryMessage.userInputMessage?.userInputMessageContext) {
             triggerPayload.toolResults = fixedHistoryMessage.userInputMessage.userInputMessageContext.toolResults
         }
+        triggerPayload.chatHistory = chatHistory.getHistory()
         const request = triggerPayloadToChatRequest(triggerPayload)
         const conversationId = chatHistory.getConversationId() || randomUUID()
         chatHistory.setConversationId(conversationId)
@@ -1417,8 +1434,13 @@ export class ChatController {
                 } metadata: ${inspect(response.$metadata, { depth: 12 })}`
             )
             await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload, chatHistory)
+
+            // Turn off AgentLoop flag after sending the AI response
+            this.sessionStorage.setAgentLoopInProgress(tabID, false)
         } catch (e: any) {
             this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)
+            // Turn off AgentLoop flag in case of exception
+            this.sessionStorage.setAgentLoopInProgress(tabID, false)
             // clears session, record telemetry before this call
             this.processException(e, tabID)
         }

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

@@ -320,7 +320,7 @@ export class Messenger {
                 }
                 return true
             },
-            { timeout: 60000, truthy: true }
+            { timeout: 600000, truthy: true }
         )
             .catch((error: any) => {
                 let errorMessage = 'Error reading chat stream.'

packages/core/src/codewhispererChat/storages/chatHistory.ts

@@ -2,14 +2,7 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import {
-    ChatMessage,
-    Tool,
-    ToolResult,
-    ToolResultStatus,
-    UserInputMessage,
-    UserInputMessageContext,
-} from '@amzn/codewhisperer-streaming'
+import { ChatMessage, Tool, ToolResult, ToolResultStatus, ToolUse } from '@amzn/codewhisperer-streaming'
 import { randomUUID } from '../../shared/crypto'
 import { getLogger } from '../../shared/logger/logger'
 import { tools } from '../constants'
@@ -105,168 +98,99 @@ export class ChatHistoryManager {
      *    message is set without tool results, then the user message will have cancelled tool results.
      */
     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
-            let indexToTrim: number | undefined
+        this.trimConversationHistory()
+        this.ensureLastMessageFromAssistant()
+        return this.handleToolUses(newUserMessage)
+    }
 
-            for (let i = 1; i < this.history.length; i++) {
-                const message = this.history[i]
-                if (message.userInputMessage) {
-                    const userMessage = message.userInputMessage
-                    const ctx = userMessage.userInputMessageContext
-                    const hasNoToolResults = ctx && (!ctx.toolResults || ctx.toolResults.length === 0)
-                    if (hasNoToolResults && userMessage.content !== '') {
-                        indexToTrim = i
-                        break
-                    }
-                }
-            }
-            if (indexToTrim !== undefined) {
-                this.logger.debug(`Removing the first ${indexToTrim} elements in the history`)
-                this.history.splice(0, indexToTrim)
-            } else {
-                this.logger.debug('No valid starting user message found in the history, clearing')
-                this.history = []
-            }
+    private trimConversationHistory(): void {
+        if (this.history.length <= MaxConversationHistoryLength) {
+            return
         }
 
-        // Ensure the last message is from the assistant
-        if (this.history.length > 0 && this.history[this.history.length - 1].userInputMessage !== undefined) {
-            this.logger.debug('Last message in history is from the user, dropping')
-            this.history.pop()
+        const indexToTrim = this.findIndexToTrim()
+        if (indexToTrim !== undefined) {
+            this.logger.debug(`Removing the first ${indexToTrim} elements in the history`)
+            this.history.splice(0, indexToTrim)
+        } else {
+            this.logger.debug('No valid starting user message found in the history, clearing')
+            this.history = []
         }
+    }
 
-        // 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) &&
-            newUserMessage
-        ) {
-            const toolUses = lastHistoryMessage.assistantResponseMessage.toolUses
-
-            if (toolUses && toolUses.length > 0) {
-                if (newUserMessage.userInputMessage) {
-                    if (newUserMessage.userInputMessage.userInputMessageContext) {
-                        const ctx = newUserMessage.userInputMessage.userInputMessageContext
-
-                        if (!ctx.toolResults || ctx.toolResults.length === 0) {
-                            ctx.toolResults = toolUses.map((toolUse) => ({
-                                toolUseId: toolUse.toolUseId,
-                                content: [
-                                    {
-                                        type: 'Text',
-                                        text: 'Tool use was cancelled by the user',
-                                    },
-                                ],
-                                status: ToolResultStatus.ERROR,
-                            }))
-                        }
-                    } else {
-                        const toolResults = toolUses.map((toolUse) => ({
-                            toolUseId: toolUse.toolUseId,
-                            content: [
-                                {
-                                    type: 'Text',
-                                    text: 'Tool use was cancelled by the user',
-                                },
-                            ],
-                            status: ToolResultStatus.ERROR,
-                        }))
-
-                        newUserMessage.userInputMessage.userInputMessageContext = {
-                            shellState: undefined,
-                            envState: undefined,
-                            toolResults: toolResults,
-                            tools: this.tools.length === 0 ? undefined : [...this.tools],
-                        }
-
-                        return newUserMessage
-                    }
-                }
+    private findIndexToTrim(): number | undefined {
+        for (let i = 1; i < this.history.length; i++) {
+            const message = this.history[i]
+            if (this.isValidUserMessageWithoutToolResults(message)) {
+                return i
             }
         }
-
-        // Always return the message to fix the TypeScript error
-        return newUserMessage
+        return undefined
     }
 
-    /**
-     * Adds tool results to the conversation.
-     */
-    addToolResults(toolResults: ToolResult[]): void {
-        const userInputMessageContext: UserInputMessageContext = {
-            shellState: undefined,
-            envState: undefined,
-            toolResults: toolResults,
-            tools: this.tools.length === 0 ? undefined : [...this.tools],
-        }
-
-        const msg: UserInputMessage = {
-            content: '',
-            userInputMessageContext: userInputMessageContext,
+    private isValidUserMessageWithoutToolResults(message: ChatMessage): boolean {
+        if (!message.userInputMessage) {
+            return false
         }
+        const ctx = message.userInputMessage.userInputMessageContext
+        return Boolean(
+            ctx && (!ctx.toolResults || ctx.toolResults.length === 0) && message.userInputMessage.content !== ''
+        )
+    }
 
-        if (this.lastUserMessage?.userInputMessage) {
-            this.lastUserMessage.userInputMessage = msg
+    private ensureLastMessageFromAssistant(): void {
+        if (this.history.length > 0 && this.history[this.history.length - 1].userInputMessage !== undefined) {
+            this.logger.debug('Last message in history is from the user, dropping')
+            this.history.pop()
         }
     }
 
-    /**
-     * 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
+    private handleToolUses(newUserMessage: ChatMessage): ChatMessage {
+        const lastHistoryMessage = this.history[this.history.length - 1]
+        if (!lastHistoryMessage || !lastHistoryMessage.assistantResponseMessage || !newUserMessage) {
+            return newUserMessage
         }
 
-        const lastMessage = this.history[this.history.length - 1]
-
-        if (lastMessage.assistantResponseMessage) {
-            const toolUses = lastMessage.assistantResponseMessage.toolUses
+        const toolUses = lastHistoryMessage.assistantResponseMessage.toolUses
+        if (!toolUses || toolUses.length === 0) {
+            return newUserMessage
+        }
 
-            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,
-                }))
+        return this.addToolResultsToUserMessage(newUserMessage, toolUses)
+    }
 
-                // 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],
-                }
+    private addToolResultsToUserMessage(newUserMessage: ChatMessage, toolUses: ToolUse[]): ChatMessage {
+        if (!newUserMessage.userInputMessage) {
+            return newUserMessage
+        }
 
-                const userMessage: ChatMessage = {
-                    userInputMessage: {
-                        content: '',
-                        userInputMessageContext: userInputMessageContext,
-                    },
-                }
+        const toolResults = this.createToolResults(toolUses)
 
-                this.history.push(this.formatChatHistoryMessage(userMessage))
-                this.logger.debug('Added user message with cancelled tool results')
+        if (newUserMessage.userInputMessage.userInputMessageContext) {
+            newUserMessage.userInputMessage.userInputMessageContext.toolResults = toolResults
+        } else {
+            newUserMessage.userInputMessage.userInputMessageContext = {
+                shellState: undefined,
+                envState: undefined,
+                toolResults: toolResults,
+                tools: this.tools.length === 0 ? undefined : [...this.tools],
             }
         }
+
+        return newUserMessage
+    }
+
+    private createToolResults(toolUses: ToolUse[]): ToolResult[] {
+        return toolUses.map((toolUse) => ({
+            toolUseId: toolUse.toolUseId,
+            content: [
+                {
+                    type: 'Text',
+                    text: 'Tool use was cancelled by the user',
+                },
+            ],
+            status: ToolResultStatus.ERROR,
+        }))
     }
 
     private formatChatHistoryMessage(message: ChatMessage): ChatMessage {
@@ -283,4 +207,18 @@ export class ChatHistoryManager {
         }
         return message
     }
+
+    public clearRecentHistory(): void {
+        if (this.history.length === 0) {
+            return
+        }
+
+        const lastHistoryMessage = this.history[this.history.length - 1]
+
+        if (lastHistoryMessage.userInputMessage?.userInputMessageContext) {
+            this.history.pop()
+        } else if (lastHistoryMessage.assistantResponseMessage) {
+            this.history.splice(-2)
+        }
+    }
 }