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.
This commit is contained in:
Akhilesh Kumar
2026-03-11 20:43:25 +00:00
parent eb5d22848c
commit 5a020e7720
3 changed files with 50 additions and 66 deletions
+13 -51
View File
@@ -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<typeof mockConfig.getMcpClientManager>);
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();
});
});
});
+29 -15
View File
@@ -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<TOutput extends z.ZodTypeAny> {
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<TOutput extends z.ZodTypeAny> {
);
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<TOutput extends z.ZodTypeAny> {
} 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<TOutput extends z.ZodTypeAny> {
parentPromptId: string | undefined,
parentCallId: string | undefined,
onActivity?: ActivityCallback,
mcpClientManager?: McpClientManager,
) {
this.definition = definition;
this.context = context;
@@ -250,7 +268,6 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
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<TOutput extends z.ZodTypeAny> {
throw error; // Re-throw other errors or external aborts.
} finally {
deadlineTimer.abort();
if (this.mcpClientManager) {
await this.mcpClientManager.stop();
}
logAgentFinish(
this.config,
new AgentFinishEvent(
+8
View File
@@ -565,6 +565,14 @@ export class ToolRegistry {
return;
}
if (
this.isMainRegistry &&
tool instanceof DiscoveredMCPTool &&
tool.serverName.startsWith('__agent__')
) {
return;
}
if (
mainAgentTools &&
!mainAgentTools.includes(toolName) &&