Skip to content

Commit 44ef040

Browse files
feat(tools): Centralize shell tool summarization (google-gemini#4009)
1 parent 09a3b7d commit 44ef040

File tree

6 files changed

+51
-58
lines changed

6 files changed

+51
-58
lines changed

packages/core/src/core/coreToolScheduler.ts

Lines changed: 2 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -656,39 +656,15 @@ export class CoreToolScheduler {
656656
return;
657657
}
658658

659-
let resultForDisplay: ToolResult = toolResult;
660-
let summary: string | undefined;
661-
if (scheduledCall.tool.summarizer) {
662-
try {
663-
const toolSignal = new AbortController();
664-
summary = await scheduledCall.tool.summarizer(
665-
toolResult,
666-
this.config.getGeminiClient(),
667-
toolSignal.signal,
668-
);
669-
if (toolSignal.signal.aborted) {
670-
console.debug('aborted summarizing tool result');
671-
return;
672-
}
673-
if (scheduledCall.tool?.shouldSummarizeDisplay) {
674-
resultForDisplay = {
675-
...toolResult,
676-
returnDisplay: summary,
677-
};
678-
}
679-
} catch (e) {
680-
console.error('Error summarizing tool result:', e);
681-
}
682-
}
683659
const response = convertToFunctionResponse(
684660
toolName,
685661
callId,
686-
summary ? [summary] : toolResult.llmContent,
662+
toolResult.llmContent,
687663
);
688664
const successResponse: ToolCallResponseInfo = {
689665
callId,
690666
responseParts: response,
691-
resultDisplay: resultForDisplay.returnDisplay,
667+
resultDisplay: toolResult.returnDisplay,
692668
error: undefined,
693669
};
694670

packages/core/src/core/nonInteractiveToolExecutor.ts

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -68,17 +68,9 @@ export async function executeToolCall(
6868
// No live output callback for non-interactive mode
6969
);
7070

71-
const tool_output = tool.summarizer
72-
? await tool.summarizer(
73-
toolResult,
74-
config.getGeminiClient(),
75-
effectiveAbortSignal,
76-
)
77-
: toolResult.llmContent;
71+
const tool_output = toolResult.llmContent;
7872

79-
const tool_display = tool.shouldSummarizeDisplay
80-
? (tool_output as string)
81-
: toolResult.returnDisplay;
73+
const tool_display = toolResult.returnDisplay;
8274

8375
const durationMs = Date.now() - startTime;
8476
logToolCall(config, {

packages/core/src/tools/shell.test.ts

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7-
import { expect, describe, it } from 'vitest';
7+
import { expect, describe, it, vi, beforeEach } from 'vitest';
88
import { ShellTool } from './shell.js';
99
import { Config } from '../config/config.js';
10+
import * as summarizer from '../utils/summarizer.js';
11+
import { GeminiClient } from '../core/client.js';
1012

1113
describe('ShellTool', () => {
1214
it('should allow a command if no restrictions are provided', async () => {
@@ -396,3 +398,35 @@ describe('ShellTool', () => {
396398
);
397399
});
398400
});
401+
402+
describe('ShellTool Bug Reproduction', () => {
403+
let shellTool: ShellTool;
404+
let config: Config;
405+
406+
beforeEach(() => {
407+
config = {
408+
getCoreTools: () => undefined,
409+
getExcludeTools: () => undefined,
410+
getDebugMode: () => false,
411+
getGeminiClient: () => ({}) as GeminiClient,
412+
getTargetDir: () => '.',
413+
} as unknown as Config;
414+
shellTool = new ShellTool(config);
415+
});
416+
417+
it('should not let the summarizer override the return display', async () => {
418+
const summarizeSpy = vi
419+
.spyOn(summarizer, 'summarizeToolOutput')
420+
.mockResolvedValue('summarized output');
421+
422+
const abortSignal = new AbortController().signal;
423+
const result = await shellTool.execute(
424+
{ command: 'echo "hello"' },
425+
abortSignal,
426+
);
427+
428+
expect(result.returnDisplay).toBe('hello\n');
429+
expect(result.llmContent).toBe('summarized output');
430+
expect(summarizeSpy).toHaveBeenCalled();
431+
});
432+
});

packages/core/src/tools/shell.ts

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ export interface ShellToolParams {
2727
directory?: string;
2828
}
2929
import { spawn } from 'child_process';
30-
import { llmSummarizer } from '../utils/summarizer.js';
30+
import { summarizeToolOutput } from '../utils/summarizer.js';
3131

3232
const OUTPUT_UPDATE_INTERVAL_MS = 1000;
3333

@@ -74,8 +74,6 @@ Process Group PGID: Process group started or \`(none)\``,
7474
},
7575
false, // output is not markdown
7676
true, // output can be updated
77-
llmSummarizer,
78-
true, // should summarize display output
7977
);
8078
}
8179

@@ -490,6 +488,16 @@ Process Group PGID: Process group started or \`(none)\``,
490488
// returnDisplayMessage will remain empty, which is fine.
491489
}
492490
}
493-
return { llmContent, returnDisplay: returnDisplayMessage };
491+
492+
const summary = await summarizeToolOutput(
493+
llmContent,
494+
this.config.getGeminiClient(),
495+
abortSignal,
496+
);
497+
498+
return {
499+
llmContent: summary,
500+
returnDisplay: returnDisplayMessage,
501+
};
494502
}
495503
}

packages/core/src/tools/tool-registry.ts

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import { spawn } from 'node:child_process';
1111
import { StringDecoder } from 'node:string_decoder';
1212
import { discoverMcpTools } from './mcp-client.js';
1313
import { DiscoveredMCPTool } from './mcp-tool.js';
14-
import { defaultSummarizer } from '../utils/summarizer.js';
1514
import { parse } from 'shell-quote';
1615

1716
type ToolParams = Record<string, unknown>;
@@ -48,7 +47,6 @@ Signal: Signal number or \`(none)\` if no signal was received.
4847
parameterSchema,
4948
false, // isOutputMarkdown
5049
false, // canUpdateOutput
51-
defaultSummarizer,
5250
);
5351
}
5452

packages/core/src/tools/tools.ts

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
*/
66

77
import { FunctionDeclaration, PartListUnion, Schema } from '@google/genai';
8-
import { Summarizer, defaultSummarizer } from '../utils/summarizer.js';
98

109
/**
1110
* Interface representing the base Tool functionality
@@ -44,16 +43,6 @@ export interface Tool<
4443
*/
4544
canUpdateOutput: boolean;
4645

47-
/**
48-
* A function that summarizes the result of the tool execution.
49-
*/
50-
summarizer?: Summarizer;
51-
52-
/**
53-
* Whether the tool's display output should be summarized
54-
*/
55-
shouldSummarizeDisplay?: boolean;
56-
5746
/**
5847
* Validates the parameters for the tool
5948
* Should be called from both `shouldConfirmExecute` and `execute`
@@ -109,8 +98,6 @@ export abstract class BaseTool<
10998
* @param isOutputMarkdown Whether the tool's output should be rendered as markdown
11099
* @param canUpdateOutput Whether the tool supports live (streaming) output
111100
* @param parameterSchema JSON Schema defining the parameters
112-
* @param summarizer Function to summarize the tool's output
113-
* @param shouldSummarizeDisplay Whether the tool's display output should be summarized
114101
*/
115102
constructor(
116103
readonly name: string,
@@ -119,8 +106,6 @@ export abstract class BaseTool<
119106
readonly parameterSchema: Schema,
120107
readonly isOutputMarkdown: boolean = true,
121108
readonly canUpdateOutput: boolean = false,
122-
readonly summarizer: Summarizer = defaultSummarizer,
123-
readonly shouldSummarizeDisplay: boolean = false,
124109
) {}
125110

126111
/**

0 commit comments

Comments
 (0)