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

View File

@@ -4,7 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
import type React from 'react';
import React from 'react';
import { Box, Text } from 'ink';
import type { AnsiLine, AnsiOutput, AnsiToken } from '@google/gemini-cli-core';
@@ -53,23 +53,26 @@ export const AnsiOutputText: React.FC<AnsiOutputProps> = ({
);
};
export const AnsiLineText: React.FC<{ line: AnsiLine }> = ({ line }) => (
<Text>
{line.length > 0
? line.map((token: AnsiToken, tokenIndex: number) => (
<Text
key={tokenIndex}
color={token.fg}
backgroundColor={token.bg}
inverse={token.inverse}
dimColor={token.dim}
bold={token.bold}
italic={token.italic}
underline={token.underline}
>
{token.text}
</Text>
))
: null}
</Text>
export const AnsiLineText = React.memo<{ line: AnsiLine }>(
({ line }: { line: AnsiLine }) => (
<Text>
{line.length > 0
? line.map((token: AnsiToken, tokenIndex: number) => (
<Text
key={tokenIndex}
color={token.fg}
backgroundColor={token.bg}
inverse={token.inverse}
dimColor={token.dim}
bold={token.bold}
italic={token.italic}
underline={token.underline}
>
{token.text}
</Text>
))
: null}
</Text>
),
);
AnsiLineText.displayName = 'AnsiLineText';

View File

@@ -144,6 +144,11 @@ describe('useExecutionLifecycle', () => {
enableEnvironmentVariableRedaction: false,
},
}),
getTruncateToolOutputThreshold: () => 40000,
storage: {
getProjectTempDir: () => '/tmp/project',
},
getSessionId: () => 'test-session',
} as unknown as Config;
mockGeminiClient = { addHistory: vi.fn() } as unknown as GeminiClient;
@@ -155,6 +160,16 @@ describe('useExecutionLifecycle', () => {
mockIsBinary.mockReturnValue(false);
vi.mocked(fs.existsSync).mockReturnValue(false);
vi.mocked(fs.createWriteStream).mockReturnValue({
write: vi.fn(),
end: vi.fn().mockImplementation((cb: () => void) => {
if (cb) cb();
}),
destroy: vi.fn(),
bytesWritten: 0,
closed: false,
} as unknown as fs.WriteStream);
mockShellExecutionService.mockImplementation((_cmd, _cwd, callback) => {
mockShellOutputCallback = callback;
return Promise.resolve({

View File

@@ -20,6 +20,7 @@ import {
ShellExecutionService,
ExecutionLifecycleService,
CoreToolCallStatus,
moveToolOutputToFile,
} from '@google/gemini-cli-core';
import { type PartListUnion } from '@google/genai';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
@@ -39,16 +40,15 @@ export { type BackgroundTask };
export const OUTPUT_UPDATE_INTERVAL_MS = 1000;
const RESTORE_VISIBILITY_DELAY_MS = 300;
const MAX_OUTPUT_LENGTH = 10000;
function addShellCommandToGeminiHistory(
geminiClient: GeminiClient,
rawQuery: string,
resultText: string,
maxOutputLength: number,
) {
const modelContent =
resultText.length > MAX_OUTPUT_LENGTH
? resultText.substring(0, MAX_OUTPUT_LENGTH) + '\n... (truncated)'
maxOutputLength > 0 && resultText.length > maxOutputLength
? resultText.substring(0, maxOutputLength) + '\n... (truncated)'
: resultText;
// Escape backticks to prevent prompt injection breakouts
@@ -378,6 +378,11 @@ export const useExecutionLifecycle = (
let cumulativeStdout: string | AnsiOutput = '';
let isBinaryStream = false;
let binaryBytesReceived = 0;
let totalBytesWritten = 0;
const outputFileName = `gemini_shell_output_${crypto.randomBytes(6).toString('hex')}.log`;
const outputFilePath = path.join(os.tmpdir(), outputFileName);
const outputStream = fs.createWriteStream(outputFilePath);
const initialToolDisplay: IndividualToolCallDisplay = {
callId,
@@ -395,6 +400,7 @@ export const useExecutionLifecycle = (
});
let executionPid: number | undefined;
let fullOutputReturned = false;
const abortHandler = () => {
onDebugMessage(
@@ -424,11 +430,23 @@ export const useExecutionLifecycle = (
let shouldUpdate = false;
switch (event.type) {
case 'raw_data':
if (!isBinaryStream) {
outputStream.write(event.chunk);
totalBytesWritten += Buffer.byteLength(event.chunk);
}
break;
case 'data':
if (isBinaryStream) break;
if (typeof event.chunk === 'string') {
if (typeof cumulativeStdout === 'string') {
cumulativeStdout += event.chunk;
// Keep a small buffer for the UI to prevent memory spikes and Ink lagging
const MAX_UI_LENGTH = 100000; // 100KB
if (cumulativeStdout.length > MAX_UI_LENGTH) {
cumulativeStdout =
cumulativeStdout.slice(-MAX_UI_LENGTH);
}
} else {
cumulativeStdout = event.chunk;
}
@@ -514,6 +532,9 @@ export const useExecutionLifecycle = (
}
const result = await resultPromise;
await new Promise<void>((resolve) => {
outputStream.end(resolve);
});
setPendingHistoryItem(null);
if (result.backgrounded && result.pid) {
@@ -533,6 +554,26 @@ export const useExecutionLifecycle = (
} else {
mainContent =
result.output.trim() || '(Command produced no output)';
const threshold = config.getTruncateToolOutputThreshold();
if (threshold > 0 && totalBytesWritten >= threshold) {
const { outputFile: savedPath } = await moveToolOutputToFile(
outputFilePath,
SHELL_COMMAND_NAME,
callId,
config.storage.getProjectTempDir(),
config.getSessionId(),
);
const warning = `[Full command output saved to: ${savedPath}]`;
mainContent = mainContent.includes(
'[GEMINI_CLI_WARNING: Output truncated.',
)
? mainContent.replace(
/\[GEMINI_CLI_WARNING: Output truncated\..*?\]/,
warning,
)
: `${mainContent}\n\n${warning}`;
fullOutputReturned = true;
}
}
let finalOutput: string | AnsiOutput =
@@ -617,7 +658,12 @@ export const useExecutionLifecycle = (
);
}
addShellCommandToGeminiHistory(geminiClient, rawQuery, mainContent);
addShellCommandToGeminiHistory(
geminiClient,
rawQuery,
mainContent,
config.getTruncateToolOutputThreshold(),
);
} catch (err) {
setPendingHistoryItem(null);
const errorMessage = err instanceof Error ? err.message : String(err);
@@ -630,12 +676,23 @@ export const useExecutionLifecycle = (
);
} finally {
abortSignal.removeEventListener('abort', abortHandler);
if (!outputStream.closed) {
outputStream.destroy();
}
if (pwdFilePath && fs.existsSync(pwdFilePath)) {
fs.unlinkSync(pwdFilePath);
}
dispatch({ type: 'SET_ACTIVE_PTY', pid: null });
setShellInputFocused(false);
if (!fullOutputReturned && fs.existsSync(outputFilePath)) {
try {
fs.unlinkSync(outputFilePath);
} catch {
// Ignore errors during unlink
}
}
}
};