fix mcp timeouts and missing description on mcp errors (#868)
This commit is contained in:
parent
a3d11e8fef
commit
a2fee6bdd3
|
@ -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<
|
||||||
|
|
|
@ -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,
|
||||||
|
|
|
@ -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,
|
||||||
),
|
),
|
||||||
);
|
);
|
||||||
|
|
|
@ -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();
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue