From d2ae869bb406ace0498de3a3b5e11a3ef191bfde Mon Sep 17 00:00:00 2001
From: Jacob MacDonald <jakemac@google.com>
Date: Wed, 3 Sep 2025 10:34:32 -0700
Subject: [PATCH] Simplify MCP server timeout configuration (#7661)

---
 packages/core/src/tools/mcp-client.ts    | 17 +++--------------
 packages/core/src/tools/mcp-tool.test.ts | 18 ------------------
 packages/core/src/tools/mcp-tool.ts      |  4 ----
 3 files changed, 3 insertions(+), 36 deletions(-)

diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts
index d34c413065..e434ead39e 100644
--- a/packages/core/src/tools/mcp-client.ts
+++ b/packages/core/src/tools/mcp-client.ts
@@ -644,7 +644,9 @@ export async function discoverTools(
   cliConfig: Config,
 ): Promise<DiscoveredMCPTool[]> {
   try {
-    const mcpCallableTool = mcpToTool(mcpClient);
+    const mcpCallableTool = mcpToTool(mcpClient, {
+      timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
+    });
     const tool = await mcpCallableTool.tool();
 
     if (!Array.isArray(tool.functionDeclarations)) {
@@ -675,7 +677,6 @@ export async function discoverTools(
             funcDecl.name!,
             funcDecl.description ?? '',
             funcDecl.parametersJsonSchema ?? { type: 'object', properties: {} },
-            mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
             mcpServerConfig.trust,
             undefined,
             cliConfig,
@@ -871,18 +872,6 @@ export async function connectToMcpServer(
     unlistenDirectories = undefined;
   };
 
-  // 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 {
     const transport = await createTransport(
       mcpServerName,
diff --git a/packages/core/src/tools/mcp-tool.test.ts b/packages/core/src/tools/mcp-tool.test.ts
index 0e9396b040..9fb155e4c8 100644
--- a/packages/core/src/tools/mcp-tool.test.ts
+++ b/packages/core/src/tools/mcp-tool.test.ts
@@ -98,20 +98,6 @@ describe('DiscoveredMCPTool', () => {
       expect(tool.schema.parameters).toBeUndefined();
       expect(tool.schema.parametersJsonSchema).toEqual(inputSchema);
       expect(tool.serverToolName).toBe(serverToolName);
-      expect(tool.timeout).toBeUndefined();
-    });
-
-    it('should accept and store a custom timeout', () => {
-      const customTimeout = 5000;
-      const toolWithTimeout = new DiscoveredMCPTool(
-        mockCallableToolInstance,
-        serverName,
-        serverToolName,
-        baseDescription,
-        inputSchema,
-        customTimeout,
-      );
-      expect(toolWithTimeout.timeout).toBe(customTimeout);
     });
   });
 
@@ -596,7 +582,6 @@ describe('DiscoveredMCPTool', () => {
         serverToolName,
         baseDescription,
         inputSchema,
-        undefined,
         true,
         undefined,
         { isTrustedFolder: () => true } as any,
@@ -761,7 +746,6 @@ describe('DiscoveredMCPTool', () => {
         serverToolName,
         baseDescription,
         inputSchema,
-        undefined,
         true, // trust = true
         undefined,
         mockConfig(true) as any, // isTrustedFolder = true
@@ -779,7 +763,6 @@ describe('DiscoveredMCPTool', () => {
         serverToolName,
         baseDescription,
         inputSchema,
-        undefined,
         true, // trust = true
         undefined,
         mockConfig(false) as any, // isTrustedFolder = false
@@ -799,7 +782,6 @@ describe('DiscoveredMCPTool', () => {
         serverToolName,
         baseDescription,
         inputSchema,
-        undefined,
         false, // trust = false
         undefined,
         mockConfig(true) as any, // isTrustedFolder = true
diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts
index 66eee800bb..b53487799b 100644
--- a/packages/core/src/tools/mcp-tool.ts
+++ b/packages/core/src/tools/mcp-tool.ts
@@ -68,7 +68,6 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation<
     readonly serverName: string,
     readonly serverToolName: string,
     readonly displayName: string,
-    readonly timeout?: number,
     readonly trust?: boolean,
     params: ToolParams = {},
     private readonly cliConfig?: Config,
@@ -182,7 +181,6 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
     readonly serverToolName: string,
     description: string,
     override readonly parameterSchema: unknown,
-    readonly timeout?: number,
     readonly trust?: boolean,
     nameOverride?: string,
     private readonly cliConfig?: Config,
@@ -205,7 +203,6 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
       this.serverToolName,
       this.description,
       this.parameterSchema,
-      this.timeout,
       this.trust,
       `${this.serverName}__${this.serverToolName}`,
       this.cliConfig,
@@ -220,7 +217,6 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool<
       this.serverName,
       this.serverToolName,
       this.displayName,
-      this.timeout,
       this.trust,
       params,
       this.cliConfig,