Checkpoint of shell optimization

fix(cli): Write shell command output to a file and limit memory buffered in UI

Fixes.

Checkpoint.

fix(core, cli): await outputStream.end() to prevent race conditions

This commit fixes a critical race condition where
was called synchronously without being awaited. This led to potential file
truncation or EBUSY errors on Windows when attempting to manipulate the file
immediately after the  call.

Additionally, this change removes fixed wait times (`setTimeout`) that
were previously used in test files as a band-aid.

fix(core): stream processed xterm output to file to remove spurious escape codes

test(core): update shell regression tests to use file_data events
This commit is contained in:
jacob314
2026-02-19 11:12:13 -08:00
committed by Spencer
parent 82e8d67a78
commit 954835123f
16 changed files with 963 additions and 63 deletions
+1
View File
@@ -119,6 +119,7 @@ describe('ShellTool', () => {
getDebugMode: vi.fn().mockReturnValue(false),
getTargetDir: vi.fn().mockReturnValue(tempRootDir),
getSummarizeToolOutputConfig: vi.fn().mockReturnValue(undefined),
getTruncateToolOutputThreshold: vi.fn().mockReturnValue(40000),
getWorkspaceContext: vi
.fn()
.mockReturnValue(new WorkspaceContext(tempRootDir)),
+75 -19
View File
@@ -9,8 +9,12 @@ import fs from 'node:fs';
import path from 'node:path';
import os from 'node:os';
import crypto from 'node:crypto';
import { debugLogger } from '../index.js';
import { type SandboxPermissions } from '../services/sandboxManager.js';
import type { Config } from '../config/config.js';
import { debugLogger } from '../utils/debugLogger.js';
import {
type SandboxPermissions,
getPathIdentity,
} from '../services/sandboxManager.js';
import { ToolErrorType } from './tool-error.js';
import {
BaseDeclarativeTool,
@@ -459,6 +463,12 @@ export class ShellToolInvocation extends BaseToolInvocation<
const onAbort = () => combinedController.abort();
const outputFileName = `gemini_shell_output_${crypto.randomBytes(6).toString('hex')}.log`;
const outputFilePath = path.join(os.tmpdir(), outputFileName);
const outputStream = fs.createWriteStream(outputFilePath);
let fullOutputReturned = false;
try {
// pgrep is not available on Windows, so we can't get background PIDs
const commandToExecute = this.wrapCommandForPgrep(
@@ -485,6 +495,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
let cumulativeOutput: string | AnsiOutput = '';
let lastUpdateTime = Date.now();
let isBinaryStream = false;
let totalBytesWritten = 0;
const resetTimeout = () => {
if (timeoutMs <= 0) {
@@ -510,31 +521,46 @@ export class ShellToolInvocation extends BaseToolInvocation<
cwd,
(event: ShellOutputEvent) => {
resetTimeout(); // Reset timeout on any event
if (!updateOutput) {
return;
}
let shouldUpdate = false;
switch (event.type) {
case 'raw_data':
// We do not write raw data to the file to avoid spurious escape codes.
// We rely on 'file_data' for the clean output stream.
break;
case 'file_data':
if (!isBinaryStream) {
totalBytesWritten += Buffer.byteLength(event.chunk);
outputStream.write(event.chunk);
}
break;
case 'data':
if (isBinaryStream) break;
cumulativeOutput = event.chunk;
shouldUpdate = true;
if (updateOutput && !this.params.is_background) {
updateOutput(cumulativeOutput);
lastUpdateTime = Date.now();
}
break;
case 'binary_detected':
isBinaryStream = true;
cumulativeOutput =
'[Binary output detected. Halting stream...]';
shouldUpdate = true;
if (updateOutput && !this.params.is_background) {
updateOutput(cumulativeOutput);
}
break;
case 'binary_progress':
isBinaryStream = true;
cumulativeOutput = `[Receiving binary output... ${formatBytes(
event.bytesReceived,
)} received]`;
if (Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS) {
shouldUpdate = true;
if (
updateOutput &&
!this.params.is_background &&
Date.now() - lastUpdateTime > OUTPUT_UPDATE_INTERVAL_MS
) {
updateOutput(cumulativeOutput);
lastUpdateTime = Date.now();
}
break;
case 'exit':
@@ -543,11 +569,6 @@ export class ShellToolInvocation extends BaseToolInvocation<
throw new Error('An unhandled ShellOutputEvent was found.');
}
}
if (shouldUpdate && !this.params.is_background) {
updateOutput(cumulativeOutput);
lastUpdateTime = Date.now();
}
},
combinedController.signal,
this.context.config.getEnableInteractiveShell(),
@@ -620,6 +641,9 @@ export class ShellToolInvocation extends BaseToolInvocation<
}
const result = await resultPromise;
await new Promise<void>((resolve) => {
outputStream.end(resolve);
});
const backgroundPIDs: number[] = [];
if (os.platform() !== 'win32') {
@@ -913,21 +937,46 @@ export class ShellToolInvocation extends BaseToolInvocation<
this.context.geminiClient,
signal,
);
return {
const threshold = this.config.getTruncateToolOutputThreshold();
const fullOutputFilePath =
threshold > 0 && totalBytesWritten >= threshold
? outputFilePath
: undefined;
const toolResult: ToolResult = {
llmContent: summary,
returnDisplay,
returnDisplay: typeof returnDisplayMessage !== 'undefined' ? returnDisplayMessage : returnDisplay,
fullOutputFilePath,
...executionError,
};
if (toolResult.fullOutputFilePath) {
fullOutputReturned = true;
}
return toolResult;
}
return {
const threshold = this.config.getTruncateToolOutputThreshold();
const fullOutputFilePath =
threshold > 0 && totalBytesWritten >= threshold
? outputFilePath
: undefined;
const toolResult: ToolResult = {
llmContent,
returnDisplay,
data,
fullOutputFilePath,
...executionError,
};
if (toolResult.fullOutputFilePath) {
fullOutputReturned = true;
}
return toolResult;
} finally {
if (timeoutTimer) clearTimeout(timeoutTimer);
if (!outputStream.closed) {
outputStream.destroy();
}
signal.removeEventListener('abort', onAbort);
timeoutController.signal.removeEventListener('abort', onAbort);
try {
@@ -935,6 +984,13 @@ export class ShellToolInvocation extends BaseToolInvocation<
} catch {
// Ignore errors during unlink
}
if (!fullOutputReturned) {
try {
await fsPromises.unlink(outputFilePath);
} catch {
// Ignore errors during unlink
}
}
}
}
}
+7
View File
@@ -776,6 +776,13 @@ export interface ToolResult {
name: string;
args: Record<string, unknown>;
};
/**
* Optional path to a file containing the full, non-truncated output of the tool.
* If provided, the scheduler may use this file for long-term storage and
* reference it in the conversation history if the output is truncated.
*/
fullOutputFilePath?: string;
}
/**