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,
);
renderOutputAsMarkdown = currentToolInstance.isOutputMarkdown;
}
if (trackedCall.status === 'error') {
description = '';
} else if ('request' in trackedCall && 'args' in trackedCall.request) {
description = JSON.stringify(trackedCall.request.args);
}
const baseDisplayProperties: Omit<

View File

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

View File

@ -13,6 +13,8 @@ import { DiscoveredMCPTool } from './mcp-tool.js';
import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai';
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
*/
@ -149,11 +151,24 @@ async function connectAndDiscover(
const mcpClient = new Client({
name: 'gemini-cli-mcp-client',
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 {
await mcpClient.connect(transport);
await mcpClient.connect(transport, {
timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
});
// Connection successful
updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED);
} catch (error) {
@ -242,7 +257,7 @@ async function connectAndDiscover(
funcDecl.description ?? '',
parameterSchema,
funcDecl.name,
mcpServerConfig.timeout,
mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
mcpServerConfig.trust,
),
);

View File

@ -15,8 +15,6 @@ import { CallableTool, Part, FunctionCall } from '@google/genai';
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> {
private static readonly allowlist: Set<string> = new Set();