From 5a020e7720d7048196b80a8824c648ded06e16de Mon Sep 17 00:00:00 2001 From: Akhilesh Kumar Date: Wed, 11 Mar 2026 20:43:25 +0000 Subject: [PATCH] fix(core): avoid restarting subagent MCP servers I've tactically refactored the `LocalAgentExecutor` so that it avoids shutting down and restarting subagent MCP servers for every agent execution, which mitigates the performance overhead caused by long startup times. 1. Leveraging the Global McpClientManager: Instead of instantiating an entirely new `McpClientManager` instance within the `LocalAgentExecutor` per execution (and shutting it down in its `finally` block), we now use the single global `McpClientManager` available on `context.config`. Since the global manager deduplicates connection attempts by checking if the server is already active, subagent MCP servers will now naturally stay alive after their initial initialization. 2. Prefixing to Avoid Polluting the Global Namespace: To isolate the agent-specific tools, we now register the subagent's MCP servers with a unique prefix: `__agent__${definition.name}__${name}`. 3. Strict Filtering for True Isolation (ToolRegistry): - Main CLI context: Added a block in the global `ToolRegistry.getFunctionDeclarations()` that strictly hides any tool belonging to a server prefixed with `__agent__` if the registry `isMainRegistry`. This prevents internal subagent tools from leaking to the main agent. - Subagent context (`LocalAgentExecutor`): When inheriting tools from the parent registry (the fallback when an agent doesn't explicitly define `tools: []`), the agent now ignores `__agent__` prefixed tools that belong to *other* agents, ensuring strict tool isolation while keeping the actual underlying server processes alive and reusable. --- .../core/src/agents/local-executor.test.ts | 64 ++++--------------- packages/core/src/agents/local-executor.ts | 44 ++++++++----- packages/core/src/tools/tool-registry.ts | 8 +++ 3 files changed, 50 insertions(+), 66 deletions(-) diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index c2949b10b6..653b44c294 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -23,7 +23,7 @@ const { mockStopMcp, } = vi.hoisted(() => ({ mockSendMessageStream: vi.fn().mockResolvedValue({ - async *[Symbol.asyncIterator] () { + async *[Symbol.asyncIterator]() { yield { type: 'chunk', value: { candidates: [] }, @@ -38,11 +38,11 @@ const { })); vi.mock('../tools/mcp-client-manager.js', () => ({ - McpClientManager: class { - maybeDiscoverMcpServer = mockMaybeDiscoverMcpServer; - stop = mockStopMcp; - }, - })); + McpClientManager: class { + maybeDiscoverMcpServer = mockMaybeDiscoverMcpServer; + stop = mockStopMcp; + }, +})); import { debugLogger } from '../utils/debugLogger.js'; import { LocalAgentExecutor, type ActivityCallback } from './local-executor.js'; @@ -2498,55 +2498,17 @@ describe('LocalAgentExecutor', () => { mcpServers, }; + vi.spyOn(mockConfig, 'getMcpClientManager').mockReturnValue({ + maybeDiscoverMcpServer: mockMaybeDiscoverMcpServer, + } as unknown as ReturnType); + await LocalAgentExecutor.create(definition, mockConfig); - expect(mockMaybeDiscoverMcpServer).toHaveBeenCalledWith( - 'test-server', + const mcpManager = mockConfig.getMcpClientManager(); + expect(mcpManager?.maybeDiscoverMcpServer).toHaveBeenCalledWith( + '__agent__TestAgent__test-server', mcpServers['test-server'], ); }); - - it('should stop McpClientManager when agent execution finishes', async () => { - const { MCPServerConfig } = await import('../config/config.js'); - const mcpServers = { - 'test-server': new MCPServerConfig('node', ['server.js']), - }; - - const definition = { - ...createTestDefinition(), - mcpServers, - }; - - const executor = await LocalAgentExecutor.create(definition, mockConfig); - - mockSendMessageStream.mockResolvedValueOnce({ - async *[Symbol.asyncIterator] () { - yield { - type: 'chunk', - value: { - candidates: [ - { - content: { - parts: [ - { - functionCall: { - name: TASK_COMPLETE_TOOL_NAME, - args: { result: 'Done' }, - id: 't1', - }, - }, - ], - }, - }, - ], - }, - }; - }, - }); - - await executor.run({ goal: 'test' }, signal); - - expect(mockStopMcp).toHaveBeenCalled(); - }); }); }); diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 0eed4f2888..ae62cdfa9b 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -18,7 +18,7 @@ import { } from '@google/genai'; import { ToolRegistry } from '../tools/tool-registry.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; -import { McpClientManager } from '../tools/mcp-client-manager.js'; +import type { McpClientManager } from '../tools/mcp-client-manager.js'; import { CompressionStatus } from '../core/turn.js'; import { type ToolCallRequestInfo } from '../scheduler/types.js'; import { type Message } from '../confirmation-bus/types.js'; @@ -95,7 +95,6 @@ export class LocalAgentExecutor { private readonly agentId: string; private readonly toolRegistry: ToolRegistry; private readonly context: AgentLoopContext; - private readonly mcpClientManager?: McpClientManager; private readonly onActivity?: ActivityCallback; private readonly compressionService: ChatCompressionService; private readonly parentCallId?: string; @@ -145,14 +144,12 @@ export class LocalAgentExecutor { ); let mcpClientManager: McpClientManager | undefined; if (definition.mcpServers) { - mcpClientManager = new McpClientManager( - await getVersion(), - agentToolRegistry, - context.config, - ); - - for (const [name, config] of Object.entries(definition.mcpServers)) { - await mcpClientManager.maybeDiscoverMcpServer(name, config); + const globalMcpManager = context.config.getMcpClientManager(); + if (globalMcpManager) { + for (const [name, config] of Object.entries(definition.mcpServers)) { + const prefixedName = `__agent__${definition.name}__${name}`; + await globalMcpManager.maybeDiscoverMcpServer(prefixedName, config); + } } } @@ -205,10 +202,32 @@ export class LocalAgentExecutor { } else { // If no tools are explicitly configured, default to all available tools. for (const toolName of parentToolRegistry.getAllToolNames()) { + const tool = parentToolRegistry.getTool(toolName); + if ( + tool instanceof DiscoveredMCPTool && + tool.serverName.startsWith('__agent__') + ) { + if (!tool.serverName.startsWith(`__agent__${definition.name}__`)) { + continue; // Skip other agents' MCP tools + } + } registerToolByName(toolName); } } + // Always ensure this agent's own MCP servers are included, even if toolConfig is restricted + parentToolRegistry.getActiveTools().forEach((tool) => { + if ( + tool instanceof DiscoveredMCPTool && + tool.serverName.startsWith(`__agent__${definition.name}__`) + ) { + const qualifiedName = tool.asFullyQualifiedTool().name; + if (!agentToolRegistry.getTool(qualifiedName)) { + registerToolByName(qualifiedName); + } + } + }); + agentToolRegistry.sortTools(); // Get the parent prompt ID from context @@ -242,7 +261,6 @@ export class LocalAgentExecutor { parentPromptId: string | undefined, parentCallId: string | undefined, onActivity?: ActivityCallback, - mcpClientManager?: McpClientManager, ) { this.definition = definition; this.context = context; @@ -250,7 +268,6 @@ export class LocalAgentExecutor { this.onActivity = onActivity; this.compressionService = new ChatCompressionService(); this.parentCallId = parentCallId; - this.mcpClientManager = mcpClientManager; const randomIdPart = Math.random().toString(36).slice(2, 8); // parentPromptId will be undefined if this agent is invoked directly @@ -697,9 +714,6 @@ export class LocalAgentExecutor { throw error; // Re-throw other errors or external aborts. } finally { deadlineTimer.abort(); - if (this.mcpClientManager) { - await this.mcpClientManager.stop(); - } logAgentFinish( this.config, new AgentFinishEvent( diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index c7a46e9700..2b498baf7a 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -565,6 +565,14 @@ export class ToolRegistry { return; } + if ( + this.isMainRegistry && + tool instanceof DiscoveredMCPTool && + tool.serverName.startsWith('__agent__') + ) { + return; + } + if ( mainAgentTools && !mainAgentTools.includes(toolName) &&