jcarr
/
gemini-cli
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

allow aborting of shell mode (!) commands, similar to shell tool commands. fix bug that prevented aborts after first abort. more robust killing logic (#616)

This commit is contained in:
Olcan 2025-05-30 00:46:43 -07:00 committed by GitHub
parent b0aeeb53b1
commit 8935a248f6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 92 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
packages/cli/src/ui/hooks/shellCommandProcessor.ts
View File

@ -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(

4
packages/cli/src/ui/hooks/useGeminiStream.ts
View File

@ -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(

51
packages/server/src/tools/shell.ts
View File

@ -25,8 +25,6 @@ export interface ShellToolParams {
}
import { spawn } from 'child_process';
const OUTPUT_UPDATE_INTERVAL_MS = 1000;
export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
static Name: string = 'execute_bash_command';
private whitelist: Set<string> = new Set();
@ -126,7 +124,7 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
async execute(
params: ShellToolParams,
abortSignal: AbortSignal,
updateOutput?: (output: string) => void,
onOutputChunk?: (chunk: string) => void,
): Promise<ToolResult> {
const validationError = this.validateToolParams(params);
if (validationError) {
@ -157,19 +155,6 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
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<ShellToolParams, ToolResult> {
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<ShellToolParams, ToolResult> {
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<ShellToolParams, ToolResult> {
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}`);
}
}
}