From 7bfe6ac418f6f0b0e7b6fc15d70bce8cb2cc3e84 Mon Sep 17 00:00:00 2001 From: AK Date: Tue, 17 Mar 2026 19:34:44 -0700 Subject: [PATCH] feat(core): subagent local execution and tool isolation (#22718) --- packages/cli/src/test-utils/AppRig.tsx | 10 +- .../core/src/agents/agent-scheduler.test.ts | 6 + packages/core/src/agents/agent-scheduler.ts | 11 +- .../core/src/agents/local-executor.test.ts | 108 +++++++++++++++--- packages/core/src/agents/local-executor.ts | 107 ++++++++++++----- .../core/src/config/agent-loop-context.ts | 8 ++ packages/core/src/config/config.ts | 28 ++++- 7 files changed, 222 insertions(+), 56 deletions(-) diff --git a/packages/cli/src/test-utils/AppRig.tsx b/packages/cli/src/test-utils/AppRig.tsx index 8c62592bc6..6043c7f8cc 100644 --- a/packages/cli/src/test-utils/AppRig.tsx +++ b/packages/cli/src/test-utils/AppRig.tsx @@ -280,14 +280,14 @@ export class AppRig { } private stubRefreshAuth() { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-assignment + // eslint-disable-next-line @typescript-eslint/no-explicit-any const gcConfig = this.config as any; gcConfig.refreshAuth = async (authMethod: AuthType) => { gcConfig.modelAvailabilityService.reset(); const newContentGeneratorConfig = { authType: authMethod, - // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + proxy: gcConfig.getProxy(), apiKey: process.env['GEMINI_API_KEY'] || 'test-api-key', }; @@ -456,7 +456,7 @@ export class AppRig { const actualToolName = toolName === '*' ? undefined : toolName; this.config .getPolicyEngine() - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + .removeRulesForTool(actualToolName as string, source); this.breakpointTools.delete(toolName); } @@ -729,7 +729,7 @@ export class AppRig { .getGeminiClient() ?.getChatRecordingService(); if (recordingService) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-type-assertion + // eslint-disable-next-line @typescript-eslint/no-explicit-any (recordingService as any).conversationFile = null; } } @@ -749,7 +749,7 @@ export class AppRig { MockShellExecutionService.reset(); ideContextStore.clear(); // Forcefully clear IdeClient singleton promise - // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-type-assertion + // eslint-disable-next-line @typescript-eslint/no-explicit-any (IdeClient as any).instancePromise = null; vi.clearAllMocks(); diff --git a/packages/core/src/agents/agent-scheduler.test.ts b/packages/core/src/agents/agent-scheduler.test.ts index 2be2f033d9..5d5b6569af 100644 --- a/packages/core/src/agents/agent-scheduler.test.ts +++ b/packages/core/src/agents/agent-scheduler.test.ts @@ -42,6 +42,8 @@ describe('agent-scheduler', () => { it('should create a scheduler with agent-specific config', async () => { const mockConfig = { + getPromptRegistry: vi.fn(), + getResourceRegistry: vi.fn(), messageBus: mockMessageBus, toolRegistry: mockToolRegistry, } as unknown as Mocked; @@ -91,6 +93,8 @@ describe('agent-scheduler', () => { } as unknown as Mocked; const config = { + getPromptRegistry: vi.fn(), + getResourceRegistry: vi.fn(), messageBus: mockMessageBus, } as unknown as Mocked; Object.defineProperty(config, 'toolRegistry', { @@ -123,6 +127,8 @@ describe('agent-scheduler', () => { it('should create an AgentLoopContext that has a defined .config property', async () => { const mockConfig = { + getPromptRegistry: vi.fn(), + getResourceRegistry: vi.fn(), messageBus: mockMessageBus, toolRegistry: mockToolRegistry, promptId: 'test-prompt', diff --git a/packages/core/src/agents/agent-scheduler.ts b/packages/core/src/agents/agent-scheduler.ts index 852e25b4c1..8bed1de00b 100644 --- a/packages/core/src/agents/agent-scheduler.ts +++ b/packages/core/src/agents/agent-scheduler.ts @@ -11,6 +11,8 @@ import type { CompletedToolCall, } from '../scheduler/types.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; +import type { PromptRegistry } from '../prompts/prompt-registry.js'; +import type { ResourceRegistry } from '../resources/resource-registry.js'; import type { EditorType } from '../utils/editor.js'; /** @@ -25,6 +27,10 @@ export interface AgentSchedulingOptions { parentCallId?: string; /** The tool registry specific to this agent. */ toolRegistry: ToolRegistry; + /** The prompt registry specific to this agent. */ + promptRegistry?: PromptRegistry; + /** The resource registry specific to this agent. */ + resourceRegistry?: ResourceRegistry; /** AbortSignal for cancellation. */ signal: AbortSignal; /** Optional function to get the preferred editor for tool modifications. */ @@ -51,16 +57,19 @@ export async function scheduleAgentTools( subagent, parentCallId, toolRegistry, + promptRegistry, + resourceRegistry, signal, getPreferredEditor, onWaitingForConfirmation, } = options; - // Create a proxy/override of the config to provide the agent-specific tool registry. const schedulerContext = { config, promptId: config.promptId, toolRegistry, + promptRegistry: promptRegistry ?? config.getPromptRegistry(), + resourceRegistry: resourceRegistry ?? config.getResourceRegistry(), messageBus: toolRegistry.messageBus, geminiClient: config.geminiClient, sandboxManager: config.sandboxManager, diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index 3ae273cf2f..f0afa73e6a 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -13,10 +13,43 @@ import { afterEach, type Mock, } from 'vitest'; + +const { + mockSendMessageStream, + mockScheduleAgentTools, + mockSetSystemInstruction, + mockCompress, + mockMaybeDiscoverMcpServer, + mockStopMcp, +} = vi.hoisted(() => ({ + mockSendMessageStream: vi.fn().mockResolvedValue({ + async *[Symbol.asyncIterator]() { + yield { + type: 'chunk', + value: { candidates: [] }, + }; + }, + }), + mockScheduleAgentTools: vi.fn(), + mockSetSystemInstruction: vi.fn(), + mockCompress: vi.fn(), + mockMaybeDiscoverMcpServer: vi.fn().mockResolvedValue(undefined), + mockStopMcp: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock('../tools/mcp-client-manager.js', () => ({ + McpClientManager: class { + maybeDiscoverMcpServer = mockMaybeDiscoverMcpServer; + stop = mockStopMcp; + }, +})); + import { debugLogger } from '../utils/debugLogger.js'; import { LocalAgentExecutor, type ActivityCallback } from './local-executor.js'; import { makeFakeConfig } from '../test-utils/config.js'; import { ToolRegistry } from '../tools/tool-registry.js'; +import { PromptRegistry } from '../prompts/prompt-registry.js'; +import { ResourceRegistry } from '../resources/resource-registry.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; import { LSTool } from '../tools/ls.js'; import { LS_TOOL_NAME, READ_FILE_TOOL_NAME } from '../tools/tool-names.js'; @@ -70,18 +103,6 @@ import type { import { getModelConfigAlias, type AgentRegistry } from './registry.js'; import type { ModelRouterService } from '../routing/modelRouterService.js'; -const { - mockSendMessageStream, - mockScheduleAgentTools, - mockSetSystemInstruction, - mockCompress, -} = vi.hoisted(() => ({ - mockSendMessageStream: vi.fn(), - mockScheduleAgentTools: vi.fn(), - mockSetSystemInstruction: vi.fn(), - mockCompress: vi.fn(), -})); - let mockChatHistory: Content[] = []; const mockSetHistory = vi.fn((newHistory: Content[]) => { mockChatHistory = newHistory; @@ -2722,6 +2743,67 @@ describe('LocalAgentExecutor', () => { }); }); + describe('MCP Isolation', () => { + it('should initialize McpClientManager when mcpServers are defined', async () => { + const { MCPServerConfig } = await import('../config/config.js'); + const mcpServers = { + 'test-server': new MCPServerConfig('node', ['server.js']), + }; + + const definition = { + ...createTestDefinition(), + mcpServers, + }; + + vi.spyOn(mockConfig, 'getMcpClientManager').mockReturnValue({ + maybeDiscoverMcpServer: mockMaybeDiscoverMcpServer, + } as unknown as ReturnType); + + await LocalAgentExecutor.create(definition, mockConfig); + + const mcpManager = mockConfig.getMcpClientManager(); + expect(mcpManager?.maybeDiscoverMcpServer).toHaveBeenCalledWith( + 'test-server', + mcpServers['test-server'], + expect.objectContaining({ + toolRegistry: expect.any(ToolRegistry), + promptRegistry: expect.any(PromptRegistry), + resourceRegistry: expect.any(ResourceRegistry), + }), + ); + }); + + it('should inherit main registry tools', async () => { + const parentMcpTool = new DiscoveredMCPTool( + {} as unknown as CallableTool, + 'main-server', + 'tool1', + 'desc1', + {}, + mockConfig.getMessageBus(), + ); + + parentToolRegistry.registerTool(parentMcpTool); + + const definition = createTestDefinition(); + definition.toolConfig = undefined; // trigger inheritance + + vi.spyOn(mockConfig, 'getMcpClientManager').mockReturnValue({ + maybeDiscoverMcpServer: vi.fn(), + } as unknown as ReturnType); + const executor = await LocalAgentExecutor.create( + definition, + mockConfig, + onActivity, + ); + const agentTools = ( + executor as unknown as { toolRegistry: ToolRegistry } + ).toolRegistry.getAllToolNames(); + + expect(agentTools).toContain(parentMcpTool.name); + }); + }); + describe('DeclarativeTool instance tools (browser agent pattern)', () => { /** * The browser agent passes DeclarativeTool instances (not string names) in @@ -2827,13 +2909,11 @@ describe('LocalAgentExecutor', () => { const navTool = new MockTool({ name: 'navigate_page' }); const definition = createInstanceToolDefinition([clickTool, navTool]); - const executor = await LocalAgentExecutor.create( definition, mockConfig, onActivity, ); - const registry = executor['toolRegistry']; expect(registry.getTool('click')).toBeDefined(); expect(registry.getTool('navigate_page')).toBeDefined(); diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index a177012850..a9adeb2e2d 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -4,7 +4,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { Config } from '../config/config.js'; import { type AgentLoopContext } from '../config/agent-loop-context.js'; import { reportError } from '../utils/errorReporting.js'; import { GeminiChat, StreamEventType } from '../core/geminiChat.js'; @@ -17,6 +16,8 @@ import { type Schema, } from '@google/genai'; import { ToolRegistry } from '../tools/tool-registry.js'; +import { PromptRegistry } from '../prompts/prompt-registry.js'; +import { ResourceRegistry } from '../resources/resource-registry.js'; import { type AnyDeclarativeTool } from '../tools/tools.js'; import { DiscoveredMCPTool, @@ -102,14 +103,22 @@ export class LocalAgentExecutor { private readonly agentId: string; private readonly toolRegistry: ToolRegistry; + private readonly promptRegistry: PromptRegistry; + private readonly resourceRegistry: ResourceRegistry; private readonly context: AgentLoopContext; private readonly onActivity?: ActivityCallback; private readonly compressionService: ChatCompressionService; private readonly parentCallId?: string; private hasFailedCompressionAttempt = false; - private get config(): Config { - return this.context.config; + private get executionContext(): AgentLoopContext { + return { + ...this.context, + toolRegistry: this.toolRegistry, + promptRegistry: this.promptRegistry, + resourceRegistry: this.resourceRegistry, + messageBus: this.toolRegistry.getMessageBus(), + }; } /** @@ -133,11 +142,27 @@ export class LocalAgentExecutor { // Create an override object to inject the subagent name into tool confirmation requests const subagentMessageBus = parentMessageBus.derive(definition.name); - // Create an isolated tool registry for this agent instance. + // Create isolated registries for this agent instance. const agentToolRegistry = new ToolRegistry( context.config, subagentMessageBus, ); + const agentPromptRegistry = new PromptRegistry(); + const agentResourceRegistry = new ResourceRegistry(); + + if (definition.mcpServers) { + const globalMcpManager = context.config.getMcpClientManager(); + if (globalMcpManager) { + for (const [name, config] of Object.entries(definition.mcpServers)) { + await globalMcpManager.maybeDiscoverMcpServer(name, config, { + toolRegistry: agentToolRegistry, + promptRegistry: agentPromptRegistry, + resourceRegistry: agentResourceRegistry, + }); + } + } + } + const parentToolRegistry = context.toolRegistry; const allAgentNames = new Set( context.config.getAgentRegistry().getAllAgentNames(), @@ -153,7 +178,9 @@ export class LocalAgentExecutor { return; } - agentToolRegistry.registerTool(tool); + // Clone the tool, so it gets its own state and subagent messageBus + const clonedTool = tool.clone(subagentMessageBus); + agentToolRegistry.registerTool(clonedTool); }; const registerToolByName = (toolName: string) => { @@ -228,10 +255,12 @@ export class LocalAgentExecutor { return new LocalAgentExecutor( definition, context, - agentToolRegistry, parentPromptId, - parentCallId, + agentToolRegistry, + agentPromptRegistry, + agentResourceRegistry, onActivity, + parentCallId, ); } @@ -244,14 +273,18 @@ export class LocalAgentExecutor { private constructor( definition: LocalAgentDefinition, context: AgentLoopContext, - toolRegistry: ToolRegistry, parentPromptId: string | undefined, - parentCallId: string | undefined, + toolRegistry: ToolRegistry, + promptRegistry: PromptRegistry, + resourceRegistry: ResourceRegistry, onActivity?: ActivityCallback, + parentCallId?: string, ) { this.definition = definition; this.context = context; this.toolRegistry = toolRegistry; + this.promptRegistry = promptRegistry; + this.resourceRegistry = resourceRegistry; this.onActivity = onActivity; this.compressionService = new ChatCompressionService(); this.parentCallId = parentCallId; @@ -447,7 +480,7 @@ export class LocalAgentExecutor { } finally { clearTimeout(graceTimeoutId); logRecoveryAttempt( - this.config, + this.context.config, new RecoveryAttemptEvent( this.agentId, this.definition.name, @@ -495,7 +528,7 @@ export class LocalAgentExecutor { const combinedSignal = AbortSignal.any([signal, deadlineTimer.signal]); logAgentStart( - this.config, + this.context.config, new AgentStartEvent(this.agentId, this.definition.name), ); @@ -506,7 +539,7 @@ export class LocalAgentExecutor { const augmentedInputs = { ...inputs, cliVersion: await getVersion(), - activeModel: this.config.getActiveModel(), + activeModel: this.context.config.getActiveModel(), today: new Date().toLocaleDateString(), }; @@ -528,14 +561,16 @@ export class LocalAgentExecutor { // Capture the index of the last hint before starting to avoid re-injecting old hints. // NOTE: Hints added AFTER this point will be broadcast to all currently running // local agents via the listener below. - const startIndex = this.config.injectionService.getLatestInjectionIndex(); - this.config.injectionService.onInjection(injectionListener); + const startIndex = + this.context.config.injectionService.getLatestInjectionIndex(); + this.context.config.injectionService.onInjection(injectionListener); try { - const initialHints = this.config.injectionService.getInjectionsAfter( - startIndex, - 'user_steering', - ); + const initialHints = + this.context.config.injectionService.getInjectionsAfter( + startIndex, + 'user_steering', + ); const formattedInitialHints = formatUserHintsForModel(initialHints); let currentMessage: Content = formattedInitialHints @@ -606,7 +641,16 @@ export class LocalAgentExecutor { } } } finally { - this.config.injectionService.offInjection(injectionListener); + this.context.config.injectionService.offInjection(injectionListener); + + const globalMcpManager = this.context.config.getMcpClientManager(); + if (globalMcpManager) { + globalMcpManager.removeRegistries({ + toolRegistry: this.toolRegistry, + promptRegistry: this.promptRegistry, + resourceRegistry: this.resourceRegistry, + }); + } } // === UNIFIED RECOVERY BLOCK === @@ -719,7 +763,7 @@ export class LocalAgentExecutor { } finally { deadlineTimer.abort(); logAgentFinish( - this.config, + this.context.config, new AgentFinishEvent( this.agentId, this.definition.name, @@ -742,7 +786,7 @@ export class LocalAgentExecutor { prompt_id, false, model, - this.config, + this.context.config, this.hasFailedCompressionAttempt, ); @@ -780,10 +824,11 @@ export class LocalAgentExecutor { const modelConfigAlias = getModelConfigAlias(this.definition); // Resolve the model config early to get the concrete model string (which may be `auto`). - const resolvedConfig = this.config.modelConfigService.getResolvedConfig({ - model: modelConfigAlias, - overrideScope: this.definition.name, - }); + const resolvedConfig = + this.context.config.modelConfigService.getResolvedConfig({ + model: modelConfigAlias, + overrideScope: this.definition.name, + }); const requestedModel = resolvedConfig.model; let modelToUse: string; @@ -800,7 +845,7 @@ export class LocalAgentExecutor { signal, requestedModel, }; - const router = this.config.getModelRouterService(); + const router = this.context.config.getModelRouterService(); const decision = await router.route(routingContext); modelToUse = decision.model; } catch (error) { @@ -888,7 +933,7 @@ export class LocalAgentExecutor { try { return new GeminiChat( - this.config, + this.executionContext, systemInstruction, [{ functionDeclarations: tools }], startHistory, @@ -1136,13 +1181,15 @@ export class LocalAgentExecutor { // Execute standard tool calls using the new scheduler if (toolRequests.length > 0) { const completedCalls = await scheduleAgentTools( - this.config, + this.context.config, toolRequests, { - schedulerId: this.agentId, + schedulerId: promptId, subagent: this.definition.name, parentCallId: this.parentCallId, toolRegistry: this.toolRegistry, + promptRegistry: this.promptRegistry, + resourceRegistry: this.resourceRegistry, signal, onWaitingForConfirmation, }, @@ -1277,7 +1324,7 @@ export class LocalAgentExecutor { let finalPrompt = templateString(promptConfig.systemPrompt, inputs); // Append environment context (CWD and folder structure). - const dirContext = await getDirectoryContextString(this.config); + const dirContext = await getDirectoryContextString(this.context.config); finalPrompt += `\n\n# Environment Context\n${dirContext}`; // Append standard rules for non-interactive execution. diff --git a/packages/core/src/config/agent-loop-context.ts b/packages/core/src/config/agent-loop-context.ts index 0a879d9c93..b16326a7ce 100644 --- a/packages/core/src/config/agent-loop-context.ts +++ b/packages/core/src/config/agent-loop-context.ts @@ -7,6 +7,8 @@ import type { GeminiClient } from '../core/client.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; +import type { PromptRegistry } from '../prompts/prompt-registry.js'; +import type { ResourceRegistry } from '../resources/resource-registry.js'; import type { SandboxManager } from '../services/sandboxManager.js'; import type { Config } from './config.js'; @@ -24,6 +26,12 @@ export interface AgentLoopContext { /** The registry of tools available to the agent in this context. */ readonly toolRegistry: ToolRegistry; + /** The registry of prompts available to the agent in this context. */ + readonly promptRegistry: PromptRegistry; + + /** The registry of resources available to the agent in this context. */ + readonly resourceRegistry: ResourceRegistry; + /** The bus for user confirmations and messages in this context. */ readonly messageBus: MessageBus; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index fcb6613756..aa3e9aa5b6 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -660,8 +660,8 @@ export class Config implements McpContext, AgentLoopContext { private allowedEnvironmentVariables: string[]; private blockedEnvironmentVariables: string[]; private readonly enableEnvironmentVariableRedaction: boolean; - private promptRegistry!: PromptRegistry; - private resourceRegistry!: ResourceRegistry; + private _promptRegistry!: PromptRegistry; + private _resourceRegistry!: ResourceRegistry; private agentRegistry!: AgentRegistry; private readonly acknowledgedAgentsService: AcknowledgedAgentsService; private skillManager!: SkillManager; @@ -1245,8 +1245,8 @@ export class Config implements McpContext, AgentLoopContext { if (this.getCheckpointingEnabled()) { await this.getGitService(); } - this.promptRegistry = new PromptRegistry(); - this.resourceRegistry = new ResourceRegistry(); + this._promptRegistry = new PromptRegistry(); + this._resourceRegistry = new ResourceRegistry(); this.agentRegistry = new AgentRegistry(this); await this.agentRegistry.initialize(); @@ -1482,6 +1482,22 @@ export class Config implements McpContext, AgentLoopContext { return this._toolRegistry; } + /** + * @deprecated Do not access directly on Config. + * Use the injected AgentLoopContext instead. + */ + get promptRegistry(): PromptRegistry { + return this._promptRegistry; + } + + /** + * @deprecated Do not access directly on Config. + * Use the injected AgentLoopContext instead. + */ + get resourceRegistry(): ResourceRegistry { + return this._resourceRegistry; + } + /** * @deprecated Do not access directly on Config. * Use the injected AgentLoopContext instead. @@ -1794,7 +1810,7 @@ export class Config implements McpContext, AgentLoopContext { } getPromptRegistry(): PromptRegistry { - return this.promptRegistry; + return this._promptRegistry; } getSkillManager(): SkillManager { @@ -1802,7 +1818,7 @@ export class Config implements McpContext, AgentLoopContext { } getResourceRegistry(): ResourceRegistry { - return this.resourceRegistry; + return this._resourceRegistry; } getDebugMode(): boolean {