mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-03 16:34:31 -07:00
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:
@@ -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)),
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user