From 8935a248f67d70c1f3892e7eb05aee38c8a11a4c Mon Sep 17 00:00:00 2001 From: Olcan Date: Fri, 30 May 2025 00:46:43 -0700 Subject: [PATCH] allow aborting of shell mode (!) commands, similar to shell tool commands. fix bug that prevented aborts after first abort. more robust killing logic (#616) --- .../cli/src/ui/hooks/shellCommandProcessor.ts | 68 +++++++++++++++++-- packages/cli/src/ui/hooks/useGeminiStream.ts | 4 +- packages/server/src/tools/shell.ts | 51 +++++++------- 3 files changed, 92 insertions(+), 31 deletions(-) diff --git a/packages/cli/src/ui/hooks/shellCommandProcessor.ts b/packages/cli/src/ui/hooks/shellCommandProcessor.ts index 59e337b4..d6c837d1 100644 --- a/packages/cli/src/ui/hooks/shellCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/shellCommandProcessor.ts @@ -37,7 +37,7 @@ export const useShellCommandProcessor = ( * @returns True if the query was handled as a shell command, false otherwise. */ const handleShellCommand = useCallback( - (rawQuery: PartListUnion): boolean => { + (rawQuery: PartListUnion, signal?: AbortSignal): boolean => { if (typeof rawQuery !== 'string') { return false; } @@ -120,6 +120,7 @@ export const useShellCommandProcessor = ( const child = spawn('bash', ['-c', commandToExecute], { cwd: targetDir, stdio: ['ignore', 'pipe', 'pipe'], + detached: true, // Important for process group killing }); let exited = false; @@ -148,18 +149,75 @@ export const useShellCommandProcessor = ( error = err; }); - child.on('exit', (code, signal) => { + const abortHandler = () => { + if (child.pid && !exited) { + onDebugMessage( + `Aborting shell command (PID: ${child.pid}) due to signal.`, + ); + try { + // attempt to SIGTERM process group (negative PID) + // if SIGTERM fails after 200ms, attempt SIGKILL + process.kill(-child.pid, 'SIGTERM'); + setTimeout(() => { + // if SIGTERM fails, attempt SIGKILL + try { + if (child.pid && !exited) { + process.kill(-child.pid, 'SIGKILL'); + } + } catch (_e) { + console.error( + `failed to kill shell process ${child.pid}: ${_e}`, + ); + } + }, 200); + } catch (_e) { + // if group kill fails, fall back to killing just the main process + try { + if (child.pid) { + child.kill('SIGKILL'); + } + } catch (_e) { + console.error( + `failed to kill shell process ${child.pid}: ${_e}`, + ); + } + } + } + }; + + if (signal) { + if (signal.aborted) { + abortHandler(); + // No need to add listener if already aborted + } else { + signal.addEventListener('abort', abortHandler, { once: true }); + } + } + + child.on('exit', (code, processSignal) => { exited = true; + if (signal) { + signal.removeEventListener('abort', abortHandler); + } setPendingHistoryItem(null); output = output.trim() || '(Command produced no output)'; if (error) { const text = `${error.message.replace(commandToExecute, rawQuery)}\n${output}`; addItemToHistory({ type: 'error', text }, userMessageTimestamp); - } else if (code !== 0) { + } else if (code !== null && code !== 0) { const text = `Command exited with code ${code}\n${output}`; addItemToHistory({ type: 'error', text }, userMessageTimestamp); - } else if (signal) { - const text = `Command terminated with signal ${signal}\n${output}`; + } else if (signal?.aborted) { + addItemToHistory( + { + type: 'info', + text: `Command was cancelled.\n${output}`, + }, + userMessageTimestamp, + ); + } else if (processSignal) { + const text = `Command terminated with signal ${processSignal} +${output}`; addItemToHistory({ type: 'error', text }, userMessageTimestamp); } else { addItemToHistory( diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index afaf0ccd..66fb317a 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -199,7 +199,7 @@ export const useGeminiStream = ( return { queryToSend: null, shouldProceed: false }; // Handled by scheduling the tool } - if (shellModeActive && handleShellCommand(trimmedQuery)) { + if (shellModeActive && handleShellCommand(trimmedQuery, signal)) { return { queryToSend: null, shouldProceed: false }; } @@ -492,7 +492,7 @@ export const useGeminiStream = ( const userMessageTimestamp = Date.now(); setShowHelp(false); - abortControllerRef.current ??= new AbortController(); + abortControllerRef.current = new AbortController(); const signal = abortControllerRef.current.signal; const { queryToSend, shouldProceed } = await prepareQueryForGemini( diff --git a/packages/server/src/tools/shell.ts b/packages/server/src/tools/shell.ts index 38848c4f..a4634a3e 100644 --- a/packages/server/src/tools/shell.ts +++ b/packages/server/src/tools/shell.ts @@ -25,8 +25,6 @@ export interface ShellToolParams { } import { spawn } from 'child_process'; -const OUTPUT_UPDATE_INTERVAL_MS = 1000; - export class ShellTool extends BaseTool { static Name: string = 'execute_bash_command'; private whitelist: Set = new Set(); @@ -126,7 +124,7 @@ export class ShellTool extends BaseTool { async execute( params: ShellToolParams, abortSignal: AbortSignal, - updateOutput?: (output: string) => void, + onOutputChunk?: (chunk: string) => void, ): Promise { const validationError = this.validateToolParams(params); if (validationError) { @@ -157,19 +155,6 @@ export class ShellTool extends BaseTool { let exited = false; let stdout = ''; let output = ''; - let lastUpdateTime = Date.now(); - - const appendOutput = (str: string) => { - output += str; - if ( - updateOutput && - Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS - ) { - updateOutput(output); - lastUpdateTime = Date.now(); - } - }; - shell.stdout.on('data', (data: Buffer) => { // continue to consume post-exit for background processes // removing listeners can overflow OS buffer and block subprocesses @@ -177,7 +162,10 @@ export class ShellTool extends BaseTool { if (!exited) { const str = data.toString(); stdout += str; - appendOutput(str); + output += str; + if (onOutputChunk) { + onOutputChunk(str); + } } }); @@ -186,7 +174,10 @@ export class ShellTool extends BaseTool { if (!exited) { const str = data.toString(); stderr += str; - appendOutput(str); + output += str; + if (onOutputChunk) { + onOutputChunk(str); + } } }); @@ -210,16 +201,28 @@ export class ShellTool extends BaseTool { shell.on('exit', exitHandler); const abortHandler = () => { - if (shell.pid) { + if (shell.pid && !exited) { try { - // Kill the entire process group + // attempt to SIGTERM process group (negative PID) + // if SIGTERM fails after 200ms, attempt SIGKILL process.kill(-shell.pid, 'SIGTERM'); + setTimeout(() => { + try { + if (shell.pid && !exited) { + process.kill(-shell.pid, 'SIGKILL'); + } + } catch (_e) { + console.error(`failed to kill shell process ${shell.pid}: ${_e}`); + } + }, 200); } catch (_e) { - // Fallback to killing the main process if group kill fails + // if group kill fails, fall back to killing just the main process try { - shell.kill('SIGKILL'); // or 'SIGTERM' - } catch (_killError) { - // Ignore errors if the process is already dead + if (shell.pid) { + shell.kill('SIGKILL'); + } + } catch (_e) { + console.error(`failed to kill shell process ${shell.pid}: ${_e}`); } } }