diff --git a/src/core/assistant-message/NativeToolCallParser.ts b/src/core/assistant-message/NativeToolCallParser.ts index e88c26a913..58bb17a486 100644 --- a/src/core/assistant-message/NativeToolCallParser.ts +++ b/src/core/assistant-message/NativeToolCallParser.ts @@ -12,6 +12,7 @@ import type { ApiStreamToolCallDeltaChunk, ApiStreamToolCallEndChunk, } from "../../api/transform/stream" +import { McpToolRegistry } from "../../services/mcp/McpToolRegistry" /** * Helper type to extract properly typed native arguments for a given tool. @@ -779,12 +780,12 @@ export class NativeToolCallParser { } /** - * Parse dynamic MCP tools (named mcp_serverName_toolName). - * These are generated dynamically by getMcpServerTools() and are returned - * as McpToolUse objects that preserve the original tool name. + * Parse dynamic MCP tools (named mcp_0, mcp_1, etc.). + * These are generated dynamically by getMcpServerTools() using McpToolRegistry + * and are returned as McpToolUse objects that preserve the original tool name. * * In native mode, MCP tools are NOT converted to use_mcp_tool - they keep - * their original name so it appears correctly in API conversation history. + * their registry name so it appears correctly in API conversation history. * The use_mcp_tool wrapper is only used in XML mode. */ public static parseDynamicMcpTool(toolCall: { id: string; name: string; arguments: string }): McpToolUse | null { @@ -792,27 +793,19 @@ export class NativeToolCallParser { // Parse the arguments - these are the actual tool arguments passed directly const args = JSON.parse(toolCall.arguments || "{}") - // Extract server_name and tool_name from the tool name itself - // Format: mcp_serverName_toolName - const nameParts = toolCall.name.split("_") - if (nameParts.length < 3 || nameParts[0] !== "mcp") { - console.error(`Invalid dynamic MCP tool name format: ${toolCall.name}`) + // Look up the original server/tool names from the registry + const entry = McpToolRegistry.lookup(toolCall.name) + if (!entry) { + console.error(`Unknown MCP tool (not in registry): ${toolCall.name}`) return null } - // Server name is the second part, tool name is everything after - const serverName = nameParts[1] - const toolName = nameParts.slice(2).join("_") - - if (!serverName || !toolName) { - console.error(`Could not extract server_name or tool_name from: ${toolCall.name}`) - return null - } + const { serverName, toolName } = entry const result: McpToolUse = { type: "mcp_tool_use" as const, id: toolCall.id, - // Keep the original tool name (e.g., "mcp_serverName_toolName") for API history + // Keep the API name (e.g., "mcp_0") for API history name: toolCall.name, serverName, toolName, diff --git a/src/core/prompts/tools/native-tools/mcp_server.ts b/src/core/prompts/tools/native-tools/mcp_server.ts index 36efdf83fb..437ca546f5 100644 --- a/src/core/prompts/tools/native-tools/mcp_server.ts +++ b/src/core/prompts/tools/native-tools/mcp_server.ts @@ -1,5 +1,6 @@ import type OpenAI from "openai" import { McpHub } from "../../../../services/mcp/McpHub" +import { McpToolRegistry } from "../../../../services/mcp/McpToolRegistry" /** * Dynamically generates native tool definitions for all enabled tools across connected MCP servers. @@ -30,8 +31,8 @@ export function getMcpServerTools(mcpHub?: McpHub): OpenAI.Chat.ChatCompletionTo const toolInputRequired = (originalSchema?.required ?? []) as string[] // Build parameters directly from the tool's input schema. - // The server_name and tool_name are encoded in the function name itself - // (e.g., mcp_serverName_toolName), so they don't need to be in the arguments. + // The server_name and tool_name are registered in McpToolRegistry, + // which returns an API-compatible name (e.g., mcp_0, mcp_1). const parameters: OpenAI.FunctionParameters = { type: "object", properties: toolInputProps, @@ -43,11 +44,13 @@ export function getMcpServerTools(mcpHub?: McpHub): OpenAI.Chat.ChatCompletionTo parameters.required = toolInputRequired } - // Use mcp_ prefix to identify dynamic MCP tools + // Register tool in McpToolRegistry to get an API-compatible name. + // This avoids issues with dots in MCP tool names (e.g., "agent.describe") + // which violate Bedrock's [a-zA-Z0-9_-]+ constraint. const toolDefinition: OpenAI.Chat.ChatCompletionTool = { type: "function", function: { - name: `mcp_${server.name}_${tool.name}`, + name: McpToolRegistry.register(server.name, tool.name), description: tool.description, parameters: parameters, }, diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index f7e4946f53..feffae5970 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -76,6 +76,7 @@ import { UrlContentFetcher } from "../../services/browser/UrlContentFetcher" import { BrowserSession } from "../../services/browser/BrowserSession" import { McpHub } from "../../services/mcp/McpHub" import { McpServerManager } from "../../services/mcp/McpServerManager" +import { McpToolRegistry } from "../../services/mcp/McpToolRegistry" import { RepoPerTaskCheckpointService } from "../../services/checkpoints" // integrations @@ -2389,6 +2390,8 @@ export class Task extends EventEmitter implements TaskLike { // Clear any leftover streaming tool call state from previous interrupted streams NativeToolCallParser.clearAllStreamingToolCalls() NativeToolCallParser.clearRawChunkState() + // Clear MCP tool registry to ensure fresh mappings for this request + McpToolRegistry.clear() await this.diffViewProvider.reset() diff --git a/src/services/mcp/McpToolRegistry.ts b/src/services/mcp/McpToolRegistry.ts new file mode 100644 index 0000000000..2525df8433 --- /dev/null +++ b/src/services/mcp/McpToolRegistry.ts @@ -0,0 +1,108 @@ +/** + * Registry for mapping MCP tool names to API-compatible identifiers. + * + * Problem: AWS Bedrock (and potentially other providers) only accept tool names + * matching [a-zA-Z0-9_-]+, but MCP tools can have dots in their names + * (e.g., "agent-block.describe", "debug.state.get_full"). + * + * Solution: This registry assigns simple numeric identifiers (mcp_0, mcp_1, etc.) + * to each unique server+tool combination, allowing bidirectional lookup. + * + * Lifecycle: The registry should be cleared at the start of each API request, + * as tools are re-registered when building the prompt. This is safe because + * the response parsing happens in the same request cycle. + * + * Performance: All operations are O(1) Map lookups - no string manipulation. + */ + +export interface McpToolEntry { + serverName: string + toolName: string +} + +/** + * Static registry for MCP tool name mapping. + * Use register() when building tools, lookup() when parsing responses. + */ +export class McpToolRegistry { + /** Maps API name (e.g., "mcp_0") to original server/tool names */ + private static registry = new Map() + + /** Maps "serverName:toolName" to API name for idempotent registration */ + private static reverseIndex = new Map() + + /** Counter for generating unique API names */ + private static nextId = 0 + + /** + * Register an MCP tool and get back an API-compatible name. + * Idempotent: returns existing name if already registered. + * + * @param serverName - The MCP server name + * @param toolName - The MCP tool name (may contain dots) + * @returns API-compatible name (e.g., "mcp_0") + */ + public static register(serverName: string, toolName: string): string { + const key = `${serverName}:${toolName}` + + // Return existing if already registered + const existing = this.reverseIndex.get(key) + if (existing) { + return existing + } + + // Create new API-compatible name + const apiName = `mcp_${this.nextId++}` + + this.registry.set(apiName, { serverName, toolName }) + this.reverseIndex.set(key, apiName) + + return apiName + } + + /** + * Look up the original server/tool names from an API name. + * + * @param apiName - The API name (e.g., "mcp_0") + * @returns The original entry, or undefined if not found + */ + public static lookup(apiName: string): McpToolEntry | undefined { + return this.registry.get(apiName) + } + + /** + * Check if an API name is a registered MCP tool. + * + * @param apiName - The name to check + * @returns true if this is a registered MCP tool + */ + public static isRegistered(apiName: string): boolean { + return this.registry.has(apiName) + } + + /** + * Clear all registrations. + * Should be called at the start of each API request to prevent stale mappings. + */ + public static clear(): void { + this.registry.clear() + this.reverseIndex.clear() + this.nextId = 0 + } + + /** + * Get the number of registered tools. + * Useful for debugging and testing. + */ + public static size(): number { + return this.registry.size + } + + /** + * Get all registered entries for debugging. + * Returns a shallow copy to prevent external mutation. + */ + public static getAll(): Map { + return new Map(this.registry) + } +} diff --git a/src/services/mcp/__tests__/McpToolRegistry.spec.ts b/src/services/mcp/__tests__/McpToolRegistry.spec.ts new file mode 100644 index 0000000000..9adcbd984f --- /dev/null +++ b/src/services/mcp/__tests__/McpToolRegistry.spec.ts @@ -0,0 +1,246 @@ +import { McpToolRegistry } from "../McpToolRegistry" + +describe("McpToolRegistry", () => { + beforeEach(() => { + // Clear the registry before each test to ensure isolation + McpToolRegistry.clear() + }) + + describe("register", () => { + it("should register a tool and return an API-compatible name", () => { + const apiName = McpToolRegistry.register("agent-block", "describe") + + expect(apiName).toBe("mcp_0") + }) + + it("should return different names for different tools", () => { + const name1 = McpToolRegistry.register("agent-block", "describe") + const name2 = McpToolRegistry.register("agent-block", "execute_task") + const name3 = McpToolRegistry.register("debug", "state.get_full") + + expect(name1).toBe("mcp_0") + expect(name2).toBe("mcp_1") + expect(name3).toBe("mcp_2") + }) + + it("should be idempotent - return same name for same server/tool combination", () => { + const name1 = McpToolRegistry.register("agent-block", "describe") + const name2 = McpToolRegistry.register("agent-block", "describe") + const name3 = McpToolRegistry.register("agent-block", "describe") + + expect(name1).toBe("mcp_0") + expect(name2).toBe("mcp_0") + expect(name3).toBe("mcp_0") + + // Should still be only 1 entry + expect(McpToolRegistry.size()).toBe(1) + }) + + it("should handle tool names with dots (the original problem)", () => { + // These names would fail Bedrock's [a-zA-Z0-9_-]+ constraint + const name1 = McpToolRegistry.register("agent-block-debug", "state.get_full") + const name2 = McpToolRegistry.register("agent-block-debug", "logs.query") + const name3 = McpToolRegistry.register("agent-block-debug", "cdp.passthrough") + + // API names are simple numeric identifiers that pass the constraint + expect(name1).toMatch(/^mcp_\d+$/) + expect(name2).toMatch(/^mcp_\d+$/) + expect(name3).toMatch(/^mcp_\d+$/) + + // All should be different + expect(new Set([name1, name2, name3]).size).toBe(3) + }) + + it("should handle server names with special characters", () => { + const name = McpToolRegistry.register("my-special_server", "tool.with.dots") + + expect(name).toBe("mcp_0") + + const entry = McpToolRegistry.lookup(name) + expect(entry).toEqual({ + serverName: "my-special_server", + toolName: "tool.with.dots", + }) + }) + }) + + describe("lookup", () => { + it("should return the original server and tool names", () => { + const apiName = McpToolRegistry.register("agent-block", "describe") + + const entry = McpToolRegistry.lookup(apiName) + + expect(entry).toEqual({ + serverName: "agent-block", + toolName: "describe", + }) + }) + + it("should return undefined for unregistered names", () => { + const entry = McpToolRegistry.lookup("mcp_999") + + expect(entry).toBeUndefined() + }) + + it("should return undefined for non-mcp names", () => { + const entry = McpToolRegistry.lookup("read_file") + + expect(entry).toBeUndefined() + }) + + it("should preserve original tool name with dots", () => { + const apiName = McpToolRegistry.register("debug", "state.get_full") + + const entry = McpToolRegistry.lookup(apiName) + + expect(entry?.toolName).toBe("state.get_full") + }) + }) + + describe("isRegistered", () => { + it("should return true for registered names", () => { + const apiName = McpToolRegistry.register("server", "tool") + + expect(McpToolRegistry.isRegistered(apiName)).toBe(true) + }) + + it("should return false for unregistered names", () => { + expect(McpToolRegistry.isRegistered("mcp_0")).toBe(false) + expect(McpToolRegistry.isRegistered("read_file")).toBe(false) + }) + }) + + describe("clear", () => { + it("should remove all registrations", () => { + McpToolRegistry.register("server1", "tool1") + McpToolRegistry.register("server2", "tool2") + expect(McpToolRegistry.size()).toBe(2) + + McpToolRegistry.clear() + + expect(McpToolRegistry.size()).toBe(0) + expect(McpToolRegistry.lookup("mcp_0")).toBeUndefined() + expect(McpToolRegistry.lookup("mcp_1")).toBeUndefined() + }) + + it("should reset the ID counter", () => { + McpToolRegistry.register("server", "tool") + expect(McpToolRegistry.lookup("mcp_0")).toBeDefined() + + McpToolRegistry.clear() + + // After clear, next registration should start from 0 again + const apiName = McpToolRegistry.register("new-server", "new-tool") + expect(apiName).toBe("mcp_0") + }) + + it("should allow re-registration with same names after clear", () => { + const apiName1 = McpToolRegistry.register("server", "tool") + expect(apiName1).toBe("mcp_0") + + McpToolRegistry.clear() + + // Same server/tool should get mcp_0 again after clear + const apiName2 = McpToolRegistry.register("server", "tool") + expect(apiName2).toBe("mcp_0") + }) + }) + + describe("size", () => { + it("should return 0 for empty registry", () => { + expect(McpToolRegistry.size()).toBe(0) + }) + + it("should return correct count", () => { + McpToolRegistry.register("s1", "t1") + expect(McpToolRegistry.size()).toBe(1) + + McpToolRegistry.register("s1", "t2") + expect(McpToolRegistry.size()).toBe(2) + + // Idempotent - same registration shouldn't increase count + McpToolRegistry.register("s1", "t1") + expect(McpToolRegistry.size()).toBe(2) + }) + }) + + describe("getAll", () => { + it("should return empty map for empty registry", () => { + const all = McpToolRegistry.getAll() + + expect(all.size).toBe(0) + }) + + it("should return all registrations", () => { + McpToolRegistry.register("server1", "tool1") + McpToolRegistry.register("server2", "tool2") + + const all = McpToolRegistry.getAll() + + expect(all.size).toBe(2) + expect(all.get("mcp_0")).toEqual({ serverName: "server1", toolName: "tool1" }) + expect(all.get("mcp_1")).toEqual({ serverName: "server2", toolName: "tool2" }) + }) + + it("should return a copy that doesn't affect internal state", () => { + McpToolRegistry.register("server", "tool") + + const all = McpToolRegistry.getAll() + all.clear() // Modify the returned map + + // Internal state should be unaffected + expect(McpToolRegistry.size()).toBe(1) + expect(McpToolRegistry.lookup("mcp_0")).toBeDefined() + }) + }) + + describe("integration scenarios", () => { + it("should handle typical request cycle", () => { + // Simulate building tools for API request + const tools = [ + { server: "agent-block", tool: "agent-block.describe" }, + { server: "agent-block", tool: "agent-block.execute_task" }, + { server: "debug", tool: "debug.state.get_full" }, + ] + + // Register tools (simulating getMcpServerTools) + const registeredNames = tools.map((t) => McpToolRegistry.register(t.server, t.tool)) + + expect(registeredNames).toEqual(["mcp_0", "mcp_1", "mcp_2"]) + + // Simulate parsing tool response (simulating parseDynamicMcpTool) + const entry = McpToolRegistry.lookup("mcp_1") + expect(entry).toEqual({ + serverName: "agent-block", + toolName: "agent-block.execute_task", + }) + + // Simulate new request cycle + McpToolRegistry.clear() + + // Tools get re-registered (order might differ) + const newTools = [ + { server: "debug", tool: "debug.logs.query" }, + { server: "agent-block", tool: "agent-block.describe" }, + ] + + const newNames = newTools.map((t) => McpToolRegistry.register(t.server, t.tool)) + expect(newNames).toEqual(["mcp_0", "mcp_1"]) + + // Lookup works with new mappings + expect(McpToolRegistry.lookup("mcp_0")).toEqual({ + serverName: "debug", + toolName: "debug.logs.query", + }) + }) + + it("should handle edge case of empty server or tool name", () => { + // While unlikely, the registry should handle edge cases gracefully + const apiName = McpToolRegistry.register("", "tool") + expect(apiName).toBe("mcp_0") + + const entry = McpToolRegistry.lookup(apiName) + expect(entry).toEqual({ serverName: "", toolName: "tool" }) + }) + }) +})