diff --git a/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts index 3b19598adeb..993d7b17699 100644 --- 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> = new Map() // tabId -> messageId -> ChatItem + private tabIdToTriggerId: Map = 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 index a9711c3ab02..1666ddc801d 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -786,11 +786,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 @@ -808,6 +804,10 @@ 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, @@ -815,6 +815,10 @@ export class ChatController { }) } } else { + if (this.isTriggerCancelled(triggerID)) { + getLogger().debug(`Tool execution cancelled before invoke for tabID: ${tabID}`) + return + } const toolResult: ToolResult = result toolResults.push(toolResult) } @@ -1632,7 +1636,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 index 77390ebdb7d..829e190037e 100644 --- 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 { @@ -232,18 +246,56 @@ export class ExecuteBash { } } - public async invoke(updates?: Writable, cancellationToken?: vscode.CancellationToken): Promise { + /** + * 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 { 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 index 300d357a9fa..1a286b74777 100644 --- 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 { - // Check if cancelled before executing - if (cancellationToken?.isCancellationRequested) { - throw new Error('Tool execution cancelled') - } - + static async invoke(tool: Tool, updates?: Writable, triggerId?: string): Promise { 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 index 31f2ec238f3..96245437684 100644 --- 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) => {