diff options
Diffstat (limited to 'packages/core/src/tools/shell.ts')
| -rw-r--r-- | packages/core/src/tools/shell.ts | 347 |
1 files changed, 158 insertions, 189 deletions
diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index d90ae678..02fcbb7f 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -20,22 +20,25 @@ import { import { Type } from '@google/genai'; import { SchemaValidator } from '../utils/schemaValidator.js'; import { getErrorMessage } from '../utils/errors.js'; -import stripAnsi from 'strip-ansi'; +import { summarizeToolOutput } from '../utils/summarizer.js'; +import { + ShellExecutionService, + ShellOutputEvent, +} from '../services/shellExecutionService.js'; +import { formatMemoryUsage } from '../utils/formatters.js'; import { getCommandRoots, isCommandAllowed, stripShellWrapper, } from '../utils/shell-utils.js'; +export const OUTPUT_UPDATE_INTERVAL_MS = 1000; + export interface ShellToolParams { command: string; description?: string; directory?: string; } -import { spawn } from 'child_process'; -import { summarizeToolOutput } from '../utils/summarizer.js'; - -const OUTPUT_UPDATE_INTERVAL_MS = 1000; export class ShellTool extends BaseTool<ShellToolParams, ToolResult> { static Name: string = 'run_shell_command'; @@ -196,216 +199,182 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> { .toString('hex')}.tmp`; const tempFilePath = path.join(os.tmpdir(), tempFileName); - // pgrep is not available on Windows, so we can't get background PIDs - const commandToExecute = isWindows - ? strippedCommand - : (() => { - // wrap command to append subprocess pids (via pgrep) to temporary file - let command = strippedCommand.trim(); - if (!command.endsWith('&')) command += ';'; - return `{ ${command} }; __code=$?; pgrep -g 0 >${tempFilePath} 2>&1; exit $__code;`; - })(); - - // spawn command in specified directory (or project root if not specified) - const shell = isWindows - ? spawn('cmd.exe', ['/c', commandToExecute], { - stdio: ['ignore', 'pipe', 'pipe'], - // detached: true, // ensure subprocess starts its own process group (esp. in Linux) - cwd: path.resolve(this.config.getTargetDir(), params.directory || ''), - env: { - ...process.env, - GEMINI_CLI: '1', - }, - }) - : spawn('bash', ['-c', commandToExecute], { - stdio: ['ignore', 'pipe', 'pipe'], - detached: true, // ensure subprocess starts its own process group (esp. in Linux) - cwd: path.resolve(this.config.getTargetDir(), params.directory || ''), - env: { - ...process.env, - GEMINI_CLI: '1', - }, - }); + try { + // pgrep is not available on Windows, so we can't get background PIDs + const commandToExecute = isWindows + ? strippedCommand + : (() => { + // wrap command to append subprocess pids (via pgrep) to temporary file + let command = strippedCommand.trim(); + if (!command.endsWith('&')) command += ';'; + return `{ ${command} }; __code=$?; pgrep -g 0 >${tempFilePath} 2>&1; exit $__code;`; + })(); - let exited = false; - let stdout = ''; - let output = ''; - let lastUpdateTime = Date.now(); + const cwd = path.resolve( + this.config.getTargetDir(), + params.directory || '', + ); - const appendOutput = (str: string) => { - output += str; - if ( - updateOutput && - Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS - ) { - updateOutput(output); - lastUpdateTime = Date.now(); - } - }; + let cumulativeStdout = ''; + let cumulativeStderr = ''; - shell.stdout.on('data', (data: Buffer) => { - // continue to consume post-exit for background processes - // removing listeners can overflow OS buffer and block subprocesses - // destroying (e.g. shell.stdout.destroy()) can terminate subprocesses via SIGPIPE - if (!exited) { - const str = stripAnsi(data.toString()); - stdout += str; - appendOutput(str); - } - }); + let lastUpdateTime = Date.now(); + let isBinaryStream = false; - let stderr = ''; - shell.stderr.on('data', (data: Buffer) => { - if (!exited) { - const str = stripAnsi(data.toString()); - stderr += str; - appendOutput(str); - } - }); - - let error: Error | null = null; - shell.on('error', (err: Error) => { - error = err; - // remove wrapper from user's command in error message - error.message = error.message.replace(commandToExecute, params.command); - }); + const { result: resultPromise } = ShellExecutionService.execute( + commandToExecute, + cwd, + (event: ShellOutputEvent) => { + if (!updateOutput) { + return; + } - let code: number | null = null; - let processSignal: NodeJS.Signals | null = null; - const exitHandler = ( - _code: number | null, - _signal: NodeJS.Signals | null, - ) => { - exited = true; - code = _code; - processSignal = _signal; - }; - shell.on('exit', exitHandler); + let currentDisplayOutput = ''; + let shouldUpdate = false; - const abortHandler = async () => { - if (shell.pid && !exited) { - if (os.platform() === 'win32') { - // For Windows, use taskkill to kill the process tree - spawn('taskkill', ['/pid', shell.pid.toString(), '/f', '/t']); - } else { - try { - // attempt to SIGTERM process group (negative PID) - // fall back to SIGKILL (to group) after 200ms - process.kill(-shell.pid, 'SIGTERM'); - await new Promise((resolve) => setTimeout(resolve, 200)); - if (shell.pid && !exited) { - process.kill(-shell.pid, 'SIGKILL'); - } - } catch (_e) { - // if group kill fails, fall back to killing just the main process - try { - if (shell.pid) { - shell.kill('SIGKILL'); + switch (event.type) { + case 'data': + if (isBinaryStream) break; // Don't process text if we are in binary mode + if (event.stream === 'stdout') { + cumulativeStdout += event.chunk; + } else { + cumulativeStderr += event.chunk; + } + currentDisplayOutput = + cumulativeStdout + + (cumulativeStderr ? `\n${cumulativeStderr}` : ''); + if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) { + shouldUpdate = true; + } + break; + case 'binary_detected': + isBinaryStream = true; + currentDisplayOutput = + '[Binary output detected. Halting stream...]'; + shouldUpdate = true; + break; + case 'binary_progress': + isBinaryStream = true; + currentDisplayOutput = `[Receiving binary output... ${formatMemoryUsage( + event.bytesReceived, + )} received]`; + if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) { + shouldUpdate = true; } - } catch (_e) { - console.error(`failed to kill shell process ${shell.pid}: ${_e}`); + break; + default: { + throw new Error('An unhandled ShellOutputEvent was found.'); } } - } - } - }; - signal.addEventListener('abort', abortHandler); - // wait for the shell to exit - try { - await new Promise((resolve) => shell.on('exit', resolve)); - } finally { - signal.removeEventListener('abort', abortHandler); - } + if (shouldUpdate) { + updateOutput(currentDisplayOutput); + lastUpdateTime = Date.now(); + } + }, + signal, + ); - // parse pids (pgrep output) from temporary file and remove it - const backgroundPIDs: number[] = []; - if (os.platform() !== 'win32') { - if (fs.existsSync(tempFilePath)) { - const pgrepLines = fs - .readFileSync(tempFilePath, 'utf8') - .split('\n') - .filter(Boolean); - for (const line of pgrepLines) { - if (!/^\d+$/.test(line)) { - console.error(`pgrep: ${line}`); + const result = await resultPromise; + + const backgroundPIDs: number[] = []; + if (os.platform() !== 'win32') { + if (fs.existsSync(tempFilePath)) { + const pgrepLines = fs + .readFileSync(tempFilePath, 'utf8') + .split('\n') + .filter(Boolean); + for (const line of pgrepLines) { + if (!/^\d+$/.test(line)) { + console.error(`pgrep: ${line}`); + } + const pid = Number(line); + if (pid !== result.pid) { + backgroundPIDs.push(pid); + } } - const pid = Number(line); - // exclude the shell subprocess pid - if (pid !== shell.pid) { - backgroundPIDs.push(pid); + } else { + if (!signal.aborted) { + console.error('missing pgrep output'); } } - fs.unlinkSync(tempFilePath); - } else { - if (!signal.aborted) { - console.error('missing pgrep output'); - } } - } - let llmContent = ''; - if (signal.aborted) { - llmContent = 'Command was cancelled by user before it could complete.'; - if (output.trim()) { - llmContent += ` Below is the output (on stdout and stderr) before it was cancelled:\n${output}`; + let llmContent = ''; + if (result.aborted) { + llmContent = 'Command was cancelled by user before it could complete.'; + if (result.output.trim()) { + llmContent += ` Below is the output (on stdout and stderr) before it was cancelled:\n${result.output}`; + } else { + llmContent += ' There was no output before it was cancelled.'; + } } else { - llmContent += ' There was no output before it was cancelled.'; + // Create a formatted error string for display, replacing the wrapper command + // with the user-facing command. + const finalError = result.error + ? result.error.message.replace(commandToExecute, params.command) + : '(none)'; + + llmContent = [ + `Command: ${params.command}`, + `Directory: ${params.directory || '(root)'}`, + `Stdout: ${result.stdout || '(empty)'}`, + `Stderr: ${result.stderr || '(empty)'}`, + `Error: ${finalError}`, // Use the cleaned error string. + `Exit Code: ${result.exitCode ?? '(none)'}`, + `Signal: ${result.signal ?? '(none)'}`, + `Background PIDs: ${ + backgroundPIDs.length ? backgroundPIDs.join(', ') : '(none)' + }`, + `Process Group PGID: ${result.pid ?? '(none)'}`, + ].join('\n'); } - } else { - llmContent = [ - `Command: ${params.command}`, - `Directory: ${params.directory || '(root)'}`, - `Stdout: ${stdout || '(empty)'}`, - `Stderr: ${stderr || '(empty)'}`, - `Error: ${error ?? '(none)'}`, - `Exit Code: ${code ?? '(none)'}`, - `Signal: ${processSignal ?? '(none)'}`, - `Background PIDs: ${backgroundPIDs.length ? backgroundPIDs.join(', ') : '(none)'}`, - `Process Group PGID: ${shell.pid ?? '(none)'}`, - ].join('\n'); - } - let returnDisplayMessage = ''; - if (this.config.getDebugMode()) { - returnDisplayMessage = llmContent; - } else { - if (output.trim()) { - returnDisplayMessage = output; + let returnDisplayMessage = ''; + if (this.config.getDebugMode()) { + returnDisplayMessage = llmContent; } else { - // Output is empty, let's provide a reason if the command failed or was cancelled - if (signal.aborted) { - returnDisplayMessage = 'Command cancelled by user.'; - } else if (processSignal) { - returnDisplayMessage = `Command terminated by signal: ${processSignal}`; - } else if (error) { - // If error is not null, it's an Error object (or other truthy value) - returnDisplayMessage = `Command failed: ${getErrorMessage(error)}`; - } else if (code !== null && code !== 0) { - returnDisplayMessage = `Command exited with code: ${code}`; + if (result.output.trim()) { + returnDisplayMessage = result.output; + } else { + if (result.aborted) { + returnDisplayMessage = 'Command cancelled by user.'; + } else if (result.signal) { + returnDisplayMessage = `Command terminated by signal: ${result.signal}`; + } else if (result.error) { + returnDisplayMessage = `Command failed: ${getErrorMessage( + result.error, + )}`; + } else if (result.exitCode !== null && result.exitCode !== 0) { + returnDisplayMessage = `Command exited with code: ${result.exitCode}`; + } + // If output is empty and command succeeded (code 0, no error/signal/abort), + // returnDisplayMessage will remain empty, which is fine. } - // If output is empty and command succeeded (code 0, no error/signal/abort), - // returnDisplayMessage will remain empty, which is fine. } - } - const summarizeConfig = this.config.getSummarizeToolOutputConfig(); - if (summarizeConfig && summarizeConfig[this.name]) { - const summary = await summarizeToolOutput( - llmContent, - this.config.getGeminiClient(), - signal, - summarizeConfig[this.name].tokenBudget, - ); + const summarizeConfig = this.config.getSummarizeToolOutputConfig(); + if (summarizeConfig && summarizeConfig[this.name]) { + const summary = await summarizeToolOutput( + llmContent, + this.config.getGeminiClient(), + signal, + summarizeConfig[this.name].tokenBudget, + ); + return { + llmContent: summary, + returnDisplay: returnDisplayMessage, + }; + } + return { - llmContent: summary, + llmContent, returnDisplay: returnDisplayMessage, }; + } finally { + if (fs.existsSync(tempFilePath)) { + fs.unlinkSync(tempFilePath); + } } - - return { - llmContent, - returnDisplay: returnDisplayMessage, - }; } } |