diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts index 0d77204f4e..d5f92806f5 100644 --- a/packages/core/src/scheduler/tool-executor.test.ts +++ b/packages/core/src/scheduler/tool-executor.test.ts @@ -16,6 +16,8 @@ import { MockTool } from '../test-utils/mock-tool.js'; import type { ScheduledToolCall } from './types.js'; import { CoreToolCallStatus } from './types.js'; import { SHELL_TOOL_NAME } from '../tools/tool-names.js'; +import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; +import type { CallableTool } from '@google/genai'; import * as fileUtils from '../utils/fileUtils.js'; import * as coreToolHookTriggers from '../core/coreToolHookTriggers.js'; import { ShellToolInvocation } from '../tools/shell.js'; @@ -312,6 +314,162 @@ describe('ToolExecutor', () => { } }); + 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 + 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'); + } + }); + + it('should not truncate MCP tool output with multiple Parts', async () => { + vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10); + + 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 long text that exceeds the threshold.'; + + // Part[] with multiple parts — should NOT be truncated + vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({ + llmContent: [{ text: longText }, { text: 'second part' }], + returnDisplay: longText, + }); + + const scheduledCall: ScheduledToolCall = { + status: CoreToolCallStatus.Scheduled, + request: { + callId: 'call-mcp-multi', + name: 'get_big_text', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-mcp-multi', + }, + tool: mcpTool, + invocation: invocation as unknown as AnyToolInvocation, + startTime: Date.now(), + }; + + const result = await executor.execute({ + call: scheduledCall, + signal: new AbortController().signal, + onUpdateToolCall: vi.fn(), + }); + + // Should NOT have been truncated + expect(fileUtils.saveTruncatedToolOutput).not.toHaveBeenCalled(); + expect(fileUtils.formatTruncatedToolOutput).not.toHaveBeenCalled(); + expect(result.status).toBe(CoreToolCallStatus.Success); + }); + + it('should not truncate MCP tool output when text is below threshold', async () => { + vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10000); + + const messageBus = createMockMessageBus(); + const mcpTool = new DiscoveredMCPTool( + {} as CallableTool, + 'my-server', + 'get_big_text', + 'A test MCP tool', + {}, + messageBus, + ); + const invocation = mcpTool.build({}); + + vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({ + llmContent: [{ text: 'short' }], + returnDisplay: 'short', + }); + + const scheduledCall: ScheduledToolCall = { + status: CoreToolCallStatus.Scheduled, + request: { + callId: 'call-mcp-short', + name: 'get_big_text', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-mcp-short', + }, + tool: mcpTool, + invocation: invocation as unknown as AnyToolInvocation, + startTime: Date.now(), + }; + + const result = await executor.execute({ + call: scheduledCall, + signal: new AbortController().signal, + onUpdateToolCall: vi.fn(), + }); + + expect(fileUtils.saveTruncatedToolOutput).not.toHaveBeenCalled(); + expect(result.status).toBe(CoreToolCallStatus.Success); + }); + it('should report PID updates for shell tools', async () => { // 1. Setup ShellToolInvocation const messageBus = createMockMessageBus(); diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index 7903266fe1..d37c49624c 100644 --- a/packages/core/src/scheduler/tool-executor.ts +++ b/packages/core/src/scheduler/tool-executor.ts @@ -18,6 +18,7 @@ import { runInDevTraceSpan, } from '../index.js'; import { SHELL_TOOL_NAME } from '../tools/tool-names.js'; +import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; import { ShellToolInvocation } from '../tools/shell.js'; import { executeToolWithHooks } from '../core/coreToolHookTriggers.js'; import { @@ -253,6 +254,45 @@ export class ToolExecutor { }), ); } + } else if ( + Array.isArray(content) && + content.length === 1 && + 'tool' in call && + call.tool instanceof DiscoveredMCPTool + ) { + const firstPart = content[0]; + if (typeof firstPart === 'object' && typeof firstPart.text === 'string') { + const textContent = firstPart.text; + const threshold = this.config.getTruncateToolOutputThreshold(); + + if (threshold > 0 && textContent.length > threshold) { + const originalContentLength = textContent.length; + const { outputFile: savedPath } = await saveTruncatedToolOutput( + textContent, + toolName, + callId, + this.config.storage.getProjectTempDir(), + this.config.getSessionId(), + ); + outputFile = savedPath; + const truncatedText = formatTruncatedToolOutput( + textContent, + outputFile, + threshold, + ); + content[0] = { ...firstPart, text: truncatedText }; + + logToolOutputTruncated( + this.config, + new ToolOutputTruncatedEvent(call.request.prompt_id, { + toolName, + originalContentLength, + truncatedContentLength: truncatedText.length, + threshold, + }), + ); + } + } } const response = convertToFunctionResponse(