mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-28 05:55:17 -07:00
refactor(core): Extract and integrate ToolExecutor (#15900)
This commit is contained in:
Abhi
@@ -0,0 +1,299 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { ToolExecutor } from './tool-executor.js';
|
||||
import type { Config } from '../index.js';
|
||||
import type { ToolResult } from '../tools/tools.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import { MockTool } from '../test-utils/mock-tool.js';
|
||||
import type { ScheduledToolCall } from './types.js';
|
||||
import type { AnyToolInvocation } from '../index.js';
|
||||
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
|
||||
import * as fileUtils from '../utils/fileUtils.js';
|
||||
import * as coreToolHookTriggers from '../core/coreToolHookTriggers.js';
|
||||
import { ShellToolInvocation } from '../tools/shell.js';
|
||||
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
|
||||
|
||||
// Mock file utils
|
||||
vi.mock('../utils/fileUtils.js', () => ({
|
||||
saveTruncatedContent: vi.fn(),
|
||||
}));
|
||||
|
||||
// Mock executeToolWithHooks
|
||||
vi.mock('../core/coreToolHookTriggers.js', () => ({
|
||||
executeToolWithHooks: vi.fn(),
|
||||
}));
|
||||
|
||||
describe('ToolExecutor', () => {
|
||||
let config: Config;
|
||||
let executor: ToolExecutor;
|
||||
|
||||
beforeEach(() => {
|
||||
// Use the standard fake config factory
|
||||
config = makeFakeConfig();
|
||||
executor = new ToolExecutor(config);
|
||||
|
||||
// Reset mocks
|
||||
vi.resetAllMocks();
|
||||
|
||||
// Default mock implementation for saveTruncatedContent
|
||||
vi.mocked(fileUtils.saveTruncatedContent).mockImplementation(
|
||||
async (_content, _callId, _tempDir, _threshold, _lines) => ({
|
||||
content: 'TruncatedContent...',
|
||||
outputFile: '/tmp/truncated_output.txt',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('should execute a tool successfully', async () => {
|
||||
const mockTool = new MockTool({
|
||||
name: 'testTool',
|
||||
execute: async () => ({
|
||||
llmContent: 'Tool output',
|
||||
returnDisplay: 'Tool output',
|
||||
}),
|
||||
});
|
||||
const invocation = mockTool.build({});
|
||||
|
||||
// Mock executeToolWithHooks to return success
|
||||
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
|
||||
llmContent: 'Tool output',
|
||||
returnDisplay: 'Tool output',
|
||||
} as ToolResult);
|
||||
|
||||
const scheduledCall: ScheduledToolCall = {
|
||||
status: 'scheduled',
|
||||
request: {
|
||||
callId: 'call-1',
|
||||
name: 'testTool',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'prompt-1',
|
||||
},
|
||||
tool: mockTool,
|
||||
invocation: invocation as unknown as AnyToolInvocation,
|
||||
startTime: Date.now(),
|
||||
};
|
||||
|
||||
const onUpdateToolCall = vi.fn();
|
||||
const result = await executor.execute({
|
||||
call: scheduledCall,
|
||||
signal: new AbortController().signal,
|
||||
onUpdateToolCall,
|
||||
});
|
||||
|
||||
expect(result.status).toBe('success');
|
||||
if (result.status === 'success') {
|
||||
const response = result.response.responseParts[0]?.functionResponse
|
||||
?.response as Record<string, unknown>;
|
||||
expect(response).toEqual({ output: 'Tool output' });
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle execution errors', async () => {
|
||||
const mockTool = new MockTool({
|
||||
name: 'failTool',
|
||||
});
|
||||
const invocation = mockTool.build({});
|
||||
|
||||
// Mock executeToolWithHooks to throw
|
||||
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockRejectedValue(
|
||||
new Error('Tool Failed'),
|
||||
);
|
||||
|
||||
const scheduledCall: ScheduledToolCall = {
|
||||
status: 'scheduled',
|
||||
request: {
|
||||
callId: 'call-2',
|
||||
name: 'failTool',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'prompt-2',
|
||||
},
|
||||
tool: mockTool,
|
||||
invocation: invocation as unknown as AnyToolInvocation,
|
||||
startTime: Date.now(),
|
||||
};
|
||||
|
||||
const result = await executor.execute({
|
||||
call: scheduledCall,
|
||||
signal: new AbortController().signal,
|
||||
onUpdateToolCall: vi.fn(),
|
||||
});
|
||||
|
||||
expect(result.status).toBe('error');
|
||||
if (result.status === 'error') {
|
||||
expect(result.response.error?.message).toBe('Tool Failed');
|
||||
}
|
||||
});
|
||||
|
||||
it('should return cancelled result when signal is aborted', async () => {
|
||||
const mockTool = new MockTool({
|
||||
name: 'slowTool',
|
||||
});
|
||||
const invocation = mockTool.build({});
|
||||
|
||||
// Mock executeToolWithHooks to simulate slow execution or cancellation check
|
||||
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockImplementation(
|
||||
async () => {
|
||||
await new Promise((r) => setTimeout(r, 100));
|
||||
return { llmContent: 'Done', returnDisplay: 'Done' };
|
||||
},
|
||||
);
|
||||
|
||||
const scheduledCall: ScheduledToolCall = {
|
||||
status: 'scheduled',
|
||||
request: {
|
||||
callId: 'call-3',
|
||||
name: 'slowTool',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'prompt-3',
|
||||
},
|
||||
tool: mockTool,
|
||||
invocation: invocation as unknown as AnyToolInvocation,
|
||||
startTime: Date.now(),
|
||||
};
|
||||
|
||||
const controller = new AbortController();
|
||||
const promise = executor.execute({
|
||||
call: scheduledCall,
|
||||
signal: controller.signal,
|
||||
onUpdateToolCall: vi.fn(),
|
||||
});
|
||||
|
||||
controller.abort();
|
||||
const result = await promise;
|
||||
|
||||
expect(result.status).toBe('cancelled');
|
||||
});
|
||||
|
||||
it('should truncate large shell output', async () => {
|
||||
// 1. Setup Config for Truncation
|
||||
vi.spyOn(config, 'getEnableToolOutputTruncation').mockReturnValue(true);
|
||||
vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10);
|
||||
vi.spyOn(config, 'getTruncateToolOutputLines').mockReturnValue(5);
|
||||
|
||||
const mockTool = new MockTool({ name: SHELL_TOOL_NAME });
|
||||
const invocation = mockTool.build({});
|
||||
const longOutput = 'This is a very long output that should be truncated.';
|
||||
|
||||
// 2. Mock execution returning long content
|
||||
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
|
||||
llmContent: longOutput,
|
||||
returnDisplay: longOutput,
|
||||
});
|
||||
|
||||
const scheduledCall: ScheduledToolCall = {
|
||||
status: 'scheduled',
|
||||
request: {
|
||||
callId: 'call-trunc',
|
||||
name: SHELL_TOOL_NAME,
|
||||
args: { command: 'echo long' },
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'prompt-trunc',
|
||||
},
|
||||
tool: mockTool,
|
||||
invocation: invocation as unknown as AnyToolInvocation,
|
||||
startTime: Date.now(),
|
||||
};
|
||||
|
||||
// 3. Execute
|
||||
const result = await executor.execute({
|
||||
call: scheduledCall,
|
||||
signal: new AbortController().signal,
|
||||
onUpdateToolCall: vi.fn(),
|
||||
});
|
||||
|
||||
// 4. Verify Truncation Logic
|
||||
expect(fileUtils.saveTruncatedContent).toHaveBeenCalledWith(
|
||||
longOutput,
|
||||
'call-trunc',
|
||||
expect.any(String), // temp dir
|
||||
10, // threshold
|
||||
5, // lines
|
||||
);
|
||||
|
||||
expect(result.status).toBe('success');
|
||||
if (result.status === 'success') {
|
||||
const response = result.response.responseParts[0]?.functionResponse
|
||||
?.response as Record<string, unknown>;
|
||||
// The content should be the *truncated* version returned by the mock saveTruncatedContent
|
||||
expect(response).toEqual({ output: 'TruncatedContent...' });
|
||||
expect(result.response.outputFile).toBe('/tmp/truncated_output.txt');
|
||||
}
|
||||
});
|
||||
|
||||
it('should report PID updates for shell tools', async () => {
|
||||
// 1. Setup ShellToolInvocation
|
||||
const messageBus = createMockMessageBus();
|
||||
const shellInvocation = new ShellToolInvocation(
|
||||
config,
|
||||
{ command: 'sleep 10' },
|
||||
messageBus,
|
||||
);
|
||||
// We need a dummy tool that matches the invocation just for structure
|
||||
const mockTool = new MockTool({ name: SHELL_TOOL_NAME });
|
||||
|
||||
// 2. Mock executeToolWithHooks to trigger the PID callback
|
||||
const testPid = 12345;
|
||||
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockImplementation(
|
||||
async (
|
||||
_inv,
|
||||
_name,
|
||||
_sig,
|
||||
_bus,
|
||||
_hooks,
|
||||
_tool,
|
||||
_liveCb,
|
||||
_shellCfg,
|
||||
setPidCallback,
|
||||
) => {
|
||||
// Simulate the shell tool reporting a PID
|
||||
if (setPidCallback) {
|
||||
setPidCallback(testPid);
|
||||
}
|
||||
return { llmContent: 'done', returnDisplay: 'done' };
|
||||
},
|
||||
);
|
||||
|
||||
const scheduledCall: ScheduledToolCall = {
|
||||
status: 'scheduled',
|
||||
request: {
|
||||
callId: 'call-pid',
|
||||
name: SHELL_TOOL_NAME,
|
||||
args: { command: 'sleep 10' },
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'prompt-pid',
|
||||
},
|
||||
tool: mockTool,
|
||||
invocation: shellInvocation,
|
||||
startTime: Date.now(),
|
||||
};
|
||||
|
||||
const onUpdateToolCall = vi.fn();
|
||||
|
||||
// 3. Execute
|
||||
await executor.execute({
|
||||
call: scheduledCall,
|
||||
signal: new AbortController().signal,
|
||||
onUpdateToolCall,
|
||||
});
|
||||
|
||||
// 4. Verify PID was reported
|
||||
expect(onUpdateToolCall).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
status: 'executing',
|
||||
pid: testPid,
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,310 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type {
|
||||
ToolCallRequestInfo,
|
||||
ToolCallResponseInfo,
|
||||
ToolResult,
|
||||
Config,
|
||||
AnsiOutput,
|
||||
} from '../index.js';
|
||||
import {
|
||||
ToolErrorType,
|
||||
ToolOutputTruncatedEvent,
|
||||
logToolOutputTruncated,
|
||||
runInDevTraceSpan,
|
||||
} from '../index.js';
|
||||
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
|
||||
import { ShellToolInvocation } from '../tools/shell.js';
|
||||
import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
|
||||
import { saveTruncatedContent } from '../utils/fileUtils.js';
|
||||
import { convertToFunctionResponse } from '../utils/generateContentResponseUtilities.js';
|
||||
import type {
|
||||
CompletedToolCall,
|
||||
ToolCall,
|
||||
ExecutingToolCall,
|
||||
ErroredToolCall,
|
||||
SuccessfulToolCall,
|
||||
CancelledToolCall,
|
||||
} from './types.js';
|
||||
|
||||
export interface ToolExecutionContext {
|
||||
call: ToolCall;
|
||||
signal: AbortSignal;
|
||||
outputUpdateHandler?: (callId: string, output: string | AnsiOutput) => void;
|
||||
onUpdateToolCall: (updatedCall: ToolCall) => void;
|
||||
}
|
||||
|
||||
export class ToolExecutor {
|
||||
constructor(private readonly config: Config) {}
|
||||
|
||||
async execute(context: ToolExecutionContext): Promise<CompletedToolCall> {
|
||||
const { call, signal, outputUpdateHandler, onUpdateToolCall } = context;
|
||||
const { request } = call;
|
||||
const toolName = request.name;
|
||||
const callId = request.callId;
|
||||
|
||||
if (!('tool' in call) || !call.tool || !('invocation' in call)) {
|
||||
throw new Error(
|
||||
`Cannot execute tool call ${callId}: Tool or Invocation missing.`,
|
||||
);
|
||||
}
|
||||
const { tool, invocation } = call;
|
||||
|
||||
// Setup live output handling
|
||||
const liveOutputCallback =
|
||||
tool.canUpdateOutput && outputUpdateHandler
|
||||
? (outputChunk: string | AnsiOutput) => {
|
||||
outputUpdateHandler(callId, outputChunk);
|
||||
}
|
||||
: undefined;
|
||||
|
||||
const shellExecutionConfig = this.config.getShellExecutionConfig();
|
||||
const hooksEnabled = this.config.getEnableHooks();
|
||||
const messageBus = this.config.getMessageBus();
|
||||
|
||||
return runInDevTraceSpan(
|
||||
{
|
||||
name: tool.name,
|
||||
attributes: { type: 'tool-call' },
|
||||
},
|
||||
async ({ metadata: spanMetadata }) => {
|
||||
spanMetadata.input = { request };
|
||||
|
||||
try {
|
||||
let promise: Promise<ToolResult>;
|
||||
if (invocation instanceof ShellToolInvocation) {
|
||||
const setPidCallback = (pid: number) => {
|
||||
const executingCall: ExecutingToolCall = {
|
||||
...call,
|
||||
status: 'executing',
|
||||
tool,
|
||||
invocation,
|
||||
pid,
|
||||
startTime: 'startTime' in call ? call.startTime : undefined,
|
||||
};
|
||||
onUpdateToolCall(executingCall);
|
||||
};
|
||||
promise = executeToolWithHooks(
|
||||
invocation,
|
||||
toolName,
|
||||
signal,
|
||||
messageBus,
|
||||
hooksEnabled,
|
||||
tool,
|
||||
liveOutputCallback,
|
||||
shellExecutionConfig,
|
||||
setPidCallback,
|
||||
);
|
||||
} else {
|
||||
promise = executeToolWithHooks(
|
||||
invocation,
|
||||
toolName,
|
||||
signal,
|
||||
messageBus,
|
||||
hooksEnabled,
|
||||
tool,
|
||||
liveOutputCallback,
|
||||
shellExecutionConfig,
|
||||
);
|
||||
}
|
||||
|
||||
const toolResult: ToolResult = await promise;
|
||||
spanMetadata.output = toolResult;
|
||||
|
||||
if (signal.aborted) {
|
||||
return this.createCancelledResult(
|
||||
call,
|
||||
'User cancelled tool execution.',
|
||||
);
|
||||
} else if (toolResult.error === undefined) {
|
||||
return await this.createSuccessResult(call, toolResult);
|
||||
} else {
|
||||
return this.createErrorResult(
|
||||
call,
|
||||
new Error(toolResult.error.message),
|
||||
toolResult.error.type,
|
||||
);
|
||||
}
|
||||
} catch (executionError: unknown) {
|
||||
spanMetadata.error = executionError;
|
||||
if (signal.aborted) {
|
||||
return this.createCancelledResult(
|
||||
call,
|
||||
'User cancelled tool execution.',
|
||||
);
|
||||
}
|
||||
const error =
|
||||
executionError instanceof Error
|
||||
? executionError
|
||||
: new Error(String(executionError));
|
||||
return this.createErrorResult(
|
||||
call,
|
||||
error,
|
||||
ToolErrorType.UNHANDLED_EXCEPTION,
|
||||
);
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
private createCancelledResult(
|
||||
call: ToolCall,
|
||||
reason: string,
|
||||
): CancelledToolCall {
|
||||
const errorMessage = `[Operation Cancelled] ${reason}`;
|
||||
const startTime = 'startTime' in call ? call.startTime : undefined;
|
||||
|
||||
if (!('tool' in call) || !('invocation' in call)) {
|
||||
// This should effectively never happen in execution phase, but we handle
|
||||
// it safely
|
||||
throw new Error('Cancelled tool call missing tool/invocation references');
|
||||
}
|
||||
|
||||
return {
|
||||
status: 'cancelled',
|
||||
request: call.request,
|
||||
response: {
|
||||
callId: call.request.callId,
|
||||
responseParts: [
|
||||
{
|
||||
functionResponse: {
|
||||
id: call.request.callId,
|
||||
name: call.request.name,
|
||||
response: { error: errorMessage },
|
||||
},
|
||||
},
|
||||
],
|
||||
resultDisplay: undefined,
|
||||
error: undefined,
|
||||
errorType: undefined,
|
||||
contentLength: errorMessage.length,
|
||||
},
|
||||
tool: call.tool,
|
||||
invocation: call.invocation,
|
||||
durationMs: startTime ? Date.now() - startTime : undefined,
|
||||
outcome: call.outcome,
|
||||
};
|
||||
}
|
||||
|
||||
private async createSuccessResult(
|
||||
call: ToolCall,
|
||||
toolResult: ToolResult,
|
||||
): Promise<SuccessfulToolCall> {
|
||||
let content = toolResult.llmContent;
|
||||
let outputFile: string | undefined;
|
||||
const toolName = call.request.name;
|
||||
const callId = call.request.callId;
|
||||
|
||||
if (
|
||||
typeof content === 'string' &&
|
||||
toolName === SHELL_TOOL_NAME &&
|
||||
this.config.getEnableToolOutputTruncation() &&
|
||||
this.config.getTruncateToolOutputThreshold() > 0 &&
|
||||
this.config.getTruncateToolOutputLines() > 0
|
||||
) {
|
||||
const originalContentLength = content.length;
|
||||
const threshold = this.config.getTruncateToolOutputThreshold();
|
||||
const lines = this.config.getTruncateToolOutputLines();
|
||||
const truncatedResult = await saveTruncatedContent(
|
||||
content,
|
||||
callId,
|
||||
this.config.storage.getProjectTempDir(),
|
||||
threshold,
|
||||
lines,
|
||||
);
|
||||
content = truncatedResult.content;
|
||||
outputFile = truncatedResult.outputFile;
|
||||
|
||||
if (outputFile) {
|
||||
logToolOutputTruncated(
|
||||
this.config,
|
||||
new ToolOutputTruncatedEvent(call.request.prompt_id, {
|
||||
toolName,
|
||||
originalContentLength,
|
||||
truncatedContentLength: content.length,
|
||||
threshold,
|
||||
lines,
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const response = convertToFunctionResponse(
|
||||
toolName,
|
||||
callId,
|
||||
content,
|
||||
this.config.getActiveModel(),
|
||||
);
|
||||
|
||||
const successResponse: ToolCallResponseInfo = {
|
||||
callId,
|
||||
responseParts: response,
|
||||
resultDisplay: toolResult.returnDisplay,
|
||||
error: undefined,
|
||||
errorType: undefined,
|
||||
outputFile,
|
||||
contentLength: typeof content === 'string' ? content.length : undefined,
|
||||
};
|
||||
|
||||
const startTime = 'startTime' in call ? call.startTime : undefined;
|
||||
// Ensure we have tool and invocation
|
||||
if (!('tool' in call) || !('invocation' in call)) {
|
||||
throw new Error('Successful tool call missing tool or invocation');
|
||||
}
|
||||
|
||||
return {
|
||||
status: 'success',
|
||||
request: call.request,
|
||||
tool: call.tool,
|
||||
response: successResponse,
|
||||
invocation: call.invocation,
|
||||
durationMs: startTime ? Date.now() - startTime : undefined,
|
||||
outcome: call.outcome,
|
||||
};
|
||||
}
|
||||
|
||||
private createErrorResult(
|
||||
call: ToolCall,
|
||||
error: Error,
|
||||
errorType?: ToolErrorType,
|
||||
): ErroredToolCall {
|
||||
const response = this.createErrorResponse(call.request, error, errorType);
|
||||
const startTime = 'startTime' in call ? call.startTime : undefined;
|
||||
|
||||
return {
|
||||
status: 'error',
|
||||
request: call.request,
|
||||
response,
|
||||
tool: call.tool,
|
||||
durationMs: startTime ? Date.now() - startTime : undefined,
|
||||
outcome: call.outcome,
|
||||
};
|
||||
}
|
||||
|
||||
private createErrorResponse(
|
||||
request: ToolCallRequestInfo,
|
||||
error: Error,
|
||||
errorType: ToolErrorType | undefined,
|
||||
): ToolCallResponseInfo {
|
||||
return {
|
||||
callId: request.callId,
|
||||
error,
|
||||
responseParts: [
|
||||
{
|
||||
functionResponse: {
|
||||
id: request.callId,
|
||||
name: request.name,
|
||||
response: { error: error.message },
|
||||
},
|
||||
},
|
||||
],
|
||||
resultDisplay: error.message,
|
||||
errorType,
|
||||
contentLength: error.message.length,
|
||||
};
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user