From 9ba412853409d1075018f43b4eb42883beaff708 Mon Sep 17 00:00:00 2001 From: bensynapse <118375461+bensynapse@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:33:19 +0300 Subject: [PATCH] fix(sdk): repair tool schemas via cross-tool $defs pool with transitive closure The hosted MCP server's upload_design_md outputSchema references #/$defs/ScreenInstance without shipping $defs (#367), and strict MCP clients drop the entire tool list when AJV fails to compile it. The existing repair only injected hardcoded stub definitions that had drifted from the backend (wrong SelectedScreenInstance shape, missing ScreenInstance members, unrepresentable File -> UserFeedback ref). - collectDefPool() harvests every $defs entry across the tools/list response so repair can inject the backend's real definitions - repairSchema() resolves from the pool first, falls back to stubs, and iterates to a fixpoint so defs introducing new refs (e.g. File -> UserFeedback) are fully repaired - Fallback stubs synced with the live backend shapes (2026-08-21) Verified against a live tools/list capture: all 15 tools fully resolve after repair. 198/198 package tests pass. --- packages/sdk/src/index.ts | 6 +- packages/sdk/src/schema-repair.ts | 156 +++++++++-- packages/sdk/test/unit/schema-repair.test.ts | 269 ++++++++++++++++++- 3 files changed, 405 insertions(+), 26 deletions(-) diff --git a/packages/sdk/src/index.ts b/packages/sdk/src/index.ts index 46b509c..bc6decb 100644 --- a/packages/sdk/src/index.ts +++ b/packages/sdk/src/index.ts @@ -21,7 +21,11 @@ export { DesignSystem } from "../generated/src/designsystem.js"; // Infrastructure (handwritten) export { StitchToolClient } from "./client.js"; export { StitchProxy } from "./proxy/core.js"; -export { repairToolSchemas, repairSchema } from "./schema-repair.js"; +export { + repairToolSchemas, + repairSchema, + collectDefPool, +} from "./schema-repair.js"; // Virtual Tools export { downloadAssetsTool } from "./proxy/virtual-tools.js"; diff --git a/packages/sdk/src/schema-repair.ts b/packages/sdk/src/schema-repair.ts index af3b661..604f7b7 100644 --- a/packages/sdk/src/schema-repair.ts +++ b/packages/sdk/src/schema-repair.ts @@ -16,17 +16,19 @@ import type { Tool } from "@modelcontextprotocol/sdk/types.js"; /** * Well-known $defs definitions that the Stitch backend may reference via - * $ref but omit from the schema's $defs block. When the MCP SDK's - * AJV validator tries to compile these schemas, the missing references - * cause a hard crash (`MissingRefError`). + * $ref but omit from a schema's $defs block. When the MCP SDK's AJV + * validator tries to compile these schemas, the missing references cause a + * hard crash (`MissingRefError`). * - * This registry lets us inject stub definitions *before* AJV ever sees - * the schema, making the repair order-independent of the MCP SDK version. + * These are FALLBACK stubs, used only when a referenced definition cannot + * be harvested from another schema in the same tools/list response (see + * collectDefPool). They were captured from the live Stitch tools/list on + * 2026-08-21 so the fallback shape stays faithful to the backend. */ const WELL_KNOWN_DEFS: Record = { ScreenInstance: { type: "object", - description: "An instance of a screen on the project.", + description: "An instance of a screen on the project. Next ID: 18", properties: { groupId: { type: "string" }, groupName: { type: "string" }, @@ -34,9 +36,12 @@ const WELL_KNOWN_DEFS: Record = { hidden: { type: "boolean" }, id: { type: "string" }, isFavourite: { type: "boolean" }, + isResized: { type: "boolean" }, label: { type: "string" }, + needsLayout: { type: "boolean" }, sourceAsset: { type: "string" }, sourceScreen: { type: "string" }, + textContent: { type: "string" }, type: { type: "string", enum: [ @@ -44,8 +49,13 @@ const WELL_KNOWN_DEFS: Record = { "SCREEN_INSTANCE", "DESIGN_SYSTEM_INSTANCE", "GROUP_INSTANCE", + "TEXT_INSTANCE", ], }, + variantScreenInstance: { + $ref: "#/$defs/ScreenInstance", + description: "Optional. The variant Screen Instance.", + }, width: { type: "integer", format: "int32" }, x: { type: "integer", format: "int32" }, y: { type: "integer", format: "int32" }, @@ -54,11 +64,13 @@ const WELL_KNOWN_DEFS: Record = { SelectedScreenInstance: { type: "object", - description: "A selected screen instance reference.", + description: + "A screen instance to be edited by the agent, selected by the user.", properties: { - screenId: { type: "string" }, - instanceId: { type: "string" }, + id: { type: "string" }, + sourceScreen: { type: "string" }, }, + required: ["id", "sourceScreen"], }, File: { @@ -66,10 +78,40 @@ const WELL_KNOWN_DEFS: Record = { description: "A File resource.", properties: { downloadUrl: { type: "string" }, - fileContentBase64: { type: "string" }, + fileContentBase64: { type: "string", writeOnly: true }, mimeType: { type: "string" }, name: { type: "string" }, uploadBlobId: { type: "string" }, + userFeedback: { + $ref: "#/$defs/UserFeedback", + description: "Output only. The latest feedback submitted for the file.", + readOnly: true, + }, + }, + }, + + UserFeedback: { + type: "object", + description: "User feedback for a given interaction.", + properties: { + comment: { type: "string" }, + designFeedbackReason: { + type: "string", + enum: [ + "DESIGN_FEEDBACK_REASON_UNSPECIFIED", + "DESIGN_DOESNT_MATCH_PROMPT", + "EDIT_DOESNT_MATCH_PROMPT", + "DESCRIPTION_DOESNT_MATCH", + "COMPONENT_ISSUE", + "INCORRECT_THEME", + "FIGMA_EXPORT_FAILED", + "OTHER", + ], + }, + rating: { + type: "string", + enum: ["RATING_UNSPECIFIED", "POSITIVE", "NEGATIVE"], + }, }, }, }; @@ -102,26 +144,92 @@ function collectRefTargets( return refs; } +function isPlainObject(value: unknown): value is Record { + return value !== null && typeof value === "object" && !Array.isArray(value); +} + /** - * Repair a single JSON Schema by injecting any missing well-known $defs - * that are referenced via $ref but not present. + * Harvest every `$defs` entry found across the tools/list response into a + * name → definition pool. + * + * The Stitch backend usually defines each entity (ScreenInstance, File, …) + * under `$defs` in at least one tool's schema, even when it omits them from + * other schemas that reference them. Pooling those real definitions lets + * repairSchema inject the backend's actual shapes instead of the fallback + * stubs in WELL_KNOWN_DEFS. + */ +export function collectDefPool(tools: Tool[]): Record { + const pool: Record = {}; + + const harvest = (schema: unknown) => { + if (!isPlainObject(schema)) return; + const defs = schema.$defs; + if (!isPlainObject(defs)) return; + for (const [name, def] of Object.entries(defs)) { + if (!(name in pool) && isPlainObject(def)) { + pool[name] = def; + } + } + }; + + for (const tool of tools) { + harvest(tool.inputSchema); + harvest((tool as any).outputSchema); + } + + return pool; +} + +/** + * Bound on repair passes. Injected definitions can introduce new $refs + * (e.g. the backend's File def references UserFeedback), so repair iterates + * to a fixpoint; real def chains are 1–2 deep, so 8 passes is generous. + */ +const MAX_REPAIR_PASSES = 8; + +/** + * Repair a single JSON Schema by injecting any missing $defs that are + * referenced via $ref but not present. + * + * Definitions are resolved from `defPool` first (the backend's real shapes, + * harvested from sibling tool schemas), falling back to WELL_KNOWN_DEFS + * stubs. Injected definitions are deep-cloned so schemas never share + * mutable state, and injection iterates until every transitive reference + * resolves. * * Mutates the schema in place and returns it for convenience. */ -export function repairSchema(schema: Record): Record { +export function repairSchema( + schema: Record, + defPool: Record = {}, +): Record { if (!schema || typeof schema !== "object") return schema; - const referencedDefs = collectRefTargets(schema); - if (referencedDefs.size === 0) return schema; + const unresolved = new Set(); + + for (let pass = 0; pass < MAX_REPAIR_PASSES; pass++) { + const referencedDefs = collectRefTargets(schema); + if (referencedDefs.size === 0) return schema; - // Ensure $defs block exists - schema.$defs = schema.$defs || {}; + // Ensure $defs block exists + schema.$defs = schema.$defs || {}; - for (const defName of referencedDefs) { - // Only inject if: (a) the def is missing, and (b) we have a well-known stub - if (!schema.$defs[defName] && WELL_KNOWN_DEFS[defName]) { - schema.$defs[defName] = { ...WELL_KNOWN_DEFS[defName] }; + let injected = false; + for (const defName of referencedDefs) { + // Only inject if the def is missing and has not proven unresolvable + if (schema.$defs[defName] || unresolved.has(defName)) continue; + + const source = defPool[defName] ?? WELL_KNOWN_DEFS[defName]; + if (!source) { + unresolved.add(defName); + continue; + } + + schema.$defs[defName] = JSON.parse(JSON.stringify(source)); + injected = true; } + + if (!injected) return schema; } return schema; @@ -134,15 +242,17 @@ export function repairSchema(schema: Record): Record { * Mutates tools in place. */ export function repairToolSchemas(tools: Tool[]): void { + const defPool = collectDefPool(tools); + for (const tool of tools) { if (tool.inputSchema && typeof tool.inputSchema === "object") { - repairSchema(tool.inputSchema as Record); + repairSchema(tool.inputSchema as Record, defPool); } // outputSchema was added in MCP SDK ≥1.27 and is the primary crash vector: // Client.cacheToolMetadata() eagerly compiles outputSchema with AJV. const anyTool = tool as any; if (anyTool.outputSchema && typeof anyTool.outputSchema === "object") { - repairSchema(anyTool.outputSchema); + repairSchema(anyTool.outputSchema, defPool); } } } diff --git a/packages/sdk/test/unit/schema-repair.test.ts b/packages/sdk/test/unit/schema-repair.test.ts index 5397b6d..69f6715 100644 --- a/packages/sdk/test/unit/schema-repair.test.ts +++ b/packages/sdk/test/unit/schema-repair.test.ts @@ -13,7 +13,11 @@ // limitations under the License. import { describe, it, expect } from "vitest"; -import { repairSchema, repairToolSchemas } from "../../src/schema-repair.js"; +import { + repairSchema, + repairToolSchemas, + collectDefPool, +} from "../../src/schema-repair.js"; import type { Tool } from "@modelcontextprotocol/sdk/types.js"; describe("repairSchema", () => { @@ -66,9 +70,13 @@ describe("repairSchema", () => { repairSchema(schema); expect(schema.$defs.SelectedScreenInstance).toBeDefined(); - expect(schema.$defs.SelectedScreenInstance.properties.screenId).toEqual({ + expect(schema.$defs.SelectedScreenInstance.properties.id).toEqual({ type: "string", }); + expect(schema.$defs.SelectedScreenInstance.required).toEqual([ + "id", + "sourceScreen", + ]); }); it("should NOT overwrite existing $defs", () => { @@ -236,3 +244,260 @@ describe("repairToolSchemas", () => { expect(() => repairToolSchemas([])).not.toThrow(); }); }); + +/** Collect every local `#/$defs/` ref in a schema. */ +function collectLocalRefs(node: any, out: string[] = []): string[] { + if (!node || typeof node !== "object") return out; + if (Array.isArray(node)) { + for (const v of node) collectLocalRefs(v, out); + return out; + } + if (typeof node.$ref === "string" && node.$ref.startsWith("#/$defs/")) { + out.push(node.$ref.slice("#/$defs/".length)); + } + for (const v of Object.values(node)) collectLocalRefs(v, out); + return out; +} + +/** Assert every local $ref in the schema resolves against its own $defs. */ +function expectAllRefsResolve(schema: any) { + const defs = schema.$defs || {}; + const refs = collectLocalRefs(schema); + expect(refs.length).toBeGreaterThan(0); + for (const ref of refs) { + expect(defs, `expected $defs.${ref} to be present`).toHaveProperty(ref); + } +} + +describe("collectDefPool", () => { + it("should harvest $defs from both inputSchema and outputSchema", () => { + const tools: any[] = [ + { + name: "a", + inputSchema: { + type: "object", + $defs: { FromInput: { type: "object" } }, + }, + outputSchema: { + type: "object", + $defs: { FromOutput: { type: "object" } }, + }, + }, + ]; + + const pool = collectDefPool(tools); + + expect(pool.FromInput).toBeDefined(); + expect(pool.FromOutput).toBeDefined(); + }); + + it("should keep the first definition seen for a name", () => { + const first = { + type: "object", + properties: { first: { type: "boolean" } }, + }; + const second = { + type: "object", + properties: { second: { type: "boolean" } }, + }; + const tools: any[] = [ + { name: "a", inputSchema: { type: "object", $defs: { Dup: first } } }, + { name: "b", inputSchema: { type: "object", $defs: { Dup: second } } }, + ]; + + const pool = collectDefPool(tools); + + expect(pool.Dup).toBe(first); + }); + + it("should ignore tools and schemas without $defs", () => { + const tools: any[] = [ + { name: "a", inputSchema: { type: "object" } }, + { name: "b", inputSchema: { type: "object", $defs: null } }, + ]; + + expect(collectDefPool(tools)).toEqual({}); + }); +}); + +describe("pool-based repair", () => { + it("should prefer the backend's real definition over the fallback stub", () => { + const backendScreenInstance = { + type: "object", + properties: { backendOnlyMarker: { type: "boolean" } }, + }; + const tools: any[] = [ + { + name: "list_projects", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + $defs: { ScreenInstance: backendScreenInstance }, + }, + }, + { + name: "upload_design_md", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + properties: { + variantScreenInstance: { $ref: "#/$defs/ScreenInstance" }, + }, + }, + }, + ]; + + repairToolSchemas(tools); + + const repaired = tools[1].outputSchema.$defs.ScreenInstance; + expect(repaired.properties.backendOnlyMarker).toBeDefined(); + // Injected defs must be deep copies, not shared references + expect(repaired).not.toBe(backendScreenInstance); + }); + + it("should resolve transitive refs introduced by injected defs (File -> UserFeedback)", () => { + const backendFile = { + type: "object", + properties: { + name: { type: "string" }, + userFeedback: { $ref: "#/$defs/UserFeedback" }, + }, + }; + const backendUserFeedback = { + type: "object", + properties: { rating: { type: "string" } }, + }; + const tools: any[] = [ + { + name: "get_file", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + $defs: { File: backendFile, UserFeedback: backendUserFeedback }, + }, + }, + { + name: "upload_thing", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + properties: { file: { $ref: "#/$defs/File" } }, + }, + }, + ]; + + repairToolSchemas(tools); + + expectAllRefsResolve(tools[1].outputSchema); + expect(tools[1].outputSchema.$defs.UserFeedback).toBeDefined(); + }); + + it("should fall back to well-known stubs when the pool lacks a definition", () => { + const tools: any[] = [ + { + name: "lonely_tool", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + properties: { + screens: { + type: "array", + items: { $ref: "#/$defs/ScreenInstance" }, + }, + }, + }, + }, + ]; + + repairToolSchemas(tools); + + expectAllRefsResolve(tools[0].outputSchema); + expect( + tools[0].outputSchema.$defs.ScreenInstance.properties.id, + ).toBeDefined(); + }); + + it("should not overwrite a def already present in the target schema", () => { + const own = { type: "object", properties: { own: { type: "boolean" } } }; + const fromPool = { + type: "object", + properties: { pooled: { type: "boolean" } }, + }; + const tools: any[] = [ + { + name: "definer", + inputSchema: { type: "object", $defs: { ScreenInstance: fromPool } }, + }, + { + name: "consumer", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + $defs: { ScreenInstance: own }, + properties: { + variantScreenInstance: { $ref: "#/$defs/ScreenInstance" }, + }, + }, + }, + ]; + + repairToolSchemas(tools); + + expect(tools[1].outputSchema.$defs.ScreenInstance).toBe(own); + }); + + it("should repair the real upload_design_md shape so every $ref resolves", () => { + // Regression test for https://github.com/google-labs-code/stitch-sdk/issues/367 + // The live backend emits upload_design_md's outputSchema as an inlined + // ScreenInstance object that retains a $ref to "#/$defs/ScreenInstance" + // but ships no $defs block of its own. + const recursiveScreenInstance = { + type: "object", + description: "An instance of a screen on the project.", + properties: { + id: { type: "string" }, + type: { type: "string", enum: ["SCREEN_INSTANCE", "TEXT_INSTANCE"] }, + variantScreenInstance: { $ref: "#/$defs/ScreenInstance" }, + }, + }; + const tools: any[] = [ + { + name: "create_project", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + $defs: { ScreenInstance: recursiveScreenInstance }, + properties: { + screenInstances: { + type: "array", + items: { $ref: "#/$defs/ScreenInstance" }, + }, + }, + }, + }, + { + name: "upload_design_md", + inputSchema: { type: "object" }, + outputSchema: { + type: "object", + description: "An instance of a screen on the project.", + properties: { + id: { type: "string" }, + variantScreenInstance: { $ref: "#/$defs/ScreenInstance" }, + }, + // NOTE: no $defs — the dangling reference that crashed clients + }, + }, + ]; + + repairToolSchemas(tools); + + expectAllRefsResolve(tools[0].outputSchema); + expectAllRefsResolve(tools[1].outputSchema); + // The recursive def must survive injection intact + expect( + tools[1].outputSchema.$defs.ScreenInstance.properties + .variantScreenInstance.$ref, + ).toBe("#/$defs/ScreenInstance"); + }); +});