From ca7d3b6b233abf7086e99ceb49f855aa1d0464ce Mon Sep 17 00:00:00 2001 From: Dennis Jeong <3719829+w0nche0l@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:43:29 -0500 Subject: [PATCH] fix(agent): send an MCP tool's own JSON Schema to the model wrapMcpTool converted the server's inputSchema to Zod and the wire parameters were regenerated from that Zod schema, which added a pattern per string format, safe-integer bounds, and additionalProperties, and inlined $ref/$defs. Keep both: serialize the server's schema, validate with the Zod one. Linear: DEV-1380 --- .changeset/mcp-wire-input-schema.md | 9 ++ packages/agent/src/lib/tool-executor.ts | 3 +- packages/agent/src/lib/tool-types.ts | 19 +-- packages/agent/src/lib/tool.ts | 14 +- packages/agent/src/mcp/tool-wrapper.ts | 21 +-- .../tests/unit/mcp/wire-input-schema.test.ts | 132 ++++++++++++++++++ 6 files changed, 177 insertions(+), 21 deletions(-) create mode 100644 .changeset/mcp-wire-input-schema.md create mode 100644 packages/agent/tests/unit/mcp/wire-input-schema.test.ts diff --git a/.changeset/mcp-wire-input-schema.md b/.changeset/mcp-wire-input-schema.md new file mode 100644 index 00000000..74118eb8 --- /dev/null +++ b/.changeset/mcp-wire-input-schema.md @@ -0,0 +1,9 @@ +--- +'@openrouter/agent': minor +--- + +Send an MCP tool's own JSON Schema to the model instead of one regenerated from Zod. + +`wrapMcpTool` converts the server's `inputSchema` to Zod to validate call arguments, and the model-facing `parameters` were then regenerated from that Zod schema. The round trip rewrote the schema: every string `format` gained a `pattern` (about 400 characters for `date-time`), every `integer` gained safe-integer bounds, `$ref`/`$defs` were inlined, and open objects gained `additionalProperties: {}`. + +MCP tools now keep both. The server's schema is stored as `function.wireInputSchema` and serialized verbatim (sanitized of `~`-prefixed keys); the Zod `inputSchema` still validates the call. `wireInputSchema` moves from `ManualToolFunction` to `BaseToolFunction`, and `markMcp` accepts it as an option. A server schema whose root is not an object still serializes the Zod-derived object schema. diff --git a/packages/agent/src/lib/tool-executor.ts b/packages/agent/src/lib/tool-executor.ts index b3ed0c4b..780917c9 100644 --- a/packages/agent/src/lib/tool-executor.ts +++ b/packages/agent/src/lib/tool-executor.ts @@ -22,7 +22,6 @@ import { isDeferredHandle, isGeneratorTool, isHITLTool, - isManualTool, isMcpTool, isRegularExecuteTool, isServerTool, @@ -131,7 +130,7 @@ export function convertToolsToAPIFormat( description: tool.function.description || null, strict: tool.function.strict ?? null, parameters: - isManualTool(tool) && tool.function.wireInputSchema !== undefined + tool.function.wireInputSchema !== undefined ? sanitizeJsonSchema(tool.function.wireInputSchema) : convertZodToJsonSchema(tool.function.inputSchema), }; diff --git a/packages/agent/src/lib/tool-types.ts b/packages/agent/src/lib/tool-types.ts index 90ce6671..41217cf8 100644 --- a/packages/agent/src/lib/tool-types.ts +++ b/packages/agent/src/lib/tool-types.ts @@ -440,6 +440,16 @@ export interface BaseToolFunction< name: TName; description?: string; inputSchema: TInput; + /** + * JSON Schema to serialize for this tool instead of regenerating one from + * `inputSchema`. A tool whose definition already is a JSON Schema (a manual + * tool, an MCP tool) forwards it verbatim, sanitized of `~`-prefixed + * metadata keys. Regenerating it from Zod costs CPU/allocations, cannot + * represent every construct, and adds constraints the source never declared + * (a `pattern` per string `format`, safe-integer bounds). `inputSchema` + * still validates the call's arguments. + */ + readonly wireInputSchema?: Readonly>; /** * Whether providers should enforce strict schema adherence when generating * this tool's call arguments (OpenAI structured-outputs style). Serialized @@ -575,15 +585,6 @@ export interface ManualToolFunction< TCtx extends $ZodObject<$ZodShape> = $ZodObject<$ZodShape>, TName extends string = string, > extends BaseToolFunction { - /** - * JSON Schema to serialize for this tool instead of regenerating one from - * `inputSchema`. Manual tools are never executed or input-validated by the - * SDK, so a caller that already owns a JSON Schema can forward it verbatim - * (sanitized of `~`-prefixed metadata keys) and skip the round trip through - * Zod, which both costs CPU/allocations and cannot represent constructs like - * `anyOf` / `oneOf`. - */ - readonly wireInputSchema?: Readonly>; outputSchema?: TOutput; } diff --git a/packages/agent/src/lib/tool.ts b/packages/agent/src/lib/tool.ts index e4916a32..cf095fea 100644 --- a/packages/agent/src/lib/tool.ts +++ b/packages/agent/src/lib/tool.ts @@ -1100,21 +1100,31 @@ export function serverTool( toolToMark: T, options?: { loopKey?: ToolLoopKey>; + wireInputSchema?: Readonly>; }, ): McpBranded { const marked = { ...toolToMark, _mcp: true as const, }; - if (options?.loopKey !== undefined && isClientTool(marked)) { + if (options !== undefined && isClientTool(marked)) { marked.function = { ...marked.function, - loopKey: options.loopKey, + ...(options.loopKey !== undefined && { + loopKey: options.loopKey, + }), + ...(options.wireInputSchema !== undefined && { + wireInputSchema: options.wireInputSchema, + }), }; } return marked; diff --git a/packages/agent/src/mcp/tool-wrapper.ts b/packages/agent/src/mcp/tool-wrapper.ts index 95b76bc2..e690d85e 100644 --- a/packages/agent/src/mcp/tool-wrapper.ts +++ b/packages/agent/src/mcp/tool-wrapper.ts @@ -87,8 +87,9 @@ async function invokeTool(invokeArgs: InvokeToolArgs): Promise { * Wrap one discovered MCP tool as an `@openrouter/agent` tool. Emits a * generator tool (streaming progress events) when `emitProgress` is on, * otherwise a regular tool. The input schema (and output schema, when declared) - * are converted to Zod so the model sees faithful parameters. The abort signal, - * if supplied, is threaded into the underlying `callTool`. + * are converted to Zod to validate the call; the model is sent the server's + * own input schema, not one regenerated from Zod. The abort signal, if + * supplied, is threaded into the underlying `callTool`. */ export function wrapMcpTool(def: McpToolDef, options: WrapToolOptions): McpBranded { const name = `${options.namePrefix ?? ''}${def.name}`; @@ -103,12 +104,16 @@ export function wrapMcpTool(def: McpToolDef, options: WrapToolOptions): McpBrand options.loopKeys !== undefined && Object.hasOwn(options.loopKeys, def.name) ? options.loopKeys[def.name] : def.loopKey; - const markOptions = - loopKey !== undefined - ? { - loopKey, - } - : undefined; + const markOptions = { + ...(loopKey !== undefined && { + loopKey, + }), + // A non-object root is not a usable tool schema; the Zod conversion + // degrades it to an object, so that one is serialized instead. + ...(def.inputSchema['type'] === 'object' && { + wireInputSchema: def.inputSchema, + }), + }; if (options.emitProgress === true) { return markMcp( diff --git a/packages/agent/tests/unit/mcp/wire-input-schema.test.ts b/packages/agent/tests/unit/mcp/wire-input-schema.test.ts new file mode 100644 index 00000000..01deca09 --- /dev/null +++ b/packages/agent/tests/unit/mcp/wire-input-schema.test.ts @@ -0,0 +1,132 @@ +import type { Client } from '@modelcontextprotocol/client'; +import { describe, expect, it } from 'vitest'; +import { convertToolsToAPIFormat, validateToolInput } from '../../../src/lib/tool-executor.js'; +import { isClientTool } from '../../../src/lib/tool-types.js'; +import type { McpToolDef } from '../../../src/mcp/tool-wrapper.js'; +import { wrapMcpTool } from '../../../src/mcp/tool-wrapper.js'; + +// wrapMcpTool only stores the client for the execute closure, which these +// tests never invoke. +function fakeClient(): Client { + return {} as never; +} + +function wireParameters(def: McpToolDef): unknown { + const [emitted] = convertToolsToAPIFormat([ + wrapMcpTool(def, { + client: fakeClient(), + }), + ]); + return emitted !== undefined && 'parameters' in emitted ? emitted.parameters : undefined; +} + +describe('MCP tool wire input schema', () => { + it("sends the server's schema to the model unchanged", () => { + // Each of these is rewritten by a JSON Schema -> Zod -> JSON Schema round + // trip: `format` gains a `pattern`, `integer` gains safe-integer bounds, + // `$ref`/`$defs` are inlined, `prefixItems` becomes draft-07 `items`, and + // an open object gains `additionalProperties: {}`. + const inputSchema = { + $schema: 'https://json-schema.org/draft/2020-12/schema', + type: 'object', + properties: { + when: { + type: 'string', + format: 'date-time', + }, + count: { + type: 'integer', + minimum: 1, + }, + point: { + type: 'array', + prefixItems: [ + { + type: 'number', + }, + { + type: 'number', + }, + ], + items: false, + }, + filter: { + $ref: '#/$defs/Filter', + }, + }, + required: [ + 'when', + ], + $defs: { + Filter: { + type: 'object', + properties: { + tag: { + type: 'string', + }, + }, + additionalProperties: false, + }, + }, + }; + + expect( + wireParameters({ + name: 'search', + inputSchema, + }), + ).toEqual(inputSchema); + }); + + it('still validates call arguments against the converted Zod schema', () => { + const wrapped = wrapMcpTool( + { + name: 'search', + inputSchema: { + type: 'object', + properties: { + count: { + type: 'integer', + }, + }, + required: [ + 'count', + ], + }, + }, + { + client: fakeClient(), + }, + ); + if (!isClientTool(wrapped)) { + throw new Error('expected a client tool'); + } + const { inputSchema } = wrapped.function; + + expect( + validateToolInput(inputSchema, { + count: 2, + }), + ).toEqual({ + count: 2, + }); + expect(() => + validateToolInput(inputSchema, { + count: 'two', + }), + ).toThrow(); + }); + + it('serializes the Zod-derived object schema when the server schema is not an object', () => { + expect( + wireParameters({ + name: 'odd', + inputSchema: { + type: 'string', + }, + }), + ).toMatchObject({ + type: 'object', + }); + }); +});