fix(core): ensure subagents use qualified MCP tool names (#20801)

This commit is contained in:
Abhi
2026-03-02 16:12:13 -05:00
committed by GitHub
parent 1502e5cbc3
commit b7a8f0d1f9
4 changed files with 121 additions and 34 deletions

View File

@@ -501,7 +501,7 @@ describe('LocalAgentExecutor', () => {
expect(agentRegistry.getTool(subAgentName)).toBeUndefined();
});
it('should enforce qualified names for MCP tools in agent definitions', async () => {
it('should automatically qualify MCP tools in agent definitions', async () => {
const serverName = 'mcp-server';
const toolName = 'mcp-tool';
const qualifiedName = `${serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${toolName}`;
@@ -530,7 +530,7 @@ describe('LocalAgentExecutor', () => {
return undefined;
});
// 1. Qualified name works and registers the tool (using short name per status quo)
// 1. Qualified name works and registers the tool (using qualified name)
const definition = createTestDefinition([qualifiedName]);
const executor = await LocalAgentExecutor.create(
definition,
@@ -539,14 +539,18 @@ describe('LocalAgentExecutor', () => {
);
const agentRegistry = executor['toolRegistry'];
// Registry shortening logic means it's registered as 'mcp-tool' internally
expect(agentRegistry.getTool(toolName)).toBeDefined();
// It should be registered as the qualified name
expect(agentRegistry.getTool(qualifiedName)).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/);
// 2. Unqualified name for MCP tool now also works (and gets upgraded to qualified)
const definition2 = createTestDefinition([toolName]);
const executor2 = await LocalAgentExecutor.create(
definition2,
mockConfig,
onActivity,
);
const agentRegistry2 = executor2['toolRegistry'];
expect(agentRegistry2.getTool(qualifiedName)).toBeDefined();
getToolSpy.mockRestore();
});

View File

@@ -16,10 +16,7 @@ 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 { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { CompressionStatus } from '../core/turn.js';
import { type ToolCallRequestInfo } from '../scheduler/types.js';
import { ChatCompressionService } from '../services/chatCompressionService.js';
@@ -142,15 +139,14 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// 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}'.`,
);
if (tool instanceof DiscoveredMCPTool) {
// Subagents MUST use fully qualified names for MCP tools to ensure
// unambiguous tool calls and to comply with policy requirements.
// We automatically "upgrade" any MCP tool to its qualified version.
agentToolRegistry.registerTool(tool.asFullyQualifiedTool());
} else {
agentToolRegistry.registerTool(tool);
}
agentToolRegistry.registerTool(tool);
}
};

View File

@@ -380,20 +380,36 @@ describe('ToolRegistry', () => {
});
describe('getAllToolNames', () => {
it('should return all registered tool names', () => {
it('should return all registered tool names with qualified names for MCP tools', () => {
// Register tools with displayNames in non-alphabetical order
const toolC = new MockTool({ name: 'c-tool', displayName: 'Tool C' });
const toolA = new MockTool({ name: 'a-tool', displayName: 'Tool A' });
const toolB = new MockTool({ name: 'b-tool', displayName: 'Tool B' });
const mcpTool = createMCPTool('my-server', 'my-tool', 'desc');
toolRegistry.registerTool(toolC);
toolRegistry.registerTool(toolA);
toolRegistry.registerTool(toolB);
toolRegistry.registerTool(mcpTool);
const toolNames = toolRegistry.getAllToolNames();
// Assert that the returned array contains all tool names
expect(toolNames).toEqual(['c-tool', 'a-tool', 'b-tool']);
// Assert that the returned array contains all tool names, with MCP qualified
expect(toolNames).toContain('c-tool');
expect(toolNames).toContain('a-tool');
expect(toolNames).toContain('my-server__my-tool');
expect(toolNames).toHaveLength(3);
});
it('should deduplicate tool names', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'desc');
// Register same MCP tool twice (one as alias, one as qualified)
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());
const toolNames = toolRegistry.getAllToolNames();
expect(toolNames).toEqual([`${serverName}__${toolName}`]);
});
});
@@ -465,8 +481,8 @@ describe('ToolRegistry', () => {
'builtin-1',
'builtin-2',
DISCOVERED_TOOL_PREFIX + 'discovered-1',
'mcp-apple',
'mcp-zebra',
'apple-server__mcp-apple',
'zebra-server__mcp-zebra',
]);
});
});
@@ -659,6 +675,34 @@ describe('ToolRegistry', () => {
});
});
describe('getFunctionDeclarations', () => {
it('should use fully qualified names for MCP tools in declarations', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');
toolRegistry.registerTool(mcpTool);
const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);
expect(declarations[0].name).toBe(`${serverName}__${toolName}`);
});
it('should deduplicate MCP tools in declarations', () => {
const serverName = 'my-server';
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');
// Register both alias and qualified
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());
const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);
expect(declarations[0].name).toBe(`${serverName}__${toolName}`);
});
});
describe('plan mode', () => {
it('should only return policy-allowed tools in plan mode', () => {
// Register several tools

View File

@@ -539,11 +539,32 @@ export class ToolRegistry {
const plansDir = this.config.storage.getPlansDir();
const declarations: FunctionDeclaration[] = [];
const seenNames = new Set<string>();
this.getActiveTools().forEach((tool) => {
const toolName =
tool instanceof DiscoveredMCPTool
? tool.getFullyQualifiedName()
: tool.name;
if (seenNames.has(toolName)) {
return;
}
seenNames.add(toolName);
let schema = tool.getSchema(modelId);
// Ensure the schema name matches the qualified name for MCP tools
if (tool instanceof DiscoveredMCPTool) {
schema = {
...schema,
name: toolName,
};
}
if (
isPlanMode &&
(tool.name === WRITE_FILE_TOOL_NAME || tool.name === EDIT_TOOL_NAME)
(toolName === WRITE_FILE_TOOL_NAME || toolName === EDIT_TOOL_NAME)
) {
schema = {
...schema,
@@ -576,20 +597,42 @@ export class ToolRegistry {
}
/**
* Returns an array of all registered and discovered tool names which are not
* excluded via configuration.
* Returns an array of names for all active tools.
* For MCP tools, this returns their fully qualified names.
* The list is deduplicated.
*/
getAllToolNames(): string[] {
return this.getActiveTools().map((tool) => tool.name);
const names = new Set<string>();
for (const tool of this.getActiveTools()) {
if (tool instanceof DiscoveredMCPTool) {
names.add(tool.getFullyQualifiedName());
} else {
names.add(tool.name);
}
}
return Array.from(names);
}
/**
* Returns an array of all registered and discovered tool instances.
*/
getAllTools(): AnyDeclarativeTool[] {
return this.getActiveTools().sort((a, b) =>
const seen = new Set<string>();
const tools: AnyDeclarativeTool[] = [];
for (const tool of this.getActiveTools().sort((a, b) =>
a.displayName.localeCompare(b.displayName),
);
)) {
const name =
tool instanceof DiscoveredMCPTool
? tool.getFullyQualifiedName()
: tool.name;
if (!seen.has(name)) {
seen.add(name);
tools.push(tool);
}
}
return tools;
}
/**