mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-21 01:07:46 +00:00
refactor(core): Centralize shell logic into ShellExecutionService (#4823)
This commit is contained in:
357
packages/core/src/services/shellExecutionService.test.ts
Normal file
357
packages/core/src/services/shellExecutionService.test.ts
Normal 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 }),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
229
packages/core/src/services/shellExecutionService.ts
Normal file
229
packages/core/src/services/shellExecutionService.ts
Normal 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 };
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user