Enable tool summarization only when explicitly set in settings.json (#4140)

Co-authored-by: matt korwel <matt.korwel@gmail.com>

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
+    }
+  }
 }
 ```
 

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,
   });
 }

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<string, SummarizeToolOutputSettings>;
+
   // Add other settings here.
   ideMode?: boolean;
 }

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<string, SummarizeToolOutputSettings>;
   ideMode?: boolean;
 }
 
@@ -191,6 +196,9 @@ export class Config {
   private readonly _activeExtensions: ActiveExtension[];
   flashFallbackHandler?: FlashFallbackHandler;
   private quotaErrorOccurred: boolean = false;
+  private readonly summarizeToolOutput:
+    | Record<string, SummarizeToolOutputSettings>
+    | 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<string, SummarizeToolOutputSettings>
+    | undefined {
+    return this.summarizeToolOutput;
+  }
+
   getIdeMode(): boolean {
     return this.ideMode;
   }

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,
+    );
+  });
 });

packages/core/src/tools/shell.ts

@@ -489,15 +489,23 @@ Process Group PGID: Process group started or \`(none)\``,
       }
     }
 
+    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,
+      returnDisplay: returnDisplayMessage,
+    };
+  }
 }

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.

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<string> {
-  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) {