diff options
| author | Gal Zahavi <[email protected]> | 2025-08-14 13:40:12 -0700 |
|---|---|---|
| committer | GitHub <[email protected]> | 2025-08-14 20:40:12 +0000 |
| commit | 980091cbc2a809690dbd401c16ec3ac34da56083 (patch) | |
| tree | 5bdadbdbebcaf6471f753ef31ef3fdc6a7716ae3 /packages/core/src | |
| parent | 48af0456c1883834a83ae74281f0c871129779d8 (diff) | |
feat(core): refactor shell execution to use node-pty (#6088)
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 | 25 |
7 files changed, 290 insertions, 387 deletions
diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 71b2d64c..9d7d45ea 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -107,6 +107,7 @@ describe('CoreToolScheduler', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', + getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -176,6 +177,7 @@ describe('CoreToolScheduler with payload', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', + getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -453,6 +455,7 @@ describe('CoreToolScheduler edit cancellation', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', + getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -548,6 +551,7 @@ describe('CoreToolScheduler YOLO mode', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', + getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -633,6 +637,7 @@ describe('CoreToolScheduler request queueing', () => { onAllToolCallsComplete, onToolCallsUpdate, getPreferredEditor: () => 'vscode', + getTerminalSize: () => ({ columns: 80, rows: 24 }), onEditorClose: vi.fn(), }); @@ -742,6 +747,7 @@ 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 00ff5c55..aac8f9a6 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -232,6 +232,7 @@ interface CoreToolSchedulerOptions { onToolCallsUpdate?: ToolCallsUpdateHandler; getPreferredEditor: () => EditorType | undefined; config: Config; + getTerminalSize: () => { columns: number; rows: number }; onEditorClose: () => void; } @@ -243,6 +244,7 @@ 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; @@ -260,6 +262,7 @@ export class CoreToolScheduler { this.onAllToolCallsComplete = options.onAllToolCallsComplete; this.onToolCallsUpdate = options.onToolCallsUpdate; this.getPreferredEditor = options.getPreferredEditor; + this.getTerminalSize = options.getTerminalSize; this.onEditorClose = options.onEditorClose; } @@ -820,8 +823,14 @@ export class CoreToolScheduler { } : undefined; + const terminalSize = this.getTerminalSize(); invocation - .execute(signal, liveOutputCallback) + .execute( + signal, + liveOutputCallback, + terminalSize.columns, + terminalSize.rows, + ) .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 2fe51a5e..06092bd9 100644 --- a/packages/core/src/services/shellExecutionService.test.ts +++ b/packages/core/src/services/shellExecutionService.test.ts @@ -5,14 +5,12 @@ */ import { vi, describe, it, expect, beforeEach, type Mock } from 'vitest'; -const mockSpawn = vi.hoisted(() => vi.fn()); -vi.mock('child_process', () => ({ - spawn: mockSpawn, +const mockPtySpawn = vi.hoisted(() => vi.fn()); +vi.mock('@lydell/node-pty', () => ({ + spawn: mockPtySpawn, })); import EventEmitter from 'events'; -import { Readable } from 'stream'; -import { type ChildProcess } from 'child_process'; import { ShellExecutionService, ShellOutputEvent, @@ -31,12 +29,13 @@ vi.mock('os', () => ({ platform: mockPlatform, })); -const mockProcessKill = vi - .spyOn(process, 'kill') - .mockImplementation(() => true); - describe('ShellExecutionService', () => { - let mockChildProcess: EventEmitter & Partial<ChildProcess>; + let mockPtyProcess: EventEmitter & { + pid: number; + kill: Mock; + onData: Mock; + onExit: Mock; + }; let onOutputEventMock: Mock<(event: ShellOutputEvent) => void>; beforeEach(() => { @@ -47,26 +46,27 @@ describe('ShellExecutionService', () => { onOutputEventMock = 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(); + 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(); - // FIX: Use Object.defineProperty to set the readonly 'pid' property. - Object.defineProperty(mockChildProcess, 'pid', { - value: 12345, - configurable: true, - }); - - mockSpawn.mockReturnValue(mockChildProcess); + mockPtySpawn.mockReturnValue(mockPtyProcess); }); // Helper function to run a standard execution simulation const simulateExecution = async ( command: string, - simulation: (cp: typeof mockChildProcess, ac: AbortController) => void, + simulation: ( + ptyProcess: typeof mockPtyProcess, + ac: AbortController, + ) => void, ) => { const abortController = new AbortController(); const handle = ShellExecutionService.execute( @@ -77,215 +77,123 @@ describe('ShellExecutionService', () => { ); await new Promise((resolve) => setImmediate(resolve)); - simulation(mockChildProcess, abortController); + simulation(mockPtyProcess, abortController); const result = await handle.result; return { result, handle, abortController }; }; describe('Successful Execution', () => { - 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); + 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 }); }); - expect(mockSpawn).toHaveBeenCalledWith( - 'ls -l', - [], - expect.objectContaining({ shell: 'bash' }), + expect(mockPtySpawn).toHaveBeenCalledWith( + 'bash', + ['-c', 'ls -l'], + expect.any(Object), ); expect(result.exitCode).toBe(0); expect(result.signal).toBeNull(); expect(result.error).toBeNull(); expect(result.aborted).toBe(false); - expect(result.stdout).toBe('file1.txt\n'); - expect(result.stderr).toBe('a warning'); - expect(result.output).toBe('file1.txt\n\na warning'); + expect(result.output).toBe('file1.txt'); expect(handle.pid).toBe(12345); expect(onOutputEventMock).toHaveBeenCalledWith({ type: 'data', - stream: 'stdout', - chunk: 'file1.txt\n', - }); - expect(onOutputEventMock).toHaveBeenCalledWith({ - type: 'data', - stream: 'stderr', - chunk: 'a warning', + chunk: 'file1.txt', }); }); it('should strip ANSI codes from output', async () => { - const { result } = await simulateExecution('ls --color=auto', (cp) => { - cp.stdout?.emit('data', Buffer.from('a\u001b[31mred\u001b[0mword')); - cp.emit('exit', 0, null); + 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 }); }); - expect(result.stdout).toBe('aredword'); + expect(result.output).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 "你好"', (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); + 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 }); }); - expect(result.stdout).toBe('你好'); + expect(result.output).toBe('你好'); }); it('should handle commands with no output', async () => { - const { result } = await simulateExecution('touch file', (cp) => { - cp.emit('exit', 0, null); + const { result } = await simulateExecution('touch file', (pty) => { + pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: 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 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); + 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 }); }); expect(result.exitCode).toBe(127); - expect(result.stderr).toBe('command not found'); - expect(result.stdout).toBe(''); - expect(result.output).toBe('\ncommand not found'); + expect(result.output).toBe('command not found'); expect(result.error).toBeNull(); }); it('should capture a termination signal', async () => { - const { result } = await simulateExecution('long-process', (cp) => { - cp.emit('exit', null, 'SIGTERM'); - }); - - expect(result.exitCode).toBeNull(); - expect(result.signal).toBe('SIGTERM'); - }); - - 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 { result } = await simulateExecution('long-process', (pty) => { + pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: 15 }); }); - expect(result.error).toBe(spawnError); - expect(result.exitCode).toBe(1); + expect(result.exitCode).toBe(0); + expect(result.signal).toBe(15); }); - 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. + it('should handle a synchronous spawn error', async () => { + const spawnError = new Error('spawn ENOENT'); + mockPtySpawn.mockImplementation(() => { + throw spawnError; }); - expect(result.error).toBe(error); - expect(result.exitCode).toBe(1); - }); - }); - - describe('Aborting Commands', () => { - 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', + 'any-command', '/test/dir', onOutputEventMock, - abortController.signal, - ); - - abortController.abort(); - - // Check the first kill signal - expect(mockProcessKill).toHaveBeenCalledWith( - -mockChildProcess.pid!, - 'SIGTERM', + new AbortController().signal, ); + const result = await handle.result; - // Now, advance time past the timeout - await vi.advanceTimersByTimeAsync(250); + expect(result.error).toBe(spawnError); + expect(result.exitCode).toBe(1); + expect(result.output).toBe(''); + expect(handle.pid).toBeUndefined(); + }); + }); - // Check the second kill signal - expect(mockProcessKill).toHaveBeenCalledWith( - -mockChildProcess.pid!, - 'SIGKILL', + 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 }); + }, ); - // 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(result.signal).toBe('SIGKILL'); - // The individual kill calls were already asserted above. - expect(mockProcessKill).toHaveBeenCalledTimes(2); + expect(mockPtyProcess.kill).toHaveBeenCalled(); }); }); @@ -295,10 +203,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', (cp) => { - cp.stdout?.emit('data', binaryChunk1); - cp.stdout?.emit('data', binaryChunk2); - cp.emit('exit', 0, null); + 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 }); }); expect(result.rawOutput).toEqual( @@ -321,14 +229,13 @@ describe('ShellExecutionService', () => { it('should not emit data events after binary is detected', async () => { mockIsBinary.mockImplementation((buffer) => buffer.includes(0x00)); - 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); + 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 }); }); - // FIX: Provide explicit type for the 'call' parameter in the map function. const eventTypes = onOutputEventMock.mock.calls.map( (call: [ShellOutputEvent]) => call[0].type, ); @@ -344,31 +251,27 @@ describe('ShellExecutionService', () => { describe('Platform-Specific Behavior', () => { it('should use cmd.exe on Windows', async () => { mockPlatform.mockReturnValue('win32'); - await simulateExecution('dir "foo bar"', (cp) => - cp.emit('exit', 0, null), + await simulateExecution('dir "foo bar"', (pty) => + pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }), ); - expect(mockSpawn).toHaveBeenCalledWith( - 'dir "foo bar"', - [], - expect.objectContaining({ - shell: true, - detached: false, - }), + expect(mockPtySpawn).toHaveBeenCalledWith( + 'cmd.exe', + ['/c', 'dir "foo bar"'], + expect.any(Object), ); }); - it('should use bash and detached process group on Linux', async () => { + it('should use bash on Linux', async () => { mockPlatform.mockReturnValue('linux'); - await simulateExecution('ls "foo bar"', (cp) => cp.emit('exit', 0, null)); + await simulateExecution('ls "foo bar"', (pty) => + pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: null }), + ); - expect(mockSpawn).toHaveBeenCalledWith( - 'ls "foo bar"', - [], - expect.objectContaining({ - shell: 'bash', - detached: true, - }), + expect(mockPtySpawn).toHaveBeenCalledWith( + 'bash', + ['-c', 'ls "foo bar"'], + expect.any(Object), ); }); }); diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts index 3749fcf6..26d884b4 100644 --- a/packages/core/src/services/shellExecutionService.ts +++ b/packages/core/src/services/shellExecutionService.ts @@ -4,29 +4,35 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { spawn } from 'child_process'; +import * as pty from '@lydell/node-pty'; 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; -const SIGKILL_TIMEOUT_MS = 200; +// @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(); +}; /** A structured result from a shell command execution. */ export interface ShellExecutionResult { /** The raw, unprocessed output buffer. */ rawOutput: Buffer; - /** The combined, decoded stdout and stderr as a string. */ + /** The combined, decoded output 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: NodeJS.Signals | null; + signal: number | null; /** An error object if the process failed to spawn. */ error: Error | null; /** A boolean indicating if the command was aborted by the user. */ @@ -50,8 +56,6 @@ 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; } @@ -73,7 +77,7 @@ export type ShellOutputEvent = */ export class ShellExecutionService { /** - * Executes a shell command using `spawn`, capturing all output and lifecycle events. + * Executes a shell command using `node-pty`, capturing all output and lifecycle events. * * @param commandToExecute The exact command string to run. * @param cwd The working directory to execute the command in. @@ -87,167 +91,150 @@ 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]; - 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', - }, - }); + 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 result = new Promise<ShellExecutionResult>((resolve) => { - // 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 headlessTerminal = new Terminal({ + allowProposedApi: true, + cols: terminalColumns ?? 200, + rows: terminalRows ?? 20, + }); + let processingChain = Promise.resolve(); + let decoder: TextDecoder | null = null; + let output = ''; const outputChunks: Buffer[] = []; - let error: Error | null = null; + const error: Error | null = null; let exited = false; let isStreamingRawContent = true; const MAX_SNIFF_SIZE = 4096; let sniffedBytes = 0; - 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); - - // 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; + 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'); + } + } - if (isBinary(sniffBuffer)) { - // Change state to stop streaming raw content. - isStreamingRawContent = false; - onOutputEvent({ type: 'binary_detected' }); - } - } + outputChunks.push(data); - const decodedChunk = - stream === 'stdout' - ? stdoutDecoder.decode(data, { stream: true }) - : stderrDecoder.decode(data, { stream: true }); - const strippedChunk = stripAnsi(decodedChunk); + // 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; - if (stream === 'stdout') { - stdout += strippedChunk; - } else { - stderr += strippedChunk; - } + if (isBinary(sniffBuffer)) { + isStreamingRawContent = false; + onOutputEvent({ type: 'binary_detected' }); + } + } - 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 }); - } + // 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(); + } + }), + ); }; - 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, - }); + ptyProcess.onData((data) => { + const bufferData = Buffer.from(data, 'utf-8'); + handleOutput(bufferData); }); - const abortHandler = async () => { - 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 }); + ptyProcess.onExit(({ exitCode, signal }) => { + exited = true; + abortSignal.removeEventListener('abort', abortHandler); - child.on('exit', (code: number, signal: NodeJS.Signals) => { - const { stdout, stderr, finalBuffer } = cleanup(); + processingChain.then(() => { + const finalBuffer = Buffer.concat(outputChunks); - resolve({ - rawOutput: finalBuffer, - output: stdout + (stderr ? `\n${stderr}` : ''), - stdout, - stderr, - exitCode: code, - signal, - error, - aborted: abortSignal.aborted, - pid: child.pid, + resolve({ + rawOutput: finalBuffer, + output, + exitCode, + signal: signal ?? null, + error, + aborted: abortSignal.aborted, + pid: ptyProcess.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 abortHandler = async () => { + if (ptyProcess.pid && !exited) { + ptyProcess.kill('SIGHUP'); } + }; - const finalBuffer = Buffer.concat(outputChunks); - - return { stdout, stderr, finalBuffer }; - } + abortSignal.addEventListener('abort', abortHandler, { once: true }); }); - return { pid: child.pid, result }; + return { pid: ptyProcess.pid, result }; } } diff --git a/packages/core/src/tools/shell.test.ts b/packages/core/src/tools/shell.test.ts index 96ff49a1..c968af26 100644 --- a/packages/core/src/tools/shell.test.ts +++ b/packages/core/src/tools/shell.test.ts @@ -66,7 +66,6 @@ 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 { @@ -123,8 +122,6 @@ describe('ShellTool', () => { const fullResult: ShellExecutionResult = { rawOutput: Buffer.from(result.output || ''), output: 'Success', - stdout: 'Success', - stderr: '', exitCode: 0, signal: null, error: null, @@ -141,7 +138,7 @@ describe('ShellTool', () => { resolveShellExecution({ pid: 54321 }); vi.mocked(fs.existsSync).mockReturnValue(true); - vi.mocked(fs.readFileSync).mockReturnValue('54321\n54322\n'); // Service PID and background PID + vi.mocked(fs.readFileSync).mockReturnValue('54321\n54322\n'); const result = await promise; @@ -152,6 +149,8 @@ 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); @@ -164,8 +163,6 @@ describe('ShellTool', () => { resolveShellExecution({ rawOutput: Buffer.from(''), output: '', - stdout: '', - stderr: '', exitCode: 0, signal: null, error: null, @@ -178,6 +175,8 @@ describe('ShellTool', () => { expect.any(String), expect.any(Function), mockAbortSignal, + undefined, + undefined, ); }); @@ -189,16 +188,13 @@ 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'); }); @@ -231,8 +227,6 @@ describe('ShellTool', () => { resolveExecutionPromise({ output: 'long output', rawOutput: Buffer.from('long output'), - stdout: 'long output', - stderr: '', exitCode: 0, signal: null, error: null, @@ -257,7 +251,7 @@ describe('ShellTool', () => { mockShellExecutionService.mockImplementation(() => { throw error; }); - vi.mocked(fs.existsSync).mockReturnValue(true); // Pretend the file exists + vi.mocked(fs.existsSync).mockReturnValue(true); const invocation = shellTool.build({ command: 'a-command' }); await expect(invocation.execute(mockAbortSignal)).rejects.toThrow(error); @@ -280,33 +274,26 @@ 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', - stream: 'stderr', - chunk: 'world', + chunk: 'hello world', }); // It should have been called once now with the combined output. expect(updateOutputMock).toHaveBeenCalledOnce(); - expect(updateOutputMock).toHaveBeenCalledWith('hello \nworld'); + expect(updateOutputMock).toHaveBeenCalledWith('hello world'); resolveExecutionPromise({ rawOutput: Buffer.from(''), output: '', - stdout: '', - stderr: '', exitCode: 0, signal: null, error: null, @@ -332,16 +319,13 @@ 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]', @@ -350,8 +334,6 @@ 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 5b01a82f..621ff90f 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -97,6 +97,8 @@ class ShellToolInvocation extends BaseToolInvocation< async execute( signal: AbortSignal, updateOutput?: (output: string) => void, + terminalColumns?: number, + terminalRows?: number, ): Promise<ToolResult> { const strippedCommand = stripShellWrapper(this.params.command); @@ -129,9 +131,7 @@ class ShellToolInvocation extends BaseToolInvocation< this.params.directory || '', ); - let cumulativeStdout = ''; - let cumulativeStderr = ''; - + let cumulativeOutput = ''; let lastUpdateTime = Date.now(); let isBinaryStream = false; @@ -148,15 +148,9 @@ class ShellToolInvocation extends BaseToolInvocation< 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 (isBinaryStream) break; + cumulativeOutput = event.chunk; + currentDisplayOutput = cumulativeOutput; if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) { shouldUpdate = true; } @@ -187,6 +181,8 @@ class ShellToolInvocation extends BaseToolInvocation< } }, signal, + terminalColumns, + terminalRows, ); const result = await resultPromise; @@ -218,7 +214,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 (on stdout and stderr) before it was cancelled:\n${result.output}`; + llmContent += ` Below is the output before it was cancelled:\n${result.output}`; } else { llmContent += ' There was no output before it was cancelled.'; } @@ -232,8 +228,7 @@ class ShellToolInvocation extends BaseToolInvocation< llmContent = [ `Command: ${this.params.command}`, `Directory: ${this.params.directory || '(root)'}`, - `Stdout: ${result.stdout || '(empty)'}`, - `Stderr: ${result.stderr || '(empty)'}`, + `Output: ${result.output || '(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 00f2a842..ee8b830b 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -50,6 +50,8 @@ export interface ToolInvocation< execute( signal: AbortSignal, updateOutput?: (output: string) => void, + terminalColumns?: number, + terminalRows?: number, ): Promise<TResult>; } @@ -78,6 +80,8 @@ export abstract class BaseToolInvocation< abstract execute( signal: AbortSignal, updateOutput?: (output: string) => void, + terminalColumns?: number, + terminalRows?: number, ): Promise<TResult>; } @@ -117,8 +121,16 @@ export class LegacyToolInvocation< execute( signal: AbortSignal, updateOutput?: (output: string) => void, + terminalColumns?: number, + terminalRows?: number, ): Promise<TResult> { - return this.legacyTool.execute(this.params, signal, updateOutput); + return this.legacyTool.execute( + this.params, + signal, + updateOutput, + terminalColumns, + terminalRows, + ); } } @@ -232,9 +244,16 @@ 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); + return invocation.execute( + signal, + updateOutput, + terminalColumns, + terminalRows, + ); } } @@ -373,6 +392,8 @@ export abstract class BaseTool< params: TParams, signal: AbortSignal, updateOutput?: (output: string) => void, + terminalColumns?: number, + terminalRows?: number, ): Promise<TResult>; } |
