refactor(core): Centralize shell logic into ShellExecutionService (#4823)

This commit is contained in:
Abhi
2025-07-25 21:56:49 -04:00
committed by GitHub
parent ad2ef080aa
commit ca5dd28ab6
11 changed files with 1751 additions and 812 deletions

View File

@@ -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 }),
);
});
});
});

View File

@@ -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 };
}
}