feat(chat): show stop button for agentic loop (aws#6947)

## Problem

Agentic chat should have a stop button

## Solution

Implement `asyncEventProgressMessage` for CWChat

---

- 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: ctlai95

packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts

@@ -27,13 +27,21 @@ export interface ConnectorProps extends BaseConnectorProps {
         description?: string
     ) => void
     onChatAnswerUpdated?: (tabID: string, message: ChatItem) => void
+    onAsyncEventProgress: (
+        tabID: string,
+        inProgress: boolean,
+        message: string,
+        messageId: string | undefined,
+        enableStopAction: boolean
+    ) => void
 }
 
 export class Connector extends BaseConnector {
     private readonly onCWCContextCommandMessage
     private readonly onContextCommandDataReceived
     private readonly onShowCustomForm
     private readonly onChatAnswerUpdated
+    private readonly onAsyncEventProgress
     private chatItems: Map<string, Map<string, ChatItem>> = new Map() // tabId -> messageId -> ChatItem
 
     override getTabType(): TabType {
@@ -46,6 +54,7 @@ export class Connector extends BaseConnector {
         this.onContextCommandDataReceived = props.onContextCommandDataReceived
         this.onShowCustomForm = props.onShowCustomForm
         this.onChatAnswerUpdated = props.onChatAnswerUpdated
+        this.onAsyncEventProgress = props.onAsyncEventProgress
     }
 
     onSourceLinkClick = (tabID: string, messageId: string, link: string): void => {
@@ -168,35 +177,6 @@ export class Connector extends BaseConnector {
         }
     }
 
-    private processToolMessage = async (messageData: any): Promise<void> => {
-        if (this.onChatAnswerUpdated === undefined) {
-            return
-        }
-        const answer: CWCChatItem = {
-            type: messageData.messageType,
-            messageId: messageData.messageID ?? messageData.triggerID,
-            body: messageData.message,
-            followUp: messageData.followUps,
-            canBeVoted: messageData.canBeVoted ?? false,
-            codeReference: messageData.codeReference,
-            userIntent: messageData.contextList,
-            codeBlockLanguage: messageData.codeBlockLanguage,
-            contextList: messageData.contextList,
-            title: messageData.title,
-            buttons: messageData.buttons,
-            fileList: messageData.fileList,
-            header: messageData.header ?? undefined,
-            padding: messageData.padding ?? undefined,
-            fullWidth: messageData.fullWidth ?? undefined,
-            codeBlockActions: messageData.codeBlockActions ?? undefined,
-        }
-        if (answer.messageId) {
-            this.storeChatItem(messageData.tabID, answer.messageId, answer)
-        }
-        this.onChatAnswerUpdated(messageData.tabID, answer)
-        return
-    }
-
     private storeChatItem(tabId: string, messageId: string, item: ChatItem): void {
         if (!this.chatItems.has(tabId)) {
             this.chatItems.set(tabId, new Map())
@@ -253,11 +233,6 @@ export class Connector extends BaseConnector {
             return
         }
 
-        if (messageData.type === 'toolMessage') {
-            await this.processToolMessage(messageData)
-            return
-        }
-
         if (messageData.type === 'editorContextCommandMessage') {
             await this.processEditorContextCommandMessage(messageData)
             return
@@ -276,6 +251,18 @@ export class Connector extends BaseConnector {
             this.onCustomFormAction(messageData.tabID, messageData.messageId, messageData.action)
             return
         }
+
+        if (messageData.type === 'asyncEventProgressMessage') {
+            const enableStopAction = true
+            this.onAsyncEventProgress(
+                messageData.tabID,
+                messageData.inProgress,
+                messageData.message ?? undefined,
+                messageData.messageId ?? undefined,
+                enableStopAction
+            )
+            return
+        }
         // For other message types, call the base class handleMessageReceive
         await this.baseHandleMessageReceive(messageData)
     }

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

@@ -22,14 +22,12 @@ export class ChatSession {
      * _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
      * _context = Additional context to be passed to the LLM for generating the response
-     * _messageIdToUpdate = messageId of a chat message to be updated, used for reducing consecutive tool messages
      */
     private _readFiles: string[] = []
     private _toolUse: ToolUse | undefined
     private _showDiffOnFileWrite: boolean = false
     private _context: PromptMessage['context']
     private _pairProgrammingModeOn: boolean = true
-    private _messageIdToUpdate: string | undefined
 
     contexts: Map<string, { first: number; second: number }[]> = new Map()
     // TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -63,14 +61,6 @@ export class ChatSession {
         this._context = context
     }
 
-    public get messageIdToUpdate(): string | undefined {
-        return this._messageIdToUpdate
-    }
-
-    public setMessageIdToUpdate(messageId: string | undefined) {
-        this._messageIdToUpdate = messageId
-    }
-
     public tokenSource!: vscode.CancellationTokenSource
 
     constructor() {

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

@@ -716,6 +716,7 @@ export class ChatController {
                     type: 'chat_message',
                     context,
                 })
+                this.messenger.sendAsyncEventProgress(tabID, true, '')
                 const session = this.sessionStorage.getSession(tabID)
                 const toolUse = session.toolUse
                 if (!toolUse || !toolUse.input) {
@@ -734,16 +735,9 @@ export class ChatController {
                     try {
                         await ToolUtils.validate(tool)
 
-                        const chatStream = new ChatStream(
-                            this.messenger,
-                            tabID,
-                            triggerID,
-                            // Pass in a different toolUseId so that the output does not overwrite
-                            // any previous messages
-                            { ...toolUse, toolUseId: `${toolUse.toolUseId}-output` },
-                            { requiresAcceptance: false },
-                            undefined
-                        )
+                        const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
+                            requiresAcceptance: false,
+                        })
                         const output = await ToolUtils.invoke(tool, chatStream)
                         if (output.output.content.length > maxToolOutputCharacterLength) {
                             throw Error(
@@ -1185,6 +1179,8 @@ export class ChatController {
                     context,
                 })
 
+                this.messenger.sendAsyncEventProgress(message.tabID, true, '')
+
                 // Save the context for the agentic loop
                 session.setContext(message.context)
 

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

@@ -14,7 +14,6 @@ import {
     OpenSettingsMessage,
     QuickActionMessage,
     ShowCustomFormMessage,
-    ToolMessage,
 } from '../../../view/connector/connector'
 import { EditorContextCommandType } from '../../../commands/registerCommands'
 import { ChatResponseStream as qdevChatResponseStream } from '@amzn/amazon-q-developer-streaming-client'
@@ -49,6 +48,7 @@ import { extractErrorInfo } from '../../../../shared/utilities/messageUtil'
 import { noWriteTools, tools } from '../../../constants'
 import { Change } from 'diff'
 import { FsWriteParams } from '../../../tools/fsWrite'
+import { AsyncEventProgressMessage } from '../../../../amazonq/commons/connector/connectorMessages'
 
 export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
 
@@ -238,36 +238,10 @@ export class Messenger {
                                 session.setShowDiffOnFileWrite(true)
                                 changeList = await tool.tool.getDiffChanges()
                             }
-                            if (
-                                tool.type === ToolType.FsWrite ||
-                                tool.type === ToolType.ExecuteBash ||
-                                eventCounts.has('assistantResponseEvent')
-                            ) {
-                                // FsWrite and ExecuteBash should never replace older messages
-                                // If the current stream also has assistantResponseEvent then reset this as well.
-                                session.setMessageIdToUpdate(undefined)
-                            }
                             const validation = ToolUtils.requiresAcceptance(tool)
-
-                            const chatStream = new ChatStream(
-                                this,
-                                tabID,
-                                triggerID,
-                                toolUse,
-                                validation,
-                                session.messageIdToUpdate,
-                                changeList
-                            )
+                            const chatStream = new ChatStream(this, tabID, triggerID, toolUse, validation, changeList)
                             await ToolUtils.queueDescription(tool, chatStream)
 
-                            if (
-                                session.messageIdToUpdate === undefined &&
-                                (tool.type === ToolType.FsRead || tool.type === ToolType.ListDirectory)
-                            ) {
-                                // Store the first messageId in a chain of tool uses
-                                session.setMessageIdToUpdate(toolUse.toolUseId)
-                            }
-
                             if (!validation.requiresAcceptance) {
                                 // Need separate id for read tool and safe bash command execution as 'confirm-tool-use' id is required to change button status from `Confirm` to `Confirmed` state in cwChatConnector.ts which will impact generic tool execution.
                                 this.dispatcher.sendCustomFormActionMessage(
@@ -469,33 +443,12 @@ export class Messenger {
         )
     }
 
-    public sendInitialToolMessage(tabID: string, triggerID: string, toolUseId: string | undefined) {
-        this.dispatcher.sendChatMessage(
-            new ChatMessage(
-                {
-                    message: '',
-                    messageType: 'answer',
-                    followUps: undefined,
-                    followUpsHeader: undefined,
-                    relatedSuggestions: undefined,
-                    triggerID,
-                    messageID: toolUseId ?? 'toolUse',
-                    userIntent: undefined,
-                    codeBlockLanguage: undefined,
-                    contextList: undefined,
-                },
-                tabID
-            )
-        )
-    }
-
     public sendPartialToolLog(
         message: string,
         tabID: string,
         triggerID: string,
         toolUse: ToolUse | undefined,
         validation: CommandValidation,
-        messageIdToUpdate: string | undefined,
         changeList?: Change[]
     ) {
         const buttons: ChatItemButton[] = []
@@ -550,16 +503,16 @@ export class Messenger {
             })
         }
 
-        this.dispatcher.sendToolMessage(
-            new ToolMessage(
+        this.dispatcher.sendChatMessage(
+            new ChatMessage(
                 {
                     message: message,
                     messageType: 'answer-part',
                     followUps: undefined,
                     followUpsHeader: undefined,
                     relatedSuggestions: undefined,
                     triggerID,
-                    messageID: messageIdToUpdate ?? toolUse?.toolUseId ?? '',
+                    messageID: toolUse?.toolUseId ?? '',
                     userIntent: undefined,
                     codeBlockLanguage: undefined,
                     contextList: undefined,
@@ -738,4 +691,8 @@ export class Messenger {
             new ShowCustomFormMessage(tabID, formItems, buttons, title, description)
         )
     }
+
+    public sendAsyncEventProgress(tabID: string, inProgress: boolean, message: string | undefined) {
+        this.dispatcher.sendAsyncEventProgress(new AsyncEventProgressMessage(tabID, 'CWChat', inProgress, message))
+    }
 }

packages/core/src/codewhispererChat/tools/chatStream.ts

@@ -23,17 +23,12 @@ export class ChatStream extends Writable {
         private readonly triggerID: string,
         private readonly toolUse: ToolUse | undefined,
         private readonly validation: CommandValidation,
-        private readonly messageIdToUpdate: string | undefined,
         private readonly changeList?: Change[],
         private readonly logger = getLogger('chatStream')
     ) {
         super()
         this.logger.debug(`ChatStream created for tabID: ${tabID}, triggerID: ${triggerID}`)
-        if (!messageIdToUpdate) {
-            // If messageIdToUpdate is undefined, we need to first create an empty message
-            // with messageId so it can be updated later
-            this.messenger.sendInitialToolMessage(tabID, triggerID, toolUse?.toolUseId)
-        }
+        this.messenger.sendInitalStream(tabID, triggerID, undefined)
     }
 
     override _write(chunk: Buffer, encoding: BufferEncoding, callback: (error?: Error | null) => void): void {
@@ -46,7 +41,6 @@ export class ChatStream extends Writable {
             this.triggerID,
             this.toolUse,
             this.validation,
-            this.messageIdToUpdate,
             this.changeList
         )
         callback()

packages/core/src/codewhispererChat/view/connector/connector.ts

@@ -19,6 +19,7 @@ import {
     Status,
 } from '@aws/mynah-ui'
 import { DocumentReference } from '../../controllers/chat/model'
+import { AsyncEventProgressMessage } from '../../../amazonq/commons/connector/connectorMessages'
 
 class UiMessage {
     readonly time: number = Date.now()
@@ -286,10 +287,6 @@ export class ChatMessage extends UiMessage {
     }
 }
 
-export class ToolMessage extends ChatMessage {
-    override type = 'toolMessage'
-}
-
 export interface FollowUp {
     readonly type: string
     readonly pillText: string
@@ -344,10 +341,6 @@ export class AppToWebViewMessageDispatcher {
         this.appsToWebViewMessagePublisher.publish(message)
     }
 
-    public sendToolMessage(message: ToolMessage) {
-        this.appsToWebViewMessagePublisher.publish(message)
-    }
-
     public sendEditorContextCommandMessage(message: EditorContextCommandMessage) {
         this.appsToWebViewMessagePublisher.publish(message)
     }
@@ -375,4 +368,8 @@ export class AppToWebViewMessageDispatcher {
     public sendCustomFormActionMessage(message: CustomFormActionMessage) {
         this.appsToWebViewMessagePublisher.publish(message)
     }
+
+    public sendAsyncEventProgress(message: AsyncEventProgressMessage) {
+        this.appsToWebViewMessagePublisher.publish(message)
+    }
 }