MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-18 09:11:55 -07:00
Code Issues Packages Projects Releases Wiki Activity

chore: fix truncation logic and test duplications

Browse Source
This commit is contained in:
Spencer
2026-04-07 19:33:15 +00:00
parent ea4dd5ae35
commit 7cdfaaa6bd
6 changed files with 103 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
packages/core/.geminiignore Normal file
View File

0
packages/core/.gitignore vendored Normal file
View File

134
packages/core/src/scheduler/tool-executor.test.ts
View File

@@ -5,6 +5,7 @@
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import fsPromises from 'node:fs/promises';
import { ToolExecutor } from './tool-executor.js';
import {
type Config,
@@ -32,6 +33,7 @@ import {
vi.mock('../utils/fileUtils.js', () => ({
saveTruncatedToolOutput: vi.fn(),
formatTruncatedToolOutput: vi.fn(),
moveToolOutputToFile: vi.fn(),
}));
// Mock executeToolWithHooks
@@ -431,68 +433,6 @@ describe('ToolExecutor', () => {
10, // threshold (maxChars)
);
expect(result.status).toBe(CoreToolCallStatus.Success);
if (result.status === CoreToolCallStatus.Success) {
const response = result.response.responseParts[0]?.functionResponse
?.response as Record<string, unknown>;
// The content should be the *truncated* version returned by the mock formatTruncatedToolOutput
expect(response).toEqual({ output: 'TruncatedContent...' });
expect(result.response.outputFile).toBe('/tmp/truncated_output.txt');
}
});
it('should truncate large output and move file when fullOutputFilePath is provided', async () => {
// 1. Setup Config for Truncation
vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10);
vi.spyOn(config.storage, 'getProjectTempDir').mockReturnValue('/tmp');
const mockTool = new MockTool({ name: SHELL_TOOL_NAME });
const invocation = mockTool.build({});
const longOutput = 'This is a very long output that should be truncated.';
// 2. Mock execution returning long content AND fullOutputFilePath
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
llmContent: longOutput,
returnDisplay: longOutput,
fullOutputFilePath: '/tmp/temp_full_output.txt',
});
const scheduledCall: ScheduledToolCall = {
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-trunc-full',
name: SHELL_TOOL_NAME,
args: { command: 'echo long' },
isClientInitiated: false,
prompt_id: 'prompt-trunc-full',
},
tool: mockTool,
invocation: invocation as unknown as AnyToolInvocation,
startTime: Date.now(),
};
// 3. Execute
const result = await executor.execute({
call: scheduledCall,
signal: new AbortController().signal,
onUpdateToolCall: vi.fn(),
});
// 4. Verify Truncation Logic
expect(fileUtils.moveToolOutputToFile).toHaveBeenCalledWith(
'/tmp/temp_full_output.txt',
SHELL_TOOL_NAME,
'call-trunc-full',
expect.any(String), // temp dir
'test-session-id', // session id from makeFakeConfig
);
expect(fileUtils.formatTruncatedToolOutput).toHaveBeenCalledWith(
longOutput,
'/tmp/moved_output.txt',
10, // threshold (maxChars)
);
expect(result.status).toBe(CoreToolCallStatus.Success);
if (result.status === CoreToolCallStatus.Success) {
const response = result.response.responseParts[0]?.functionResponse
@@ -500,76 +440,8 @@ describe('ToolExecutor', () => {
// The content should be the *truncated* version returned by the mock formatTruncatedToolOutput
expect(response).toEqual({
output: 'TruncatedContent...',
outputFile: '/tmp/moved_output.txt',
outputFile: '/tmp/truncated_output.txt',
});
expect(result.response.outputFile).toBe('/tmp/moved_output.txt');
}
});
it('should truncate large MCP tool output with single text Part', async () => {
// 1. Setup Config for Truncation
vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10);
vi.spyOn(config.storage, 'getProjectTempDir').mockReturnValue('/tmp');
const mcpToolName = 'get_big_text';
const messageBus = createMockMessageBus();
const mcpTool = new DiscoveredMCPTool(
{} as CallableTool,
'my-server',
'get_big_text',
'A test MCP tool',
{},
messageBus,
);
const invocation = mcpTool.build({});
const longText = 'This is a very long MCP output that should be truncated.';
// 2. Mock execution returning Part[] with single text Part
// We do NOT provide fullOutputFilePath here because we want to test the path
// that uses saveTruncatedToolOutput for MCP tools.
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
llmContent: [{ text: longText }],
returnDisplay: longText,
});
const scheduledCall: ScheduledToolCall = {
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-mcp-trunc',
name: mcpToolName,
args: { query: 'test' },
isClientInitiated: false,
prompt_id: 'prompt-mcp-trunc',
},
tool: mcpTool,
invocation: invocation as unknown as AnyToolInvocation,
startTime: Date.now(),
};
// 3. Execute
const result = await executor.execute({
call: scheduledCall,
signal: new AbortController().signal,
onUpdateToolCall: vi.fn(),
});
// 4. Verify Truncation Logic
expect(fileUtils.saveTruncatedToolOutput).toHaveBeenCalledWith(
longText,
mcpToolName,
'call-mcp-trunc',
expect.any(String),
'test-session-id',
);
expect(fileUtils.formatTruncatedToolOutput).toHaveBeenCalledWith(
longText,
'/tmp/truncated_output.txt',
10,
);
expect(result.status).toBe(CoreToolCallStatus.Success);
if (result.status === CoreToolCallStatus.Success) {
expect(result.response.outputFile).toBe('/tmp/truncated_output.txt');
}
});

145
packages/core/src/scheduler/tool-executor.ts
View File

@@ -26,6 +26,7 @@ import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
import {
saveTruncatedToolOutput,
formatTruncatedToolOutput,
moveToolOutputToFile,
} from '../utils/fileUtils.js';
import { convertToFunctionResponse } from '../utils/generateContentResponseUtilities.js';
import {
@@ -161,6 +162,16 @@ export class ToolExecutor {
toolResult,
);
} else {
if (toolResult.fullOutputFilePath) {
await fsPromises
.unlink(toolResult.fullOutputFilePath)
.catch((error) => {
debugLogger.warn(
`Failed to delete temporary tool output file on error: ${toolResult.fullOutputFilePath}`,
error,
);
});
}
const displayText =
typeof toolResult.returnDisplay === 'string'
? toolResult.returnDisplay
@@ -209,6 +220,7 @@ export class ToolExecutor {
private async truncateOutputIfNeeded(
call: ToolCall,
content: PartListUnion,
fullOutputFilePath?: string,
): Promise<{ truncatedContent: PartListUnion; outputFile?: string }> {
if (this.config.isContextManagementEnabled()) {
const distiller = new ToolOutputDistillationService(
@@ -216,7 +228,22 @@ export class ToolExecutor {
this.context.geminiClient,
this.context.promptId,
);
return distiller.distill(call.request.name, call.request.callId, content);
const result = await distiller.distill(
call.request.name,
call.request.callId,
content,
);
if (fullOutputFilePath && !result.outputFile) {
try {
await fsPromises.unlink(fullOutputFilePath);
} catch (error) {
debugLogger.warn(
`Failed to delete temporary tool output file: ${fullOutputFilePath}`,
error,
);
}
}
return result;
}
const toolName = call.request.name;
@@ -228,13 +255,28 @@ export class ToolExecutor {
if (threshold > 0 && content.length > threshold) {
const originalContentLength = content.length;
const { outputFile: savedPath } = await saveTruncatedToolOutput(
content,
toolName,
callId,
this.config.storage.getProjectTempDir(),
this.context.promptId,
);
let savedPath: string;
if (fullOutputFilePath) {
const { outputFile: movedPath } = await moveToolOutputToFile(
fullOutputFilePath,
toolName,
callId,
this.config.storage.getProjectTempDir(),
this.config.getSessionId(),
);
savedPath = movedPath;
} else {
const { outputFile: writtenPath } = await saveTruncatedToolOutput(
content,
toolName,
callId,
this.config.storage.getProjectTempDir(),
this.context.promptId,
);
savedPath = writtenPath;
}
outputFile = savedPath;
const truncatedContent = formatTruncatedToolOutput(
content,
@@ -297,11 +339,33 @@ export class ToolExecutor {
}),
);
if (fullOutputFilePath) {
try {
await fsPromises.unlink(fullOutputFilePath);
} catch (error) {
debugLogger.warn(
`Failed to delete temporary tool output file: ${fullOutputFilePath}`,
error,
);
}
}
return { truncatedContent, outputFile };
}
}
}
if (fullOutputFilePath && !outputFile) {
try {
await fsPromises.unlink(fullOutputFilePath);
} catch (error) {
debugLogger.warn(
`Failed to delete temporary tool output file: ${fullOutputFilePath}`,
error,
);
}
}
return { truncatedContent: content, outputFile };
}
@@ -326,7 +390,11 @@ export class ToolExecutor {
// Attempt to truncate and save output if we have content, even in cancellation case
// This is to handle cases where the tool may have produced output before cancellation
const { truncatedContent: output, outputFile: truncatedOutputFile } =
await this.truncateOutputIfNeeded(call, toolResult?.llmContent);
await this.truncateOutputIfNeeded(
call,
toolResult.llmContent,
toolResult.fullOutputFilePath,
);
outputFile = truncatedOutputFile;
responseParts = convertToFunctionResponse(
@@ -335,6 +403,7 @@ export class ToolExecutor {
output,
this.config.getActiveModel(),
this.config,
outputFile,
);
// Inject the cancellation error into the response object
@@ -344,6 +413,16 @@ export class ToolExecutor {
respObj['error'] = errorMessage;
}
} else {
if (toolResult?.fullOutputFilePath) {
try {
await fsPromises.unlink(toolResult.fullOutputFilePath);
} catch (error) {
debugLogger.warn(
`Failed to delete temporary tool output file: ${toolResult.fullOutputFilePath}`,
error,
);
}
}
responseParts = [
{
functionResponse: {
@@ -380,56 +459,22 @@ export class ToolExecutor {
call: ToolCall,
toolResult: ToolResult,
): Promise<SuccessfulToolCall> {
let { truncatedContent: content, outputFile } =
await this.truncateOutputIfNeeded(call, toolResult.llmContent);
const { truncatedContent: content, outputFile } =
await this.truncateOutputIfNeeded(
call,
toolResult.llmContent,
toolResult.fullOutputFilePath,
);
const toolName = call.request.originalRequestName || call.request.name;
const callId = call.request.callId;
if (toolResult.fullOutputFilePath) {
const threshold = this.config.getTruncateToolOutputThreshold();
if (
threshold > 0 &&
typeof content === 'string' &&
content.length > threshold
) {
const { outputFile: savedPath } = await moveToolOutputToFile(
toolResult.fullOutputFilePath,
toolName,
callId,
this.config.storage.getProjectTempDir(),
this.config.getSessionId(),
);
outputFile = savedPath;
content = formatTruncatedToolOutput(content, outputFile, threshold);
logToolOutputTruncated(
this.config,
new ToolOutputTruncatedEvent(call.request.prompt_id, {
toolName,
originalContentLength: content.length, // approximation
truncatedContentLength: content.length,
threshold,
}),
);
} else {
// If the content is not truncated, we don't need the temporary file.
try {
await fsPromises.unlink(toolResult.fullOutputFilePath);
} catch (error) {
debugLogger.warn(
`Failed to delete temporary tool output file: ${toolResult.fullOutputFilePath}`,
error,
);
}
}
}
const response = convertToFunctionResponse(
toolName,
callId,
content,
this.config.getActiveModel(),
this.config,
outputFile,
);
const successResponse: ToolCallResponseInfo = {

3
packages/core/src/services/shellExecutionService.ts
View File

@@ -1182,7 +1182,8 @@ export class ShellExecutionService {
let lastYdisp = 0;
let hasReachedMax = false;
const scrollbackLimit = shellExecutionConfig.scrollback ?? SCROLLBACK_LIMIT;
const scrollbackLimit =
shellExecutionConfig.scrollback ?? SCROLLBACK_LIMIT;
headlessTerminal.onScroll((ydisp) => {
if (!isWriting) {

7
packages/core/src/tools/shell.ts
View File

@@ -9,7 +9,6 @@ import fs from 'node:fs';
import path from 'node:path';
import os from 'node:os';
import crypto from 'node:crypto';
import type { Config } from '../config/config.js';
import { debugLogger } from '../utils/debugLogger.js';
import {
type SandboxPermissions,
@@ -937,7 +936,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
this.context.geminiClient,
signal,
);
const threshold = this.config.getTruncateToolOutputThreshold();
const threshold = this.context.config.getTruncateToolOutputThreshold();
const fullOutputFilePath =
threshold > 0 && totalBytesWritten >= threshold
? outputFilePath
@@ -945,7 +944,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
const toolResult: ToolResult = {
llmContent: summary,
returnDisplay: typeof returnDisplayMessage !== 'undefined' ? returnDisplayMessage : returnDisplay,
returnDisplay,
fullOutputFilePath,
...executionError,
};
@@ -955,7 +954,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
return toolResult;
}
const threshold = this.config.getTruncateToolOutputThreshold();
const threshold = this.context.config.getTruncateToolOutputThreshold();
const fullOutputFilePath =
threshold > 0 && totalBytesWritten >= threshold
? outputFilePath