From e9862093a35c34c239b189b57a430eec9ebd9af6 Mon Sep 17 00:00:00 2001 From: Your Name Date: Thu, 19 Mar 2026 16:38:08 +0000 Subject: [PATCH] guard with flag --- packages/cli/src/config/config.ts | 1 + packages/cli/src/config/settingsSchema.ts | 9 ++ packages/core/src/config/config.ts | 7 ++ packages/core/src/scheduler/tool-executor.ts | 104 +++++++++++++++++-- packages/core/src/tools/web-fetch.ts | 35 ++++++- 5 files changed, 145 insertions(+), 11 deletions(-) diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 80c1e19443..940ca1dcb0 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -840,6 +840,7 @@ export async function loadCliConfig( skillsSupport: settings.skills?.enabled ?? true, disabledSkills: settings.skills?.disabled, experimentalJitContext: settings.experimental?.jitContext, + autoDistillation: settings.experimental?.autoDistillation, modelSteering: settings.experimental?.modelSteering, topicUpdateNarration: settings.experimental?.topicUpdateNarration, toolOutputMasking: settings.experimental?.toolOutputMasking, diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 8a107c4d47..6406a0af4f 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1870,6 +1870,15 @@ const SETTINGS_SCHEMA = { description: 'Enable local and remote subagents.', showInDialog: false, }, + autoDistillation: { + type: 'boolean', + label: 'Auto Distillation', + category: 'Experimental', + requiresRestart: true, + default: false, + description: 'Enable automatic distillation for large tool outputs.', + showInDialog: false, + }, extensionManagement: { type: 'boolean', label: 'Extension Management', diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index d89e694103..8d07435883 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -630,6 +630,7 @@ export interface ConfigParameters { disabledSkills?: string[]; adminSkillsEnabled?: boolean; experimentalJitContext?: boolean; + autoDistillation?: boolean; topicUpdateNarration?: boolean; toolOutputMasking?: Partial; disableLLMCorrection?: boolean; @@ -855,6 +856,7 @@ export class Config implements McpContext, AgentLoopContext { private readonly adminSkillsEnabled: boolean; private readonly experimentalJitContext: boolean; + private readonly autoDistillation: boolean; private readonly topicUpdateNarration: boolean; private readonly disableLLMCorrection: boolean; private readonly planEnabled: boolean; @@ -1016,6 +1018,7 @@ export class Config implements McpContext, AgentLoopContext { ); this.experimentalJitContext = params.experimentalJitContext ?? true; + this.autoDistillation = params.autoDistillation ?? false; this.topicUpdateNarration = params.topicUpdateNarration ?? false; this.modelSteering = params.modelSteering ?? false; this.injectionService = new InjectionService(() => @@ -2166,6 +2169,10 @@ export class Config implements McpContext, AgentLoopContext { return this.experimentalJitContext; } + isAutoDistillationEnabled(): boolean { + return this.autoDistillation; + } + isTopicUpdateNarrationEnabled(): boolean { return this.topicUpdateNarration; } diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index 97ad740bca..fcfb9776ab 100644 --- a/packages/core/src/scheduler/tool-executor.ts +++ b/packages/core/src/scheduler/tool-executor.ts @@ -7,6 +7,8 @@ import { ToolErrorType, runInDevTraceSpan, + ToolOutputTruncatedEvent, + logToolOutputTruncated, type ToolCallRequestInfo, type ToolCallResponseInfo, type ToolResult, @@ -18,8 +20,11 @@ import { isAbortError } from '../utils/errors.js'; import { SHELL_TOOL_NAME } from '../tools/tool-names.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; import { ToolOutputDistillationService } from '../services/toolDistillationService.js'; -import { ShellToolInvocation } from '../tools/shell.js'; import { executeToolWithHooks } from '../core/coreToolHookTriggers.js'; +import { + saveTruncatedToolOutput, + formatTruncatedToolOutput, +} from '../utils/fileUtils.js'; import { convertToFunctionResponse } from '../utils/generateContentResponseUtilities.js'; import { @@ -177,12 +182,97 @@ export class ToolExecutor { call: ToolCall, content: PartListUnion, ): Promise<{ truncatedContent: PartListUnion; outputFile?: string }> { - const distiller = new ToolOutputDistillationService( - this.config, - this.context.geminiClient, - this.context.promptId, - ); - return distiller.distill(call.request.name, call.request.callId, content); + if (this.config.isAutoDistillationEnabled()) { + const distiller = new ToolOutputDistillationService( + this.config, + this.context.geminiClient, + this.context.promptId, + ); + return distiller.distill(call.request.name, call.request.callId, content); + } + + const toolName = call.request.name; + const callId = call.request.callId; + let outputFile: string | undefined; + + if (typeof content === 'string' && toolName === SHELL_TOOL_NAME) { + const threshold = this.config.getTruncateToolOutputThreshold(); + + 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, + ); + outputFile = savedPath; + const truncatedContent = formatTruncatedToolOutput( + content, + outputFile, + threshold, + ); + + logToolOutputTruncated( + this.config, + new ToolOutputTruncatedEvent(call.request.prompt_id, { + toolName, + originalContentLength, + truncatedContentLength: truncatedContent.length, + threshold, + }), + ); + + return { truncatedContent, outputFile }; + } + } 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.context.promptId, + ); + outputFile = savedPath; + const truncatedText = formatTruncatedToolOutput( + textContent, + outputFile, + threshold, + ); + + const truncatedContent: Part[] = [ + { ...firstPart, text: truncatedText }, + ]; + + logToolOutputTruncated( + this.config, + new ToolOutputTruncatedEvent(call.request.prompt_id, { + toolName, + originalContentLength, + truncatedContentLength: truncatedText.length, + threshold, + }), + ); + + return { truncatedContent, outputFile }; + } + } + } + + return { truncatedContent: content, outputFile }; } private async createCancelledResult( diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts index 3c34387131..92ea78645c 100644 --- a/packages/core/src/tools/web-fetch.ts +++ b/packages/core/src/tools/web-fetch.ts @@ -22,6 +22,7 @@ import { ApprovalMode } from '../policy/types.js'; import { getResponseText } from '../utils/partUtils.js'; import { fetchWithTimeout, isPrivateIp } from '../utils/fetch.js'; import { convert } from 'html-to-text'; +import { truncateString } from '../utils/textUtils.js'; import { logWebFetchFallbackAttempt, WebFetchFallbackAttemptEvent, @@ -41,6 +42,8 @@ import type { AgentLoopContext } from '../config/agent-loop-context.js'; const URL_FETCH_TIMEOUT_MS = 10000; const MAX_EXPERIMENTAL_FETCH_SIZE = 10 * 1024 * 1024; // 10MB +const MAX_CONTENT_LENGTH = 100000; +const TRUNCATION_WARNING = '\n\n... [Content truncated due to size limit] ...'; const USER_AGENT = 'Mozilla/5.0 (compatible; Google-Gemini-CLI/1.0; +https://github.com/google-gemini/gemini-cli)'; @@ -329,6 +332,10 @@ class WebFetchToolInvocation extends BaseToolInvocation< textContent = rawContent; } + if (!this.context.config.isAutoDistillationEnabled()) { + return truncateString(textContent, MAX_CONTENT_LENGTH, TRUNCATION_WARNING); + } + return textContent; } @@ -634,7 +641,14 @@ ${aggregatedContent} ); if (status >= 400) { - const rawResponseText = bodyBuffer.toString('utf8'); + let rawResponseText = bodyBuffer.toString('utf8'); + if (!this.context.config.isAutoDistillationEnabled()) { + rawResponseText = truncateString( + rawResponseText, + 10000, + '\n\n... [Error response truncated] ...', + ); + } const headers: Record = {}; response.headers.forEach((value, key) => { headers[key] = value; @@ -657,7 +671,10 @@ Response: ${rawResponseText}`; lowContentType.includes('text/plain') || lowContentType.includes('application/json') ) { - const text = bodyBuffer.toString('utf8'); + let text = bodyBuffer.toString('utf8'); + if (!this.context.config.isAutoDistillationEnabled()) { + text = truncateString(text, MAX_CONTENT_LENGTH, TRUNCATION_WARNING); + } return { llmContent: text, returnDisplay: `Fetched ${contentType} content from ${url}`, @@ -666,12 +683,19 @@ Response: ${rawResponseText}`; if (lowContentType.includes('text/html')) { const html = bodyBuffer.toString('utf8'); - const textContent = convert(html, { + let textContent = convert(html, { wordwrap: false, selectors: [ { selector: 'a', options: { ignoreHref: false, baseUrl: url } }, ], }); + if (!this.context.config.isAutoDistillationEnabled()) { + textContent = truncateString( + textContent, + MAX_CONTENT_LENGTH, + TRUNCATION_WARNING, + ); + } return { llmContent: textContent, returnDisplay: `Fetched and converted HTML content from ${url}`, @@ -696,7 +720,10 @@ Response: ${rawResponseText}`; } // Fallback for unknown types - try as text - const text = bodyBuffer.toString('utf8'); + let text = bodyBuffer.toString('utf8'); + if (!this.context.config.isAutoDistillationEnabled()) { + text = truncateString(text, MAX_CONTENT_LENGTH, TRUNCATION_WARNING); + } return { llmContent: text, returnDisplay: `Fetched ${contentType || 'unknown'} content from ${url}`,