From a2fee6bdd3927a2a08520ab87dabc8fecdb6075b Mon Sep 17 00:00:00 2001
From: Olcan <olcans@gmail.com>
Date: Sun, 8 Jun 2025 21:52:11 -0700
Subject: fix mcp timeouts and missing description on mcp errors (#868)

---
 packages/core/src/tools/mcp-client.ts | 21 ++++++++++++++++++---
 packages/core/src/tools/mcp-tool.ts   |  2 --
 2 files changed, 18 insertions(+), 5 deletions(-)

(limited to 'packages/core/src')

diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts
index 9a02df0c..1b797ba4 100644
--- a/packages/core/src/tools/mcp-client.ts
+++ b/packages/core/src/tools/mcp-client.ts
@@ -13,6 +13,8 @@ import { DiscoveredMCPTool } from './mcp-tool.js';
 import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
 import { ToolRegistry } from './tool-registry.js';
 
+export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
+
 /**
  * Enum representing the connection status of an MCP server
  */
@@ -149,11 +151,24 @@ async function connectAndDiscover(
   const mcpClient = new Client({
     name: 'gemini-cli-mcp-client',
     version: '0.0.1',
-    timeout: mcpServerConfig.timeout,
   });
 
+  // patch Client.callTool to use request timeout as genai McpCallTool.callTool does not do it
+  // TODO: remove this hack once GenAI SDK does callTool with request options
+  if ('callTool' in mcpClient) {
+    const origCallTool = mcpClient.callTool.bind(mcpClient);
+    mcpClient.callTool = function (params, resultSchema, options) {
+      return origCallTool(params, resultSchema, {
+        ...options,
+        timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
+      });
+    };
+  }
+
   try {
-    await mcpClient.connect(transport);
+    await mcpClient.connect(transport, {
+      timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
+    });
     // Connection successful
     updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED);
   } catch (error) {
@@ -242,7 +257,7 @@ async function connectAndDiscover(
           funcDecl.description ?? '',
           parameterSchema,
           funcDecl.name,
-          mcpServerConfig.timeout,
+          mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
           mcpServerConfig.trust,
         ),
       );
diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts
index d5a8ccb5..086db763 100644
--- a/packages/core/src/tools/mcp-tool.ts
+++ b/packages/core/src/tools/mcp-tool.ts
@@ -15,8 +15,6 @@ import { CallableTool, Part, FunctionCall } from '@google/genai';
 
 type ToolParams = Record<string, unknown>;
 
-export const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
-
 export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
   private static readonly allowlist: Set<string> = new Set();
 
-- 
cgit v1.2.3