From a01d411c5a7140e289f1f3b003d9d4bb9bfc071c Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Tue, 19 Aug 2025 16:27:15 -0700 Subject: [PATCH] Get ToolRegistry from config instead of passing it (#6592) --- packages/cli/src/nonInteractiveCli.test.ts | 1 - packages/cli/src/nonInteractiveCli.ts | 3 -- .../core/nonInteractiveToolExecutor.test.ts | 28 ++++++++----------- .../src/core/nonInteractiveToolExecutor.ts | 4 +-- packages/core/src/core/subagent.test.ts | 3 +- packages/core/src/core/subagent.ts | 4 --- 6 files changed, 14 insertions(+), 29 deletions(-) diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index b08862f5..d7f7ad70 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -137,7 +137,6 @@ describe('runNonInteractive', () => { expect(mockCoreExecuteToolCall).toHaveBeenCalledWith( mockConfig, expect.objectContaining({ name: 'testTool' }), - mockToolRegistry, expect.any(AbortSignal), ); expect(mockGeminiClient.sendMessageStream).toHaveBeenNthCalledWith( diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index e6af6665..6aec2754 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -8,7 +8,6 @@ import { Config, ToolCallRequestInfo, executeToolCall, - ToolRegistry, shutdownTelemetry, isTelemetrySdkInitialized, GeminiEventType, @@ -39,7 +38,6 @@ export async function runNonInteractive( }); const geminiClient = config.getGeminiClient(); - const toolRegistry: ToolRegistry = config.getToolRegistry(); const abortController = new AbortController(); let currentMessages: Content[] = [ @@ -100,7 +98,6 @@ export async function runNonInteractive( const toolResponse = await executeToolCall( config, requestInfo, - toolRegistry, abortController.signal, ); diff --git a/packages/core/src/core/nonInteractiveToolExecutor.test.ts b/packages/core/src/core/nonInteractiveToolExecutor.test.ts index 0c1164ea..38afa697 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.test.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.test.ts @@ -16,20 +16,11 @@ import { import { Part } from '@google/genai'; import { MockTool } from '../test-utils/tools.js'; -const mockConfig = { - getSessionId: () => 'test-session-id', - getUsageStatisticsEnabled: () => true, - getDebugMode: () => false, - getContentGeneratorConfig: () => ({ - model: 'test-model', - authType: 'oauth-personal', - }), -} as unknown as Config; - describe('executeToolCall', () => { let mockToolRegistry: ToolRegistry; let mockTool: MockTool; let abortController: AbortController; + let mockConfig: Config; beforeEach(() => { mockTool = new MockTool(); @@ -39,6 +30,17 @@ describe('executeToolCall', () => { // Add other ToolRegistry methods if needed, or use a more complete mock } as unknown as ToolRegistry; + mockConfig = { + getSessionId: () => 'test-session-id', + getUsageStatisticsEnabled: () => true, + getDebugMode: () => false, + getContentGeneratorConfig: () => ({ + model: 'test-model', + authType: 'oauth-personal', + }), + getToolRegistry: () => mockToolRegistry, + } as unknown as Config; + abortController = new AbortController(); }); @@ -60,7 +62,6 @@ describe('executeToolCall', () => { const response = await executeToolCall( mockConfig, request, - mockToolRegistry, abortController.signal, ); @@ -94,7 +95,6 @@ describe('executeToolCall', () => { const response = await executeToolCall( mockConfig, request, - mockToolRegistry, abortController.signal, ); @@ -141,7 +141,6 @@ describe('executeToolCall', () => { const response = await executeToolCall( mockConfig, request, - mockToolRegistry, abortController.signal, ); expect(response).toStrictEqual({ @@ -185,7 +184,6 @@ describe('executeToolCall', () => { const response = await executeToolCall( mockConfig, request, - mockToolRegistry, abortController.signal, ); expect(response).toStrictEqual({ @@ -222,7 +220,6 @@ describe('executeToolCall', () => { const response = await executeToolCall( mockConfig, request, - mockToolRegistry, abortController.signal, ); @@ -262,7 +259,6 @@ describe('executeToolCall', () => { const response = await executeToolCall( mockConfig, request, - mockToolRegistry, abortController.signal, ); diff --git a/packages/core/src/core/nonInteractiveToolExecutor.ts b/packages/core/src/core/nonInteractiveToolExecutor.ts index 3849d52a..c116ca33 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.ts @@ -10,7 +10,6 @@ import { ToolCallRequestInfo, ToolCallResponseInfo, ToolErrorType, - ToolRegistry, ToolResult, } from '../index.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; @@ -25,10 +24,9 @@ import { ToolCallDecision } from '../telemetry/tool-call-decision.js'; export async function executeToolCall( config: Config, toolCallRequest: ToolCallRequestInfo, - toolRegistry: ToolRegistry, abortSignal?: AbortSignal, ): Promise { - const tool = toolRegistry.getTool(toolCallRequest.name); + const tool = config.getToolRegistry().getTool(toolCallRequest.name); const startTime = Date.now(); if (!tool) { diff --git a/packages/core/src/core/subagent.test.ts b/packages/core/src/core/subagent.test.ts index 43b656c2..978a686b 100644 --- a/packages/core/src/core/subagent.test.ts +++ b/packages/core/src/core/subagent.test.ts @@ -534,7 +534,7 @@ describe('subagent.ts', () => { parameters: { type: Type.OBJECT, properties: {} }, }; - const { config, toolRegistry } = await createMockConfig({ + const { config } = await createMockConfig({ getFunctionDeclarationsFiltered: vi .fn() .mockReturnValue([listFilesToolDef]), @@ -580,7 +580,6 @@ describe('subagent.ts', () => { expect(executeToolCall).toHaveBeenCalledWith( config, expect.objectContaining({ name: 'list_files', args: { path: '.' } }), - toolRegistry, expect.any(AbortSignal), ); diff --git a/packages/core/src/core/subagent.ts b/packages/core/src/core/subagent.ts index 10776f0e..3abe4816 100644 --- a/packages/core/src/core/subagent.ts +++ b/packages/core/src/core/subagent.ts @@ -5,7 +5,6 @@ */ import { reportError } from '../utils/errorReporting.js'; -import { ToolRegistry } from '../tools/tool-registry.js'; import { Config } from '../config/config.js'; import { ToolCallRequestInfo } from './turn.js'; import { executeToolCall } from './nonInteractiveToolExecutor.js'; @@ -422,7 +421,6 @@ export class SubAgentScope { if (functionCalls.length > 0) { currentMessages = await this.processFunctionCalls( functionCalls, - toolRegistry, abortController, promptId, ); @@ -479,7 +477,6 @@ export class SubAgentScope { */ private async processFunctionCalls( functionCalls: FunctionCall[], - toolRegistry: ToolRegistry, abortController: AbortController, promptId: string, ): Promise { @@ -513,7 +510,6 @@ export class SubAgentScope { toolResponse = await executeToolCall( this.runtimeContext, requestInfo, - toolRegistry, abortController.signal, ); }