Skip to content
15 changes: 14 additions & 1 deletion packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -351,7 +364,7 @@ export class Connector extends BaseConnector {
formSelectedValues: action.formItemValues,
tabType: this.getTabType(),
tabID: tabId,
triggerId: triggerId,
triggerId: this.getTriggerId(tabId),
})

if (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -808,13 +804,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)
}
Expand Down Expand Up @@ -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) {
Expand Down
128 changes: 110 additions & 18 deletions packages/core/src/codewhispererChat/tools/executeBash.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -114,6 +115,7 @@ export interface ExecuteBashParams {
command: string
cwd?: string
explanation?: string
triggerId?: string
}

export interface CommandValidation {
Expand All @@ -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> {
Expand Down Expand Up @@ -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}) &
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of inlining the command literally, it would avoid problems with "escaping" to write the command to a file, then execute the file here (bash foo.sh).


# 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[] = []
Expand Down Expand Up @@ -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)
Expand All @@ -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
}
Expand All @@ -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
Expand Down Expand Up @@ -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}`)
Expand Down
18 changes: 6 additions & 12 deletions packages/core/src/codewhispererChat/tools/toolUtils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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)
}
Expand Down
43 changes: 41 additions & 2 deletions packages/core/src/shared/utilities/processUtils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand All @@ -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) => {
Expand Down
Loading