Sanitize MCP FunctionDeclarations to workaround Vertex bug (#1330)

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();
+  });
+});

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<string, unknown> =
         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);
+    }
+  }
+}