From a0fe672ca2812c6dcccbc56ecdfcb26b03829a20 Mon Sep 17 00:00:00 2001 From: "Miroslav Drbal [ApoC]" Date: Tue, 3 Mar 2026 22:52:41 +0100 Subject: [PATCH] fix: dereference MCP tool schema $ref/$defs and suppress Ajv format warnings (#276) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add dereferenceJsonSchema() that inlines local $ref pointers and strips $defs/definitions from MCP tool schemas before they reach LLM providers. Previously, Anthropic's convertTools() extracted only properties/required, dropping $defs and leaving dangling $ref — the LLM never saw the actual type definitions (e.g. SourceAnchorInput enum values from nucleus). - Silence Ajv logger (logger: false) on all three instances that use strict: false. MCP servers may declare non-standard format keywords (e.g. "uint") that caused console.warn() to corrupt TUI output. - Cache compiled Ajv validators per schema object identity in validation.ts, eliminating redundant recompilation on every tool call. Co-authored-by: Miroslav Drbal --- packages/ai/CHANGELOG.md | 5 + packages/ai/src/utils/schema/dereference.ts | 93 +++++++ packages/ai/src/utils/schema/index.ts | 1 + .../ai/src/utils/schema/sanitize-google.ts | 7 +- packages/ai/src/utils/validation.ts | 19 +- packages/ai/test/schema-dereference.test.ts | 255 ++++++++++++++++++ packages/coding-agent/src/task/executor.ts | 2 +- .../coding-agent/src/tools/submit-result.ts | 55 +--- 8 files changed, 382 insertions(+), 55 deletions(-) create mode 100644 packages/ai/src/utils/schema/dereference.ts create mode 100644 packages/ai/test/schema-dereference.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index b32572fa1..aec52fb7e 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Fixed + +- MCP tool schemas with `$ref`/`$defs` are now dereferenced before being sent to LLM providers, fixing dangling references that left models without type definitions +- Ajv schema validation no longer emits `console.warn()` for non-standard format keywords (e.g. `"uint"`) from MCP servers, preventing TUI corruption +- Tool schema compilation is now cached per schema identity, eliminating redundant recompilation on every tool call ## [13.6.0] - 2026-03-03 ### Added diff --git a/packages/ai/src/utils/schema/dereference.ts b/packages/ai/src/utils/schema/dereference.ts new file mode 100644 index 000000000..0f2716b4e --- /dev/null +++ b/packages/ai/src/utils/schema/dereference.ts @@ -0,0 +1,93 @@ +/** + * Inline `$ref` / `$defs` in a JSON Schema so every consumer sees + * the full definition without needing a resolver. + * + * Handles: + * - Local `$ref` pointers (`#/$defs/Foo`, `#/definitions/Foo`) + * - Nested `$defs` / `definitions` blocks + * - Circular references (breaks the cycle by emitting `{}`) + * + * After dereferencing, `$defs` and `definitions` are stripped from the root. + */ +import { isJsonObject, type JsonObject } from "./types"; + +/** + * Resolve a JSON-pointer-style `$ref` against the root schema's `$defs` + * or `definitions` block. Returns `undefined` for external or unresolvable refs. + */ +function resolveLocalRef(ref: string, root: JsonObject): JsonObject | undefined { + // Only handle local refs: #/$defs/Name or #/definitions/Name + const match = /^#\/(\$defs|definitions)\/(.+)$/.exec(ref); + if (!match) return undefined; + + const [, defsKey, name] = match; + const defs = root[defsKey!]; + if (!isJsonObject(defs)) return undefined; + + const resolved = defs[name!]; + return isJsonObject(resolved) ? resolved : undefined; +} + +/** + * Recursively dereference a JSON Schema node, inlining all local `$ref` pointers. + */ +function dereferenceNode(node: unknown, root: JsonObject, visiting: Set): unknown { + if (!isJsonObject(node)) return node; + if (Array.isArray(node)) return node.map(item => dereferenceNode(item, root, visiting)); + + const ref = node.$ref; + if (typeof ref === "string") { + // Break circular references + if (visiting.has(ref)) return {}; + const resolved = resolveLocalRef(ref, root); + if (!resolved) return node; // External ref — leave as-is + visiting.add(ref); + const inlined = dereferenceNode(resolved, root, visiting); + visiting.delete(ref); + + // Merge sibling keywords (e.g. description, default) from the + // referencing node. In draft 2020-12 these are valid alongside $ref. + const hasSiblings = Object.keys(node).some(k => k !== "$ref"); + if (!hasSiblings || !isJsonObject(inlined)) return inlined; + const merged: JsonObject = { ...inlined }; + for (const [key, value] of Object.entries(node)) { + if (key !== "$ref") merged[key] = value; + } + return merged; + } + + const result: JsonObject = {}; + for (const [key, value] of Object.entries(node)) { + // Skip $defs/definitions — they get inlined into consumers + if (key === "$defs" || key === "definitions") continue; + + if (Array.isArray(value)) { + result[key] = value.map(item => dereferenceNode(item, root, visiting)); + } else if (isJsonObject(value)) { + result[key] = dereferenceNode(value, root, visiting); + } else { + result[key] = value; + } + } + return result; +} + +/** + * Dereference all local `$ref` pointers in a JSON Schema, inlining definitions + * from `$defs` / `definitions`. The `$defs` block is stripped from the output. + * + * Non-local refs (e.g. `http://...`) are left untouched. + * Circular references are broken with `{}`. + * + * @returns A new schema object with all local refs inlined, or the input unchanged + * if it's not an object or has no `$defs`/`definitions`. + */ +export function dereferenceJsonSchema(schema: unknown): unknown { + if (!isJsonObject(schema)) return schema; + + // Fast path: nothing to dereference + const hasDefs = schema.$defs !== undefined || schema.definitions !== undefined; + if (!hasDefs) return schema; + + return dereferenceNode(schema, schema, new Set()); +} diff --git a/packages/ai/src/utils/schema/index.ts b/packages/ai/src/utils/schema/index.ts index be356e59f..3771a5ef0 100644 --- a/packages/ai/src/utils/schema/index.ts +++ b/packages/ai/src/utils/schema/index.ts @@ -1,5 +1,6 @@ export * from "./adapt"; export * from "./compatibility"; +export * from "./dereference"; export * from "./equality"; export * from "./fields"; export * from "./normalize-cca"; diff --git a/packages/ai/src/utils/schema/sanitize-google.ts b/packages/ai/src/utils/schema/sanitize-google.ts index d7a0def32..7cc445667 100644 --- a/packages/ai/src/utils/schema/sanitize-google.ts +++ b/packages/ai/src/utils/schema/sanitize-google.ts @@ -1,3 +1,4 @@ +import { dereferenceJsonSchema } from "./dereference"; import { areJsonValuesEqual } from "./equality"; import { UNSUPPORTED_SCHEMA_FIELDS } from "./fields"; @@ -197,7 +198,11 @@ const MCP_UNSUPPORTED_SCHEMA_FIELDS = new Set(["$schema"]); * (`pattern`, `format`, `additionalProperties`, etc.) and `$ref`/`$defs`. */ export function sanitizeSchemaForMCP(value: unknown): unknown { - return sanitizeSchemaImpl(value, { + // Dereference $ref/$defs first — MCP servers emit standard JSON Schema + // with $defs, but providers (Anthropic, Google) only forward `properties` + // and `required`, dropping $defs and leaving dangling $ref pointers. + const dereferenced = dereferenceJsonSchema(value); + return sanitizeSchemaImpl(dereferenced, { insideProperties: false, normalizeTypeArrayToNullable: false, stripNullableKeyword: true, diff --git a/packages/ai/src/utils/validation.ts b/packages/ai/src/utils/validation.ts index 9599db34e..702f80790 100644 --- a/packages/ai/src/utils/validation.ts +++ b/packages/ai/src/utils/validation.ts @@ -451,12 +451,28 @@ function coerceArgsFromErrors( // Create a singleton AJV instance with formats (only if not in browser extension) // AJV requires 'unsafe-eval' CSP which is not allowed in Manifest V3 +// +// Silent logger: MCP servers may declare non-standard format keywords (e.g. "uint") +// which cause Ajv to emit console.warn() with strict:false — corrupting TUI output. const ajv = new Ajv({ allErrors: true, strict: false, + logger: false, }); addFormats(ajv); +// Cache compiled validators by schema object identity to avoid +// re-compiling the same tool schema on every call. +const compiledSchemaCache = new WeakMap(); +function compileSchema(schema: object): import("ajv").ValidateFunction { + let validate = compiledSchemaCache.get(schema); + if (!validate) { + validate = ajv.compile(schema); + compiledSchemaCache.set(schema, validate); + } + return validate; +} + const MAX_TYPE_COERCION_PASSES = 5; /** @@ -484,8 +500,7 @@ export function validateToolCall(tools: Tool[], toolCall: ToolCall): any { export function validateToolArguments(tool: Tool, toolCall: ToolCall): any { const originalArgs = toolCall.arguments; - // Compile the schema - const validate = ajv.compile(tool.parameters); + const validate = compileSchema(tool.parameters); // Validate the arguments if (validate(originalArgs)) { diff --git a/packages/ai/test/schema-dereference.test.ts b/packages/ai/test/schema-dereference.test.ts new file mode 100644 index 000000000..c2dbbd8bc --- /dev/null +++ b/packages/ai/test/schema-dereference.test.ts @@ -0,0 +1,255 @@ +import { describe, expect, it } from "bun:test"; +import { dereferenceJsonSchema } from "@oh-my-pi/pi-ai/utils/schema"; + +describe("dereferenceJsonSchema", () => { + it("returns non-object input unchanged", () => { + expect(dereferenceJsonSchema(null)).toBe(null); + expect(dereferenceJsonSchema("string")).toBe("string"); + expect(dereferenceJsonSchema(42)).toBe(42); + }); + + it("returns schema without $defs unchanged", () => { + const schema = { + type: "object", + properties: { name: { type: "string" } }, + required: ["name"], + }; + expect(dereferenceJsonSchema(schema)).toBe(schema); + }); + + it("inlines a simple $ref from $defs", () => { + const schema = { + type: "object", + properties: { + anchor: { $ref: "#/$defs/Anchor" }, + }, + $defs: { + Anchor: { + type: "object", + properties: { + path: { type: "string" }, + line: { type: "integer", minimum: 0 }, + }, + }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.$defs).toBeUndefined(); + expect(result.properties.anchor).toEqual({ + type: "object", + properties: { + path: { type: "string" }, + line: { type: "integer", minimum: 0 }, + }, + }); + }); + + it("inlines $ref from definitions (legacy keyword)", () => { + const schema = { + type: "object", + properties: { + item: { $ref: "#/definitions/Item" }, + }, + definitions: { + Item: { type: "string", enum: ["a", "b"] }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.definitions).toBeUndefined(); + expect(result.properties.item).toEqual({ type: "string", enum: ["a", "b"] }); + }); + + it("inlines nested $ref inside arrays (items, anyOf, oneOf)", () => { + const schema = { + type: "object", + properties: { + anchors: { + type: "array", + items: { $ref: "#/$defs/SourceAnchorInput" }, + }, + value: { + anyOf: [{ $ref: "#/$defs/Str" }, { type: "null" }], + }, + }, + $defs: { + SourceAnchorInput: { + type: "object", + properties: { + anchor_type: { type: "string", enum: ["file", "symbol", "pattern"] }, + path: { type: "string" }, + }, + required: ["anchor_type", "path"], + }, + Str: { type: "string" }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.$defs).toBeUndefined(); + expect(result.properties.anchors.items).toEqual({ + type: "object", + properties: { + anchor_type: { type: "string", enum: ["file", "symbol", "pattern"] }, + path: { type: "string" }, + }, + required: ["anchor_type", "path"], + }); + expect(result.properties.value.anyOf[0]).toEqual({ type: "string" }); + expect(result.properties.value.anyOf[1]).toEqual({ type: "null" }); + }); + + it("handles $ref inside a definition pointing to another definition", () => { + const schema = { + type: "object", + properties: { + wrapper: { $ref: "#/$defs/Wrapper" }, + }, + $defs: { + Inner: { type: "integer", minimum: 0 }, + Wrapper: { + type: "object", + properties: { + value: { $ref: "#/$defs/Inner" }, + }, + }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.$defs).toBeUndefined(); + expect(result.properties.wrapper).toEqual({ + type: "object", + properties: { + value: { type: "integer", minimum: 0 }, + }, + }); + }); + + it("breaks circular $ref with empty object", () => { + const schema = { + type: "object", + properties: { + node: { $ref: "#/$defs/TreeNode" }, + }, + $defs: { + TreeNode: { + type: "object", + properties: { + name: { type: "string" }, + children: { + type: "array", + items: { $ref: "#/$defs/TreeNode" }, + }, + }, + }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.$defs).toBeUndefined(); + const node = result.properties.node; + expect(node.properties.name).toEqual({ type: "string" }); + // Circular ref breaks to {} + expect(node.properties.children.items).toEqual({}); + }); + + it("leaves external $ref untouched", () => { + const schema = { + type: "object", + properties: { + ext: { $ref: "https://example.com/schema.json#/Foo" }, + }, + $defs: {}, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.properties.ext).toEqual({ $ref: "https://example.com/schema.json#/Foo" }); + }); + + it("preserves sibling keywords alongside $ref", () => { + const schema = { + type: "object", + properties: { + anchor_type: { + $ref: "#/$defs/AnchorType", + description: "Field-level description", + default: "file", + }, + }, + $defs: { + AnchorType: { + type: "string", + enum: ["file", "symbol", "pattern"], + description: "Type-level description", + }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.$defs).toBeUndefined(); + // Sibling description overrides the definition's description + expect(result.properties.anchor_type).toEqual({ + type: "string", + enum: ["file", "symbol", "pattern"], + description: "Field-level description", + default: "file", + }); + }); + + it("handles multiple properties referencing the same $def", () => { + const schema = { + type: "object", + properties: { + a: { $ref: "#/$defs/Shared" }, + b: { $ref: "#/$defs/Shared" }, + }, + $defs: { + Shared: { type: "string", maxLength: 100 }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + expect(result.$defs).toBeUndefined(); + expect(result.properties.a).toEqual({ type: "string", maxLength: 100 }); + expect(result.properties.b).toEqual({ type: "string", maxLength: 100 }); + }); + + it("reproduces the nucleus write_memory schema pattern", () => { + // Simplified version of the actual nucleus tool schema + const schema = { + type: "object", + properties: { + kind: { type: "string", description: "Memory kind" }, + content: { type: "string", description: "Memory content" }, + anchors: { + description: "Source anchors", + items: { $ref: "#/$defs/SourceAnchorInput" }, + type: "array", + }, + }, + required: ["kind", "content", "anchors"], + $defs: { + SourceAnchorInput: { + type: "object", + properties: { + anchor_type: { + description: "Anchor type", + enum: ["file", "symbol", "pattern"], + type: "string", + }, + path: { type: "string" }, + role: { type: "string" }, + symbol: { type: "string" }, + }, + required: ["anchor_type", "path"], + }, + }, + }; + const result = dereferenceJsonSchema(schema) as any; + + // $defs stripped + expect(result.$defs).toBeUndefined(); + + // anchors items fully inlined + expect(result.properties.anchors.items.properties.anchor_type.enum).toEqual(["file", "symbol", "pattern"]); + expect(result.properties.anchors.items.required).toEqual(["anchor_type", "path"]); + + // Other properties preserved + expect(result.properties.kind).toEqual({ type: "string", description: "Memory kind" }); + expect(result.required).toEqual(["kind", "content", "anchors"]); + }); +}); diff --git a/packages/coding-agent/src/task/executor.ts b/packages/coding-agent/src/task/executor.ts index a34af06da..e75e537aa 100644 --- a/packages/coding-agent/src/task/executor.ts +++ b/packages/coding-agent/src/task/executor.ts @@ -41,7 +41,7 @@ import { } from "./types"; const MCP_CALL_TIMEOUT_MS = 60_000; -const ajv = new Ajv({ allErrors: true, strict: false }); +const ajv = new Ajv({ allErrors: true, strict: false, logger: false }); /** Agent event types to forward for progress tracking. */ const agentEventTypes = new Set([ diff --git a/packages/coding-agent/src/tools/submit-result.ts b/packages/coding-agent/src/tools/submit-result.ts index cde31fe7f..4dac49761 100644 --- a/packages/coding-agent/src/tools/submit-result.ts +++ b/packages/coding-agent/src/tools/submit-result.ts @@ -4,7 +4,7 @@ * Subagents must call this tool to finish and return structured JSON output. */ import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core"; -import { sanitizeSchemaForStrictMode } from "@oh-my-pi/pi-ai/utils/schema"; +import { dereferenceJsonSchema, sanitizeSchemaForStrictMode } from "@oh-my-pi/pi-ai/utils/schema"; import type { Static, TSchema } from "@sinclair/typebox"; import { Type } from "@sinclair/typebox"; import Ajv, { type ErrorObject, type ValidateFunction } from "ajv"; @@ -18,7 +18,7 @@ export interface SubmitResultDetails { error?: string; } -const ajv = new Ajv({ allErrors: true, strict: false }); +const ajv = new Ajv({ allErrors: true, strict: false, logger: false }); function normalizeSchema(schema: unknown): { normalized?: unknown; error?: string } { if (schema === undefined || schema === null) return {}; @@ -52,53 +52,6 @@ function formatAjvErrors(errors: ErrorObject[] | null | undefined): string { .join("; "); } -/** - * Resolve all $ref references in a JSON Schema by inlining definitions. - * Handles $defs and definitions at any nesting level. - * Removes $defs/definitions from the output since all refs are inlined. - */ -function resolveSchemaRefs(schema: Record): Record { - const defs: Record> = {}; - const defsObj = schema.$defs ?? schema.definitions; - if (defsObj && typeof defsObj === "object" && !Array.isArray(defsObj)) { - for (const [name, def] of Object.entries(defsObj as Record)) { - if (def && typeof def === "object" && !Array.isArray(def)) { - defs[name] = def as Record; - } - } - } - if (Object.keys(defs).length === 0) return schema; - - const inlining = new Set(); - function inline(node: unknown): unknown { - if (node === null || typeof node !== "object") return node; - if (Array.isArray(node)) return node.map(inline); - const obj = node as Record; - const ref = obj.$ref; - if (typeof ref === "string") { - const match = ref.match(/^#\/(?:\$defs|definitions)\/(.+)$/); - if (match) { - const name = match[1]; - const def = defs[name]; - if (def) { - if (inlining.has(name)) return {}; - inlining.add(name); - const resolved = inline(def); - inlining.delete(name); - return resolved; - } - } - } - const result: Record = {}; - for (const [key, value] of Object.entries(obj)) { - if (key === "$defs" || key === "definitions") continue; - result[key] = inline(value); - } - return result; - } - return inline(schema) as Record; -} - export class SubmitResultTool implements AgentTool { readonly name = "submit_result"; readonly label = "Submit Result"; @@ -168,11 +121,11 @@ export class SubmitResultTool implements AgentTool : undefined; if (sanitizedSchema !== undefined) { - const resolved = resolveSchemaRefs({ + const resolved = dereferenceJsonSchema({ ...sanitizedSchema, description: schemaDescription, }); - dataSchema = Type.Unsafe(resolved); + dataSchema = Type.Unsafe(resolved as Record); } else { dataSchema = Type.Record(Type.String(), Type.Any(), { description: schemaError ? schemaDescription : "Structured JSON output (no schema specified)",