added timeout setting to mcp server config, also switched to custom config type without "stderr" field that does not make sense in settings (#410)

4 files changed, 45 insertions, 42 deletions

diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts
index fa8a545d..9f00aabb 100644
--- a/packages/cli/src/config/settings.ts
+++ b/packages/cli/src/config/settings.ts
@@ -7,7 +7,7 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import { homedir } from 'os';
-import { StdioServerParameters } from '@modelcontextprotocol/sdk/client/stdio.js';
+import { MCPServerConfig } from '@gemini-code/server/src/config/config.js';
 
 export const SETTINGS_DIRECTORY_NAME = '.gemini';
 export const USER_SETTINGS_DIR = path.join(homedir(), SETTINGS_DIRECTORY_NAME);
@@ -24,7 +24,7 @@ export interface Settings {
   toolDiscoveryCommand?: string;
   toolCallCommand?: string;
   mcpServerCommand?: string;
-  mcpServers?: Record<string, StdioServerParameters>;
+  mcpServers?: Record<string, MCPServerConfig>;
   // Add other settings here.
 }
 
@@ -69,10 +69,10 @@ export class LoadedSettings {
   setValue(
     scope: SettingScope,
     key: keyof Settings,
-    value: string | Record<string, StdioServerParameters> | undefined,
+    value: string | Record<string, MCPServerConfig> | undefined,
   ): void {
     const settingsFile = this.forScope(scope);
-    // @ts-expect-error - value can be string | Record<string, StdioServerParameters>
+    // @ts-expect-error - value can be string | Record<string, MCPServerConfig>
     settingsFile.settings[key] = value;
     this._merged = this.computeMergedSettings();
     saveSettings(settingsFile);
diff --git a/packages/server/src/config/config.ts b/packages/server/src/config/config.ts
index fdd7973e..6093b7ab 100644
--- a/packages/server/src/config/config.ts
+++ b/packages/server/src/config/config.ts
@@ -21,7 +21,16 @@ import { WebFetchTool } from '../tools/web-fetch.js';
 import { ReadManyFilesTool } from '../tools/read-many-files.js';
 import { BaseTool, ToolResult } from '../tools/tools.js';
 import { MemoryTool } from '../tools/memoryTool.js';
-import { StdioServerParameters } from '@modelcontextprotocol/sdk/client/stdio.js';
+
+export class MCPServerConfig {
+  constructor(
+    readonly command: string,
+    readonly args?: string[],
+    readonly env?: Record<string, string>,
+    readonly cwd?: string,
+    readonly timeout?: number,
+  ) {}
+}
 
 export class Config {
   private toolRegistry: ToolRegistry;
@@ -37,9 +46,7 @@ export class Config {
     private readonly toolDiscoveryCommand: string | undefined,
     private readonly toolCallCommand: string | undefined,
     private readonly mcpServerCommand: string | undefined,
-    private readonly mcpServers:
-      | Record<string, StdioServerParameters>
-      | undefined,
+    private readonly mcpServers: Record<string, MCPServerConfig> | undefined,
     private readonly userAgent: string,
     private userMemory: string = '', // Made mutable for refresh
     private geminiMdFileCount: number = 0,
@@ -92,7 +99,7 @@ export class Config {
     return this.mcpServerCommand;
   }
 
-  getMcpServers(): Record<string, StdioServerParameters> | undefined {
+  getMcpServers(): Record<string, MCPServerConfig> | undefined {
     return this.mcpServers;
   }
 
@@ -164,7 +171,7 @@ export function createServerConfig(
   toolDiscoveryCommand?: string,
   toolCallCommand?: string,
   mcpServerCommand?: string,
-  mcpServers?: Record<string, StdioServerParameters>,
+  mcpServers?: Record<string, MCPServerConfig>,
   userAgent?: string,
   userMemory?: string,
   geminiMdFileCount?: number,
diff --git a/packages/server/src/tools/tool-registry.test.ts b/packages/server/src/tools/tool-registry.test.ts
index cf699377..7f3c2e86 100644
--- a/packages/server/src/tools/tool-registry.test.ts
+++ b/packages/server/src/tools/tool-registry.test.ts
@@ -577,7 +577,6 @@ describe('DiscoveredTool', () => {
 });
 
 describe('DiscoveredMCPTool', () => {
-  let config: Config;
   let mockMcpClient: Client;
   const toolName = 'my-mcp-tool';
   const toolDescription = 'An MCP-discovered tool.';
@@ -587,21 +586,6 @@ describe('DiscoveredMCPTool', () => {
   };
 
   beforeEach(() => {
-    const mockTargetDir = '/test/dir';
-    config = new Config(
-      'test-api-key',
-      'test-model',
-      false, // sandbox
-      mockTargetDir, // targetDir
-      false, // debugMode
-      undefined, // question
-      false, // fullContext
-      undefined, // toolDiscoveryCommand
-      undefined, // toolCallCommand
-      undefined, // mcpServerCommand
-      undefined, // mcpServers
-      'TestAgent/1.0', // userAgent
-    );
     mockMcpClient = new Client({
       name: 'test-client',
       version: '0.0.0',
@@ -615,7 +599,6 @@ describe('DiscoveredMCPTool', () => {
   it('constructor should set up properties correctly and enhance description', () => {
     const tool = new DiscoveredMCPTool(
       mockMcpClient,
-      config,
       toolName,
       toolDescription,
       toolInputSchema,
@@ -631,7 +614,6 @@ describe('DiscoveredMCPTool', () => {
   it('execute should call mcpClient.callTool with correct params and return serialized result', async () => {
     const tool = new DiscoveredMCPTool(
       mockMcpClient,
-      config,
       toolName,
       toolDescription,
       toolInputSchema,
@@ -644,10 +626,16 @@ describe('DiscoveredMCPTool', () => {
 
     const result = await tool.execute(params);
 
-    expect(mockMcpClient.callTool).toHaveBeenCalledWith({
-      name: toolName,
-      arguments: params,
-    });
+    expect(mockMcpClient.callTool).toHaveBeenCalledWith(
+      {
+        name: toolName,
+        arguments: params,
+      },
+      undefined,
+      {
+        timeout: 10 * 60 * 1000,
+      },
+    );
     const expectedOutput = JSON.stringify(mcpResult, null, 2);
     expect(result.llmContent).toBe(expectedOutput);
     expect(result.returnDisplay).toBe(expectedOutput);
diff --git a/packages/server/src/tools/tool-registry.ts b/packages/server/src/tools/tool-registry.ts
index d0bdfba0..a544c8c1 100644
--- a/packages/server/src/tools/tool-registry.ts
+++ b/packages/server/src/tools/tool-registry.ts
@@ -14,6 +14,8 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
 type ToolParams = Record<string, unknown>;
 
+const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
+
 export class DiscoveredTool extends BaseTool<ToolParams, ToolResult> {
   constructor(
     private readonly config: Config,
@@ -95,11 +97,11 @@ Signal: Signal number or \`(none)\` if no signal was received.
 export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
   constructor(
     private readonly mcpClient: Client,
-    private readonly config: Config,
     readonly name: string,
     readonly description: string,
     readonly parameterSchema: Record<string, unknown>,
     readonly serverToolName: string,
+    readonly timeout?: number,
   ) {
     description += `
 
@@ -112,10 +114,16 @@ Returns the MCP server response as a json string.
   }
 
   async execute(params: ToolParams): Promise<ToolResult> {
-    const result = await this.mcpClient.callTool({
-      name: this.serverToolName,
-      arguments: params,
-    });
+    const result = await this.mcpClient.callTool(
+      {
+        name: this.serverToolName,
+        arguments: params,
+      },
+      undefined, // skip resultSchema to specify options (RequestOptions)
+      {
+        timeout: this.timeout ?? MCP_TOOL_DEFAULT_TIMEOUT_MSEC,
+      },
+    );
     return {
       llmContent: JSON.stringify(result, null, 2),
       returnDisplay: JSON.stringify(result, null, 2),
@@ -189,17 +197,17 @@ export class ToolRegistry {
         args: args.slice(1),
       };
     }
-    for (const [mcpServerName, mcpServer] of Object.entries(mcpServers)) {
+    for (const [mcpServerName, mcpServerConfig] of Object.entries(mcpServers)) {
       (async () => {
         const mcpClient = new Client({
           name: 'mcp-client',
           version: '0.0.1',
         });
         const transport = new StdioClientTransport({
-          ...mcpServer,
+          ...mcpServerConfig,
           env: {
             ...process.env,
-            ...(mcpServer.env || {}),
+            ...(mcpServerConfig.env || {}),
           } as Record<string, string>,
           stderr: 'pipe',
         });
@@ -208,7 +216,7 @@ export class ToolRegistry {
         } catch (error) {
           console.error(
             `failed to start or connect to MCP server '${mcpServerName}' ` +
-              `${JSON.stringify(mcpServer)}; \n${error}`,
+              `${JSON.stringify(mcpServerConfig)}; \n${error}`,
           );
           // Do not re-throw, let other MCP servers be discovered.
           return; // Exit this async IIFE if connection failed
@@ -246,13 +254,13 @@ export class ToolRegistry {
           this.registerTool(
             new DiscoveredMCPTool(
               mcpClient,
-              this.config,
               Object.keys(mcpServers).length > 1
                 ? mcpServerName + '__' + tool.name
                 : tool.name,
               tool.description ?? '',
               tool.inputSchema,
               tool.name,
+              mcpServerConfig.timeout,
             ),
           );
         }