From d3ee9de3c3b976ad45fe1a13eb49271d17a32e37 Mon Sep 17 00:00:00 2001 From: anj-s <32556631+anj-s@users.noreply.github.com> Date: Tue, 15 Jul 2025 10:22:31 -0700 Subject: [PATCH] Enable tool summarization only when explicitly set in settings.json (#4140) Co-authored-by: matt korwel --- docs/cli/configuration.md | 20 ++++- packages/cli/src/config/config.ts | 1 + packages/cli/src/config/settings.ts | 7 ++ packages/core/src/config/config.ts | 15 ++++ packages/core/src/tools/shell.test.ts | 85 ++++++++++++++++++++++ packages/core/src/tools/shell.ts | 20 +++-- packages/core/src/utils/summarizer.test.ts | 2 +- packages/core/src/utils/summarizer.ts | 23 +++--- 8 files changed, 153 insertions(+), 20 deletions(-) diff --git a/docs/cli/configuration.md b/docs/cli/configuration.md index f200c5df..e6a9ee72 100644 --- a/docs/cli/configuration.md +++ b/docs/cli/configuration.md @@ -206,6 +206,19 @@ In addition to a project settings file, a project's `.gemini` directory can cont "maxSessionTurns": 10 ``` +- **`summarizeToolOutput`** (object): + - **Description:** Enables or disables the summarization of tool output. You can specify the token budget for the summarization using the `tokenBudget` setting. + - Note: Currently only the `run_shell_command` tool is supported. + - **Default:** `{}` (Disabled by default) + - **Example:** + ```json + "summarizeToolOutput": { + "run_shell_command": { + "tokenBudget": 2000 + } + } + ``` + ### Example `settings.json`: ```json @@ -232,7 +245,12 @@ In addition to a project settings file, a project's `.gemini` directory can cont "usageStatisticsEnabled": true, "hideTips": false, "hideBanner": false, - "maxSessionTurns": 10 + "maxSessionTurns": 10, + "summarizeToolOutput": { + "run_shell_command": { + "tokenBudget": 100 + } + } } ``` diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index a5eb327d..d116bc67 100644 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -363,6 +363,7 @@ export async function loadCliConfig( version: e.config.version, })), noBrowser: !!process.env.NO_BROWSER, + summarizeToolOutput: settings.summarizeToolOutput, ideMode, }); } diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index 50c07cf1..f0258db3 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -46,6 +46,10 @@ export interface CheckpointingSettings { enabled?: boolean; } +export interface SummarizeToolOutputSettings { + tokenBudget?: number; +} + export interface AccessibilitySettings { disableLoadingPhrases?: boolean; } @@ -84,6 +88,9 @@ export interface Settings { // Setting for setting maximum number of user/model/tool turns in a session. maxSessionTurns?: number; + // A map of tool names to their summarization settings. + summarizeToolOutput?: Record; + // Add other settings here. ideMode?: boolean; } diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index fead0a9c..268871ca 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -60,6 +60,10 @@ export interface BugCommandSettings { urlTemplate: string; } +export interface SummarizeToolOutputSettings { + tokenBudget?: number; +} + export interface TelemetrySettings { enabled?: boolean; target?: TelemetryTarget; @@ -144,6 +148,7 @@ export interface ConfigParameters { listExtensions?: boolean; activeExtensions?: ActiveExtension[]; noBrowser?: boolean; + summarizeToolOutput?: Record; ideMode?: boolean; } @@ -191,6 +196,9 @@ export class Config { private readonly _activeExtensions: ActiveExtension[]; flashFallbackHandler?: FlashFallbackHandler; private quotaErrorOccurred: boolean = false; + private readonly summarizeToolOutput: + | Record + | undefined; constructor(params: ConfigParameters) { this.sessionId = params.sessionId; @@ -236,6 +244,7 @@ export class Config { this.listExtensions = params.listExtensions ?? false; this._activeExtensions = params.activeExtensions ?? []; this.noBrowser = params.noBrowser ?? false; + this.summarizeToolOutput = params.summarizeToolOutput; this.ideMode = params.ideMode ?? false; if (params.contextFileName) { @@ -497,6 +506,12 @@ export class Config { return this.noBrowser; } + getSummarizeToolOutputConfig(): + | Record + | undefined { + return this.summarizeToolOutput; + } + getIdeMode(): boolean { return this.ideMode; } diff --git a/packages/core/src/tools/shell.test.ts b/packages/core/src/tools/shell.test.ts index a4d56e22..f358f972 100644 --- a/packages/core/src/tools/shell.test.ts +++ b/packages/core/src/tools/shell.test.ts @@ -410,6 +410,9 @@ describe('ShellTool Bug Reproduction', () => { getDebugMode: () => false, getGeminiClient: () => ({}) as GeminiClient, getTargetDir: () => '.', + getSummarizeToolOutputConfig: () => ({ + [shellTool.name]: {}, + }), } as unknown as Config; shellTool = new ShellTool(config); }); @@ -429,4 +432,86 @@ describe('ShellTool Bug Reproduction', () => { expect(result.llmContent).toBe('summarized output'); expect(summarizeSpy).toHaveBeenCalled(); }); + + it('should not call summarizer if disabled in config', async () => { + config = { + getCoreTools: () => undefined, + getExcludeTools: () => undefined, + getDebugMode: () => false, + getGeminiClient: () => ({}) as GeminiClient, + getTargetDir: () => '.', + getSummarizeToolOutputConfig: () => ({}), + } as unknown as Config; + shellTool = new ShellTool(config); + + const summarizeSpy = vi + .spyOn(summarizer, 'summarizeToolOutput') + .mockResolvedValue('summarized output'); + + const abortSignal = new AbortController().signal; + const result = await shellTool.execute( + { command: 'echo "hello"' }, + abortSignal, + ); + + expect(result.returnDisplay).toBe('hello\n'); + expect(result.llmContent).not.toBe('summarized output'); + expect(summarizeSpy).not.toHaveBeenCalled(); + }); + + it('should pass token budget to summarizer', async () => { + config = { + getCoreTools: () => undefined, + getExcludeTools: () => undefined, + getDebugMode: () => false, + getGeminiClient: () => ({}) as GeminiClient, + getTargetDir: () => '.', + getSummarizeToolOutputConfig: () => ({ + [shellTool.name]: { tokenBudget: 1000 }, + }), + } as unknown as Config; + shellTool = new ShellTool(config); + + const summarizeSpy = vi + .spyOn(summarizer, 'summarizeToolOutput') + .mockResolvedValue('summarized output'); + + const abortSignal = new AbortController().signal; + await shellTool.execute({ command: 'echo "hello"' }, abortSignal); + + expect(summarizeSpy).toHaveBeenCalledWith( + expect.any(String), + expect.any(Object), + expect.any(Object), + 1000, + ); + }); + + it('should use default token budget if not specified', async () => { + config = { + getCoreTools: () => undefined, + getExcludeTools: () => undefined, + getDebugMode: () => false, + getGeminiClient: () => ({}) as GeminiClient, + getTargetDir: () => '.', + getSummarizeToolOutputConfig: () => ({ + [shellTool.name]: {}, + }), + } as unknown as Config; + shellTool = new ShellTool(config); + + const summarizeSpy = vi + .spyOn(summarizer, 'summarizeToolOutput') + .mockResolvedValue('summarized output'); + + const abortSignal = new AbortController().signal; + await shellTool.execute({ command: 'echo "hello"' }, abortSignal); + + expect(summarizeSpy).toHaveBeenCalledWith( + expect.any(String), + expect.any(Object), + expect.any(Object), + undefined, + ); + }); }); diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 7e79c717..3dc3d0a6 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -489,14 +489,22 @@ Process Group PGID: Process group started or \`(none)\``, } } - const summary = await summarizeToolOutput( - llmContent, - this.config.getGeminiClient(), - abortSignal, - ); + const summarizeConfig = this.config.getSummarizeToolOutputConfig(); + if (summarizeConfig && summarizeConfig[this.name]) { + const summary = await summarizeToolOutput( + llmContent, + this.config.getGeminiClient(), + abortSignal, + summarizeConfig[this.name].tokenBudget, + ); + return { + llmContent: summary, + returnDisplay: returnDisplayMessage, + }; + } return { - llmContent: summary, + llmContent, returnDisplay: returnDisplayMessage, }; } diff --git a/packages/core/src/utils/summarizer.test.ts b/packages/core/src/utils/summarizer.test.ts index 85f09cee..c3a521b7 100644 --- a/packages/core/src/utils/summarizer.test.ts +++ b/packages/core/src/utils/summarizer.test.ts @@ -121,7 +121,7 @@ describe('summarizers', () => { await summarizeToolOutput(longText, mockGeminiClient, abortSignal, 1000); - const expectedPrompt = `Summarize the following tool output to be a maximum of 1000 characters. The summary should be concise and capture the main points of the tool output. + const expectedPrompt = `Summarize the following tool output to be a maximum of 1000 tokens. The summary should be concise and capture the main points of the tool output. The summarization should be done based on the content that is provided. Here are the basic rules to follow: 1. If the text is a directory listing or any output that is structural, use the history of the conversation to understand the context. Using this context try to understand what information we need from the tool output and return that as a response. diff --git a/packages/core/src/utils/summarizer.ts b/packages/core/src/utils/summarizer.ts index 57d58a61..3f0fd52b 100644 --- a/packages/core/src/utils/summarizer.ts +++ b/packages/core/src/utils/summarizer.ts @@ -74,12 +74,7 @@ function getResponseText(response: GenerateContentResponse): string | null { return null; } -const toolOutputSummarizerModel = DEFAULT_GEMINI_FLASH_MODEL; -const toolOutputSummarizerConfig: GenerateContentConfig = { - maxOutputTokens: 2000, -}; - -const SUMMARIZE_TOOL_OUTPUT_PROMPT = `Summarize the following tool output to be a maximum of {maxLength} characters. The summary should be concise and capture the main points of the tool output. +const SUMMARIZE_TOOL_OUTPUT_PROMPT = `Summarize the following tool output to be a maximum of {maxOutputTokens} tokens. The summary should be concise and capture the main points of the tool output. The summarization should be done based on the content that is provided. Here are the basic rules to follow: 1. If the text is a directory listing or any output that is structural, use the history of the conversation to understand the context. Using this context try to understand what information we need from the tool output and return that as a response. @@ -104,24 +99,28 @@ export async function summarizeToolOutput( textToSummarize: string, geminiClient: GeminiClient, abortSignal: AbortSignal, - maxLength: number = 2000, + maxOutputTokens: number = 2000, ): Promise { - if (!textToSummarize || textToSummarize.length < maxLength) { + // There is going to be a slight difference here since we are comparing length of string with maxOutputTokens. + // This is meant to be a ballpark estimation of if we need to summarize the tool output. + if (!textToSummarize || textToSummarize.length < maxOutputTokens) { return textToSummarize; } const prompt = SUMMARIZE_TOOL_OUTPUT_PROMPT.replace( - '{maxLength}', - String(maxLength), + '{maxOutputTokens}', + String(maxOutputTokens), ).replace('{textToSummarize}', textToSummarize); const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }]; - + const toolOutputSummarizerConfig: GenerateContentConfig = { + maxOutputTokens, + }; try { const parsedResponse = (await geminiClient.generateContent( contents, toolOutputSummarizerConfig, abortSignal, - toolOutputSummarizerModel, + DEFAULT_GEMINI_FLASH_MODEL, )) as unknown as GenerateContentResponse; return getResponseText(parsedResponse) || textToSummarize; } catch (error) {