diff options
| author | Abhi <[email protected]> | 2025-07-25 21:56:49 -0400 |
|---|---|---|
| committer | GitHub <[email protected]> | 2025-07-26 01:56:49 +0000 |
| commit | ca5dd28ab60d78f42150460cbe9d4ed58d40afe4 (patch) | |
| tree | 6f3b3b373133a6c0bf8bb9cee8a10cc06ac83c9a /packages/core/src/services | |
| parent | ad2ef080aae2e21bf04ad8e922719ceaa81f1e5f (diff) | |
refactor(core): Centralize shell logic into ShellExecutionService (#4823)
Diffstat (limited to 'packages/core/src/services')
| -rw-r--r-- | packages/core/src/services/shellExecutionService.test.ts | 357 | ||||
| -rw-r--r-- | packages/core/src/services/shellExecutionService.ts | 229 |
2 files changed, 586 insertions, 0 deletions
diff --git a/packages/core/src/services/shellExecutionService.test.ts b/packages/core/src/services/shellExecutionService.test.ts new file mode 100644 index 00000000..4d1655a2 --- /dev/null +++ b/packages/core/src/services/shellExecutionService.test.ts @@ -0,0 +1,357 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi, describe, it, expect, beforeEach, type Mock } from 'vitest'; +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, +} from './shellExecutionService.js'; + +const mockIsBinary = vi.hoisted(() => vi.fn()); +vi.mock('../utils/textUtils.js', () => ({ + isBinary: mockIsBinary, +})); + +const mockPlatform = vi.hoisted(() => vi.fn()); +vi.mock('os', () => ({ + default: { + platform: mockPlatform, + }, + platform: mockPlatform, +})); + +const mockProcessKill = vi + .spyOn(process, 'kill') + .mockImplementation(() => true); + +describe('ShellExecutionService', () => { + let mockChildProcess: EventEmitter & Partial<ChildProcess>; + let onOutputEventMock: Mock<(event: ShellOutputEvent) => void>; + + beforeEach(() => { + vi.clearAllMocks(); + + mockIsBinary.mockReturnValue(false); + mockPlatform.mockReturnValue('linux'); + + 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(); + + // 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: (cp: typeof mockChildProcess, ac: AbortController) => void, + ) => { + const abortController = new AbortController(); + const handle = ShellExecutionService.execute( + command, + '/test/dir', + onOutputEventMock, + abortController.signal, + ); + + await new Promise((resolve) => setImmediate(resolve)); + simulation(mockChildProcess, 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); + }); + + expect(mockSpawn).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(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', + }); + }); + + 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); + }); + + 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 "你好"', (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.stdout).toBe('你好'); + }); + + it('should handle commands with no output', async () => { + 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 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.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', (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); + }); + + expect(result.error).toBe(spawnError); + 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', + '/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(result.signal).toBe('SIGKILL'); + // The individual kill calls were already asserted above. + expect(mockProcessKill).toHaveBeenCalledTimes(2); + }); + }); + + describe('Binary Output', () => { + it('should detect binary output and switch to progress events', async () => { + mockIsBinary.mockReturnValueOnce(true); + 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); + }); + + expect(result.rawOutput).toEqual( + Buffer.concat([binaryChunk1, binaryChunk2]), + ); + expect(onOutputEventMock).toHaveBeenCalledTimes(3); + expect(onOutputEventMock.mock.calls[0][0]).toEqual({ + type: 'binary_detected', + }); + expect(onOutputEventMock.mock.calls[1][0]).toEqual({ + type: 'binary_progress', + bytesReceived: 4, + }); + expect(onOutputEventMock.mock.calls[2][0]).toEqual({ + type: 'binary_progress', + bytesReceived: 8, + }); + }); + + 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); + }); + + // FIX: Provide explicit type for the 'call' parameter in the map function. + const eventTypes = onOutputEventMock.mock.calls.map( + (call: [ShellOutputEvent]) => call[0].type, + ); + expect(eventTypes).toEqual([ + 'data', + 'binary_detected', + 'binary_progress', + 'binary_progress', + ]); + }); + }); + + describe('Platform-Specific Behavior', () => { + it('should use cmd.exe on Windows', async () => { + mockPlatform.mockReturnValue('win32'); + await simulateExecution('dir', (cp) => cp.emit('exit', 0, null)); + + expect(mockSpawn).toHaveBeenCalledWith( + 'cmd.exe', + ['/c', 'dir'], + expect.objectContaining({ detached: false }), + ); + }); + + it('should use bash and detached process group on Linux', async () => { + mockPlatform.mockReturnValue('linux'); + await simulateExecution('ls', (cp) => cp.emit('exit', 0, null)); + + expect(mockSpawn).toHaveBeenCalledWith( + 'bash', + ['-c', 'ls'], + expect.objectContaining({ detached: true }), + ); + }); + }); +}); diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts new file mode 100644 index 00000000..0f0002cd --- /dev/null +++ b/packages/core/src/services/shellExecutionService.ts @@ -0,0 +1,229 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +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'; + +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 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: 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. */ + aborted: boolean; + /** The process ID of the spawned shell. */ + pid: number | undefined; +} + +/** A handle for an ongoing shell execution. */ +export interface ShellExecutionHandle { + /** The process ID of the spawned shell. */ + pid: number | undefined; + /** A promise that resolves with the complete execution result. */ + result: Promise<ShellExecutionResult>; +} + +/** + * Describes a structured event emitted during shell command execution. + */ +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; + } + | { + /** Signals that the output stream has been identified as binary. */ + type: 'binary_detected'; + } + | { + /** Provides progress updates for a binary stream. */ + type: 'binary_progress'; + /** The total number of bytes received so far. */ + bytesReceived: number; + }; + +/** + * A centralized service for executing shell commands with robust process + * management, cross-platform compatibility, and streaming output capabilities. + * + */ +export class ShellExecutionService { + /** + * 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. + * @param onOutputEvent A callback for streaming structured events about the execution, including data chunks and status updates. + * @param abortSignal An AbortSignal to terminate the process and its children. + * @returns An object containing the process ID (pid) and a promise that + * resolves with the complete execution result. + */ + static execute( + commandToExecute: string, + cwd: string, + onOutputEvent: (event: ShellOutputEvent) => void, + abortSignal: AbortSignal, + ): ShellExecutionHandle { + const isWindows = os.platform() === 'win32'; + const shell = isWindows ? 'cmd.exe' : 'bash'; + const shellArgs = [isWindows ? '/c' : '-c', commandToExecute]; + + const child = spawn(shell, shellArgs, { + cwd, + stdio: ['ignore', 'pipe', 'pipe'], + detached: !isWindows, // Use process groups on non-Windows for robust killing + env: { + ...process.env, + GEMINI_CLI: '1', + }, + }); + + 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 outputChunks: Buffer[] = []; + let 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; + + if (isBinary(sniffBuffer)) { + // Change state to stop streaming raw content. + isStreamingRawContent = false; + onOutputEvent({ type: 'binary_detected' }); + } + } + + const decodedChunk = + stream === 'stdout' + ? stdoutDecoder.decode(data, { stream: true }) + : stderrDecoder.decode(data, { stream: true }); + const strippedChunk = stripAnsi(decodedChunk); + + if (stream === 'stdout') { + stdout += strippedChunk; + } else { + stderr += strippedChunk; + } + + 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 }); + } + }; + + child.stdout.on('data', (data) => handleOutput(data, 'stdout')); + child.stderr.on('data', (data) => handleOutput(data, 'stderr')); + child.on('error', (err) => { + error = err; + }); + + 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 }); + + child.on('exit', (code, signal) => { + exited = true; + abortSignal.removeEventListener('abort', abortHandler); + + if (stdoutDecoder) { + stdout += stripAnsi(stdoutDecoder.decode()); + } + if (stderrDecoder) { + stderr += stripAnsi(stderrDecoder.decode()); + } + + 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, + }); + }); + }); + + return { pid: child.pid, result }; + } +} |
