From 5e30969cc5f1255d44ec7a1f1ca8dbac0af5e7c4 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:23:21 +0800 Subject: [PATCH] fix(mcp): advertise self-contained tool inputs Inline reused input schemas at their projection owner and remove the superseded reference-budget fallback. Validate actual SDK input and output schemas while preserving canonical parsing, output references, and projection cache options. Claude Desktop and affected Windows connector behavior remain unverified; the SDK catalog has no schema-local input references. Co-authored-by: akushonkamen --- src/contracts/inputSchemaPresentation.ts | 73 ---------- src/contracts/toolInputReferences.test.ts | 130 ++---------------- src/contracts/toolSchemaMetadata.ts | 33 +---- tests/boundary/mcp/toolSchemaValidity.test.ts | 128 ++++++++++------- 4 files changed, 102 insertions(+), 262 deletions(-) diff --git a/src/contracts/inputSchemaPresentation.ts b/src/contracts/inputSchemaPresentation.ts index 740f1f3fae..1a9754ace7 100644 --- a/src/contracts/inputSchemaPresentation.ts +++ b/src/contracts/inputSchemaPresentation.ts @@ -246,76 +246,3 @@ export const presentInputJsonSchema = ( ].filter((text) => text !== undefined); return { ...presentation, description: description.join(" ") }; }; - -/** - * Clients that render advertised inputs into SDK declarations expand - * schema-local references with a shared per-input budget: Pi 1.1.0 code-mode - * allows 32 expansions and degrades every later reference to `unknown` - * (morluto/rea#1632), hiding the fields those references describe. The budget - * counts conservatively — every resolvable, non-cyclic reference occurrence - * in every schema position — so any reference-resolving client stays within - * its measured limit. - */ -export const advertisedReferenceExpansionBudget = 32; - -const definitionKeys = new Set(["$defs", "definitions"]); - -const countReferenceExpansions = ( - node: unknown, - root: Readonly>, - active: ReadonlySet, - count: { value: number }, -): void => { - if (Array.isArray(node)) { - node.forEach((child) => - countReferenceExpansions(child, root, active, count), - ); - return; - } - if (!isObject(node)) return; - if (typeof node.$ref === "string" && !active.has(node.$ref)) { - const target = resolveReference(root, node.$ref); - if (target !== undefined) { - count.value += 1; - countReferenceExpansions( - target, - root, - new Set([...active, node.$ref]), - count, - ); - return; - } - } - for (const [key, value] of Object.entries(node)) { - if (definitionKeys.has(key)) continue; - if (schemaMaps.has(key) && isObject(value)) { - for (const child of Object.values(value)) - countReferenceExpansions(child, root, active, count); - } else if (schemaArrays.has(key) || schemaChildren.has(key)) { - countReferenceExpansions(value, root, active, count); - } - } -}; - -/** Conservative client expansion count over every schema position. */ -export const countLocalReferenceExpansions = ( - schema: Record, -): number => { - const count = { value: 0 }; - countReferenceExpansions(schema, schema, new Set(), count); - return count.value; -}; - -/** - * Keep the shared-definition form of an advertised input while clients can - * still resolve it. Once rendering would exhaust the reference budget, fall - * back to the inline projection, which carries no shared definitions for the - * nonrecursive schemas this repository advertises, so every field renders - * with complete structure in every reference-resolving client. - */ -export const inputSchemaWithinReferenceExpansionBudget = ( - shared: Record, - inline: () => Record, - budget: number, -): Record => - countLocalReferenceExpansions(shared) <= budget ? shared : inline(); diff --git a/src/contracts/toolInputReferences.test.ts b/src/contracts/toolInputReferences.test.ts index 657b77f3cd..80775dad34 100644 --- a/src/contracts/toolInputReferences.test.ts +++ b/src/contracts/toolInputReferences.test.ts @@ -3,12 +3,7 @@ import { Ajv2020 } from "ajv/dist/2020.js"; import { describe, expect, it } from "vitest"; import { z } from "zod"; -import { - advertisedReferenceExpansionBudget, - countLocalReferenceExpansions, - inputSchemaWithinReferenceExpansionBudget, - presentInputJsonSchema, -} from "./inputSchemaPresentation.js"; +import { presentInputJsonSchema } from "./inputSchemaPresentation.js"; import { toolContract } from "./toolContracts.js"; import { toolInputSchemaWithMetadata, @@ -20,27 +15,27 @@ it.each(["draft-2020-12", "draft-07"] as const)( "honors projection options independently of cached wire transformations in %s", (target) => { const range = z.strictObject({ offset: z.number(), length: z.number() }); - const canonical = toolInputSchemaWithMetadata({ - ...toolContract("open_binary"), - inputSchema: z.strictObject({ before: range, after: range }), - }); + const canonical = z.strictObject({ before: range, after: range }); const transform = (schema: Record) => ({ ...schema, title: "Transformed range inputs", }); - const options = { target, libraryOptions: { reused: "inline" as const } }; - for (const inlineFirst of [false, true]) { + const options = { target, libraryOptions: { reused: "ref" as const } }; + for (const sharedFirst of [false, true]) { const wrapped = transformAdvertisedInputJsonSchema(canonical, transform); const project = wrapped["~standard"].jsonSchema.input; - if (inlineFirst) project(options); - const shared = project({ target }); + if (sharedFirst) project(options); + const advertised = project({ target }); + expect(JSON.stringify(advertised)).not.toContain('"$ref":'); + expect(advertised).toEqual( + transform(canonical["~standard"].jsonSchema.input({ target })), + ); + const shared = project(options); expect(JSON.stringify(shared)).toContain('"$ref":'); - const inline = project(options); - expect(inline).toEqual( + expect(shared).toEqual( transform(canonical["~standard"].jsonSchema.input(options)), ); - expect(JSON.stringify(inline)).not.toContain('"$ref":'); - expect(project({ target })).toBe(shared); + expect(project({ target })).toBe(advertised); } }, ); @@ -239,105 +234,6 @@ describe("advertised schema metadata and validation", () => { ); }); -describe("advertised reference expansion budget", () => { - it("counts references in keyword-named fields, excluding literals", () => { - const fields = ["default", "enum", "examples", "const", "$defs"]; - const schema = { - type: "object", - properties: Object.fromEntries( - fields.map((name) => [name, { $ref: "#/$defs/value" }]), - ), - $defs: { value: { type: "string" } }, - default: { $ref: "#/$defs/value" }, - examples: [{ $ref: "#/$defs/value" }], - }; - expect(countLocalReferenceExpansions(schema)).toBe(fields.length); - }); - const reusedComponent = z - .object({ kind: z.literal("rec"), value: z.string() }) - .strict(); - const inputWithReusedFields = (count: number) => - z.strictObject( - Object.fromEntries( - Array.from({ length: count }, (_, index) => [ - `field_${String(index).padStart(2, "0")}`, - reusedComponent, - ]), - ), - ); - const advertisedInput = ( - canonical: z.ZodType>>, - ) => - toolInputSchemaWithMetadata({ - ...toolContract("open_binary"), - inputSchema: canonical, - })["~standard"].jsonSchema.input({ target: "draft-2020-12" }); - - it("keeps shared definitions while every client can resolve them", () => { - const advertised = advertisedInput(inputWithReusedFields(8)); - expect(advertised.$defs).toBeDefined(); - expect(countLocalReferenceExpansions(advertised)).toBe(8); - }); - - it("keeps shared definitions exactly at the budget", () => { - const advertised = advertisedInput( - inputWithReusedFields(advertisedReferenceExpansionBudget), - ); - expect(advertised.$defs).toBeDefined(); - expect(countLocalReferenceExpansions(advertised)).toBe( - advertisedReferenceExpansionBudget, - ); - }); - - it("re-projects without shared definitions once the budget is exceeded", () => { - const advertised = advertisedInput( - inputWithReusedFields(advertisedReferenceExpansionBudget + 1), - ); - expect(advertised.$defs).toBeUndefined(); - expect(countLocalReferenceExpansions(advertised)).toBe(0); - const properties = advertised.properties as Record; - expect(properties.field_00).toMatchObject({ - type: "object", - properties: { kind: { const: "rec" }, value: { type: "string" } }, - required: ["kind", "value"], - }); - }); - - it("keeps the shared form untouched when it fits and swaps it when not", () => { - const shared = { - type: "object", - properties: { value: { $ref: "#/$defs/shared" } }, - $defs: { shared: { type: "string" } }, - }; - const inlined = { - type: "object", - properties: { value: { type: "string" } }, - }; - expect( - inputSchemaWithinReferenceExpansionBudget( - shared, - () => inlined, - advertisedReferenceExpansionBudget, - ), - ).toBe(shared); - expect( - inputSchemaWithinReferenceExpansionBudget( - { - ...shared, - properties: Object.fromEntries( - Array.from({ length: 33 }, (_, index) => [ - `field_${index}`, - { $ref: "#/$defs/shared" }, - ]), - ), - }, - () => inlined, - advertisedReferenceExpansionBudget, - ), - ).toBe(inlined); - }); -}); - describe("root input presentation", () => { it("preserves prototype-named properties in object-root presentation", () => { const fields = Object.fromEntries([ diff --git a/src/contracts/toolSchemaMetadata.ts b/src/contracts/toolSchemaMetadata.ts index a0821612db..a6e6740d9a 100644 --- a/src/contracts/toolSchemaMetadata.ts +++ b/src/contracts/toolSchemaMetadata.ts @@ -2,11 +2,7 @@ import { z } from "zod"; import type { ToolContract } from "./toolContractTypes.js"; import { advertisedInputExamples } from "./advertisedInputExamples.js"; -import { - advertisedReferenceExpansionBudget, - inputSchemaWithinReferenceExpansionBudget, - presentInputJsonSchema, -} from "./inputSchemaPresentation.js"; +import { presentInputJsonSchema } from "./inputSchemaPresentation.js"; import { digestCanonicalValue } from "../domain/canonicalDigest.js"; const PROPERTY_DESCRIPTIONS: Readonly> = { @@ -64,31 +60,16 @@ export const toolInputSchemaWithMetadata = ( contract.inputSchema, "input", (project) => (options) => { - const shared = project({ + const inline = project({ ...options, - libraryOptions: { reused: "ref", ...options.libraryOptions }, + libraryOptions: { ...options.libraryOptions, reused: "inline" }, }); const examples = advertisedInputExamples(contract); return { - // Sharing definitions only helps clients that resolve them. A client - // rendering declarations from a shared expansion budget degrades - // every reference past it to `unknown`, so an input whose shared form - // exceeds the budget is re-projected without shared definitions. - ...inputSchemaWithinReferenceExpansionBudget( - presentInputJsonSchema(shared, fallbackPropertyDescription), - () => - presentInputJsonSchema( - project({ - ...options, - libraryOptions: { - ...options.libraryOptions, - reused: "inline", - }, - }), - fallbackPropertyDescription, - ), - advertisedReferenceExpansionBudget, - ), + // Keep advertised parameters self-contained for clients that do not + // resolve local references. The registered canonical schema remains + // responsible for runtime validation. + ...presentInputJsonSchema(inline, fallbackPropertyDescription), ...(examples.length === 0 ? {} : { examples }), }; }, diff --git a/tests/boundary/mcp/toolSchemaValidity.test.ts b/tests/boundary/mcp/toolSchemaValidity.test.ts index 8716aa21a5..53008b94f6 100644 --- a/tests/boundary/mcp/toolSchemaValidity.test.ts +++ b/tests/boundary/mcp/toolSchemaValidity.test.ts @@ -240,46 +240,6 @@ function resolveReference(schema: unknown, reference: string): unknown { return value; } -function localReferenceExpansions(schema: Record): number { - // Clients that render SDK declarations from the advertised input expand - // schema-local references against a shared per-input budget; Pi 1.1.0 - // code-mode allows 32 and renders every later reference as `unknown`. - // Count conservatively: every resolvable, non-cyclic reference occurrence - // reached while walking the input, never the definitions themselves. - let count = 0; - const visit = (node: unknown, active: ReadonlySet): void => { - if (Array.isArray(node)) { - node.forEach((child) => visit(child, active)); - return; - } - if (!isRecord(node)) return; - if (typeof node.$ref === "string" && !active.has(node.$ref)) { - const target = resolveReference(schema, node.$ref); - if (target !== undefined) { - count += 1; - visit(target, new Set([...active, node.$ref])); - return; - } - } - for (const [key, value] of Object.entries(node)) { - if ( - [ - "$defs", - "definitions", - "examples", - "default", - "const", - "enum", - ].includes(key) - ) - continue; - visit(value, active); - } - }; - visit(schema, new Set()); - return count; -} - function recursiveReferences(tools: readonly ToolSchemas[]): string[] { return tools.flatMap((tool) => (["inputSchema", "outputSchema"] as const).flatMap((kind) => { @@ -312,6 +272,85 @@ function recursiveReferences(tools: readonly ToolSchemas[]): string[] { ); } +const schemaMapKeywords = new Set([ + "$defs", + "definitions", + "properties", + "patternProperties", + "dependentSchemas", + "dependencies", +]); +const schemaArrayKeywords = new Set(["allOf", "anyOf", "oneOf", "prefixItems"]); +const schemaChildKeywords = new Set([ + "items", + "additionalItems", + "contains", + "additionalProperties", + "unevaluatedProperties", + "unevaluatedItems", + "propertyNames", + "not", + "if", + "then", + "else", +]); + +function localInputReferences(tools: readonly ToolSchemas[]): string[] { + return tools.flatMap((tool) => { + const references: string[] = []; + const visit = (node: unknown, path: string): void => { + if (!isRecord(node)) return; + if (typeof node.$ref === "string") references.push(`${path}.$ref`); + for (const [key, value] of Object.entries(node)) { + if (schemaMapKeywords.has(key) && isRecord(value)) { + for (const [name, schema] of Object.entries(value)) + visit(schema, `${path}.${key}.${name}`); + } else if (schemaArrayKeywords.has(key) && Array.isArray(value)) { + value.forEach((schema, index) => + visit(schema, `${path}.${key}[${index}]`), + ); + } else if (schemaChildKeywords.has(key)) { + if (Array.isArray(value)) + value.forEach((schema, index) => + visit(schema, `${path}.${key}[${index}]`), + ); + else visit(value, `${path}.${key}`); + } + } + }; + visit(tool.inputSchema, `${tool.name}.inputSchema`); + return references; + }); +} + +it("finds references only in input schema positions, not literal data", () => { + const schema = { + type: "object", + properties: { + $ref: { type: "string" }, + default: { type: "array", items: { type: "string" } }, + }, + default: { $ref: "literal data" }, + examples: [{ $defs: { field: { $ref: "literal data" } } }], + }; + expect( + localInputReferences([{ name: "literal", inputSchema: schema }]), + ).toEqual([]); + expect( + localInputReferences([ + { + name: "nested", + inputSchema: { + type: "object", + properties: { + default: { type: "array", items: { $ref: "#/definitions/value" } }, + }, + }, + }, + ]), + ).toEqual(["nested.inputSchema.properties.default.items.$ref"]); +}); + const advertiseAndEnforceProcessEnvironmentKeyConstraint = async (): Promise => { const contract = TOOL_CONTRACTS.find( @@ -423,6 +462,7 @@ describe("MCP JSON Schema validity", () => { TOOL_CONTRACTS.map(({ name }) => name).sort(), ); expect(schemaErrors(tools)).toEqual([]); + expect(localInputReferences(tools)).toEqual([]); expect(recursiveReferences(tools)).toEqual([]); expectKnownAuthorityHints(tools); const byName = new Map(tools.map((tool) => [tool.name, tool])); @@ -737,12 +777,8 @@ it("ships nonrecursive input and output schemas in the generated catalog", () => expect(recursiveReferences(GENERATED_MCP_TOOL_CATALOG)).toEqual([]); }); -it("keeps advertised input references within the client expansion budget", () => { - for (const tool of GENERATED_MCP_TOOL_CATALOG) - expect( - localReferenceExpansions(tool.inputSchema), - `${tool.name} exhausts the shared reference-expansion budget`, - ).toBeLessThanOrEqual(32); +it("ships input schemas without local references in the generated catalog", () => { + expect(localInputReferences(GENERATED_MCP_TOOL_CATALOG)).toEqual([]); }); it("rejects unresolved output-schema references even for uncalled tools", () => {