mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-21 09:17:53 +00:00
feat(core): refactor shell execution to use node-pty (#6088)
This commit is contained in:
Gal Zahavi
GitHub
@@ -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');
|
||||
const { result } = await simulateExecution('long-process', (pty) => {
|
||||
pty.onExit.mock.calls[0][0]({ exitCode: 0, signal: 15 });
|
||||
});
|
||||
|
||||
expect(result.exitCode).toBeNull();
|
||||
expect(result.signal).toBe('SIGTERM');
|
||||
expect(result.exitCode).toBe(0);
|
||||
expect(result.signal).toBe(15);
|
||||
});
|
||||
|
||||
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);
|
||||
it('should handle a synchronous spawn error', async () => {
|
||||
const spawnError = new Error('spawn ENOENT');
|
||||
mockPtySpawn.mockImplementation(() => {
|
||||
throw spawnError;
|
||||
});
|
||||
|
||||
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);
|
||||
});
|
||||
|
||||
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);
|
||||
expect(result.output).toBe('');
|
||||
expect(handle.pid).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
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,
|
||||
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 });
|
||||
},
|
||||
);
|
||||
|
||||
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);
|
||||
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),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user