feat(tools): Centralize shell tool summarization (#4009)

This commit is contained in:
N. Taylor Mullen 2025-07-12 21:09:12 -07:00 committed by GitHub
parent 09a3b7d5e1
commit 44ef0408f3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 51 additions and 58 deletions

View File

@ -656,39 +656,15 @@ export class CoreToolScheduler {
return;
}
let resultForDisplay: ToolResult = toolResult;
let summary: string | undefined;
if (scheduledCall.tool.summarizer) {
try {
const toolSignal = new AbortController();
summary = await scheduledCall.tool.summarizer(
toolResult,
this.config.getGeminiClient(),
toolSignal.signal,
);
if (toolSignal.signal.aborted) {
console.debug('aborted summarizing tool result');
return;
}
if (scheduledCall.tool?.shouldSummarizeDisplay) {
resultForDisplay = {
...toolResult,
returnDisplay: summary,
};
}
} catch (e) {
console.error('Error summarizing tool result:', e);
}
}
const response = convertToFunctionResponse(
toolName,
callId,
summary ? [summary] : toolResult.llmContent,
toolResult.llmContent,
);
const successResponse: ToolCallResponseInfo = {
callId,
responseParts: response,
resultDisplay: resultForDisplay.returnDisplay,
resultDisplay: toolResult.returnDisplay,
error: undefined,
};

View File

@ -68,17 +68,9 @@ export async function executeToolCall(
// No live output callback for non-interactive mode
);
const tool_output = tool.summarizer
? await tool.summarizer(
toolResult,
config.getGeminiClient(),
effectiveAbortSignal,
)
: toolResult.llmContent;
const tool_output = toolResult.llmContent;
const tool_display = tool.shouldSummarizeDisplay
? (tool_output as string)
: toolResult.returnDisplay;
const tool_display = toolResult.returnDisplay;
const durationMs = Date.now() - startTime;
logToolCall(config, {

View File

@ -4,9 +4,11 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { expect, describe, it } from 'vitest';
import { expect, describe, it, vi, beforeEach } from 'vitest';
import { ShellTool } from './shell.js';
import { Config } from '../config/config.js';
import * as summarizer from '../utils/summarizer.js';
import { GeminiClient } from '../core/client.js';
describe('ShellTool', () => {
it('should allow a command if no restrictions are provided', async () => {
@ -396,3 +398,35 @@ describe('ShellTool', () => {
);
});
});
describe('ShellTool Bug Reproduction', () => {
let shellTool: ShellTool;
let config: Config;
beforeEach(() => {
config = {
getCoreTools: () => undefined,
getExcludeTools: () => undefined,
getDebugMode: () => false,
getGeminiClient: () => ({}) as GeminiClient,
getTargetDir: () => '.',
} as unknown as Config;
shellTool = new ShellTool(config);
});
it('should not let the summarizer override the return display', async () => {
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).toBe('summarized output');
expect(summarizeSpy).toHaveBeenCalled();
});
});

View File

@ -27,7 +27,7 @@ export interface ShellToolParams {
directory?: string;
}
import { spawn } from 'child_process';
import { llmSummarizer } from '../utils/summarizer.js';
import { summarizeToolOutput } from '../utils/summarizer.js';
const OUTPUT_UPDATE_INTERVAL_MS = 1000;
@ -74,8 +74,6 @@ Process Group PGID: Process group started or \`(none)\``,
},
false, // output is not markdown
true, // output can be updated
llmSummarizer,
true, // should summarize display output
);
}
@ -490,6 +488,16 @@ Process Group PGID: Process group started or \`(none)\``,
// returnDisplayMessage will remain empty, which is fine.
}
}
return { llmContent, returnDisplay: returnDisplayMessage };
const summary = await summarizeToolOutput(
llmContent,
this.config.getGeminiClient(),
abortSignal,
);
return {
llmContent: summary,
returnDisplay: returnDisplayMessage,
};
}
}

View File

@ -11,7 +11,6 @@ import { spawn } from 'node:child_process';
import { StringDecoder } from 'node:string_decoder';
import { discoverMcpTools } from './mcp-client.js';
import { DiscoveredMCPTool } from './mcp-tool.js';
import { defaultSummarizer } from '../utils/summarizer.js';
import { parse } from 'shell-quote';
type ToolParams = Record<string, unknown>;
@ -48,7 +47,6 @@ Signal: Signal number or \`(none)\` if no signal was received.
parameterSchema,
false, // isOutputMarkdown
false, // canUpdateOutput
defaultSummarizer,
);
}

View File

@ -5,7 +5,6 @@
*/
import { FunctionDeclaration, PartListUnion, Schema } from '@google/genai';
import { Summarizer, defaultSummarizer } from '../utils/summarizer.js';
/**
* Interface representing the base Tool functionality
@ -44,16 +43,6 @@ export interface Tool<
*/
canUpdateOutput: boolean;
/**
* A function that summarizes the result of the tool execution.
*/
summarizer?: Summarizer;
/**
* Whether the tool's display output should be summarized
*/
shouldSummarizeDisplay?: boolean;
/**
* Validates the parameters for the tool
* Should be called from both `shouldConfirmExecute` and `execute`
@ -109,8 +98,6 @@ export abstract class BaseTool<
* @param isOutputMarkdown Whether the tool's output should be rendered as markdown
* @param canUpdateOutput Whether the tool supports live (streaming) output
* @param parameterSchema JSON Schema defining the parameters
* @param summarizer Function to summarize the tool's output
* @param shouldSummarizeDisplay Whether the tool's display output should be summarized
*/
constructor(
readonly name: string,
@ -119,8 +106,6 @@ export abstract class BaseTool<
readonly parameterSchema: Schema,
readonly isOutputMarkdown: boolean = true,
readonly canUpdateOutput: boolean = false,
readonly summarizer: Summarizer = defaultSummarizer,
readonly shouldSummarizeDisplay: boolean = false,
) {}
/**