From b623b619de1c5bb73dd2eb85b99922d2207681cd Mon Sep 17 00:00:00 2001 From: debian Date: Fri, 30 Oct 2026 08:52:54 -0300 Subject: [PATCH] fix(core): resolve and preserve MCP tool inputSchema $defs/$ref definitions --- packages/core/src/mcp/mcp-client.ts | 2 + packages/core/src/mcp/mcp-manager.ts | 111 +++++++++++++++++++-- packages/core/src/prompt.ts | 2 + packages/core/src/tests/mcp-client.test.ts | 92 +++++++++++++++++ 4 files changed, 199 insertions(+), 8 deletions(-) diff --git a/packages/core/src/mcp/mcp-client.ts b/packages/core/src/mcp/mcp-client.ts index 4420c569..1ab06bfd 100644 --- a/packages/core/src/mcp/mcp-client.ts +++ b/packages/core/src/mcp/mcp-client.ts @@ -31,6 +31,8 @@ export type McpToolDefinition = { properties: Record; required?: string[]; additionalProperties?: boolean; + $defs?: Record; + definitions?: Record; }; }; diff --git a/packages/core/src/mcp/mcp-manager.ts b/packages/core/src/mcp/mcp-manager.ts index 6d2edc63..28ba64f5 100644 --- a/packages/core/src/mcp/mcp-manager.ts +++ b/packages/core/src/mcp/mcp-manager.ts @@ -313,6 +313,8 @@ export class McpManager { properties: Record; required?: string[]; additionalProperties?: boolean; + $defs?: Record; + definitions?: Record; }; }; }> { @@ -321,14 +323,7 @@ export class McpManager { function: { name: t.namespacedName, description: this.buildMcpToolDescription(t), - parameters: { - type: "object" as const, - properties: t.definition.inputSchema.properties, - required: t.definition.inputSchema.required, - ...(t.definition.inputSchema.additionalProperties !== undefined - ? { additionalProperties: t.definition.inputSchema.additionalProperties } - : {}), - }, + parameters: normalizeMcpInputSchema(t.definition.inputSchema), }, })); } @@ -522,3 +517,103 @@ function fitApiToolNameWithSuffix(name: string, suffix: string): string { function hashToolName(value: string): string { return createHash("sha256").update(value).digest("hex").slice(0, 8); } + +export function normalizeMcpInputSchema(inputSchema?: McpToolDefinition["inputSchema"]): { + type: "object"; + properties: Record; + required?: string[]; + additionalProperties?: boolean; + $defs?: Record; + definitions?: Record; +} { + if (!inputSchema || typeof inputSchema !== "object") { + return { type: "object", properties: {} }; + } + + const defs: Record = { + ...(isPlainObject(inputSchema.definitions) ? inputSchema.definitions : {}), + ...(isPlainObject(inputSchema.$defs) ? inputSchema.$defs : {}), + }; + + let hasUnresolvedRef = false; + + const deref = (node: unknown, seen: ReadonlySet = new Set()): unknown => { + if (Array.isArray(node)) { + return node.map((item) => deref(item, seen)); + } + if (!isPlainObject(node)) { + return node; + } + + if (typeof node.$ref === "string") { + const ref = node.$ref; + let defKey: string | null = null; + if (ref.startsWith("#/$defs/")) { + defKey = ref.slice("#/$defs/".length); + } else if (ref.startsWith("#/definitions/")) { + defKey = ref.slice("#/definitions/".length); + } + + if (defKey && Object.prototype.hasOwnProperty.call(defs, defKey) && !seen.has(defKey)) { + const nextSeen = new Set(seen); + nextSeen.add(defKey); + const resolvedTarget = deref(defs[defKey], nextSeen); + const rest = { ...node }; + delete rest.$ref; + const resolvedRest = deref(rest, seen); + if (isPlainObject(resolvedTarget) && isPlainObject(resolvedRest)) { + return { ...resolvedTarget, ...resolvedRest }; + } + } + + hasUnresolvedRef = true; + } + + const out: Record = {}; + for (const [key, value] of Object.entries(node)) { + out[key] = deref(value, seen); + } + + if (Array.isArray(out.allOf) && out.allOf.length === 1 && isPlainObject(out.allOf[0])) { + const { allOf, ...rest } = out; + return { ...allOf[0], ...rest }; + } + + return out; + }; + + const resolvedProperties = isPlainObject(inputSchema.properties) + ? (deref(inputSchema.properties) as Record) + : {}; + + const result: { + type: "object"; + properties: Record; + required?: string[]; + additionalProperties?: boolean; + $defs?: Record; + definitions?: Record; + } = { + type: "object", + properties: resolvedProperties, + ...(inputSchema.required !== undefined ? { required: inputSchema.required } : {}), + ...(inputSchema.additionalProperties !== undefined + ? { additionalProperties: inputSchema.additionalProperties } + : {}), + }; + + if (hasUnresolvedRef) { + if (isPlainObject(inputSchema.$defs)) { + result.$defs = inputSchema.$defs; + } + if (isPlainObject(inputSchema.definitions)) { + result.definitions = inputSchema.definitions; + } + } + + return result; +} + +function isPlainObject(value: unknown): value is Record { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} diff --git a/packages/core/src/prompt.ts b/packages/core/src/prompt.ts index f7f147d3..7d938c63 100644 --- a/packages/core/src/prompt.ts +++ b/packages/core/src/prompt.ts @@ -394,6 +394,8 @@ export type ToolDefinition = { properties: Record; required?: string[]; additionalProperties?: boolean; + $defs?: Record; + definitions?: Record; }; }; }; diff --git a/packages/core/src/tests/mcp-client.test.ts b/packages/core/src/tests/mcp-client.test.ts index 6a7dc016..d39f9745 100644 --- a/packages/core/src/tests/mcp-client.test.ts +++ b/packages/core/src/tests/mcp-client.test.ts @@ -4,6 +4,7 @@ import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import path from "node:path"; import { McpClient, createMcpSpawnSpec } from "../mcp/mcp-client"; +import { normalizeMcpInputSchema } from "../mcp/mcp-manager"; test("createMcpSpawnSpec keeps non-Windows MCP launches shell-free", () => { assert.deepEqual(createMcpSpawnSpec("npx", ["-y", "@playwright/mcp@latest"], "darwin"), { @@ -108,3 +109,94 @@ test("McpClient starts a PATH-resolved cmd MCP server on Windows", { skip: proce rmSync(serverDir, { recursive: true, force: true }); } }); + +test("normalizeMcpInputSchema inlines $defs and definitions $ref pointers", () => { + const schema = normalizeMcpInputSchema({ + type: "object", + properties: { + body: { + $ref: "#/$defs/CreateDelegationRequest", + type: "object", + }, + tier: { + allOf: [{ $ref: "#/definitions/VolumeType" }], + description: "Storage tier override.", + }, + }, + required: ["body"], + $defs: { + CreateDelegationRequest: { + type: "object", + additionalProperties: false, + required: ["resource"], + properties: { + resource: { + type: "string", + description: "ECR resource ARN", + }, + }, + }, + }, + definitions: { + VolumeType: { + type: "string", + description: "Default volume type.", + enum: ["STANDARD", "HIGH_PERFORMANCE"], + }, + }, + }); + + assert.deepEqual(schema, { + type: "object", + properties: { + body: { + type: "object", + additionalProperties: false, + required: ["resource"], + properties: { + resource: { + type: "string", + description: "ECR resource ARN", + }, + }, + }, + tier: { + type: "string", + description: "Storage tier override.", + enum: ["STANDARD", "HIGH_PERFORMANCE"], + }, + }, + required: ["body"], + }); +}); + +test("normalizeMcpInputSchema preserves $defs when cyclic $ref cannot be fully inlined", () => { + const schema = normalizeMcpInputSchema({ + type: "object", + properties: { + node: { + $ref: "#/$defs/TreeNode", + }, + }, + $defs: { + TreeNode: { + type: "object", + properties: { + child: { + $ref: "#/$defs/TreeNode", + }, + }, + }, + }, + }); + + assert.equal(typeof schema.$defs, "object"); + assert.deepEqual(schema.properties.node, { + type: "object", + properties: { + child: { + $ref: "#/$defs/TreeNode", + }, + }, + }); +});