diff options
Diffstat (limited to 'packages/core/src')
| -rw-r--r-- | packages/core/src/core/coreToolScheduler.test.ts | 6 | ||||
| -rw-r--r-- | packages/core/src/core/coreToolScheduler.ts | 11 | ||||
| -rw-r--r-- | packages/core/src/services/shellExecutionService.test.ts | 299 | ||||
| -rw-r--r-- | packages/core/src/services/shellExecutionService.ts | 277 | ||||
| -rw-r--r-- | packages/core/src/tools/shell.test.ts | 34 | ||||
| -rw-r--r-- | packages/core/src/tools/shell.ts | 25 | ||||
| -rw-r--r-- | packages/core/src/tools/tools.ts | 13 |
7 files changed, 386 insertions, 279 deletions
diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 6ba85b04..5cf49350 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -61,7 +61,6 @@ describe('CoreToolScheduler', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', - getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -118,7 +117,6 @@ describe('CoreToolScheduler with payload', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', - getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -415,7 +413,6 @@ describe('CoreToolScheduler edit cancellation', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', - getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -504,7 +501,6 @@ describe('CoreToolScheduler YOLO mode', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', - getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -590,7 +586,6 @@ describe('CoreToolScheduler request queueing', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', - getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -700,7 +695,6 @@ describe('CoreToolScheduler request queueing', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', - getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index bccb724a..0d442f9d 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -232,7 +232,6 @@ interface CoreToolSchedulerOptions { onToolCallsUpdate?: ToolCallsUpdateHandler; getPreferredEditor: () => EditorType | undefined; config: Config; - getTerminalSize: () => { columns: number; rows: number }; onEditorClose: () => void; } @@ -244,7 +243,6 @@ export class CoreToolScheduler { private onToolCallsUpdate?: ToolCallsUpdateHandler; private getPreferredEditor: () => EditorType | undefined; private config: Config; - private getTerminalSize: () => { columns: number; rows: number }; private onEditorClose: () => void; private isFinalizingToolCalls = false; private isScheduling = false; @@ -262,7 +260,6 @@ export class CoreToolScheduler { this.onAllToolCallsComplete = options.onAllToolCallsComplete; this.onToolCallsUpdate = options.onToolCallsUpdate; this.getPreferredEditor = options.getPreferredEditor; - this.getTerminalSize = options.getTerminalSize; this.onEditorClose = options.onEditorClose; } @@ -831,14 +828,8 @@ export class CoreToolScheduler { } : undefined; - const terminalSize = this.getTerminalSize(); invocation - .execute( - signal, - liveOutputCallback, - terminalSize.columns, - terminalSize.rows, - ) + .execute(signal, liveOutputCallback) .then(async (toolResult: ToolResult) => { if (signal.aborted) { this.setStatusInternal( diff --git a/packages/core/src/services/shellExecutionService.test.ts b/packages/core/src/services/shellExecutionService.test.ts index 06092bd9..2fe51a5e 100644 --- a/packages/core/src/services/shellExecutionService.test.ts +++ b/packages/core/src/services/shellExecutionService.test.ts @@ -5,12 +5,14 @@ */ import { vi, describe, it, expect, beforeEach, type Mock } from 'vitest'; -const mockPtySpawn = vi.hoisted(() => vi.fn()); -vi.mock('@lydell/node-pty', () => ({ - spawn: mockPtySpawn, +const mockSpawn = vi.hoisted(() => vi.fn()); +vi.mock('child_process', () => ({ + spawn: mockSpawn, })); import EventEmitter from 'events'; +import { Readable } from 'stream'; +import { type ChildProcess } from 'child_process'; import { ShellExecutionService, ShellOutputEvent, @@ -29,13 +31,12 @@ vi.mock('os', () => ({ platform: mockPlatform, })); +const mockProcessKill = vi + .spyOn(process, 'kill') + .mockImplementation(() => true); + describe('ShellExecutionService', () => { - let mockPtyProcess: EventEmitter & { - pid: number; - kill: Mock; - onData: Mock; - onExit: Mock; - }; + let mockChildProcess: EventEmitter & Partial<ChildProcess>; let onOutputEventMock: Mock<(event: ShellOutputEvent) => void>; beforeEach(() => { @@ -46,27 +47,26 @@ describe('ShellExecutionService', () => { onOutputEventMock = vi.fn(); - mockPtyProcess = new EventEmitter() as EventEmitter & { - pid: number; - kill: Mock; - onData: Mock; - onExit: Mock; - }; - mockPtyProcess.pid = 12345; - mockPtyProcess.kill = vi.fn(); - mockPtyProcess.onData = vi.fn(); - mockPtyProcess.onExit = vi.fn(); + mockChildProcess = new EventEmitter() as EventEmitter & + Partial<ChildProcess>; + // FIX: Cast simple EventEmitters to the expected stream type. + mockChildProcess.stdout = new EventEmitter() as Readable; + mockChildProcess.stderr = new EventEmitter() as Readable; + mockChildProcess.kill = vi.fn(); - mockPtySpawn.mockReturnValue(mockPtyProcess); + // FIX: Use Object.defineProperty to set the readonly 'pid' property. + Object.defineProperty(mockChildProcess, 'pid', { + value: 12345, + configurable: true, + }); + + mockSpawn.mockReturnValue(mockChildProcess); }); // Helper function to run a standard execution simulation const simulateExecution = async ( command: string, - simulation: ( - ptyProcess: typeof mockPtyProcess, - ac: AbortController, - ) => void, + simulation: (cp: typeof mockChildProcess, ac: AbortController) => void, ) => { const abortController = new AbortController(); const handle = ShellExecutionService.execute( @@ -77,123 +77,215 @@ describe('ShellExecutionService', () => { ); await new Promise((resolve) => setImmediate(resolve)); - simulation(mockPtyProcess, abortController); + simulation(mockChildProcess, abortController); const result = await handle.result; return { result, handle, abortController }; }; describe('Successful Execution', () => { - it('should execute a command and capture output', async () => { - const { result, handle } = await simulateExecution('ls -l', (pty) => { - pty.onData.mock.calls[0][0]('file1.txt\n'); - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }); + it('should execute a command and capture stdout and stderr', async () => { + const { result, handle } = await simulateExecution('ls -l', (cp) => { + cp.stdout?.emit('data', Buffer.from('file1.txt\n')); + cp.stderr?.emit('data', Buffer.from('a warning')); + cp.emit('exit', 0, null); }); - expect(mockPtySpawn).toHaveBeenCalledWith( - 'bash', - ['-c', 'ls -l'], - expect.any(Object), + expect(mockSpawn).toHaveBeenCalledWith( + 'ls -l', + [], + expect.objectContaining({ shell: 'bash' }), ); expect(result.exitCode).toBe(0); expect(result.signal).toBeNull(); expect(result.error).toBeNull(); expect(result.aborted).toBe(false); - expect(result.output).toBe('file1.txt'); + expect(result.stdout).toBe('file1.txt\n'); + expect(result.stderr).toBe('a warning'); + expect(result.output).toBe('file1.txt\n\na warning'); expect(handle.pid).toBe(12345); expect(onOutputEventMock).toHaveBeenCalledWith({ type: 'data', - chunk: 'file1.txt', + stream: 'stdout', + chunk: 'file1.txt\n', + }); + expect(onOutputEventMock).toHaveBeenCalledWith({ + type: 'data', + stream: 'stderr', + chunk: 'a warning', }); }); it('should strip ANSI codes from output', async () => { - const { result } = await simulateExecution('ls --color=auto', (pty) => { - pty.onData.mock.calls[0][0]('a\u001b[31mred\u001b[0mword'); - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }); + const { result } = await simulateExecution('ls --color=auto', (cp) => { + cp.stdout?.emit('data', Buffer.from('a\u001b[31mred\u001b[0mword')); + cp.emit('exit', 0, null); }); - expect(result.output).toBe('aredword'); + expect(result.stdout).toBe('aredword'); expect(onOutputEventMock).toHaveBeenCalledWith({ type: 'data', + stream: 'stdout', chunk: 'aredword', }); }); it('should correctly decode multi-byte characters split across chunks', async () => { - const { result } = await simulateExecution('echo "你好"', (pty) => { - const multiByteChar = '你好'; - pty.onData.mock.calls[0][0](multiByteChar.slice(0, 1)); - pty.onData.mock.calls[0][0](multiByteChar.slice(1)); - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }); + const { result } = await simulateExecution('echo "你好"', (cp) => { + const multiByteChar = Buffer.from('你好', 'utf-8'); + cp.stdout?.emit('data', multiByteChar.slice(0, 2)); + cp.stdout?.emit('data', multiByteChar.slice(2)); + cp.emit('exit', 0, null); }); - expect(result.output).toBe('你好'); + expect(result.stdout).toBe('你好'); }); it('should handle commands with no output', async () => { - const { result } = await simulateExecution('touch file', (pty) => { - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }); + const { result } = await simulateExecution('touch file', (cp) => { + cp.emit('exit', 0, null); }); + expect(result.stdout).toBe(''); + expect(result.stderr).toBe(''); expect(result.output).toBe(''); expect(onOutputEventMock).not.toHaveBeenCalled(); }); }); describe('Failed Execution', () => { - it('should capture a non-zero exit code', async () => { - const { result } = await simulateExecution('a-bad-command', (pty) => { - pty.onData.mock.calls[0][0]('command not found'); - pty.onExit.mock.calls[0][0]({ exitCode: 127, signal: null }); + it('should capture a non-zero exit code and format output correctly', async () => { + const { result } = await simulateExecution('a-bad-command', (cp) => { + cp.stderr?.emit('data', Buffer.from('command not found')); + cp.emit('exit', 127, null); }); expect(result.exitCode).toBe(127); - expect(result.output).toBe('command not found'); + expect(result.stderr).toBe('command not found'); + expect(result.stdout).toBe(''); + expect(result.output).toBe('\ncommand not found'); expect(result.error).toBeNull(); }); it('should capture a termination signal', async () => { - const { result } = await simulateExecution('long-process', (pty) => { - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: 15 }); + const { result } = await simulateExecution('long-process', (cp) => { + cp.emit('exit', null, 'SIGTERM'); }); - expect(result.exitCode).toBe(0); - expect(result.signal).toBe(15); + expect(result.exitCode).toBeNull(); + expect(result.signal).toBe('SIGTERM'); }); - it('should handle a synchronous spawn error', async () => { - const spawnError = new Error('spawn ENOENT'); - mockPtySpawn.mockImplementation(() => { - throw spawnError; + it('should handle a spawn error', async () => { + const spawnError = new Error('spawn EACCES'); + const { result } = await simulateExecution('protected-cmd', (cp) => { + cp.emit('error', spawnError); + cp.emit('exit', 1, null); }); - const handle = ShellExecutionService.execute( - 'any-command', - '/test/dir', - onOutputEventMock, - new AbortController().signal, - ); - const result = await handle.result; - expect(result.error).toBe(spawnError); expect(result.exitCode).toBe(1); - expect(result.output).toBe(''); - expect(handle.pid).toBeUndefined(); + }); + + it('handles errors that do not fire the exit event', async () => { + const error = new Error('spawn abc ENOENT'); + const { result } = await simulateExecution('touch cat.jpg', (cp) => { + cp.emit('error', error); // No exit event is fired. + }); + + expect(result.error).toBe(error); + expect(result.exitCode).toBe(1); }); }); describe('Aborting Commands', () => { - it('should abort a running process and set the aborted flag', async () => { - const { result } = await simulateExecution( - 'sleep 10', - (pty, abortController) => { - abortController.abort(); - pty.onExit.mock.calls[0][0]({ exitCode: 1, signal: null }); - }, + describe.each([ + { + platform: 'linux', + expectedSignal: 'SIGTERM', + expectedExit: { signal: 'SIGKILL' as const }, + }, + { + platform: 'win32', + expectedCommand: 'taskkill', + expectedExit: { code: 1 }, + }, + ])( + 'on $platform', + ({ platform, expectedSignal, expectedCommand, expectedExit }) => { + it('should abort a running process and set the aborted flag', async () => { + mockPlatform.mockReturnValue(platform); + + const { result } = await simulateExecution( + 'sleep 10', + (cp, abortController) => { + abortController.abort(); + if (expectedExit.signal) + cp.emit('exit', null, expectedExit.signal); + if (typeof expectedExit.code === 'number') + cp.emit('exit', expectedExit.code, null); + }, + ); + + expect(result.aborted).toBe(true); + + if (platform === 'linux') { + expect(mockProcessKill).toHaveBeenCalledWith( + -mockChildProcess.pid!, + expectedSignal, + ); + } else { + expect(mockSpawn).toHaveBeenCalledWith(expectedCommand, [ + '/pid', + String(mockChildProcess.pid), + '/f', + '/t', + ]); + } + }); + }, + ); + + it('should gracefully attempt SIGKILL on linux if SIGTERM fails', async () => { + mockPlatform.mockReturnValue('linux'); + vi.useFakeTimers(); + + // Don't await the result inside the simulation block for this specific test. + // We need to control the timeline manually. + const abortController = new AbortController(); + const handle = ShellExecutionService.execute( + 'unresponsive_process', + '/test/dir', + onOutputEventMock, + abortController.signal, ); + abortController.abort(); + + // Check the first kill signal + expect(mockProcessKill).toHaveBeenCalledWith( + -mockChildProcess.pid!, + 'SIGTERM', + ); + + // Now, advance time past the timeout + await vi.advanceTimersByTimeAsync(250); + + // Check the second kill signal + expect(mockProcessKill).toHaveBeenCalledWith( + -mockChildProcess.pid!, + 'SIGKILL', + ); + + // Finally, simulate the process exiting and await the result + mockChildProcess.emit('exit', null, 'SIGKILL'); + const result = await handle.result; + + vi.useRealTimers(); + expect(result.aborted).toBe(true); - expect(mockPtyProcess.kill).toHaveBeenCalled(); + expect(result.signal).toBe('SIGKILL'); + // The individual kill calls were already asserted above. + expect(mockProcessKill).toHaveBeenCalledTimes(2); }); }); @@ -203,10 +295,10 @@ describe('ShellExecutionService', () => { const binaryChunk1 = Buffer.from([0x89, 0x50, 0x4e, 0x47]); const binaryChunk2 = Buffer.from([0x0d, 0x0a, 0x1a, 0x0a]); - const { result } = await simulateExecution('cat image.png', (pty) => { - pty.onData.mock.calls[0][0](binaryChunk1); - pty.onData.mock.calls[0][0](binaryChunk2); - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }); + const { result } = await simulateExecution('cat image.png', (cp) => { + cp.stdout?.emit('data', binaryChunk1); + cp.stdout?.emit('data', binaryChunk2); + cp.emit('exit', 0, null); }); expect(result.rawOutput).toEqual( @@ -229,13 +321,14 @@ describe('ShellExecutionService', () => { it('should not emit data events after binary is detected', async () => { mockIsBinary.mockImplementation((buffer) => buffer.includes(0x00)); - await simulateExecution('cat mixed_file', (pty) => { - pty.onData.mock.calls[0][0](Buffer.from('some text')); - pty.onData.mock.calls[0][0](Buffer.from([0x00, 0x01, 0x02])); - pty.onData.mock.calls[0][0](Buffer.from('more text')); - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }); + await simulateExecution('cat mixed_file', (cp) => { + cp.stdout?.emit('data', Buffer.from('some text')); + cp.stdout?.emit('data', Buffer.from([0x00, 0x01, 0x02])); + cp.stdout?.emit('data', Buffer.from('more text')); + cp.emit('exit', 0, null); }); + // FIX: Provide explicit type for the 'call' parameter in the map function. const eventTypes = onOutputEventMock.mock.calls.map( (call: [ShellOutputEvent]) => call[0].type, ); @@ -251,27 +344,31 @@ describe('ShellExecutionService', () => { describe('Platform-Specific Behavior', () => { it('should use cmd.exe on Windows', async () => { mockPlatform.mockReturnValue('win32'); - await simulateExecution('dir "foo bar"', (pty) => - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }), + await simulateExecution('dir "foo bar"', (cp) => + cp.emit('exit', 0, null), ); - expect(mockPtySpawn).toHaveBeenCalledWith( - 'cmd.exe', - ['/c', 'dir "foo bar"'], - expect.any(Object), + expect(mockSpawn).toHaveBeenCalledWith( + 'dir "foo bar"', + [], + expect.objectContaining({ + shell: true, + detached: false, + }), ); }); - it('should use bash on Linux', async () => { + it('should use bash and detached process group on Linux', async () => { mockPlatform.mockReturnValue('linux'); - await simulateExecution('ls "foo bar"', (pty) => - pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }), - ); + await simulateExecution('ls "foo bar"', (cp) => cp.emit('exit', 0, null)); - expect(mockPtySpawn).toHaveBeenCalledWith( - 'bash', - ['-c', 'ls "foo bar"'], - expect.any(Object), + expect(mockSpawn).toHaveBeenCalledWith( + 'ls "foo bar"', + [], + expect.objectContaining({ + shell: 'bash', + detached: true, + }), ); }); }); diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts index 26d884b4..3749fcf6 100644 --- a/packages/core/src/services/shellExecutionService.ts +++ b/packages/core/src/services/shellExecutionService.ts @@ -4,35 +4,29 @@ * SPDX-License-Identifier: Apache-2.0 */ -import * as pty from '@lydell/node-pty'; +import { spawn } from 'child_process'; import { TextDecoder } from 'util'; import os from 'os'; +import stripAnsi from 'strip-ansi'; import { getCachedEncodingForBuffer } from '../utils/systemEncoding.js'; import { isBinary } from '../utils/textUtils.js'; -import pkg from '@xterm/headless'; -const { Terminal } = pkg; -// @ts-expect-error getFullText is not a public API. -const getFullText = (terminal: Terminal) => { - const buffer = terminal.buffer.active; - const lines: string[] = []; - for (let i = 0; i < buffer.length; i++) { - const line = buffer.getLine(i); - lines.push(line ? line.translateToString(true) : ''); - } - return lines.join('\n').trim(); -}; +const SIGKILL_TIMEOUT_MS = 200; /** A structured result from a shell command execution. */ export interface ShellExecutionResult { /** The raw, unprocessed output buffer. */ rawOutput: Buffer; - /** The combined, decoded output as a string. */ + /** The combined, decoded stdout and stderr as a string. */ output: string; + /** The decoded stdout as a string. */ + stdout: string; + /** The decoded stderr as a string. */ + stderr: string; /** The process exit code, or null if terminated by a signal. */ exitCode: number | null; /** The signal that terminated the process, if any. */ - signal: number | null; + signal: NodeJS.Signals | null; /** An error object if the process failed to spawn. */ error: Error | null; /** A boolean indicating if the command was aborted by the user. */ @@ -56,6 +50,8 @@ export type ShellOutputEvent = | { /** The event contains a chunk of output data. */ type: 'data'; + /** The stream from which the data originated. */ + stream: 'stdout' | 'stderr'; /** The decoded string chunk. */ chunk: string; } @@ -77,7 +73,7 @@ export type ShellOutputEvent = */ export class ShellExecutionService { /** - * Executes a shell command using `node-pty`, capturing all output and lifecycle events. + * Executes a shell command using `spawn`, capturing all output and lifecycle events. * * @param commandToExecute The exact command string to run. * @param cwd The working directory to execute the command in. @@ -91,150 +87,167 @@ export class ShellExecutionService { cwd: string, onOutputEvent: (event: ShellOutputEvent) => void, abortSignal: AbortSignal, - terminalColumns?: number, - terminalRows?: number, ): ShellExecutionHandle { const isWindows = os.platform() === 'win32'; - const shell = isWindows ? 'cmd.exe' : 'bash'; - const args = isWindows - ? ['/c', commandToExecute] - : ['-c', commandToExecute]; - let ptyProcess; - try { - ptyProcess = pty.spawn(shell, args, { - cwd, - name: 'xterm-color', - cols: terminalColumns ?? 200, - rows: terminalRows ?? 20, - env: { - ...process.env, - GEMINI_CLI: '1', - }, - handleFlowControl: true, - }); - } catch (e) { - const error = e as Error; - return { - pid: undefined, - result: Promise.resolve({ - rawOutput: Buffer.from(''), - output: '', - exitCode: 1, - signal: null, - error, - aborted: false, - pid: undefined, - }), - }; - } + const child = spawn(commandToExecute, [], { + cwd, + stdio: ['ignore', 'pipe', 'pipe'], + // Use bash unless in Windows (since it doesn't support bash). + // For windows, just use the default. + shell: isWindows ? true : 'bash', + // Use process groups on non-Windows for robust killing. + // Windows process termination is handled by `taskkill /t`. + detached: !isWindows, + env: { + ...process.env, + GEMINI_CLI: '1', + }, + }); const result = new Promise<ShellExecutionResult>((resolve) => { - const headlessTerminal = new Terminal({ - allowProposedApi: true, - cols: terminalColumns ?? 200, - rows: terminalRows ?? 20, - }); - let processingChain = Promise.resolve(); - let decoder: TextDecoder | null = null; - let output = ''; + // Use decoders to handle multi-byte characters safely (for streaming output). + let stdoutDecoder: TextDecoder | null = null; + let stderrDecoder: TextDecoder | null = null; + + let stdout = ''; + let stderr = ''; const outputChunks: Buffer[] = []; - const error: Error | null = null; + let error: Error | null = null; let exited = false; let isStreamingRawContent = true; const MAX_SNIFF_SIZE = 4096; let sniffedBytes = 0; - const handleOutput = (data: Buffer) => { - // NOTE: The migration from `child_process` to `node-pty` means we - // no longer have separate `stdout` and `stderr` streams. The `data` - // buffer contains the merged output. If a drop in LLM quality is - // observed after this change, we may need to revisit this and - // explore ways to re-introduce that distinction. - processingChain = processingChain.then( - () => - new Promise<void>((resolve) => { - if (!decoder) { - const encoding = getCachedEncodingForBuffer(data); - try { - decoder = new TextDecoder(encoding); - } catch { - decoder = new TextDecoder('utf-8'); - } - } + const handleOutput = (data: Buffer, stream: 'stdout' | 'stderr') => { + if (!stdoutDecoder || !stderrDecoder) { + const encoding = getCachedEncodingForBuffer(data); + try { + stdoutDecoder = new TextDecoder(encoding); + stderrDecoder = new TextDecoder(encoding); + } catch { + // If the encoding is not supported, fall back to utf-8. + // This can happen on some platforms for certain encodings like 'utf-32le'. + stdoutDecoder = new TextDecoder('utf-8'); + stderrDecoder = new TextDecoder('utf-8'); + } + } - outputChunks.push(data); + outputChunks.push(data); - // First, check if we need to switch to binary mode. - if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) { - const sniffBuffer = Buffer.concat(outputChunks.slice(0, 20)); - sniffedBytes = sniffBuffer.length; + // Binary detection logic. This only runs until we've made a determination. + if (isStreamingRawContent && sniffedBytes < MAX_SNIFF_SIZE) { + const sniffBuffer = Buffer.concat(outputChunks.slice(0, 20)); + sniffedBytes = sniffBuffer.length; - if (isBinary(sniffBuffer)) { - isStreamingRawContent = false; - onOutputEvent({ type: 'binary_detected' }); - } - } - - // Now, based on the *current* state, either process as text or binary. - if (isStreamingRawContent) { - const decodedChunk = decoder.decode(data, { stream: true }); - headlessTerminal.write(decodedChunk, () => { - const newStrippedOutput = getFullText(headlessTerminal); - output = newStrippedOutput; - onOutputEvent({ type: 'data', chunk: newStrippedOutput }); - resolve(); - }); - } else { - // Once in binary mode, we only emit progress events. - const totalBytes = outputChunks.reduce( - (sum, chunk) => sum + chunk.length, - 0, - ); - onOutputEvent({ - type: 'binary_progress', - bytesReceived: totalBytes, - }); - resolve(); - } - }), - ); - }; + if (isBinary(sniffBuffer)) { + // Change state to stop streaming raw content. + isStreamingRawContent = false; + onOutputEvent({ type: 'binary_detected' }); + } + } - ptyProcess.onData((data) => { - const bufferData = Buffer.from(data, 'utf-8'); - handleOutput(bufferData); - }); + const decodedChunk = + stream === 'stdout' + ? stdoutDecoder.decode(data, { stream: true }) + : stderrDecoder.decode(data, { stream: true }); + const strippedChunk = stripAnsi(decodedChunk); - ptyProcess.onExit(({ exitCode, signal }) => { - exited = true; - abortSignal.removeEventListener('abort', abortHandler); + if (stream === 'stdout') { + stdout += strippedChunk; + } else { + stderr += strippedChunk; + } - processingChain.then(() => { - const finalBuffer = Buffer.concat(outputChunks); + if (isStreamingRawContent) { + onOutputEvent({ type: 'data', stream, chunk: strippedChunk }); + } else { + const totalBytes = outputChunks.reduce( + (sum, chunk) => sum + chunk.length, + 0, + ); + onOutputEvent({ type: 'binary_progress', bytesReceived: totalBytes }); + } + }; - resolve({ - rawOutput: finalBuffer, - output, - exitCode, - signal: signal ?? null, - error, - aborted: abortSignal.aborted, - pid: ptyProcess.pid, - }); + child.stdout.on('data', (data) => handleOutput(data, 'stdout')); + child.stderr.on('data', (data) => handleOutput(data, 'stderr')); + child.on('error', (err) => { + const { stdout, stderr, finalBuffer } = cleanup(); + error = err; + resolve({ + error, + stdout, + stderr, + rawOutput: finalBuffer, + output: stdout + (stderr ? `\n${stderr}` : ''), + exitCode: 1, + signal: null, + aborted: false, + pid: child.pid, }); }); const abortHandler = async () => { - if (ptyProcess.pid && !exited) { - ptyProcess.kill('SIGHUP'); + if (child.pid && !exited) { + if (isWindows) { + spawn('taskkill', ['/pid', child.pid.toString(), '/f', '/t']); + } else { + try { + // Kill the entire process group (negative PID). + // SIGTERM first, then SIGKILL if it doesn't die. + process.kill(-child.pid, 'SIGTERM'); + await new Promise((res) => setTimeout(res, SIGKILL_TIMEOUT_MS)); + if (!exited) { + process.kill(-child.pid, 'SIGKILL'); + } + } catch (_e) { + // Fall back to killing just the main process if group kill fails. + if (!exited) child.kill('SIGKILL'); + } + } } }; abortSignal.addEventListener('abort', abortHandler, { once: true }); + + child.on('exit', (code: number, signal: NodeJS.Signals) => { + const { stdout, stderr, finalBuffer } = cleanup(); + + resolve({ + rawOutput: finalBuffer, + output: stdout + (stderr ? `\n${stderr}` : ''), + stdout, + stderr, + exitCode: code, + signal, + error, + aborted: abortSignal.aborted, + pid: child.pid, + }); + }); + + /** + * Cleans up a process (and it's accompanying state) that is exiting or + * erroring and returns output formatted output buffers and strings + */ + function cleanup() { + exited = true; + abortSignal.removeEventListener('abort', abortHandler); + if (stdoutDecoder) { + stdout += stripAnsi(stdoutDecoder.decode()); + } + if (stderrDecoder) { + stderr += stripAnsi(stderrDecoder.decode()); + } + + const finalBuffer = Buffer.concat(outputChunks); + + return { stdout, stderr, finalBuffer }; + } }); - return { pid: ptyProcess.pid, result }; + return { pid: child.pid, result }; } } diff --git a/packages/core/src/tools/shell.test.ts b/packages/core/src/tools/shell.test.ts index c968af26..96ff49a1 100644 --- a/packages/core/src/tools/shell.test.ts +++ b/packages/core/src/tools/shell.test.ts @@ -66,6 +66,7 @@ describe('ShellTool', () => { Buffer.from('abcdef', 'hex'), ); + // Capture the output callback to simulate streaming events from the service mockShellExecutionService.mockImplementation((_cmd, _cwd, callback) => { mockShellOutputCallback = callback; return { @@ -122,6 +123,8 @@ describe('ShellTool', () => { const fullResult: ShellExecutionResult = { rawOutput: Buffer.from(result.output || ''), output: 'Success', + stdout: 'Success', + stderr: '', exitCode: 0, signal: null, error: null, @@ -138,7 +141,7 @@ describe('ShellTool', () => { resolveShellExecution({ pid: 54321 }); vi.mocked(fs.existsSync).mockReturnValue(true); - vi.mocked(fs.readFileSync).mockReturnValue('54321\n54322\n'); + vi.mocked(fs.readFileSync).mockReturnValue('54321\n54322\n'); // Service PID and background PID const result = await promise; @@ -149,8 +152,6 @@ describe('ShellTool', () => { expect.any(String), expect.any(Function), mockAbortSignal, - undefined, - undefined, ); expect(result.llmContent).toContain('Background PIDs: 54322'); expect(vi.mocked(fs.unlinkSync)).toHaveBeenCalledWith(tmpFile); @@ -163,6 +164,8 @@ describe('ShellTool', () => { resolveShellExecution({ rawOutput: Buffer.from(''), output: '', + stdout: '', + stderr: '', exitCode: 0, signal: null, error: null, @@ -175,8 +178,6 @@ describe('ShellTool', () => { expect.any(String), expect.any(Function), mockAbortSignal, - undefined, - undefined, ); }); @@ -188,13 +189,16 @@ describe('ShellTool', () => { error, exitCode: 1, output: 'err', + stderr: 'err', rawOutput: Buffer.from('err'), + stdout: '', signal: null, aborted: false, pid: 12345, }); const result = await promise; + // The final llmContent should contain the user's command, not the wrapper expect(result.llmContent).toContain('Error: wrapped command failed'); expect(result.llmContent).not.toContain('pgrep'); }); @@ -227,6 +231,8 @@ describe('ShellTool', () => { resolveExecutionPromise({ output: 'long output', rawOutput: Buffer.from('long output'), + stdout: 'long output', + stderr: '', exitCode: 0, signal: null, error: null, @@ -251,7 +257,7 @@ describe('ShellTool', () => { mockShellExecutionService.mockImplementation(() => { throw error; }); - vi.mocked(fs.existsSync).mockReturnValue(true); + vi.mocked(fs.existsSync).mockReturnValue(true); // Pretend the file exists const invocation = shellTool.build({ command: 'a-command' }); await expect(invocation.execute(mockAbortSignal)).rejects.toThrow(error); @@ -274,26 +280,33 @@ describe('ShellTool', () => { const invocation = shellTool.build({ command: 'stream' }); const promise = invocation.execute(mockAbortSignal, updateOutputMock); + // First chunk, should be throttled. mockShellOutputCallback({ type: 'data', + stream: 'stdout', chunk: 'hello ', }); expect(updateOutputMock).not.toHaveBeenCalled(); + // Advance time past the throttle interval. await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1); + // Send a second chunk. THIS event triggers the update with the CUMULATIVE content. mockShellOutputCallback({ type: 'data', - chunk: 'hello world', + stream: 'stderr', + chunk: 'world', }); // It should have been called once now with the combined output. expect(updateOutputMock).toHaveBeenCalledOnce(); - expect(updateOutputMock).toHaveBeenCalledWith('hello world'); + expect(updateOutputMock).toHaveBeenCalledWith('hello \nworld'); resolveExecutionPromise({ rawOutput: Buffer.from(''), output: '', + stdout: '', + stderr: '', exitCode: 0, signal: null, error: null, @@ -319,13 +332,16 @@ describe('ShellTool', () => { }); expect(updateOutputMock).toHaveBeenCalledOnce(); + // Advance time past the throttle interval. await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1); + // Send a SECOND progress event. This one will trigger the flush. mockShellOutputCallback({ type: 'binary_progress', bytesReceived: 2048, }); + // Now it should be called a second time with the latest progress. expect(updateOutputMock).toHaveBeenCalledTimes(2); expect(updateOutputMock).toHaveBeenLastCalledWith( '[Receiving binary output... 2.0 KB received]', @@ -334,6 +350,8 @@ describe('ShellTool', () => { resolveExecutionPromise({ rawOutput: Buffer.from(''), output: '', + stdout: '', + stderr: '', exitCode: 0, signal: null, error: null, diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 621ff90f..5b01a82f 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -97,8 +97,6 @@ class ShellToolInvocation extends BaseToolInvocation< async execute( signal: AbortSignal, updateOutput?: (output: string) => void, - terminalColumns?: number, - terminalRows?: number, ): Promise<ToolResult> { const strippedCommand = stripShellWrapper(this.params.command); @@ -131,7 +129,9 @@ class ShellToolInvocation extends BaseToolInvocation< this.params.directory || '', ); - let cumulativeOutput = ''; + let cumulativeStdout = ''; + let cumulativeStderr = ''; + let lastUpdateTime = Date.now(); let isBinaryStream = false; @@ -148,9 +148,15 @@ class ShellToolInvocation extends BaseToolInvocation< switch (event.type) { case 'data': - if (isBinaryStream) break; - cumulativeOutput = event.chunk; - currentDisplayOutput = cumulativeOutput; + 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; } @@ -181,8 +187,6 @@ class ShellToolInvocation extends BaseToolInvocation< } }, signal, - terminalColumns, - terminalRows, ); const result = await resultPromise; @@ -214,7 +218,7 @@ class ShellToolInvocation extends BaseToolInvocation< if (result.aborted) { llmContent = 'Command was cancelled by user before it could complete.'; if (result.output.trim()) { - llmContent += ` Below is the output before it was cancelled:\n${result.output}`; + 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.'; } @@ -228,7 +232,8 @@ class ShellToolInvocation extends BaseToolInvocation< llmContent = [ `Command: ${this.params.command}`, `Directory: ${this.params.directory || '(root)'}`, - `Output: ${result.output || '(empty)'}`, + `Stdout: ${result.stdout || '(empty)'}`, + `Stderr: ${result.stderr || '(empty)'}`, `Error: ${finalError}`, // Use the cleaned error string. `Exit Code: ${result.exitCode ?? '(none)'}`, `Signal: ${result.signal ?? '(none)'}`, diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 5684e4ac..68ea00af 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -50,8 +50,6 @@ export interface ToolInvocation< execute( signal: AbortSignal, updateOutput?: (output: string) => void, - terminalColumns?: number, - terminalRows?: number, ): Promise<TResult>; } @@ -80,8 +78,6 @@ export abstract class BaseToolInvocation< abstract execute( signal: AbortSignal, updateOutput?: (output: string) => void, - terminalColumns?: number, - terminalRows?: number, ): Promise<TResult>; } @@ -200,16 +196,9 @@ export abstract class DeclarativeTool< params: TParams, signal: AbortSignal, updateOutput?: (output: string) => void, - terminalColumns?: number, - terminalRows?: number, ): Promise<TResult> { const invocation = this.build(params); - return invocation.execute( - signal, - updateOutput, - terminalColumns, - terminalRows, - ); + return invocation.execute(signal, updateOutput); } } |