summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--packages/core/src/tools/mcp-client.test.ts87
-rw-r--r--packages/core/src/tools/mcp-client.ts30
2 files changed, 115 insertions, 2 deletions
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<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);
+ }
+ }
+}