Fix mcp tool lookup in tool registry (#17054)

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

packages/core/src/tools/mcp-tool.ts

@@ -264,6 +264,10 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
     return `${this.serverName}__`;
   }
 
+  getFullyQualifiedName(): string {
+    return `${this.getFullyQualifiedPrefix()}${generateValidName(this.serverToolName)}`;
+  }
+
   asFullyQualifiedTool(): DiscoveredMCPTool {
     return new DiscoveredMCPTool(
       this.mcpTool,
@@ -273,7 +277,7 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
       this.parameterSchema,
       this.messageBus,
       this.trust,
-      `${this.getFullyQualifiedPrefix()}${this.serverToolName}`,
+      this.getFullyQualifiedName(),
       this.cliConfig,
       this.extensionName,
       this.extensionId,

packages/core/src/tools/tool-registry.test.ts

@@ -525,6 +525,51 @@ describe('ToolRegistry', () => {
     });
   });
 
+  describe('getTool', () => {
+    it('should retrieve an MCP tool by its fully qualified name even if registered with simple name', () => {
+      const serverName = 'my-server';
+      const toolName = 'my-tool';
+      const mcpTool = createMCPTool(serverName, toolName, 'description');
+
+      // Register tool (will be registered as 'my-tool' since no conflict)
+      toolRegistry.registerTool(mcpTool);
+
+      // Verify it is available as 'my-tool'
+      expect(toolRegistry.getTool('my-tool')).toBeDefined();
+      expect(toolRegistry.getTool('my-tool')?.name).toBe('my-tool');
+
+      // Verify it is available as 'my-server__my-tool'
+      const fullyQualifiedName = `${serverName}__${toolName}`;
+      const retrievedTool = toolRegistry.getTool(fullyQualifiedName);
+
+      expect(retrievedTool).toBeDefined();
+      // The returned tool object is the same, so its name property is still 'my-tool'
+      expect(retrievedTool?.name).toBe('my-tool');
+    });
+
+    it('should retrieve an MCP tool by its fully qualified name when tool name has special characters', () => {
+      const serverName = 'my-server';
+      // Use a space which is invalid and will be replaced by underscore
+      const toolName = 'my tool';
+      const validToolName = 'my_tool';
+      const mcpTool = createMCPTool(serverName, toolName, 'description');
+
+      // Register tool (will be registered as sanitized name)
+      toolRegistry.registerTool(mcpTool);
+
+      // Verify it is available as sanitized name
+      expect(toolRegistry.getTool(validToolName)).toBeDefined();
+      expect(toolRegistry.getTool(validToolName)?.name).toBe(validToolName);
+
+      // Verify it is available as 'my-server__my_tool'
+      const fullyQualifiedName = `${serverName}__${validToolName}`;
+      const retrievedTool = toolRegistry.getTool(fullyQualifiedName);
+
+      expect(retrievedTool).toBeDefined();
+      expect(retrievedTool?.name).toBe(validToolName);
+    });
+  });
+
   describe('DiscoveredToolInvocation', () => {
     it('should return the stringified params from getDescription', () => {
       const tool = new DiscoveredTool(

packages/core/src/tools/tool-registry.ts

@@ -530,7 +530,18 @@ export class ToolRegistry {
    * Get the definition of a specific tool.
    */
   getTool(name: string): AnyDeclarativeTool | undefined {
-    const tool = this.allKnownTools.get(name);
+    let tool = this.allKnownTools.get(name);
+    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;
     }