diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index abd9c58f..c9546ba3 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -14,7 +14,8 @@ import { afterEach, Mocked, } from 'vitest'; -import { discoverMcpTools } from './mcp-client.js'; +import { discoverMcpTools, sanatizeParameters } from './mcp-client.js'; +import { Schema, Type } from '@google/genai'; import { Config, MCPServerConfig } from '../config/config.js'; import { DiscoveredMCPTool } from './mcp-tool.js'; import { Client } from '@modelcontextprotocol/sdk/client/index.js'; @@ -536,3 +537,87 @@ describe('discoverMcpTools', () => { expect(lastClientInstance?.onerror).toEqual(expect.any(Function)); }); }); + +describe('sanatizeParameters', () => { + it('should do nothing for an undefined schema', () => { + const schema = undefined; + sanatizeParameters(schema); + }); + + it('should remove default when anyOf is present', () => { + const schema: Schema = { + anyOf: [{ type: Type.STRING }, { type: Type.NUMBER }], + default: 'hello', + }; + sanatizeParameters(schema); + expect(schema.default).toBeUndefined(); + }); + + it('should recursively sanatize items in anyOf', () => { + const schema: Schema = { + anyOf: [ + { + anyOf: [{ type: Type.STRING }], + default: 'world', + }, + { type: Type.NUMBER }, + ], + }; + sanatizeParameters(schema); + expect(schema.anyOf![0].default).toBeUndefined(); + }); + + it('should recursively sanatize items in items', () => { + const schema: Schema = { + items: { + anyOf: [{ type: Type.STRING }], + default: 'world', + }, + }; + sanatizeParameters(schema); + expect(schema.items!.default).toBeUndefined(); + }); + + it('should recursively sanatize items in properties', () => { + const schema: Schema = { + properties: { + prop1: { + anyOf: [{ type: Type.STRING }], + default: 'world', + }, + }, + }; + sanatizeParameters(schema); + expect(schema.properties!.prop1.default).toBeUndefined(); + }); + + it('should handle complex nested schemas', () => { + const schema: Schema = { + properties: { + prop1: { + items: { + anyOf: [{ type: Type.STRING }], + default: 'world', + }, + }, + prop2: { + anyOf: [ + { + properties: { + nestedProp: { + anyOf: [{ type: Type.NUMBER }], + default: 123, + }, + }, + }, + ], + }, + }, + }; + sanatizeParameters(schema); + expect(schema.properties!.prop1.items!.default).toBeUndefined(); + const nestedProp = + schema.properties!.prop2.anyOf![0].properties!.nestedProp; + expect(nestedProp?.default).toBeUndefined(); + }); +}); diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 7fde99b7..ceedb068 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -11,7 +11,12 @@ import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/ import { parse } from 'shell-quote'; import { MCPServerConfig } from '../config/config.js'; import { DiscoveredMCPTool } from './mcp-tool.js'; -import { CallableTool, FunctionDeclaration, mcpToTool } from '@google/genai'; +import { + CallableTool, + FunctionDeclaration, + mcpToTool, + Schema, +} from '@google/genai'; import { ToolRegistry } from './tool-registry.js'; export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes @@ -299,6 +304,8 @@ async function connectAndDiscover( toolNameForModel.slice(0, 28) + '___' + toolNameForModel.slice(-32); } + sanatizeParameters(funcDecl.parameters); + // Ensure parameters is a valid JSON schema object, default to empty if not. const parameterSchema: Record = funcDecl.parameters && typeof funcDecl.parameters === 'object' @@ -354,3 +361,24 @@ async function connectAndDiscover( } } } + +export function sanatizeParameters(schema?: Schema) { + if (!schema) { + return; + } + if (schema.anyOf) { + // Vertex AI gets confused if both anyOf and default are set. + schema.default = undefined; + for (const item of schema.anyOf) { + sanatizeParameters(item); + } + } + if (schema.items) { + sanatizeParameters(schema.items); + } + if (schema.properties) { + for (const item of Object.values(schema.properties)) { + sanatizeParameters(item); + } + } +}