Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 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
92 changes: 91 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,91 @@ 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();

await vi.advanceTimersByTimeAsync(OUTPUT_UPDATE_INTERVAL_MS + 1);
mockShellOutputCallback({ type: 'data', chunk: 'third' });

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

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

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 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 NOT call updateOutput if the command is backgrounded', async () => {
const invocation = shellTool.build({
command: 'sleep 10',
Expand Down
55 changes: 47 additions & 8 deletions packages/core/src/tools/shell.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ 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;
Expand Down Expand Up @@ -502,9 +503,37 @@ 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 = chunk.slice(-LIVE_OUTPUT_MAX_BUFFER_CHARS);
return;
}

const nextOutput = currentOutput + chunk;
cumulativeOutput =
nextOutput.length > LIVE_OUTPUT_MAX_BUFFER_CHARS
? nextOutput.slice(-LIVE_OUTPUT_MAX_BUFFER_CHARS)
: nextOutput;
};

const flushOutput = () => {
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 resetTimeout = () => {
if (timeoutMs <= 0) {
return;
Expand All @@ -529,43 +558,50 @@ 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;
} 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 +675,9 @@ export class ShellToolInvocation extends BaseToolInvocation<
}

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

const backgroundPIDs: number[] = [];
if (os.platform() !== 'win32') {
Expand Down
Loading