Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
186 changes: 185 additions & 1 deletion packages/core/src/tools/shell.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,11 @@ vi.mock('crypto');
vi.mock('../utils/summarizer.js');

import { initializeShellParsers } from '../utils/shell-utils.js';
import { ShellTool, OUTPUT_UPDATE_INTERVAL_MS } from './shell.js';
import {
ShellTool,
OUTPUT_UPDATE_INTERVAL_MS,
LIVE_OUTPUT_MAX_BUFFER_CHARS,
} from './shell.js';
import { debugLogger } from '../index.js';
import { type Config } from '../config/config.js';
import { NoopSandboxManager } from '../services/sandboxManager.js';
Expand Down Expand Up @@ -77,6 +81,7 @@ import {
} from '../confirmation-bus/types.js';
import { type MessageBus } from '../confirmation-bus/message-bus.js';
import { type SandboxManager } from '../services/sandboxManager.js';
import type { AnsiOutput } from '../utils/terminalSerializer.js';

interface TestableMockMessageBus extends MessageBus {
defaultToolDecision: 'allow' | 'deny' | 'ask_user';
Expand Down Expand Up @@ -686,6 +691,185 @@ EOF`;
await promise;
});

it('should show the first text output immediately and throttle subsequent text updates', async () => {
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});

mockShellOutputCallback({ type: 'data', chunk: 'first' });
expect(updateOutputMock).toHaveBeenCalledOnce();
expect(updateOutputMock).toHaveBeenLastCalledWith('first');

mockShellOutputCallback({ type: 'data', chunk: 'second' });
expect(updateOutputMock).toHaveBeenCalledOnce();

mockShellOutputCallback({ type: 'data', chunk: 'third' });
expect(updateOutputMock).toHaveBeenCalledOnce();

resolveShellExecution({ output: 'firstsecondthird' });
await promise;

expect(updateOutputMock).toHaveBeenCalledTimes(2);
expect(updateOutputMock).toHaveBeenLastCalledWith('firstsecondthird');
});

it('should flush trailing throttled text output when the command completes', async () => {
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});

mockShellOutputCallback({ type: 'data', chunk: 'first' });
mockShellOutputCallback({ type: 'data', chunk: 'second' });
expect(updateOutputMock).toHaveBeenCalledOnce();

resolveShellExecution({ output: 'firstsecond' });
await promise;

expect(updateOutputMock).toHaveBeenCalledTimes(2);
expect(updateOutputMock).toHaveBeenLastCalledWith('firstsecond');
});

it('should keep only a bounded text buffer for live display', async () => {
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});

mockShellOutputCallback({
type: 'data',
chunk: `older${'x'.repeat(LIVE_OUTPUT_MAX_BUFFER_CHARS)}`,
});

expect(updateOutputMock).toHaveBeenCalledOnce();
expect(updateOutputMock).toHaveBeenLastCalledWith(
'x'.repeat(LIVE_OUTPUT_MAX_BUFFER_CHARS),
);

resolveShellExecution({
output: `older${'x'.repeat(LIVE_OUTPUT_MAX_BUFFER_CHARS)}`,
});
await promise;
});

it('should not start the bounded live text buffer with a low surrogate', async () => {
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});
const emoji = '\uD83D\uDE00';

mockShellOutputCallback({
type: 'data',
chunk: `${emoji}${'x'.repeat(LIVE_OUTPUT_MAX_BUFFER_CHARS - 1)}`,
});

expect(updateOutputMock).toHaveBeenCalledOnce();
const displayedOutput = updateOutputMock.mock.calls[0][0] as string;
expect(displayedOutput.charCodeAt(0)).not.toBe(0xde00);
expect(displayedOutput).toHaveLength(LIVE_OUTPUT_MAX_BUFFER_CHARS - 1);

resolveShellExecution();
await promise;
});

it('should not throttle PTY AnsiOutput snapshots in the shell tool', async () => {
const firstAnsiOutput = [[{ text: 'first' }]] as AnsiOutput;
const secondAnsiOutput = [[{ text: 'second' }]] as AnsiOutput;
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});

mockShellOutputCallback({ type: 'data', chunk: firstAnsiOutput });
mockShellOutputCallback({ type: 'data', chunk: secondAnsiOutput });

expect(updateOutputMock).toHaveBeenCalledTimes(2);
expect(updateOutputMock).toHaveBeenNthCalledWith(1, firstAnsiOutput);
expect(updateOutputMock).toHaveBeenNthCalledWith(2, secondAnsiOutput);

resolveShellExecution({ ansiOutput: secondAnsiOutput });
await promise;
});

it('should trailing-flush throttled text output when the command goes silent', async () => {
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});

mockShellOutputCallback({ type: 'data', chunk: 'first' });
expect(updateOutputMock).toHaveBeenCalledOnce();
expect(updateOutputMock).toHaveBeenLastCalledWith('first');

mockShellOutputCallback({ type: 'data', chunk: 'second' });
expect(updateOutputMock).toHaveBeenCalledOnce();

await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1);

expect(updateOutputMock).toHaveBeenCalledTimes(2);
expect(updateOutputMock).toHaveBeenLastCalledWith('firstsecond');

resolveShellExecution({ output: 'firstsecond' });
await promise;
});

it('should trailing-flush throttled text output after only the remaining interval', async () => {
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});

mockShellOutputCallback({ type: 'data', chunk: 'first' });
expect(updateOutputMock).toHaveBeenCalledOnce();
expect(updateOutputMock).toHaveBeenLastCalledWith('first');

await vi.advanceTimersByTimeAsync(750);
mockShellOutputCallback({ type: 'data', chunk: 'second' });
expect(updateOutputMock).toHaveBeenCalledOnce();

await vi.advanceTimersByTimeAsync(249);
expect(updateOutputMock).toHaveBeenCalledOnce();

await vi.advanceTimersByTimeAsync(1);
expect(updateOutputMock).toHaveBeenCalledTimes(2);
expect(updateOutputMock).toHaveBeenLastCalledWith('firstsecond');

resolveShellExecution({ output: 'firstsecond' });
await promise;
});

it('should cancel the scheduled trailing flush when the command exits', async () => {
const invocation = shellTool.build({ command: 'printf output' });
const promise = invocation.execute({
abortSignal: mockAbortSignal,
updateOutput: updateOutputMock,
});

mockShellOutputCallback({ type: 'data', chunk: 'first' });
expect(updateOutputMock).toHaveBeenCalledOnce();

mockShellOutputCallback({ type: 'data', chunk: 'second' });
expect(updateOutputMock).toHaveBeenCalledOnce();

resolveShellExecution({ output: 'firstsecond' });
await promise;

expect(updateOutputMock).toHaveBeenCalledTimes(2);
expect(updateOutputMock).toHaveBeenLastCalledWith('firstsecond');

await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS * 5);
expect(updateOutputMock).toHaveBeenCalledTimes(2);
});

it('should NOT call updateOutput if the command is backgrounded', async () => {
const invocation = shellTool.build({
command: 'sleep 10',
Expand Down
105 changes: 97 additions & 8 deletions packages/core/src/tools/shell.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,10 +58,29 @@ import {
} from '../sandbox/utils/proactivePermissions.js';

export const OUTPUT_UPDATE_INTERVAL_MS = 1000;
export const LIVE_OUTPUT_MAX_BUFFER_CHARS = 100_000;

// Delay so user does not see the output of the process before the process is moved to the background.
const BACKGROUND_DELAY_MS = 200;
const SHOW_NL_DESCRIPTION_THRESHOLD = 150;
const LOW_SURROGATE_START = 0xdc00;
const LOW_SURROGATE_END = 0xdfff;

function trimLiveOutputBuffer(output: string): string {
if (output.length <= LIVE_OUTPUT_MAX_BUFFER_CHARS) {
return output;
}

let startIndex = output.length - LIVE_OUTPUT_MAX_BUFFER_CHARS;
const firstCodeUnit = output.charCodeAt(startIndex);
if (
firstCodeUnit >= LOW_SURROGATE_START &&
firstCodeUnit <= LOW_SURROGATE_END
) {
startIndex += 1;
}
return output.slice(startIndex);
}

export interface ShellToolParams {
command: string;
Expand Down Expand Up @@ -470,6 +489,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
const timeoutMs = this.context.config.getShellToolInactivityTimeout();
const timeoutController = new AbortController();
let timeoutTimer: NodeJS.Timeout | undefined;
let trailingFlushTimer: ReturnType<typeof setTimeout> | null = null;

// Handle signal combination manually to avoid TS issues or runtime missing features
const combinedController = new AbortController();
Expand Down Expand Up @@ -502,9 +522,61 @@ export class ShellToolInvocation extends BaseToolInvocation<
};
}
let cumulativeOutput: string | AnsiOutput = '';
let lastUpdateTime = Date.now();
let lastUpdateTime = 0;
let hasFlushedOutput = false;
let hasPendingOutput = false;
let isBinaryStream = false;

const appendToLiveOutputBuffer = (chunk: string) => {
const currentOutput =
typeof cumulativeOutput === 'string' ? cumulativeOutput : '';
if (chunk.length >= LIVE_OUTPUT_MAX_BUFFER_CHARS) {
cumulativeOutput = trimLiveOutputBuffer(chunk);
return;
}

const nextOutput = currentOutput + chunk;
cumulativeOutput = trimLiveOutputBuffer(nextOutput);
};

const cancelTrailingFlush = () => {
if (trailingFlushTimer !== null) {
clearTimeout(trailingFlushTimer);
trailingFlushTimer = null;
}
};

const flushOutput = () => {
cancelTrailingFlush();
if (!hasPendingOutput || !updateOutput || this.params.is_background) {
return;
}

updateOutput(cumulativeOutput);
hasPendingOutput = false;
hasFlushedOutput = true;
lastUpdateTime = Date.now();
};
Comment on lines +525 to +559

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

To ensure that the last chunk of output is not stuck in the buffer when a command becomes silent but remains active, we should implement a trailing edge flush using a timer. This ensures that the UI is updated even if no further output events are received within the throttle interval. When implementing this, ensure the flush emits only the incremental delta instead of the full accumulated text to avoid redundant data transfer. The check for signal.aborted in flushOutput correctly utilizes the AbortSignal for cancellation safety. Note that while this check provides safety, it is still recommended to clear the timer in the finally block of the execute method to avoid unnecessary resource usage, although that block is outside the current diff scope.

References
  1. When implementing streaming message events, emit only the incremental delta instead of the full accumulated text to avoid redundant data transfer and potential display issues, particularly when a higher-level component is responsible for content accumulation.
  2. Asynchronous operations waiting for user input via the MessageBus should rely on the provided AbortSignal for cancellation, rather than implementing a separate timeout, to maintain consistency with existing patterns.


const scheduleTrailingFlush = () => {
if (
trailingFlushTimer !== null ||
!updateOutput ||
this.params.is_background
) {
return;
}
const elapsedSinceLastUpdate = Date.now() - lastUpdateTime;
const trailingDelayMs = Math.max(
OUTPUT_UPDATE_INTERVAL_MS - elapsedSinceLastUpdate,
0,
);
trailingFlushTimer = setTimeout(() => {
trailingFlushTimer = null;
flushOutput();
}, trailingDelayMs);
};

const resetTimeout = () => {
if (timeoutMs <= 0) {
return;
Expand All @@ -529,43 +601,53 @@ export class ShellToolInvocation extends BaseToolInvocation<
cwd,
(event: ShellOutputEvent) => {
resetTimeout(); // Reset timeout on any event
if (!updateOutput) {
return;
}

let shouldUpdate = false;

switch (event.type) {
case 'data':
if (isBinaryStream) break;
cumulativeOutput = event.chunk;
shouldUpdate = true;
if (typeof event.chunk === 'string') {
appendToLiveOutputBuffer(event.chunk);
shouldUpdate =
!hasFlushedOutput ||
Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS;
if (!shouldUpdate) {
scheduleTrailingFlush();
}
} else {
cumulativeOutput = event.chunk;
shouldUpdate = true;
}
hasPendingOutput = true;
break;
case 'binary_detected':
isBinaryStream = true;
cumulativeOutput =
'[Binary output detected. Halting stream...]';
hasPendingOutput = true;
shouldUpdate = true;
break;
case 'binary_progress':
isBinaryStream = true;
cumulativeOutput = `[Receiving binary output... ${formatBytes(
event.bytesReceived,
)} received]`;
hasPendingOutput = true;
if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) {
shouldUpdate = true;
}
break;
case 'exit':
flushOutput();
break;
default: {
throw new Error('An unhandled ShellOutputEvent was found.');
}
}

if (shouldUpdate && !this.params.is_background) {
updateOutput(cumulativeOutput);
lastUpdateTime = Date.now();
flushOutput();
}
Comment on lines 649 to 651

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the data handler to schedule a trailing flush if the current event is throttled. This ensures that sporadic output is eventually displayed as an incremental delta even if the process doesn't exit immediately.

References
  1. When implementing streaming message events, emit only the incremental delta instead of the full accumulated text to avoid redundant data transfer and potential display issues, particularly when a higher-level component is responsible for content accumulation.

},
combinedController.signal,
Expand Down Expand Up @@ -639,6 +721,9 @@ export class ShellToolInvocation extends BaseToolInvocation<
}

const result = await resultPromise;
if (!result.backgrounded) {
flushOutput();
}

const backgroundPIDs: number[] = [];
if (os.platform() !== 'win32') {
Expand Down Expand Up @@ -966,6 +1051,10 @@ export class ShellToolInvocation extends BaseToolInvocation<
};
} finally {
if (timeoutTimer) clearTimeout(timeoutTimer);
if (trailingFlushTimer) {
clearTimeout(trailingFlushTimer);
trailingFlushTimer = null;
}
signal.removeEventListener('abort', onAbort);
timeoutController.signal.removeEventListener('abort', onAbort);
if (tempFilePath) {
Expand Down
Loading