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

This commit is contained in:
Olcan 2025-05-17 16:53:22 -07:00 committed by GitHub
parent 324040032a
commit 4de4822219
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 45 additions and 42 deletions

View File

@ -7,7 +7,7 @@
import * as fs from 'fs'; import * as fs from 'fs';
import * as path from 'path'; import * as path from 'path';
import { homedir } from 'os'; 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 SETTINGS_DIRECTORY_NAME = '.gemini';
export const USER_SETTINGS_DIR = path.join(homedir(), SETTINGS_DIRECTORY_NAME); export const USER_SETTINGS_DIR = path.join(homedir(), SETTINGS_DIRECTORY_NAME);
@ -24,7 +24,7 @@ export interface Settings {
toolDiscoveryCommand?: string; toolDiscoveryCommand?: string;
toolCallCommand?: string; toolCallCommand?: string;
mcpServerCommand?: string; mcpServerCommand?: string;
mcpServers?: Record<string, StdioServerParameters>; mcpServers?: Record<string, MCPServerConfig>;
// Add other settings here. // Add other settings here.
} }
@ -69,10 +69,10 @@ export class LoadedSettings {
setValue( setValue(
scope: SettingScope, scope: SettingScope,
key: keyof Settings, key: keyof Settings,
value: string | Record<string, StdioServerParameters> | undefined, value: string | Record<string, MCPServerConfig> | undefined,
): void { ): void {
const settingsFile = this.forScope(scope); 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; settingsFile.settings[key] = value;
this._merged = this.computeMergedSettings(); this._merged = this.computeMergedSettings();
saveSettings(settingsFile); saveSettings(settingsFile);

View File

@ -21,7 +21,16 @@ import { WebFetchTool } from '../tools/web-fetch.js';
import { ReadManyFilesTool } from '../tools/read-many-files.js'; import { ReadManyFilesTool } from '../tools/read-many-files.js';
import { BaseTool, ToolResult } from '../tools/tools.js'; import { BaseTool, ToolResult } from '../tools/tools.js';
import { MemoryTool } from '../tools/memoryTool.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 { export class Config {
private toolRegistry: ToolRegistry; private toolRegistry: ToolRegistry;
@ -37,9 +46,7 @@ export class Config {
private readonly toolDiscoveryCommand: string | undefined, private readonly toolDiscoveryCommand: string | undefined,
private readonly toolCallCommand: string | undefined, private readonly toolCallCommand: string | undefined,
private readonly mcpServerCommand: string | undefined, private readonly mcpServerCommand: string | undefined,
private readonly mcpServers: private readonly mcpServers: Record<string, MCPServerConfig> | undefined,
| Record<string, StdioServerParameters>
| undefined,
private readonly userAgent: string, private readonly userAgent: string,
private userMemory: string = '', // Made mutable for refresh private userMemory: string = '', // Made mutable for refresh
private geminiMdFileCount: number = 0, private geminiMdFileCount: number = 0,
@ -92,7 +99,7 @@ export class Config {
return this.mcpServerCommand; return this.mcpServerCommand;
} }
getMcpServers(): Record<string, StdioServerParameters> | undefined { getMcpServers(): Record<string, MCPServerConfig> | undefined {
return this.mcpServers; return this.mcpServers;
} }
@ -164,7 +171,7 @@ export function createServerConfig(
toolDiscoveryCommand?: string, toolDiscoveryCommand?: string,
toolCallCommand?: string, toolCallCommand?: string,
mcpServerCommand?: string, mcpServerCommand?: string,
mcpServers?: Record<string, StdioServerParameters>, mcpServers?: Record<string, MCPServerConfig>,
userAgent?: string, userAgent?: string,
userMemory?: string, userMemory?: string,
geminiMdFileCount?: number, geminiMdFileCount?: number,

View File

@ -577,7 +577,6 @@ describe('DiscoveredTool', () => {
}); });
describe('DiscoveredMCPTool', () => { describe('DiscoveredMCPTool', () => {
let config: Config;
let mockMcpClient: Client; let mockMcpClient: Client;
const toolName = 'my-mcp-tool'; const toolName = 'my-mcp-tool';
const toolDescription = 'An MCP-discovered tool.'; const toolDescription = 'An MCP-discovered tool.';
@ -587,21 +586,6 @@ describe('DiscoveredMCPTool', () => {
}; };
beforeEach(() => { 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({ mockMcpClient = new Client({
name: 'test-client', name: 'test-client',
version: '0.0.0', version: '0.0.0',
@ -615,7 +599,6 @@ describe('DiscoveredMCPTool', () => {
it('constructor should set up properties correctly and enhance description', () => { it('constructor should set up properties correctly and enhance description', () => {
const tool = new DiscoveredMCPTool( const tool = new DiscoveredMCPTool(
mockMcpClient, mockMcpClient,
config,
toolName, toolName,
toolDescription, toolDescription,
toolInputSchema, toolInputSchema,
@ -631,7 +614,6 @@ describe('DiscoveredMCPTool', () => {
it('execute should call mcpClient.callTool with correct params and return serialized result', async () => { it('execute should call mcpClient.callTool with correct params and return serialized result', async () => {
const tool = new DiscoveredMCPTool( const tool = new DiscoveredMCPTool(
mockMcpClient, mockMcpClient,
config,
toolName, toolName,
toolDescription, toolDescription,
toolInputSchema, toolInputSchema,
@ -644,10 +626,16 @@ describe('DiscoveredMCPTool', () => {
const result = await tool.execute(params); const result = await tool.execute(params);
expect(mockMcpClient.callTool).toHaveBeenCalledWith({ expect(mockMcpClient.callTool).toHaveBeenCalledWith(
name: toolName, {
arguments: params, name: toolName,
}); arguments: params,
},
undefined,
{
timeout: 10 * 60 * 1000,
},
);
const expectedOutput = JSON.stringify(mcpResult, null, 2); const expectedOutput = JSON.stringify(mcpResult, null, 2);
expect(result.llmContent).toBe(expectedOutput); expect(result.llmContent).toBe(expectedOutput);
expect(result.returnDisplay).toBe(expectedOutput); expect(result.returnDisplay).toBe(expectedOutput);

View File

@ -14,6 +14,8 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'; import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
type ToolParams = Record<string, unknown>; 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> { export class DiscoveredTool extends BaseTool<ToolParams, ToolResult> {
constructor( constructor(
private readonly config: Config, 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> { export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
constructor( constructor(
private readonly mcpClient: Client, private readonly mcpClient: Client,
private readonly config: Config,
readonly name: string, readonly name: string,
readonly description: string, readonly description: string,
readonly parameterSchema: Record<string, unknown>, readonly parameterSchema: Record<string, unknown>,
readonly serverToolName: string, readonly serverToolName: string,
readonly timeout?: number,
) { ) {
description += ` description += `
@ -112,10 +114,16 @@ Returns the MCP server response as a json string.
} }
async execute(params: ToolParams): Promise<ToolResult> { async execute(params: ToolParams): Promise<ToolResult> {
const result = await this.mcpClient.callTool({ const result = await this.mcpClient.callTool(
name: this.serverToolName, {
arguments: params, name: this.serverToolName,
}); arguments: params,
},
undefined, // skip resultSchema to specify options (RequestOptions)
{
timeout: this.timeout ?? MCP_TOOL_DEFAULT_TIMEOUT_MSEC,
},
);
return { return {
llmContent: JSON.stringify(result, null, 2), llmContent: JSON.stringify(result, null, 2),
returnDisplay: JSON.stringify(result, null, 2), returnDisplay: JSON.stringify(result, null, 2),
@ -189,17 +197,17 @@ export class ToolRegistry {
args: args.slice(1), args: args.slice(1),
}; };
} }
for (const [mcpServerName, mcpServer] of Object.entries(mcpServers)) { for (const [mcpServerName, mcpServerConfig] of Object.entries(mcpServers)) {
(async () => { (async () => {
const mcpClient = new Client({ const mcpClient = new Client({
name: 'mcp-client', name: 'mcp-client',
version: '0.0.1', version: '0.0.1',
}); });
const transport = new StdioClientTransport({ const transport = new StdioClientTransport({
...mcpServer, ...mcpServerConfig,
env: { env: {
...process.env, ...process.env,
...(mcpServer.env || {}), ...(mcpServerConfig.env || {}),
} as Record<string, string>, } as Record<string, string>,
stderr: 'pipe', stderr: 'pipe',
}); });
@ -208,7 +216,7 @@ export class ToolRegistry {
} catch (error) { } catch (error) {
console.error( console.error(
`failed to start or connect to MCP server '${mcpServerName}' ` + `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. // Do not re-throw, let other MCP servers be discovered.
return; // Exit this async IIFE if connection failed return; // Exit this async IIFE if connection failed
@ -246,13 +254,13 @@ export class ToolRegistry {
this.registerTool( this.registerTool(
new DiscoveredMCPTool( new DiscoveredMCPTool(
mcpClient, mcpClient,
this.config,
Object.keys(mcpServers).length > 1 Object.keys(mcpServers).length > 1
? mcpServerName + '__' + tool.name ? mcpServerName + '__' + tool.name
: tool.name, : tool.name,
tool.description ?? '', tool.description ?? '',
tool.inputSchema, tool.inputSchema,
tool.name, tool.name,
mcpServerConfig.timeout,
), ),
); );
} }