fix mcp timeouts and missing description on mcp errors (#868)

4 files changed, 23 insertions, 19 deletions

diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
index ae58ed38..8dbcfb87 100644
--- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts
+++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
@@ -249,10 +249,8 @@ export function mapToDisplay(
           trackedCall.request.args,
         );
         renderOutputAsMarkdown = currentToolInstance.isOutputMarkdown;
-      }
-
-      if (trackedCall.status === 'error') {
-        description = '';
+      } else if ('request' in trackedCall && 'args' in trackedCall.request) {
+        description = JSON.stringify(trackedCall.request.args);
       }
 
       const baseDisplayProperties: Omit<
diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts
index f5a3529c..8e3f139b 100644
--- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts
+++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts
@@ -924,7 +924,7 @@ describe('mapToDisplay', () => {
       expectedStatus: ToolCallStatus.Error,
       expectedResultDisplay: 'Error display tool not found',
       expectedName: baseRequest.name,
-      expectedDescription: '',
+      expectedDescription: JSON.stringify(baseRequest.args),
     },
     {
       name: 'error tool execution failed',
@@ -940,7 +940,7 @@ describe('mapToDisplay', () => {
       expectedStatus: ToolCallStatus.Error,
       expectedResultDisplay: 'Execution failed display',
       expectedName: baseTool.displayName, // Changed from baseTool.name
-      expectedDescription: '',
+      expectedDescription: baseTool.getDescription(baseRequest.args),
     },
     {
       name: 'cancelled',
@@ -986,14 +986,7 @@ describe('mapToDisplay', () => {
         expect(toolDisplay.resultDisplay).toBe(expectedResultDisplay);
 
         expect(toolDisplay.name).toBe(expectedName);
-
-        if (status === 'error' && !extraProps?.tool) {
-          expect(toolDisplay.description).toBe('');
-        } else {
-          expect(toolDisplay.description).toBe(
-            expectedDescription ?? baseTool.getDescription(baseRequest.args),
-          );
-        }
+        expect(toolDisplay.description).toBe(expectedDescription);
 
         expect(toolDisplay.renderOutputAsMarkdown).toBe(
           extraProps?.tool?.isOutputMarkdown ?? false,
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();