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 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);

View File

@ -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,

View File

@ -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);

View File

@ -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,
),
);
}