From c71d6ddc3b6dd68bde40cc2b208825ba173f3ba1 Mon Sep 17 00:00:00 2001
From: "N. Taylor Mullen" <ntaylormullen@google.com>
Date: Tue, 3 Jun 2025 00:40:51 -0700
Subject: Fix: Ensure MCP tools are discovered from slow-starting servers
 (#717)

---
 packages/core/src/tools/mcp-client.ts | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

(limited to 'packages/core/src/tools/mcp-client.ts')

diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts
index 87835219..1b7823c7 100644
--- a/packages/core/src/tools/mcp-client.ts
+++ b/packages/core/src/tools/mcp-client.ts
@@ -8,15 +8,18 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
 import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
 import { parse } from 'shell-quote';
-import { Config, MCPServerConfig } from '../config/config.js';
+import { MCPServerConfig } from '../config/config.js';
 import { DiscoveredMCPTool } from './mcp-tool.js';
 import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
+import { ToolRegistry } from './tool-registry.js';
 
-export async function discoverMcpTools(config: Config): Promise<void> {
-  const mcpServers = config.getMcpServers() || {};
-
-  if (config.getMcpServerCommand()) {
-    const cmd = config.getMcpServerCommand()!;
+export async function discoverMcpTools(
+  mcpServers: Record<string, MCPServerConfig>,
+  mcpServerCommand: string | undefined,
+  toolRegistry: ToolRegistry,
+): Promise<void> {
+  if (mcpServerCommand) {
+    const cmd = mcpServerCommand;
     const args = parse(cmd, process.env) as string[];
     if (args.some((arg) => typeof arg !== 'string')) {
       throw new Error('failed to parse mcpServerCommand: ' + cmd);
@@ -30,7 +33,7 @@ export async function discoverMcpTools(config: Config): Promise<void> {
 
   const discoveryPromises = Object.entries(mcpServers).map(
     ([mcpServerName, mcpServerConfig]) =>
-      connectAndDiscover(mcpServerName, mcpServerConfig, config),
+      connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry),
   );
   await Promise.all(discoveryPromises);
 }
@@ -38,7 +41,7 @@ export async function discoverMcpTools(config: Config): Promise<void> {
 async function connectAndDiscover(
   mcpServerName: string,
   mcpServerConfig: MCPServerConfig,
-  config: Config,
+  toolRegistry: ToolRegistry,
 ): Promise<void> {
   let transport;
   if (mcpServerConfig.url) {
@@ -90,7 +93,6 @@ async function connectAndDiscover(
     });
   }
 
-  const toolRegistry = await config.getToolRegistry();
   try {
     const mcpCallableTool: CallableTool = mcpToTool(mcpClient);
     const discoveredToolFunctions = await mcpCallableTool.tool();
-- 
cgit v1.2.3