fix for new user prompt stoppage

Author: ashishrp-aws

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

@@ -257,13 +257,12 @@ export class Connector extends BaseConnector {
         }
 
         if (messageData.type === 'asyncEventProgressMessage') {
-            const enableStopAction = true
             this.onAsyncEventProgress(
                 messageData.tabID,
                 messageData.inProgress,
                 messageData.message ?? undefined,
                 messageData.messageId ?? undefined,
-                enableStopAction,
+                messageData.inProgress,
                 false
             )
             return

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

@@ -691,6 +691,7 @@ export const createMynahUI = (
         onTabChange: connector.onTabChange,
         // TODO: update mynah-ui this type doesn't seem correct https://github.com/aws/mynah-ui/blob/3777a39eb534a91fd6b99d6cf421ce78ee5c7526/src/main.ts#L372
         onStopChatResponse: (tabID: string) => {
+            console.log('stop response started')
             mynahUI.updateStore(tabID, {
                 loadingChat: false,
                 promptInputDisabledState: false,
@@ -702,6 +703,12 @@ export const createMynahUI = (
                 return
             }
 
+            mynahUI.updateStore(tabID, {
+                loadingChat: false,
+                promptInputDisabledState: false,
+            })
+            connector.onStopChatResponse(tabID)
+
             const tabType = tabsStorage.getTab(tabID)?.type
             if (tabType === 'featuredev') {
                 mynahUI.addChatItem(tabID, {
@@ -736,6 +743,8 @@ export const createMynahUI = (
                 return
             }
 
+            console.log('prompt entered: ', prompt)
+
             /**
              * Update the tab title if coming from the welcome page
              * non cwc panels will have this updated automatically

packages/core/src/codewhisperer/client/codewhisperer.ts

@@ -31,8 +31,8 @@ export interface CodeWhispererConfig {
 }
 
 export const defaultServiceConfig: CodeWhispererConfig = {
-    region: 'us-east-1',
-    endpoint: 'https://codewhisperer.us-east-1.amazonaws.com/',
+    region: 'us-west-2',
+    endpoint: 'https://rts.alpha-us-west-2.codewhisperer.ai.aws.dev/',
 }
 
 export function getCodewhispererConfig(): CodeWhispererConfig {

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

@@ -37,6 +37,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.
      */
@@ -50,6 +52,33 @@ export class ChatSession {
         return this.sessionId
     }
 
+    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) {
+                // Ignore errors during disposal
+            }
+        }
+    }
+
     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 {
@@ -101,6 +102,7 @@ import { FsWriteParams } from '../../tools/fsWrite'
 import { tempDirPath } from '../../../shared/filesystemUtilities'
 import { Database } from '../../../shared/db/chatDb/chatDb'
 import { TabBarController } from './tabBarController'
+import { sleep } from '../../../shared'
 
 export interface ChatControllerMessagePublishers {
     readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -407,9 +409,25 @@ export class ChatController {
 
     private async processStopResponseMessage(message: StopResponseMessage) {
         const session = this.sessionStorage.getSession(message.tabID)
+        const wasInAgenticLoop = session.agenticLoopInProgress
         session.tokenSource.cancel()
+        console.log('process stop has been triggered')
+        session.setAgenticLoopInProgress(false)
+        session.setToolUseWithError(undefined)
+
+        // Mark any active triggers as completed when stopping the response
+        const triggerEvents = this.triggerEventsStorage.getTriggerEventsByTabID(message.tabID)
+        if (triggerEvents && triggerEvents.length > 0) {
+            const conversationTracker = ConversationTracker.getInstance()
+            triggerEvents.forEach((event) => {
+                conversationTracker.markTriggerCompleted(event.id)
+            })
+        }
+
+        wasInAgenticLoop && (await sleep(1000))
+
         this.messenger.sendEmptyMessage(message.tabID, '', undefined)
-        this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
+        // this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
     }
 
     private async processTriggerTabIDReceived(message: TriggerTabIDReceived) {
@@ -465,6 +483,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.chatHistoryStorage.deleteHistory(message.tabID)
         this.triggerEventsStorage.removeTabEvents(message.tabID)
@@ -680,7 +705,7 @@ export class ChatController {
         this.editorContextExtractor
             .extractContextForTrigger('ChatMessage')
             .then(async (context) => {
-                const triggerID = randomUUID()
+                const triggerID = message.triggerId ?? randomUUID()
                 this.triggerEventsStorage.addTriggerEvent({
                     id: triggerID,
                     tabID: message.tabID,
@@ -690,10 +715,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 (session.tokenSource.token.isCancellationRequested) {
+                    getLogger().debug(`Tool execution cancelled for tabID: ${tabID}`)
+                    return
+                }
+
                 const toolUseWithError = session.toolUseWithError
                 if (!toolUseWithError || !toolUseWithError.toolUse || !toolUseWithError.toolUse.input) {
                     // Turn off AgentLoop flag if there's no tool use
-                    this.sessionStorage.setAgentLoopInProgress(tabID, false)
                     return
                 }
                 session.setToolUseWithError(undefined)
@@ -716,14 +747,33 @@ export class ChatController {
                         try {
                             await ToolUtils.validate(tool)
 
-                            const chatStream = new ChatStream(this.messenger, tabID, triggerID, toolUse, {
-                                requiresAcceptance: false,
-                            })
+                            // Get the cancellation token from the session
+                            const cancellationToken = session.tokenSource.token
+
+                            // Pass the cancellation token to ChatStream
+                            const chatStream = new ChatStream(
+                                this.messenger,
+                                tabID,
+                                triggerID,
+                                toolUse,
+                                { requiresAcceptance: false },
+                                undefined,
+                                undefined,
+                                cancellationToken
+                            )
+
                             if (tool.type === ToolType.FsWrite && toolUse.toolUseId) {
                                 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 (cancellationToken.isCancellationRequested) {
+                                getLogger().debug(`Tool execution cancelled before invoke for tabID: ${tabID}`)
+                                return
+                            }
+
+                            const output = await ToolUtils.invoke(tool, chatStream, cancellationToken)
                             ToolUtils.validateOutput(output)
 
                             toolResults.push({
@@ -985,7 +1035,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()
+                triggerEvents.forEach((event) => {
+                    conversationTracker.markTriggerCompleted(event.id)
+                })
+            }
         }
 
         this.messenger.sendErrorMessage(errorMessage, tabID, requestID)
@@ -1171,22 +1231,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.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)
+                console.log('conversation tracker:', conversationTracker.getTokenForTrigger(triggerID))
+
                 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 ?? '',
@@ -1381,16 +1452,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()
 
@@ -1493,6 +1554,7 @@ export class ChatController {
                 session.setContext(triggerPayload.context)
             }
             this.messenger.sendInitalStream(tabID, triggerID)
+            this.messenger.sendAsyncEventProgress(tabID, true, '')
             this.telemetryHelper.setConversationStreamStartTime(tabID)
             if (isSsoConnection(AuthUtil.instance.conn)) {
                 const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)
@@ -1527,13 +1589,8 @@ 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)
         }