fix(chat): fix to stop bash executions for stop button (aws#7100)

ashishrp-aws · web-flow · commit a597cd2817bf · 2025-04-18T20:19:24.000-07:00
## Problem Long running ExecuteBash commands keep running after clicking on stop button. ## Solution - Added force Stop to kill bash executions - Added tabTotrigger Id mapping to maintain trigger Id consistency across agentic loop. --- - 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.
diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts
@@ -44,6 +44,7 @@ export class Connector extends BaseConnector {
     private readonly onChatAnswerUpdated
     private readonly onAsyncEventProgress
     private chatItems: Map<string, Map<string, ChatItem>> = new Map() // tabId -> messageId -> ChatItem
+    private tabIdToTriggerId: Map<string, string> = new Map() // tabId -> triggerId
 
     override getTabType(): TabType {
         return 'cwc'
@@ -135,6 +136,10 @@ export class Connector extends BaseConnector {
                 }
             }
 
+            if (messageData.tabID && messageData.triggerID && messageData.triggerID !== '') {
+                this.storetriggerId(messageData.tabID, messageData.triggerID)
+            }
+
             if (answer.messageId) {
                 this.storeChatItem(messageData.tabID, answer.messageId, answer)
             }
@@ -213,6 +218,14 @@ export class Connector extends BaseConnector {
         this.chatItems.get(tabId)?.set(messageId, { ...item })
     }
 
+    private storetriggerId(tabId: string, triggerId: string): void {
+        this.tabIdToTriggerId.set(tabId, triggerId)
+    }
+
+    private getTriggerId(tabId: string): string | undefined {
+        return this.tabIdToTriggerId.get(tabId)
+    }
+
     private getCurrentChatItem(tabId: string, messageId: string | undefined): ChatItem | undefined {
         if (!messageId) {
             return
@@ -351,7 +364,7 @@ export class Connector extends BaseConnector {
             formSelectedValues: action.formItemValues,
             tabType: this.getTabType(),
             tabID: tabId,
-            triggerId: triggerId,
+            triggerId: this.getTriggerId(tabId),
         })
 
         if (
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -788,11 +788,7 @@ export class ChatController {
                                 return
                             }
 
-                            const output = await ToolUtils.invoke(
-                                tool,
-                                chatStream,
-                                ConversationTracker.getInstance().getTokenForTrigger(triggerID)
-                            )
+                            const output = await ToolUtils.invoke(tool, chatStream, triggerID)
                             ToolUtils.validateOutput(output, tool.type)
 
                             let status: ToolResultStatus = ToolResultStatus.SUCCESS
@@ -810,13 +806,21 @@ export class ChatController {
                                 status,
                             })
                         } catch (e: any) {
+                            if (this.isTriggerCancelled(triggerID)) {
+                                getLogger().debug(`Tool execution cancelled before invoke for tabID: ${tabID}`)
+                                return
+                            }
                             toolResults.push({
                                 content: [{ text: e.message }],
                                 toolUseId: toolUse.toolUseId,
                                 status: ToolResultStatus.ERROR,
                             })
                         }
                     } else {
+                        if (this.isTriggerCancelled(triggerID)) {
+                            getLogger().debug(`Tool execution cancelled before invoke for tabID: ${tabID}`)
+                            return
+                        }
                         const toolResult: ToolResult = result
                         toolResults.push(toolResult)
                     }
@@ -1634,7 +1638,6 @@ export class ChatController {
             })}`
         )
         let response: MessengerResponseType | undefined = undefined
-        session.createNewTokenSource()
         // TODO: onProfileChanged, abort previous response?
         try {
             if (!session.context && triggerPayload.context.length) {
diff --git a/packages/core/src/codewhispererChat/tools/executeBash.ts b/packages/core/src/codewhispererChat/tools/executeBash.ts
@@ -12,6 +12,7 @@ import { split } from 'shlex'
 import path from 'path'
 import * as vscode from 'vscode'
 import { isInDirectory } from '../../shared/filesystemUtilities'
+import { ConversationTracker } from '../storages/conversationTracker'
 
 export enum CommandCategory {
     ReadOnly,
@@ -114,6 +115,7 @@ export interface ExecuteBashParams {
     command: string
     cwd?: string
     explanation?: string
+    triggerId?: string
 }
 
 export interface CommandValidation {
@@ -134,10 +136,22 @@ export class ExecuteBash {
     private readonly workingDirectory?: string
     private readonly logger = getLogger('executeBash')
     private childProcess?: ChildProcess
+    // Make triggerId writable so it can be set after construction
+    private _triggerId?: string
 
     constructor(params: ExecuteBashParams) {
         this.command = params.command
         this.workingDirectory = params.cwd ? sanitizePath(params.cwd) : fs.getUserHomeDir()
+        this._triggerId = params.triggerId
+    }
+
+    // Getter and setter for triggerId
+    get triggerId(): string | undefined {
+        return this._triggerId
+    }
+
+    set triggerId(id: string | undefined) {
+        this._triggerId = id
     }
 
     public async validate(): Promise<void> {
@@ -232,18 +246,56 @@ export class ExecuteBash {
         }
     }
 
-    public async invoke(updates?: Writable, cancellationToken?: vscode.CancellationToken): Promise<InvokeOutput> {
+    /**
+     * Check if the trigger has been cancelled using ConversationTracker
+     */
+    private isTriggerCancelled(): boolean {
+        if (!this.triggerId) {
+            return false
+        }
+        const cancellationtracker = ConversationTracker.getInstance()
+        return cancellationtracker.isTriggerCancelled(this.triggerId)
+    }
+
+    public async invoke(updates?: Writable): Promise<InvokeOutput> {
         this.logger.info(`Invoking bash command: "${this.command}" in cwd: "${this.workingDirectory}"`)
 
         return new Promise(async (resolve, reject) => {
-            // Check if cancelled before starting
-            if (cancellationToken?.isCancellationRequested) {
+            // Check if cancelled before starting using triggerId
+            if (this.isTriggerCancelled()) {
                 this.logger.debug('Bash command execution cancelled before starting')
                 reject(new Error('Command execution cancelled'))
                 return
             }
 
-            this.logger.debug(`Spawning process with command: bash -c "${this.command}" (cwd=${this.workingDirectory})`)
+            // Modify the command to make it more cancellable by using process groups
+            // This ensures that when we kill the parent process, all child processes are also terminated
+            // The trap ensures cleanup on SIGTERM/SIGINT and sends SIGTERM to the child process group
+            const modifiedCommand = `
+exec bash -c "
+  # Create a new process group
+  set -m
+  
+  # Set up trap to kill the entire process group on exit
+  trap 'kill -TERM -\\$CMD_PID 2>/dev/null || true; exit' TERM INT
+  
+  # Run the actual command in background
+  # Use '()' to create a subshell which becomes the process group leader
+  (${this.command}) &
+  
+  # Store the PID
+  CMD_PID=\\$!
+  
+  # Wait for the command to finish
+  wait \\$CMD_PID
+  exit_code=\\$?
+  exit \\$exit_code
+"
+`
+
+            this.logger.debug(
+                `Spawning process with modified command for better cancellation support (cwd=${this.workingDirectory})`
+            )
 
             const stdoutBuffer: string[] = []
             const stderrBuffer: string[] = []
@@ -282,16 +334,53 @@ export class ExecuteBash {
                 }
             }
 
+            // Setup a periodic check for trigger cancellation
+            let checkCancellationInterval: NodeJS.Timeout | undefined
+            if (this.triggerId) {
+                checkCancellationInterval = setInterval(() => {
+                    if (this.isTriggerCancelled()) {
+                        this.logger.debug('Trigger cancellation detected, killing child process')
+
+                        // Kill the process
+                        this.childProcess?.stop(false, 'SIGTERM')
+
+                        // After a short delay, force kill with SIGKILL if still running
+                        setTimeout(() => {
+                            if (this.childProcess && !this.childProcess.stopped) {
+                                this.logger.debug('Process still running after SIGTERM, sending SIGKILL')
+
+                                // Try to kill the process group with SIGKILL
+                                this.childProcess.stop(true, 'SIGKILL')
+                            }
+                        }, 500)
+
+                        if (checkCancellationInterval) {
+                            clearInterval(checkCancellationInterval)
+                        }
+
+                        // Return from the function after cancellation
+                        reject(new Error('Command execution cancelled'))
+                        return
+                    }
+                }, 500) // Check every 500ms
+            }
+
             const childProcessOptions: ChildProcessOptions = {
                 spawnOptions: {
                     cwd: this.workingDirectory,
                     stdio: ['pipe', 'pipe', 'pipe'],
+                    // Set detached to true to create a new process group
+                    // This allows us to kill the entire process group later
+                    detached: true,
+                    // On Windows, we need to create a new process group
+                    // On Unix, we need to create a new session
+                    ...(process.platform === 'win32' ? { windowsVerbatimArguments: true } : {}),
                 },
                 collect: false,
                 waitForStreams: true,
                 onStdout: async (chunk: string) => {
-                    if (cancellationToken?.isCancellationRequested) {
-                        this.logger.debug('Bash command execution cancelled during stderr processing')
+                    if (this.isTriggerCancelled()) {
+                        this.logger.debug('Bash command execution cancelled during stdout processing')
                         return
                     }
                     const isFirst = getAndSetFirstChunk(false)
@@ -305,7 +394,7 @@ export class ExecuteBash {
                     processQueue()
                 },
                 onStderr: async (chunk: string) => {
-                    if (cancellationToken?.isCancellationRequested) {
+                    if (this.isTriggerCancelled()) {
                         this.logger.debug('Bash command execution cancelled during stderr processing')
                         return
                     }
@@ -321,21 +410,19 @@ export class ExecuteBash {
                 },
             }
 
-            this.childProcess = new ChildProcess('bash', ['-c', this.command], childProcessOptions)
-
-            // Set up cancellation listener
-            if (cancellationToken) {
-                cancellationToken.onCancellationRequested(() => {
-                    this.logger.debug('Cancellation requested, killing child process')
-                    this.childProcess?.stop()
-                })
-            }
+            // Use bash directly with the modified command
+            this.childProcess = new ChildProcess('bash', ['-c', modifiedCommand], childProcessOptions)
 
             try {
                 const result = await this.childProcess.run()
 
+                // Clean up the interval if it exists
+                if (checkCancellationInterval) {
+                    clearInterval(checkCancellationInterval)
+                }
+
                 // Check if cancelled after execution
-                if (cancellationToken?.isCancellationRequested) {
+                if (this.isTriggerCancelled()) {
                     this.logger.debug('Bash command execution cancelled after completion')
                     reject(new Error('Command execution cancelled'))
                     return
@@ -368,8 +455,13 @@ export class ExecuteBash {
                     },
                 })
             } catch (err: any) {
+                // Clean up the interval if it exists
+                if (checkCancellationInterval) {
+                    clearInterval(checkCancellationInterval)
+                }
+
                 // Check if this was due to cancellation
-                if (cancellationToken?.isCancellationRequested) {
+                if (this.isTriggerCancelled()) {
                     reject(new Error('Command execution cancelled'))
                 } else {
                     this.logger.error(`Failed to execute bash command '${this.command}': ${err.message}`)
diff --git a/packages/core/src/codewhispererChat/tools/toolUtils.ts b/packages/core/src/codewhispererChat/tools/toolUtils.ts
@@ -2,7 +2,6 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import * as vscode from 'vscode'
 import { Writable } from 'stream'
 import { FsRead, FsReadParams } from './fsRead'
 import { FsWrite, FsWriteParams } from './fsWrite'
@@ -56,23 +55,18 @@ export class ToolUtils {
         }
     }
 
-    static async invoke(
-        tool: Tool,
-        updates?: Writable,
-        cancellationToken?: vscode.CancellationToken
-    ): Promise<InvokeOutput> {
-        // Check if cancelled before executing
-        if (cancellationToken?.isCancellationRequested) {
-            throw new Error('Tool execution cancelled')
-        }
-
+    static async invoke(tool: Tool, updates?: Writable, triggerId?: string): Promise<InvokeOutput> {
         switch (tool.type) {
             case ToolType.FsRead:
                 return tool.tool.invoke(updates)
             case ToolType.FsWrite:
                 return tool.tool.invoke(updates)
             case ToolType.ExecuteBash:
-                return tool.tool.invoke(updates ?? undefined, cancellationToken)
+                // If triggerId is provided, update the tool's triggerId
+                if (triggerId) {
+                    ;(tool.tool as ExecuteBash).triggerId = triggerId
+                }
+                return tool.tool.invoke(updates)
             case ToolType.ListDirectory:
                 return tool.tool.invoke(updates)
         }
diff --git a/packages/core/src/shared/utilities/processUtils.ts b/packages/core/src/shared/utilities/processUtils.ts
@@ -409,6 +409,41 @@ export class ChildProcess {
         return this.#childProcess?.exitCode ?? -1
     }
 
+    /**
+     * Attempts to kill a process group on Unix systems, falling back to regular kill if it fails
+     * @param child The child process to kill
+     * @param signal The signal to send (defaults to SIGTERM)
+     * @param checkSpawnProps Whether to check for child.spawnargs and child.spawnfile (used in first call only)
+     * @returns true if process group kill was attempted, false if regular kill was used
+     */
+    public killProcessGroup(
+        child: proc.ChildProcess,
+        signal: NodeJS.Signals = 'SIGTERM',
+        checkSpawnProps: boolean = false
+    ): boolean {
+        // First call in stop() uses the stricter condition with spawnargs/spawnfile checks
+        // Second call in the force kill uses the less restrictive condition
+        const condition = checkSpawnProps
+            ? process.platform !== 'win32' && child.spawnargs && child.spawnfile
+            : process.platform !== 'win32'
+
+        if (condition && child.pid) {
+            try {
+                // On Unix systems, negative PID kills the process group
+                this.#log.debug(`Attempting to kill process group -${child.pid} with ${signal}`)
+                process.kill(-child.pid, signal)
+                return true
+            } catch (err) {
+                this.#log.debug(`Failed to kill process group with ${signal}: ${err}, falling back to regular kill`)
+                child.kill(signal)
+                return false
+            }
+        } else {
+            child.kill(signal)
+            return false
+        }
+    }
+
     /**
      * Stops the process.
      *
@@ -426,13 +461,17 @@ export class ChildProcess {
         const command = this.#command
         const pid = this.pid()
         if (!this.stopped) {
-            child.kill(signal)
+            // Try to kill the process group if possible (Unix systems only)
+            // First call uses the stricter condition with spawnargs/spawnfile checks
+            this.killProcessGroup(child, signal || 'SIGTERM', true)
 
             if (force === true) {
                 waitUntil(async () => this.stopped, { timeout: ChildProcess.stopTimeout, interval: 200, truthy: true })
                     .then((stopped) => {
                         if (!stopped) {
-                            child.kill('SIGKILL')
+                            // Try to kill the process group with SIGKILL if possible
+                            // Second call uses the less restrictive condition
+                            this.killProcessGroup(child, 'SIGKILL', false)
                         }
                     })
                     .catch((e) => {
