diff --git a/packages/core/src/agents/agentLoader.ts b/packages/core/src/agents/agentLoader.ts index 12337c6248..e0ccba0782 100644 --- a/packages/core/src/agents/agentLoader.ts +++ b/packages/core/src/agents/agentLoader.ts @@ -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(), diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index cbc6260304..6a9dfe0151 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -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 { 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 { 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 { 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. diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index 523eac62ad..5702f88a52 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -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, diff --git a/packages/core/src/tools/tool-names.test.ts b/packages/core/src/tools/tool-names.test.ts index 8ff871986f..c631541171 100644 --- a/packages/core/src/tools/tool-names.test.ts +++ b/packages/core/src/tools/tool-names.test.ts @@ -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); }); }); diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index 38a868d665..91b0574d9e 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -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; diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts index ea560865e6..21bbb0cc71 100644 --- a/packages/core/src/tools/tool-registry.test.ts +++ b/packages/core/src/tools/tool-registry.test.ts @@ -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); diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index 69695877c2..f8542112bb 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -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; }