fix(core): resolve MCP tool FQN validation, schema export, and wildcards in subagents (#22069)

This commit is contained in:
Abhi
2026-03-12 10:17:36 -04:00
committed by GitHub
parent a38aaa47fb
commit 8432bcee75
7 changed files with 136 additions and 99 deletions

View File

@@ -107,9 +107,11 @@ const localAgentSchema = z
display_name: z.string().optional(),
tools: z
.array(
z.string().refine((val) => isValidToolName(val), {
message: 'Invalid tool name',
}),
z
.string()
.refine((val) => isValidToolName(val, { allowWildcards: true }), {
message: 'Invalid tool name',
}),
)
.optional(),
model: z.string().optional(),

View File

@@ -17,7 +17,13 @@ import {
type Schema,
} from '@google/genai';
import { ToolRegistry } from '../tools/tool-registry.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { type AnyDeclarativeTool } from '../tools/tools.js';
import {
DiscoveredMCPTool,
isMcpToolName,
parseMcpToolName,
MCP_TOOL_PREFIX,
} from '../tools/mcp-tool.js';
import { CompressionStatus } from '../core/turn.js';
import { type ToolCallRequestInfo } from '../scheduler/types.js';
import { type Message } from '../confirmation-bus/types.js';
@@ -146,28 +152,55 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
context.config.getAgentRegistry().getAllAgentNames(),
);
const registerToolByName = (toolName: string) => {
const registerToolInstance = (tool: AnyDeclarativeTool) => {
// Check if the tool is a subagent to prevent recursion.
// We do not allow agents to call other agents.
if (allAgentNames.has(toolName)) {
if (allAgentNames.has(tool.name)) {
debugLogger.warn(
`[LocalAgentExecutor] Skipping subagent tool '${toolName}' for agent '${definition.name}' to prevent recursion.`,
`[LocalAgentExecutor] Skipping subagent tool '${tool.name}' for agent '${definition.name}' to prevent recursion.`,
);
return;
}
agentToolRegistry.registerTool(tool);
};
const registerToolByName = (toolName: string) => {
// Handle global wildcard
if (toolName === '*') {
for (const tool of parentToolRegistry.getAllTools()) {
registerToolInstance(tool);
}
return;
}
// Handle MCP wildcards
if (isMcpToolName(toolName)) {
if (toolName === `${MCP_TOOL_PREFIX}*`) {
for (const tool of parentToolRegistry.getAllTools()) {
if (tool instanceof DiscoveredMCPTool) {
registerToolInstance(tool);
}
}
return;
}
const parsed = parseMcpToolName(toolName);
if (parsed.serverName && parsed.toolName === '*') {
for (const tool of parentToolRegistry.getToolsByServer(
parsed.serverName,
)) {
registerToolInstance(tool);
}
return;
}
}
// If the tool is referenced by name, retrieve it from the parent
// registry and register it with the agent's isolated registry.
const tool = parentToolRegistry.getTool(toolName);
if (tool) {
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);
}
registerToolInstance(tool);
}
};
@@ -1175,10 +1208,9 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const { toolConfig, outputConfig } = this.definition;
if (toolConfig) {
const toolNamesToLoad: string[] = [];
for (const toolRef of toolConfig.tools) {
if (typeof toolRef === 'string') {
toolNamesToLoad.push(toolRef);
// The names were already expanded and loaded into the registry during creation.
} else if (typeof toolRef === 'object' && 'schema' in toolRef) {
// Tool instance with an explicit schema property.
toolsList.push(toolRef.schema);
@@ -1187,10 +1219,8 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
toolsList.push(toolRef);
}
}
// Add schemas from tools that were registered by name.
toolsList.push(
...this.toolRegistry.getFunctionDeclarationsFiltered(toolNamesToLoad),
);
// Add schemas from tools that were explicitly registered by name or wildcard.
toolsList.push(...this.toolRegistry.getFunctionDeclarations());
}
// Always inject complete_task.

View File

@@ -58,6 +58,7 @@ export function parseMcpToolName(name: string): {
// Remove the prefix
const withoutPrefix = name.slice(MCP_TOOL_PREFIX.length);
// The first segment is the server name, the rest is the tool name
// Must be strictly `server_tool` where neither are empty
const match = withoutPrefix.match(/^([^_]+)_(.+)$/);
if (match) {
return {
@@ -390,25 +391,6 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
`${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${this.serverToolName}`,
);
}
asFullyQualifiedTool(): DiscoveredMCPTool {
return new DiscoveredMCPTool(
this.mcpTool,
this.serverName,
this.serverToolName,
this.description,
this.parameterSchema,
this.messageBus,
this.trust,
this.isReadOnly,
this.getFullyQualifiedName(),
this.cliConfig,
this.extensionName,
this.extensionId,
this._toolAnnotations,
);
}
protected createInvocation(
params: ToolParams,
messageBus: MessageBus,

View File

@@ -25,7 +25,8 @@ vi.mock('./tool-names.js', async (importOriginal) => {
...actual,
TOOL_LEGACY_ALIASES: mockedAliases,
isValidToolName: vi.fn().mockImplementation((name: string, options) => {
if (mockedAliases[name]) return true;
if (Object.prototype.hasOwnProperty.call(mockedAliases, name))
return true;
return actual.isValidToolName(name, options);
}),
getToolAliases: vi.fn().mockImplementation((name: string) => {
@@ -55,11 +56,9 @@ describe('tool-names', () => {
expect(isValidToolName(`${DISCOVERED_TOOL_PREFIX}my_tool`)).toBe(true);
});
it('should validate MCP tool names (server__tool)', () => {
expect(isValidToolName('server__tool')).toBe(true);
expect(isValidToolName('my-server__my-tool')).toBe(true);
expect(isValidToolName('my.server__my:tool')).toBe(true);
expect(isValidToolName('my-server...truncated__tool')).toBe(true);
it('should validate modern MCP FQNs (mcp_server_tool)', () => {
expect(isValidToolName('mcp_server_tool')).toBe(true);
expect(isValidToolName('mcp_my-server_my-tool')).toBe(true);
});
it('should validate legacy tool aliases', async () => {
@@ -69,28 +68,33 @@ describe('tool-names', () => {
}
});
it('should reject invalid tool names', () => {
expect(isValidToolName('')).toBe(false);
expect(isValidToolName('invalid-name')).toBe(false);
expect(isValidToolName('server__')).toBe(false);
expect(isValidToolName('__tool')).toBe(false);
expect(isValidToolName('server__tool__extra')).toBe(false);
it('should return false for invalid tool names', () => {
expect(isValidToolName('invalid-tool-name')).toBe(false);
expect(isValidToolName('mcp_server')).toBe(false);
expect(isValidToolName('mcp__tool')).toBe(false);
expect(isValidToolName('mcp_invalid server_tool')).toBe(false);
expect(isValidToolName('mcp_server_invalid tool')).toBe(false);
expect(isValidToolName('mcp_server_')).toBe(false);
});
it('should handle wildcards when allowed', () => {
// Default: not allowed
expect(isValidToolName('*')).toBe(false);
expect(isValidToolName('server__*')).toBe(false);
expect(isValidToolName('mcp_*')).toBe(false);
expect(isValidToolName('mcp_server_*')).toBe(false);
// Explicitly allowed
expect(isValidToolName('*', { allowWildcards: true })).toBe(true);
expect(isValidToolName('server__*', { allowWildcards: true })).toBe(true);
expect(isValidToolName('mcp_*', { allowWildcards: true })).toBe(true);
expect(isValidToolName('mcp_server_*', { allowWildcards: true })).toBe(
true,
);
// Invalid wildcards
expect(isValidToolName('__*', { allowWildcards: true })).toBe(false);
expect(isValidToolName('server__tool*', { allowWildcards: true })).toBe(
false,
);
expect(isValidToolName('mcp__*', { allowWildcards: true })).toBe(false);
expect(
isValidToolName('mcp_server_tool*', { allowWildcards: true }),
).toBe(false);
});
});

View File

@@ -221,6 +221,12 @@ export const DISCOVERED_TOOL_PREFIX = 'discovered_tool_';
/**
* List of all built-in tool names.
*/
import {
isMcpToolName,
parseMcpToolName,
MCP_TOOL_PREFIX,
} from './mcp-tool.js';
export const ALL_BUILTIN_TOOL_NAMES = [
GLOB_TOOL_NAME,
WRITE_TODOS_TOOL_NAME,
@@ -290,25 +296,44 @@ export function isValidToolName(
return true;
}
// MCP tools (format: server__tool)
if (name.includes('__')) {
const parts = name.split('__');
if (parts.length !== 2 || parts[0].length === 0 || parts[1].length === 0) {
// Handle standard MCP FQNs (mcp_server_tool or wildcards mcp_*, mcp_server_*)
if (isMcpToolName(name)) {
// Global wildcard: mcp_*
if (name === `${MCP_TOOL_PREFIX}*` && options.allowWildcards) {
return true;
}
// Explicitly reject names with empty server component (e.g. mcp__tool)
if (name.startsWith(`${MCP_TOOL_PREFIX}_`)) {
return false;
}
const server = parts[0];
const tool = parts[1];
const parsed = parseMcpToolName(name);
// Ensure that both components are populated. parseMcpToolName splits at the second _,
// so `mcp__tool` has serverName="", toolName="tool"
if (parsed.serverName && parsed.toolName) {
// Basic slug validation for server and tool names.
// We allow dots (.) and colons (:) as they are valid in function names and
// used for truncation markers.
const slugRegex = /^[a-z0-9_.:-]+$/i;
if (tool === '*') {
return !!options.allowWildcards;
if (!slugRegex.test(parsed.serverName)) {
return false;
}
if (parsed.toolName === '*') {
return options.allowWildcards === true;
}
// A tool name consisting only of underscores is invalid.
if (/^_*$/.test(parsed.toolName)) {
return false;
}
return slugRegex.test(parsed.toolName);
}
// Basic slug validation for server and tool names.
// We allow dots (.) and colons (:) as they are valid in function names and
// used for truncation markers.
const slugRegex = /^[a-z0-9_.:-]+$/i;
return slugRegex.test(server) && slugRegex.test(tool);
return false;
}
return false;

View File

@@ -310,13 +310,13 @@ describe('ToolRegistry', () => {
excludedTools: ['tool-a'],
},
{
name: 'should match simple MCP tool names, when qualified or unqualified',
tools: [mcpTool, mcpTool.asFullyQualifiedTool()],
name: 'should match simple MCP tool names',
tools: [mcpTool],
excludedTools: [mcpTool.name],
},
{
name: 'should match qualified MCP tool names when qualified or unqualified',
tools: [mcpTool, mcpTool.asFullyQualifiedTool()],
name: 'should match qualified MCP tool names',
tools: [mcpTool],
excludedTools: [mcpTool.name],
},
{
@@ -414,9 +414,9 @@ describe('ToolRegistry', () => {
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'desc');
// Register same MCP tool twice (one as alias, one as qualified)
// Register same MCP tool twice
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());
const toolNames = toolRegistry.getAllToolNames();
expect(toolNames).toEqual([`mcp_${serverName}_${toolName}`]);
@@ -698,9 +698,8 @@ describe('ToolRegistry', () => {
const toolName = 'my-tool';
const mcpTool = createMCPTool(serverName, toolName, 'description');
// Register both alias and qualified
toolRegistry.registerTool(mcpTool);
toolRegistry.registerTool(mcpTool.asFullyQualifiedTool());
toolRegistry.registerTool(mcpTool);
const declarations = toolRegistry.getFunctionDeclarations();
expect(declarations).toHaveLength(1);

View File

@@ -222,14 +222,10 @@ export class ToolRegistry {
*/
registerTool(tool: AnyDeclarativeTool): void {
if (this.allKnownTools.has(tool.name)) {
if (tool instanceof DiscoveredMCPTool) {
tool = tool.asFullyQualifiedTool();
} else {
// Decide on behavior: throw error, log warning, or allow overwrite
debugLogger.warn(
`Tool with name "${tool.name}" is already registered. Overwriting.`,
);
}
// Decide on behavior: throw error, log warning, or allow overwrite
debugLogger.warn(
`Tool with name "${tool.name}" is already registered. Overwriting.`,
);
}
this.allKnownTools.set(tool.name, tool);
}
@@ -594,7 +590,17 @@ export class ToolRegistry {
for (const name of toolNames) {
const tool = this.getTool(name);
if (tool) {
declarations.push(tool.getSchema(modelId));
let schema = tool.getSchema(modelId);
// Ensure the schema name matches the qualified name for MCP tools
if (tool instanceof DiscoveredMCPTool) {
schema = {
...schema,
name: tool.getFullyQualifiedName(),
};
}
declarations.push(schema);
}
}
return declarations;
@@ -670,17 +676,6 @@ export class ToolRegistry {
}
}
if (!tool && name.includes('__')) {
for (const t of this.allKnownTools.values()) {
if (t instanceof DiscoveredMCPTool) {
if (t.getFullyQualifiedName() === name) {
tool = t;
break;
}
}
}
}
if (tool && this.isActiveTool(tool)) {
return tool;
}