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; + mcpServers?: Record; // Add other settings here. } @@ -69,10 +69,10 @@ export class LoadedSettings { setValue( scope: SettingScope, key: keyof Settings, - value: string | Record | undefined, + value: string | Record | undefined, ): void { const settingsFile = this.forScope(scope); - // @ts-expect-error - value can be string | Record + // @ts-expect-error - value can be string | Record 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, + 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 - | undefined, + private readonly mcpServers: Record | 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 | undefined { + getMcpServers(): Record | undefined { return this.mcpServers; } @@ -164,7 +171,7 @@ export function createServerConfig( toolDiscoveryCommand?: string, toolCallCommand?: string, mcpServerCommand?: string, - mcpServers?: Record, + mcpServers?: Record, 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; +const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes + export class DiscoveredTool extends BaseTool { constructor( private readonly config: Config, @@ -95,11 +97,11 @@ Signal: Signal number or \`(none)\` if no signal was received. export class DiscoveredMCPTool extends BaseTool { constructor( private readonly mcpClient: Client, - private readonly config: Config, readonly name: string, readonly description: string, readonly parameterSchema: Record, 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 { - 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, 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, ), ); }