mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-25 12:34:38 -07:00
feat(core): truncate large MCP tool output (#19365)
This commit is contained in:
@@ -16,6 +16,8 @@ import { MockTool } from '../test-utils/mock-tool.js';
|
|||||||
import type { ScheduledToolCall } from './types.js';
|
import type { ScheduledToolCall } from './types.js';
|
||||||
import { CoreToolCallStatus } from './types.js';
|
import { CoreToolCallStatus } from './types.js';
|
||||||
import { SHELL_TOOL_NAME } from '../tools/tool-names.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 fileUtils from '../utils/fileUtils.js';
|
||||||
import * as coreToolHookTriggers from '../core/coreToolHookTriggers.js';
|
import * as coreToolHookTriggers from '../core/coreToolHookTriggers.js';
|
||||||
import { ShellToolInvocation } from '../tools/shell.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 () => {
|
it('should report PID updates for shell tools', async () => {
|
||||||
// 1. Setup ShellToolInvocation
|
// 1. Setup ShellToolInvocation
|
||||||
const messageBus = createMockMessageBus();
|
const messageBus = createMockMessageBus();
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ import {
|
|||||||
runInDevTraceSpan,
|
runInDevTraceSpan,
|
||||||
} from '../index.js';
|
} from '../index.js';
|
||||||
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
|
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
|
||||||
|
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
|
||||||
import { ShellToolInvocation } from '../tools/shell.js';
|
import { ShellToolInvocation } from '../tools/shell.js';
|
||||||
import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
|
import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
|
||||||
import {
|
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(
|
const response = convertToFunctionResponse(
|
||||||
|
|||||||
Reference in New Issue
Block a user