fix mcp timeouts and missing description on mcp errors (#868)

This commit is contained in:
Olcan 2025-06-08 21:52:11 -07:00 committed by GitHub
parent a3d11e8fef
commit a2fee6bdd3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 23 additions and 19 deletions

View File

@ -249,10 +249,8 @@ export function mapToDisplay(
trackedCall.request.args, trackedCall.request.args,
); );
renderOutputAsMarkdown = currentToolInstance.isOutputMarkdown; renderOutputAsMarkdown = currentToolInstance.isOutputMarkdown;
} } else if ('request' in trackedCall && 'args' in trackedCall.request) {
description = JSON.stringify(trackedCall.request.args);
if (trackedCall.status === 'error') {
description = '';
} }
const baseDisplayProperties: Omit< const baseDisplayProperties: Omit<

View File

@ -924,7 +924,7 @@ describe('mapToDisplay', () => {
expectedStatus: ToolCallStatus.Error, expectedStatus: ToolCallStatus.Error,
expectedResultDisplay: 'Error display tool not found', expectedResultDisplay: 'Error display tool not found',
expectedName: baseRequest.name, expectedName: baseRequest.name,
expectedDescription: '', expectedDescription: JSON.stringify(baseRequest.args),
}, },
{ {
name: 'error tool execution failed', name: 'error tool execution failed',
@ -940,7 +940,7 @@ describe('mapToDisplay', () => {
expectedStatus: ToolCallStatus.Error, expectedStatus: ToolCallStatus.Error,
expectedResultDisplay: 'Execution failed display', expectedResultDisplay: 'Execution failed display',
expectedName: baseTool.displayName, // Changed from baseTool.name expectedName: baseTool.displayName, // Changed from baseTool.name
expectedDescription: '', expectedDescription: baseTool.getDescription(baseRequest.args),
}, },
{ {
name: 'cancelled', name: 'cancelled',
@ -986,14 +986,7 @@ describe('mapToDisplay', () => {
expect(toolDisplay.resultDisplay).toBe(expectedResultDisplay); expect(toolDisplay.resultDisplay).toBe(expectedResultDisplay);
expect(toolDisplay.name).toBe(expectedName); expect(toolDisplay.name).toBe(expectedName);
expect(toolDisplay.description).toBe(expectedDescription);
if (status === 'error' && !extraProps?.tool) {
expect(toolDisplay.description).toBe('');
} else {
expect(toolDisplay.description).toBe(
expectedDescription ?? baseTool.getDescription(baseRequest.args),
);
}
expect(toolDisplay.renderOutputAsMarkdown).toBe( expect(toolDisplay.renderOutputAsMarkdown).toBe(
extraProps?.tool?.isOutputMarkdown ?? false, extraProps?.tool?.isOutputMarkdown ?? false,

View File

@ -13,6 +13,8 @@ import { DiscoveredMCPTool } from './mcp-tool.js';
import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai'; import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
import { ToolRegistry } from './tool-registry.js'; import { ToolRegistry } from './tool-registry.js';
export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
/** /**
* Enum representing the connection status of an MCP server * Enum representing the connection status of an MCP server
*/ */
@ -149,11 +151,24 @@ async function connectAndDiscover(
const mcpClient = new Client({ const mcpClient = new Client({
name: 'gemini-cli-mcp-client', name: 'gemini-cli-mcp-client',
version: '0.0.1', version: '0.0.1',
timeout: mcpServerConfig.timeout,
}); });
// 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 { try {
await mcpClient.connect(transport); await mcpClient.connect(transport, {
timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
});
// Connection successful // Connection successful
updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED); updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED);
} catch (error) { } catch (error) {
@ -242,7 +257,7 @@ async function connectAndDiscover(
funcDecl.description ?? '', funcDecl.description ?? '',
parameterSchema, parameterSchema,
funcDecl.name, funcDecl.name,
mcpServerConfig.timeout, mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
mcpServerConfig.trust, mcpServerConfig.trust,
), ),
); );

View File

@ -15,8 +15,6 @@ import { CallableTool, Part, FunctionCall } from '@google/genai';
type ToolParams = Record<string, unknown>; type ToolParams = Record<string, unknown>;
export const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> { export class DiscoveredMCPTool extends BaseTool<ToolParams, ToolResult> {
private static readonly allowlist: Set<string> = new Set(); private static readonly allowlist: Set<string> = new Set();