From 7cdfaaa6bdc6ccb3c1759e44d2fc8a2363f813f5 Mon Sep 17 00:00:00 2001 From: Spencer Date: Tue, 7 Apr 2026 19:33:15 +0000 Subject: [PATCH] chore: fix truncation logic and test duplications --- packages/core/.geminiignore | 0 packages/core/.gitignore | 0 .../core/src/scheduler/tool-executor.test.ts | 134 +--------------- packages/core/src/scheduler/tool-executor.ts | 145 ++++++++++++------ .../src/services/shellExecutionService.ts | 3 +- packages/core/src/tools/shell.ts | 7 +- 6 files changed, 103 insertions(+), 186 deletions(-) create mode 100644 packages/core/.geminiignore create mode 100644 packages/core/.gitignore diff --git a/packages/core/.geminiignore b/packages/core/.geminiignore new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/core/.gitignore b/packages/core/.gitignore new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts index 3007ffb9be..f3a3ff9da6 100644 --- a/packages/core/src/scheduler/tool-executor.test.ts +++ b/packages/core/src/scheduler/tool-executor.test.ts @@ -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; - // 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'); } }); diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index 4c5fdef5c0..2a407bddba 100644 --- a/packages/core/src/scheduler/tool-executor.ts +++ b/packages/core/src/scheduler/tool-executor.ts @@ -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 { - 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 = { diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts index fbb7df7e4a..ba7c33a896 100644 --- a/packages/core/src/services/shellExecutionService.ts +++ b/packages/core/src/services/shellExecutionService.ts @@ -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) { diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 096db9d674..8f1d9873a9 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -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