fix: repair tool schemas before AJV compilation to prevent MissingRefError (#355)

The MCP SDK (≥1.27) introduced Client.cacheToolMetadata() which eagerly
compiles outputSchema with AJV during listTools(). When the Stitch
backend returns schemas with $ref to missing $defs (e.g.
#/$defs/ScreenInstance), AJV throws MissingRefError before the SDK's
schema repair code can run.

Fix:
- Extract schema repair into a standalone module (schema-repair.ts)
  that scans for $ref targets and injects well-known stub definitions
- In StitchToolClient.listTools(): use raw request() instead of
  Client.listTools() to bypass AJV compilation, then apply repair
- In proxy refreshTools(): apply repair before re-serving tools to
  MCP clients whose AJV validators would also crash

The repair module:
- Scans both inputSchema and outputSchema for $ref targets
- Only injects stubs for referenced-but-missing well-known defs
- Never overwrites existing $defs
- Handles ScreenInstance, SelectedScreenInstance, and File types

Includes 13 unit tests covering injection, preservation, nesting,
unknown refs, and null safety.

Author: davideast

packages/sdk/src/client.ts

@@ -14,6 +14,7 @@
 
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import { ListToolsResultSchema } from "@modelcontextprotocol/sdk/types.js";
 import {
   StitchConfigSchema,
   StitchConfig,
@@ -23,6 +24,7 @@ import {
 import { StitchError, StitchErrorCode } from "./spec/errors.js";
 import { buildAuthHeaders as buildBaseAuthHeaders } from "./auth.js";
 import { SDK_VERSION } from "./version.js";
+import { repairToolSchemas } from "./schema-repair.js";
 
 /**
  * Authenticated tool pipe for the Stitch MCP Server.
@@ -235,63 +237,25 @@ export class StitchToolClient implements StitchToolClientSpec {
 
   async listTools() {
     if (!this.isConnected) await this.connect();
-    const remoteTools = await this.client.listTools();
+
+    // CRITICAL: We use a raw request() instead of this.client.listTools()
+    // because Client.listTools() eagerly compiles outputSchema with AJV
+    // via cacheToolMetadata(). If the Stitch backend returns schemas with
+    // $ref to missing $defs (e.g. #/$defs/ScreenInstance), AJV throws a
+    // MissingRefError BEFORE our schema repair code can run.
+    //
+    // By using request() directly, we get the raw tool list, apply schema
+    // repair to inject missing $defs, and avoid the AJV crash entirely.
+    const remoteTools = await (this.client as any).request(
+      { method: 'tools/list', params: {} },
+      ListToolsResultSchema,
+    );
 
     const tools = remoteTools.tools || [];
 
-    // Resilient Schema Repair: Dynamically patch broken $ref schemas from the Stitch backend
-    for (const tool of tools) {
-      const schema = tool.inputSchema as any;
-      if (schema && typeof schema === 'object') {
-        schema.$defs = schema.$defs || {};
-        
-        // Inject the missing ScreenInstance definition if referenced
-        if (!schema.$defs.ScreenInstance) {
-          schema.$defs.ScreenInstance = {
-            type: 'object',
-            description: 'An instance of a screen on the project.',
-            properties: {
-              groupId: { type: 'string' },
-              groupName: { type: 'string' },
-              height: { type: 'integer', format: 'int32' },
-              hidden: { type: 'boolean' },
-              id: { type: 'string' },
-              isFavourite: { type: 'boolean' },
-              label: { type: 'string' },
-              sourceAsset: { type: 'string' },
-              sourceScreen: { type: 'string' },
-              type: {
-                type: 'string',
-                enum: [
-                  'SCREEN_INSTANCE_TYPE_UNSPECIFIED',
-                  'SCREEN_INSTANCE',
-                  'DESIGN_SYSTEM_INSTANCE',
-                  'GROUP_INSTANCE'
-                ]
-              },
-              width: { type: 'integer', format: 'int32' },
-              x: { type: 'integer', format: 'int32' },
-              y: { type: 'integer', format: 'int32' }
-            }
-          };
-        }
-
-        // Inject the missing File definition if referenced
-        if (!schema.$defs.File) {
-          schema.$defs.File = {
-            type: 'object',
-            description: 'A File resource.',
-            properties: {
-              downloadUrl: { type: 'string' },
-              fileContentBase64: { type: 'string', writeOnly: true },
-              mimeType: { type: 'string' },
-              name: { type: 'string' },
-              uploadBlobId: { type: 'string' }
-            }
-          };
-        }
-      }
-    }
+    // Resilient Schema Repair: Inject missing $defs BEFORE any AJV
+    // compilation can occur. Repairs both inputSchema and outputSchema.
+    repairToolSchemas(tools);
 
     const localTools = this.localVirtualTools.map(t => ({
       name: t.name,

packages/sdk/src/index.ts

@@ -22,6 +22,7 @@ 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";
 
 // Virtual Tools
 export {

packages/sdk/src/proxy/client.ts

@@ -15,6 +15,7 @@
 import { StitchProxyConfig } from '../spec/proxy.js';
 import { buildAuthHeaders } from '../auth.js';
 import type { Tool } from '@modelcontextprotocol/sdk/types.js';
+import { repairToolSchemas } from '../schema-repair.js';
 
 /**
  * Shared state for proxy handlers.
@@ -111,10 +112,16 @@ export async function initializeStitchConnection(
 
 /**
  * Refresh the cached tools list from Stitch.
+ *
+ * Applies schema repair to inject missing $defs before the tools are
+ * re-served to MCP clients whose AJV validators would otherwise crash
+ * on unresolved $ref targets.
  */
 export async function refreshTools(ctx: ProxyContext): Promise<void> {
   const toolsResult = (await forwardToStitch(ctx.config, 'tools/list', {})) as {
     tools: Tool[];
   };
-  ctx.remoteTools = toolsResult.tools || [];
+  const tools = toolsResult.tools || [];
+  repairToolSchemas(tools);
+  ctx.remoteTools = tools;
 }

packages/sdk/src/schema-repair.ts

@@ -0,0 +1,145 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+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`).
+ *
+ * This registry lets us inject stub definitions *before* AJV ever sees
+ * the schema, making the repair order-independent of the MCP SDK version.
+ */
+const WELL_KNOWN_DEFS: Record<string, object> = {
+  ScreenInstance: {
+    type: 'object',
+    description: 'An instance of a screen on the project.',
+    properties: {
+      groupId: { type: 'string' },
+      groupName: { type: 'string' },
+      height: { type: 'integer', format: 'int32' },
+      hidden: { type: 'boolean' },
+      id: { type: 'string' },
+      isFavourite: { type: 'boolean' },
+      label: { type: 'string' },
+      sourceAsset: { type: 'string' },
+      sourceScreen: { type: 'string' },
+      type: {
+        type: 'string',
+        enum: [
+          'SCREEN_INSTANCE_TYPE_UNSPECIFIED',
+          'SCREEN_INSTANCE',
+          'DESIGN_SYSTEM_INSTANCE',
+          'GROUP_INSTANCE',
+        ],
+      },
+      width: { type: 'integer', format: 'int32' },
+      x: { type: 'integer', format: 'int32' },
+      y: { type: 'integer', format: 'int32' },
+    },
+  },
+
+  SelectedScreenInstance: {
+    type: 'object',
+    description: 'A selected screen instance reference.',
+    properties: {
+      screenId: { type: 'string' },
+      instanceId: { type: 'string' },
+    },
+  },
+
+  File: {
+    type: 'object',
+    description: 'A File resource.',
+    properties: {
+      downloadUrl: { type: 'string' },
+      fileContentBase64: { type: 'string' },
+      mimeType: { type: 'string' },
+      name: { type: 'string' },
+      uploadBlobId: { type: 'string' },
+    },
+  },
+};
+
+/**
+ * Collect every `$ref` target of the form `#/$defs/<Name>` from a schema
+ * object (recursively). Returns the set of referenced definition names.
+ */
+function collectRefTargets(obj: unknown, refs: Set<string> = new Set()): Set<string> {
+  if (obj === null || typeof obj !== 'object') return refs;
+
+  if (Array.isArray(obj)) {
+    for (const item of obj) collectRefTargets(item, refs);
+    return refs;
+  }
+
+  const record = obj as Record<string, unknown>;
+  if (typeof record.$ref === 'string') {
+    const match = record.$ref.match(/^#\/\$defs\/(.+)$/);
+    if (match) refs.add(match[1]);
+  }
+
+  for (const value of Object.values(record)) {
+    collectRefTargets(value, refs);
+  }
+
+  return refs;
+}
+
+/**
+ * Repair a single JSON Schema by injecting any missing well-known $defs
+ * that are referenced via $ref but not present.
+ *
+ * Mutates the schema in place and returns it for convenience.
+ */
+export function repairSchema(schema: Record<string, any>): Record<string, any> {
+  if (!schema || typeof schema !== 'object') return schema;
+
+  const referencedDefs = collectRefTargets(schema);
+  if (referencedDefs.size === 0) return schema;
+
+  // 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] };
+    }
+  }
+
+  return schema;
+}
+
+/**
+ * Apply schema repair to every tool's inputSchema and outputSchema.
+ *
+ * This MUST run before the MCP SDK's AJV validator sees the schemas.
+ * Mutates tools in place.
+ */
+export function repairToolSchemas(tools: Tool[]): void {
+  for (const tool of tools) {
+    if (tool.inputSchema && typeof tool.inputSchema === 'object') {
+      repairSchema(tool.inputSchema as Record<string, any>);
+    }
+    // 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);
+    }
+  }
+}