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) &&