feat(chat): agentic chat loop stop with trigger ID for tracking conversation aws#6990

## Problem
Currently stop button and new prompt only stop with UI changes but
agentic loop continues on the background and new messages keep streaming
in.

## Solution
- Correcting the setup for the agentic loop.
- This stops the agentic loop completely with multiple checks.
- Trigger Id for each new user prompt linked to a token to track an
agentic loop Conversation.
- Added multiple cancellation checks to stop the loop once we hit the
stop condition.
- Stops the Execute Bash executions once stop is clicked

Author: ashishrp-aws

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

@@ -285,19 +285,18 @@ export class Connector extends BaseConnector {
         }
 
         if (messageData.type === 'customFormActionMessage') {
-            this.onCustomFormAction(messageData.tabID, messageData.messageId, messageData.action)
+            this.onCustomFormAction(messageData.tabID, messageData.messageId, messageData.action, messageData.triggerId)
             return
         }
 
         if (messageData.type === 'asyncEventProgressMessage') {
-            const enableStopAction = false
             const isPromptInputDisabled = true
             this.onAsyncEventProgress(
                 messageData.tabID,
                 messageData.inProgress,
                 messageData.message ?? undefined,
                 messageData.messageId ?? undefined,
-                enableStopAction,
+                messageData.inProgress,
                 isPromptInputDisabled
             )
             return
@@ -335,7 +334,8 @@ export class Connector extends BaseConnector {
             id: string
             text?: string | undefined
             formItemValues?: Record<string, string> | undefined
-        }
+        },
+        triggerId: string
     ) {
         if (action === undefined) {
             return
@@ -351,6 +351,7 @@ export class Connector extends BaseConnector {
             formSelectedValues: action.formItemValues,
             tabType: this.getTabType(),
             tabID: tabId,
+            triggerId: triggerId,
         })
 
         if (

packages/core/src/amazonq/webview/ui/connector.ts

@@ -744,7 +744,7 @@ export class Connector {
                         tabType: 'cwc',
                     })
                 } else {
-                    this.cwChatConnector.onCustomFormAction(tabId, messageId ?? '', action)
+                    this.cwChatConnector.onCustomFormAction(tabId, messageId ?? '', action, messageId ?? '')
                 }
                 break
             case 'agentWalkthrough': {

packages/core/src/amazonq/webview/ui/main.ts

@@ -719,6 +719,13 @@ export const createMynahUI = (
                 return
             }
 
+            // For new user prompt stopping chat with UI changes
+            mynahUI.updateStore(tabID, {
+                loadingChat: false,
+                promptInputDisabledState: false,
+            })
+            connector.onStopChatResponse(tabID)
+
             const tabType = tabsStorage.getTab(tabID)?.type
             if (tabType === 'featuredev') {
                 mynahUI.addChatItem(tabID, {

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

@@ -17,6 +17,7 @@ import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWr
 import { DocumentReference, PromptMessage } from '../../../controllers/chat/model'
 import { FsWriteBackup } from '../../../../codewhispererChat/tools/fsWrite'
 import { randomUUID } from '../../../../shared/crypto'
+import { getLogger } from '../../../../shared/logger/logger'
 
 export type ToolUseWithError = {
     toolUse: ToolUse
@@ -38,6 +39,8 @@ export class ChatSession {
     private _context: PromptMessage['context']
     private _pairProgrammingModeOn: boolean = true
     private _fsWriteBackups: Map<string, FsWriteBackup> = new Map()
+    private _agenticLoopInProgress: boolean = false
+
     /**
      * True if messages from local history have been sent to session.
      */
@@ -68,6 +71,33 @@ export class ChatSession {
         this._messageIdToUpdateListDirectory = messageId
     }
 
+    public get agenticLoopInProgress(): boolean {
+        return this._agenticLoopInProgress
+    }
+
+    public setAgenticLoopInProgress(value: boolean) {
+        // When setting agenticLoop to false (ending the loop), dispose the current token source
+        if (this._agenticLoopInProgress === true && value === false) {
+            this.disposeTokenSource()
+            // Create a new token source for future operations
+            this.createNewTokenSource()
+        }
+        this._agenticLoopInProgress = value
+    }
+
+    /**
+     * Safely disposes the current token source if it exists
+     */
+    disposeTokenSource() {
+        if (this.tokenSource) {
+            try {
+                this.tokenSource.dispose()
+            } catch (error) {
+                getLogger().debug(`Error disposing token source: ${error}`)
+            }
+        }
+    }
+
     public get pairProgrammingModeOn(): boolean {
         return this._pairProgrammingModeOn
     }

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

@@ -6,6 +6,7 @@ import * as path from 'path'
 import * as vscode from 'vscode'
 import { Event as VSCodeEvent, Uri, workspace, window, ViewColumn, Position, Selection } from 'vscode'
 import { EditorContextExtractor } from '../../editor/context/extractor'
+import { ConversationTracker } from '../../storages/conversationTracker'
 import { ChatSessionStorage } from '../../storages/chatSession'
 import { Messenger, MessengerResponseType, StaticTextResponseType } from './messenger/messenger'
 import {
@@ -414,8 +415,19 @@ export class ChatController {
     private async processStopResponseMessage(message: StopResponseMessage) {
         const session = this.sessionStorage.getSession(message.tabID)
         session.tokenSource.cancel()
+        session.setAgenticLoopInProgress(false)
+        session.setToolUseWithError(undefined)
+
+        // Mark any active triggers as cancelled when stopping the response
+        const triggerEvents = this.triggerEventsStorage.getTriggerEventsByTabID(message.tabID)
+        if (triggerEvents && triggerEvents.length > 0) {
+            const conversationTracker = ConversationTracker.getInstance()
+            for (const event of triggerEvents) {
+                conversationTracker.cancelTrigger(event.id)
+            }
+        }
+
         this.messenger.sendEmptyMessage(message.tabID, '', undefined)
-        this.chatHistoryDb.clearRecentHistory(message.tabID)
         this.telemetryHelper.recordInteractionWithAgenticChat(AgenticChatInteractionType.StopChat, message)
     }
 
@@ -472,6 +484,13 @@ export class ChatController {
     }
 
     private async processTabCloseMessage(message: TabClosedMessage) {
+        // First cancel any active triggers to stop ongoing operations
+        const conversationTracker = ConversationTracker.getInstance()
+        conversationTracker.cancelTabTriggers(message.tabID)
+
+        // Then clear all triggers to release resources
+        conversationTracker.clearTabTriggers(message.tabID)
+
         this.sessionStorage.deleteSession(message.tabID)
         this.triggerEventsStorage.removeTabEvents(message.tabID)
         // this.telemetryHelper.recordCloseChat(message.tabID)
@@ -686,7 +705,13 @@ export class ChatController {
         this.editorContextExtractor
             .extractContextForTrigger('ChatMessage')
             .then(async (context) => {
-                const triggerID = randomUUID()
+                const triggerID = message.triggerId
+
+                // Check if this trigger has already been cancelled
+                if (this.isTriggerCancelled(triggerID)) {
+                    return
+                }
+
                 this.triggerEventsStorage.addTriggerEvent({
                     id: triggerID,
                     tabID: message.tabID,
@@ -696,13 +721,16 @@ export class ChatController {
                 })
                 this.messenger.sendAsyncEventProgress(tabID, true, '')
                 const session = this.sessionStorage.getSession(tabID)
+
+                // Check if the session has been cancelled before proceeding
+                if (this.isTriggerCancelled(triggerID)) {
+                    getLogger().debug(`Tool execution cancelled for tabID: ${tabID}`)
+                    return
+                }
+
                 const toolUseWithError = session.toolUseWithError
-                getLogger().debug(
-                    `processToolUseMessage: ${toolUseWithError?.toolUse.name}:${toolUseWithError?.toolUse.toolUseId} with error: ${toolUseWithError?.error}`
-                )
                 if (!toolUseWithError || !toolUseWithError.toolUse) {
-                    // Turn off AgentLoop flag if there's no tool use
-                    this.sessionStorage.setAgentLoopInProgress(tabID, false)
+                    session.setAgenticLoopInProgress(false)
                     return
                 }
                 session.setToolUseWithError(undefined)
@@ -742,7 +770,18 @@ export class ChatController {
                                 const backup = await tool.tool.getBackup()
                                 session.setFsWriteBackup(toolUse.toolUseId, backup)
                             }
-                            const output = await ToolUtils.invoke(tool, chatStream)
+
+                            // Check again if cancelled before invoking the tool
+                            if (this.isTriggerCancelled(triggerID)) {
+                                getLogger().debug(`Tool execution cancelled before invoke for tabID: ${tabID}`)
+                                return
+                            }
+
+                            const output = await ToolUtils.invoke(
+                                tool,
+                                chatStream,
+                                ConversationTracker.getInstance().getTokenForTrigger(triggerID)
+                            )
                             ToolUtils.validateOutput(output)
 
                             toolResults.push({
@@ -1039,7 +1078,17 @@ export class ChatController {
 
         // Turn off AgentLoop flag in case of exception
         if (tabID) {
-            this.sessionStorage.setAgentLoopInProgress(tabID, false)
+            const session = this.sessionStorage.getSession(tabID)
+            session.setAgenticLoopInProgress(false)
+
+            // Mark any active triggers as completed when there's an exception
+            const triggerEvents = this.triggerEventsStorage.getTriggerEventsByTabID(tabID)
+            if (triggerEvents && triggerEvents.length > 0) {
+                const conversationTracker = ConversationTracker.getInstance()
+                for (const event of triggerEvents) {
+                    conversationTracker.cancelTrigger(event.id)
+                }
+            }
         }
 
         this.messenger.sendErrorMessage(errorMessage, tabID, requestID, statusCode)
@@ -1226,23 +1275,33 @@ export class ChatController {
 
     private async processPromptMessageAsNewThread(message: PromptMessage) {
         const session = this.sessionStorage.getSession(message.tabID)
+        // If there's an existing conversation, ensure we dispose the previous token
+        if (session.agenticLoopInProgress) {
+            session.disposeTokenSource()
+        }
+
+        // Create a fresh token for this new conversation
+        session.createNewTokenSource()
+        session.setAgenticLoopInProgress(true)
         session.clearListOfReadFiles()
         session.clearListOfReadFolders()
         session.setShowDiffOnFileWrite(false)
         this.editorContextExtractor
             .extractContextForTrigger('ChatMessage')
             .then(async (context) => {
                 const triggerID = randomUUID()
+
+                // Register the trigger ID with the token for cancellation tracking
+                const conversationTracker = ConversationTracker.getInstance()
+                conversationTracker.registerTrigger(triggerID, session.tokenSource, message.tabID)
+
                 this.triggerEventsStorage.addTriggerEvent({
                     id: triggerID,
                     tabID: message.tabID,
                     message: message.message,
                     type: 'chat_message',
                     context,
                 })
-
-                this.messenger.sendAsyncEventProgress(message.tabID, true, '')
-
                 await this.generateResponse(
                     {
                         message: message.message ?? '',
@@ -1438,16 +1497,6 @@ 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()
 
@@ -1545,7 +1594,11 @@ export class ChatController {
                 session.setContext(triggerPayload.context)
             }
             this.messenger.sendInitalStream(tabID, triggerID)
+            this.messenger.sendAsyncEventProgress(tabID, true, '')
             this.telemetryHelper.setConversationStreamStartTime(tabID)
+            if (this.isTriggerCancelled(triggerID)) {
+                return
+            }
             if (isSsoConnection(AuthUtil.instance.conn)) {
                 const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)
                 response = {
@@ -1562,7 +1615,7 @@ export class ChatController {
             this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
             this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
 
-            if (currentMessage && session.sessionIdentifier) {
+            if (currentMessage && session.sessionIdentifier && !this.isTriggerCancelled(triggerID)) {
                 this.chatHistoryDb.addMessage(tabID, 'cwc', session.sessionIdentifier, {
                     body: triggerPayload.message,
                     type: 'prompt' as any,
@@ -1577,22 +1630,14 @@ 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,
-                this.cancelTokenSource.token
-            )
 
-            // Turn off AgentLoop flag after sending the AI response
-            this.sessionStorage.setAgentLoopInProgress(tabID, false)
+            if (this.isTriggerCancelled(triggerID)) {
+                return
+            }
+
+            await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload)
         } 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)
         }
@@ -1635,4 +1680,17 @@ export class ChatController {
             return { relativeFilePath: filePath, lineRanges: mergedRanges }
         })
     }
+
+    /**
+     * Check if a trigger has been cancelled and should not proceed
+     * @param triggerId The trigger ID to check
+     * @returns true if the trigger is cancelled and should not proceed
+     */
+    private isTriggerCancelled(triggerId: string): boolean {
+        if (!triggerId) {
+            return false
+        }
+        const conversationTracker = ConversationTracker.getInstance()
+        return conversationTracker.isTriggerCancelled(triggerId)
+    }
 }