This commit is contained in:
Gal Zahavi
2025-08-15 10:27:33 -07:00
committed by GitHub
parent ab1c483cab
commit 1a2906a8ad
16 changed files with 412 additions and 442 deletions
+26 -8
View File
@@ -66,6 +66,7 @@ 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 {
@@ -122,6 +123,8 @@ describe('ShellTool', () => {
const fullResult: ShellExecutionResult = {
rawOutput: Buffer.from(result.output || ''),
output: 'Success',
stdout: 'Success',
stderr: '',
exitCode: 0,
signal: null,
error: null,
@@ -138,7 +141,7 @@ describe('ShellTool', () => {
resolveShellExecution({ pid: 54321 });
vi.mocked(fs.existsSync).mockReturnValue(true);
vi.mocked(fs.readFileSync).mockReturnValue('54321\n54322\n');
vi.mocked(fs.readFileSync).mockReturnValue('54321\n54322\n'); // Service PID and background PID
const result = await promise;
@@ -149,8 +152,6 @@ 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);
@@ -163,6 +164,8 @@ describe('ShellTool', () => {
resolveShellExecution({
rawOutput: Buffer.from(''),
output: '',
stdout: '',
stderr: '',
exitCode: 0,
signal: null,
error: null,
@@ -175,8 +178,6 @@ describe('ShellTool', () => {
expect.any(String),
expect.any(Function),
mockAbortSignal,
undefined,
undefined,
);
});
@@ -188,13 +189,16 @@ 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');
});
@@ -227,6 +231,8 @@ describe('ShellTool', () => {
resolveExecutionPromise({
output: 'long output',
rawOutput: Buffer.from('long output'),
stdout: 'long output',
stderr: '',
exitCode: 0,
signal: null,
error: null,
@@ -251,7 +257,7 @@ describe('ShellTool', () => {
mockShellExecutionService.mockImplementation(() => {
throw error;
});
vi.mocked(fs.existsSync).mockReturnValue(true);
vi.mocked(fs.existsSync).mockReturnValue(true); // Pretend the file exists
const invocation = shellTool.build({ command: 'a-command' });
await expect(invocation.execute(mockAbortSignal)).rejects.toThrow(error);
@@ -274,26 +280,33 @@ 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',
chunk: 'hello world',
stream: 'stderr',
chunk: 'world',
});
// It should have been called once now with the combined output.
expect(updateOutputMock).toHaveBeenCalledOnce();
expect(updateOutputMock).toHaveBeenCalledWith('hello world');
expect(updateOutputMock).toHaveBeenCalledWith('hello \nworld');
resolveExecutionPromise({
rawOutput: Buffer.from(''),
output: '',
stdout: '',
stderr: '',
exitCode: 0,
signal: null,
error: null,
@@ -319,13 +332,16 @@ 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]',
@@ -334,6 +350,8 @@ describe('ShellTool', () => {
resolveExecutionPromise({
rawOutput: Buffer.from(''),
output: '',
stdout: '',
stderr: '',
exitCode: 0,
signal: null,
error: null,
+15 -10
View File
@@ -97,8 +97,6 @@ class ShellToolInvocation extends BaseToolInvocation<
async execute(
signal: AbortSignal,
updateOutput?: (output: string) => void,
terminalColumns?: number,
terminalRows?: number,
): Promise<ToolResult> {
const strippedCommand = stripShellWrapper(this.params.command);
@@ -131,7 +129,9 @@ class ShellToolInvocation extends BaseToolInvocation<
this.params.directory || '',
);
let cumulativeOutput = '';
let cumulativeStdout = '';
let cumulativeStderr = '';
let lastUpdateTime = Date.now();
let isBinaryStream = false;
@@ -148,9 +148,15 @@ class ShellToolInvocation extends BaseToolInvocation<
switch (event.type) {
case 'data':
if (isBinaryStream) break;
cumulativeOutput = event.chunk;
currentDisplayOutput = cumulativeOutput;
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 (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) {
shouldUpdate = true;
}
@@ -181,8 +187,6 @@ class ShellToolInvocation extends BaseToolInvocation<
}
},
signal,
terminalColumns,
terminalRows,
);
const result = await resultPromise;
@@ -214,7 +218,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 before it was cancelled:\n${result.output}`;
llmContent += ` Below is the output (on stdout and stderr) before it was cancelled:\n${result.output}`;
} else {
llmContent += ' There was no output before it was cancelled.';
}
@@ -228,7 +232,8 @@ class ShellToolInvocation extends BaseToolInvocation<
llmContent = [
`Command: ${this.params.command}`,
`Directory: ${this.params.directory || '(root)'}`,
`Output: ${result.output || '(empty)'}`,
`Stdout: ${result.stdout || '(empty)'}`,
`Stderr: ${result.stderr || '(empty)'}`,
`Error: ${finalError}`, // Use the cleaned error string.
`Exit Code: ${result.exitCode ?? '(none)'}`,
`Signal: ${result.signal ?? '(none)'}`,
+1 -12
View File
@@ -50,8 +50,6 @@ export interface ToolInvocation<
execute(
signal: AbortSignal,
updateOutput?: (output: string) => void,
terminalColumns?: number,
terminalRows?: number,
): Promise<TResult>;
}
@@ -80,8 +78,6 @@ export abstract class BaseToolInvocation<
abstract execute(
signal: AbortSignal,
updateOutput?: (output: string) => void,
terminalColumns?: number,
terminalRows?: number,
): Promise<TResult>;
}
@@ -200,16 +196,9 @@ 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,
terminalColumns,
terminalRows,
);
return invocation.execute(signal, updateOutput);
}
}