fix(chat): agentic chat working stop

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

@@ -32,7 +32,8 @@ export interface ConnectorProps extends BaseConnectorProps {
         inProgress: boolean,
         message: string,
         messageId: string | undefined,
-        enableStopAction: boolean
+        enableStopAction: boolean,
+        isPromptInputDisabled: boolean
     ) => void
 }
 
@@ -259,7 +260,8 @@ export class Connector extends BaseConnector {
                 messageData.inProgress,
                 messageData.message ?? undefined,
                 messageData.messageId ?? undefined,
-                enableStopAction
+                enableStopAction,
+                false
             )
             return
         }

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

@@ -287,12 +287,13 @@ export const createMynahUI = (
             inProgress: boolean,
             message: string | undefined,
             messageId: string | undefined = undefined,
-            enableStopAction: boolean = false
+            enableStopAction: boolean = false,
+            isPromptInputDisabled: boolean = true
         ) => {
             if (inProgress) {
                 mynahUI.updateStore(tabID, {
                     loadingChat: true,
-                    promptInputDisabledState: true,
+                    promptInputDisabledState: isPromptInputDisabled,
                     cancelButtonWhenLoading: enableStopAction,
                 })
 

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

@@ -373,6 +373,7 @@ export class ChatController {
     private async processStopResponseMessage(message: StopResponseMessage) {
         const session = this.sessionStorage.getSession(message.tabID)
         session.tokenSource.cancel()
+        this.messenger.sendEmptyMessage(message.tabID, '', undefined)
         this.chatHistoryStorage.getTabHistory(message.tabID).clearRecentHistory()
     }
 

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

@@ -735,4 +735,28 @@ export class Messenger {
     public sendAsyncEventProgress(tabID: string, inProgress: boolean, message: string | undefined) {
         this.dispatcher.sendAsyncEventProgress(new AsyncEventProgressMessage(tabID, 'CWChat', inProgress, message))
     }
+
+    public sendEmptyMessage(
+        tabID: string,
+        triggerId: string,
+        mergedRelevantDocuments: DocumentReference[] | undefined
+    ) {
+        this.dispatcher.sendChatMessage(
+            new ChatMessage(
+                {
+                    message: '',
+                    messageType: 'answer',
+                    followUps: undefined,
+                    followUpsHeader: undefined,
+                    relatedSuggestions: undefined,
+                    triggerID: triggerId,
+                    messageID: '',
+                    userIntent: undefined,
+                    codeBlockLanguage: undefined,
+                    contextList: undefined,
+                },
+                tabID
+            )
+        )
+    }
 }

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

@@ -80,8 +80,13 @@ export class ChatHistoryManager {
      * Push an assistant message to the history
      */
     public pushAssistantMessage(newMessage: ChatMessage): void {
-        if (newMessage !== undefined && this.lastUserMessage !== undefined) {
-            this.logger.warn('last Message should not be defined when pushing an assistant message')
+        if (newMessage !== undefined && this.lastUserMessage === undefined) {
+            this.logger.warn('first assistant response should always come after user input message')
+            return
+        }
+        // check if last message in histroy is assistant message and now replace it in that case
+        if (this.history.length > 0 && this.history.at(-1)?.assistantResponseMessage) {
+            this.history.pop()
         }
         this.history.push(newMessage)
     }
@@ -101,6 +106,18 @@ export class ChatHistoryManager {
     }
 
     private trimConversationHistory(): void {
+        // make sure the UseInputMessage is the first stored message
+        if (this.history.length === 1 && this.history[0].assistantResponseMessage) {
+            this.history = []
+        }
+
+        if (
+            this.history.at(-1)?.assistantResponseMessage?.content === '' &&
+            this.history.at(-1)?.assistantResponseMessage?.toolUses === undefined
+        ) {
+            this.clearRecentHistory()
+        }
+
         if (this.history.length <= MaxConversationHistoryLength) {
             return
         }