diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index 3cb7b188fd..6b33e0b76b 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -17,6 +17,10 @@ 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 { + DiscoveredMCPTool, + MCP_QUALIFIED_NAME_SEPARATOR, +} from '../tools/mcp-tool.js'; import { LSTool } from '../tools/ls.js'; import { LS_TOOL_NAME, READ_FILE_TOOL_NAME } from '../tools/tool-names.js'; import { @@ -31,6 +35,7 @@ import { type Content, type PartListUnion, type Tool, + type CallableTool, } from '@google/genai'; import type { Config } from '../config/config.js'; import { MockTool } from '../test-utils/mock-tool.js'; @@ -493,6 +498,56 @@ describe('LocalAgentExecutor', () => { // Should exclude subagent expect(agentRegistry.getTool(subAgentName)).toBeUndefined(); }); + + it('should enforce qualified names for MCP tools in agent definitions', async () => { + const serverName = 'mcp-server'; + const toolName = 'mcp-tool'; + const qualifiedName = `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}`; + + const mockMcpTool = { + tool: vi.fn(), + callTool: vi.fn(), + } as unknown as CallableTool; + + const mcpTool = new DiscoveredMCPTool( + mockMcpTool, + serverName, + toolName, + 'description', + {}, + mockConfig.getMessageBus(), + ); + + // Mock getTool to return our real DiscoveredMCPTool instance + const getToolSpy = vi + .spyOn(parentToolRegistry, 'getTool') + .mockImplementation((name) => { + if (name === toolName || name === qualifiedName) { + return mcpTool; + } + return undefined; + }); + + // 1. Qualified name works and registers the tool (using short name per status quo) + const definition = createTestDefinition([qualifiedName]); + const executor = await LocalAgentExecutor.create( + definition, + mockConfig, + onActivity, + ); + + const agentRegistry = executor['toolRegistry']; + // Registry shortening logic means it's registered as 'mcp-tool' internally + expect(agentRegistry.getTool(toolName)).toBeDefined(); + + // 2. Unqualified name for MCP tool THROWS + const badDefinition = createTestDefinition([toolName]); + await expect( + LocalAgentExecutor.create(badDefinition, mockConfig, onActivity), + ).rejects.toThrow(/must be requested with its server prefix/); + + getToolSpy.mockRestore(); + }); }); describe('run (Execution Loop and Logic)', () => { diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index e22143ac54..95f3ab74c8 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -16,6 +16,10 @@ import type { Schema, } from '@google/genai'; import { ToolRegistry } from '../tools/tool-registry.js'; +import { + DiscoveredMCPTool, + MCP_QUALIFIED_NAME_SEPARATOR, +} from '../tools/mcp-tool.js'; import { CompressionStatus } from '../core/turn.js'; import { type ToolCallRequestInfo } from '../scheduler/types.js'; import { ChatCompressionService } from '../services/chatCompressionService.js'; @@ -129,6 +133,14 @@ export class LocalAgentExecutor { // registry and register it with the agent's isolated registry. const tool = parentToolRegistry.getTool(toolName); if (tool) { + if ( + tool instanceof DiscoveredMCPTool && + !toolName.includes(MCP_QUALIFIED_NAME_SEPARATOR) + ) { + throw new Error( + `MCP tool '${toolName}' must be requested with its server prefix (e.g., '${tool.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}') in agent '${definition.name}'.`, + ); + } agentToolRegistry.registerTool(tool); } }; diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index e401e673f7..c096feeeee 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -23,6 +23,12 @@ import { ToolErrorType } from './tool-error.js'; import type { Config } from '../config/config.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; +/** + * The separator used to qualify MCP tool names with their server prefix. + * e.g. "server_name__tool_name" + */ +export const MCP_QUALIFIED_NAME_SEPARATOR = '__'; + type ToolParams = Record; // Discriminated union for MCP Content Blocks to ensure type safety. @@ -82,7 +88,7 @@ export class DiscoveredMCPToolInvocation extends BaseToolInvocation< super( params, messageBus, - `${serverName}__${serverToolName}`, + `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${serverToolName}`, displayName, serverName, ); @@ -261,7 +267,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< } getFullyQualifiedPrefix(): string { - return `${this.serverName}__`; + return `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}`; } getFullyQualifiedName(): string { diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts index b1be466847..df4984595a 100644 --- a/packages/core/src/tools/tool-registry.test.ts +++ b/packages/core/src/tools/tool-registry.test.ts @@ -13,7 +13,7 @@ import { ApprovalMode } from '../policy/types.js'; import { ToolRegistry, DiscoveredTool } from './tool-registry.js'; import { DISCOVERED_TOOL_PREFIX } from './tool-names.js'; -import { DiscoveredMCPTool } from './mcp-tool.js'; +import { DiscoveredMCPTool, MCP_QUALIFIED_NAME_SEPARATOR } from './mcp-tool.js'; import type { FunctionDeclaration, CallableTool } from '@google/genai'; import { mcpToTool } from '@google/genai'; import { spawn } from 'node:child_process'; @@ -568,6 +568,22 @@ describe('ToolRegistry', () => { expect(retrievedTool).toBeDefined(); expect(retrievedTool?.name).toBe(validToolName); }); + + it('should resolve qualified names in getFunctionDeclarationsFiltered', () => { + const serverName = 'my-server'; + const toolName = 'my-tool'; + const mcpTool = createMCPTool(serverName, toolName, 'description'); + + toolRegistry.registerTool(mcpTool); + + const fullyQualifiedName = `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}`; + const declarations = toolRegistry.getFunctionDeclarationsFiltered([ + fullyQualifiedName, + ]); + + expect(declarations).toHaveLength(1); + expect(declarations[0].name).toBe(toolName); + }); }); describe('DiscoveredToolInvocation', () => { diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index 45512eb8cc..f9e5f8fa8d 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -488,8 +488,8 @@ export class ToolRegistry { getFunctionDeclarationsFiltered(toolNames: string[]): FunctionDeclaration[] { const declarations: FunctionDeclaration[] = []; for (const name of toolNames) { - const tool = this.allKnownTools.get(name); - if (tool && this.isActiveTool(tool)) { + const tool = this.getTool(name); + if (tool) { declarations.push(tool.schema); } }