fix(mcp): keep advertised inputs within client reference budgets (#1749)

Squash merge of https://github.com/morluto/rea/pull/1749.

Co-authored-by: akushonkamen <akushonkamen@163.com>
Co-authored-by: GLM <noreply@z.ai>

docs/mcp-contracts.md

@@ -604,8 +604,12 @@ failures for such entries name this cause and the operation to rerun with
 
 Advertised input and output schemas contain no reachable recursive references.
 Schemas share repeated definitions through schema-local references while
-preserving complete nested fields and validation rules. Input properties retain
-their descriptions and literal examples.
+preserving complete nested fields and validation rules. Input schemas whose
+shared definitions would exhaust a client's declaration-rendering reference
+budget (Pi 1.1.0 code-mode shares 32 reference expansions across an input and
+renders every later reference as `unknown`) are re-projected without shared
+definitions, so every advertised field renders with complete structure.
+Input properties retain their descriptions and literal examples.
 
 When a host forwards input schemas to the model, their examples consume context.
 REA's advertised examples therefore never copy complete producer Evidence.
@@ -652,7 +656,7 @@ JSON Schema:
   objects or arrays.
 - The few schemas whose validation structure alone exceeds the budget
   (`compare_web_captures`, `build_reconstruction_obligation_ledger`,
-  `evaluate_reconstruction_coverage`, `project_managed_application_graph`,
+  `evaluate_reconstruction_coverage`,
   `compare_application_versions`, `capture_browser_scenario`) are advertised
   in a reduced form: property names, reference-resolved types (including union
   alternatives), first-level guidance, and the root's required-field constraints,
@@ -662,8 +666,7 @@ JSON Schema:
   instead of silent acceptance.
 
 Boundary tests pin every advertised compact schema to the budget, keep them
-valid JSON Schema that still accepts every canonical example, and hold the
-set of six reduced presentations as contracts evolve. Client verification
+valid JSON Schema that still accepts every canonical example. Client verification
 also checks provider-facing sizes and canonical examples after conversion by
 the actual Kimi Code CLI; local captures do not establish Moonshot acceptance.
 

src/contracts/inputSchemaPresentation.ts

@@ -30,6 +30,7 @@ const resolveReference = (
   root: Readonly<Record<string, unknown>>,
   reference: string,
 ): unknown => {
+  if (reference === "#") return root;
   if (!reference.startsWith("#/")) return undefined;
   let value: unknown = root;
   for (const token of reference.slice(2).split("/")) {
@@ -229,3 +230,76 @@ export const presentInputJsonSchema = (
   const described = describeSchema(flattened, flattened, describeProperty);
   return isObject(described) ? described : flattened;
 };
+
+/**
+ * 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<Record<string, unknown>>,
+  active: ReadonlySet<string>,
+  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<string, unknown>,
+): 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<string, unknown>,
+  inline: () => Record<string, unknown>,
+  budget: number,
+): Record<string, unknown> =>
+  countLocalReferenceExpansions(shared) <= budget ? shared : inline();

src/contracts/toolInputReferences.test.ts

@@ -3,7 +3,12 @@ import { Ajv2020 } from "ajv/dist/2020.js";
 import { describe, expect, it } from "vitest";
 import { z } from "zod";
 
-import { presentInputJsonSchema } from "./inputSchemaPresentation.js";
+import {
+  advertisedReferenceExpansionBudget,
+  countLocalReferenceExpansions,
+  inputSchemaWithinReferenceExpansionBudget,
+  presentInputJsonSchema,
+} from "./inputSchemaPresentation.js";
 import { toolContract } from "./toolContracts.js";
 import { toolInputSchemaWithMetadata } from "./toolSchemaMetadata.js";
 import { JAVASCRIPT_APPLICATION_EVIDENCE_EXAMPLE } from "./javascript/javascriptRuntimeReconciliationExample.js";
@@ -166,6 +171,105 @@ describe("advertised input references", () => {
   );
 });
 
+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<Readonly<Record<string, unknown>>>,
+  ) =>
+    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<string, unknown>;
+    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([

src/contracts/toolSchemaMetadata.ts

@@ -2,7 +2,11 @@ import { z } from "zod";
 
 import type { ToolContract } from "./toolContractTypes.js";
 import { advertisedInputExamples } from "./advertisedInputExamples.js";
-import { presentInputJsonSchema } from "./inputSchemaPresentation.js";
+import {
+  advertisedReferenceExpansionBudget,
+  inputSchemaWithinReferenceExpansionBudget,
+  presentInputJsonSchema,
+} from "./inputSchemaPresentation.js";
 import { digestCanonicalValue } from "../domain/canonicalDigest.js";
 
 const PROPERTY_DESCRIPTIONS: Readonly<Record<string, string>> = {
@@ -66,7 +70,25 @@ export const toolInputSchemaWithMetadata = <Contract extends ToolContract>(
       });
       const examples = advertisedInputExamples(contract);
       return {
-        ...presentInputJsonSchema(shared, fallbackPropertyDescription),
+        // 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,
+        ),
         ...(examples.length === 0 ? {} : { examples }),
       };
     },

tests/boundary/mcp/toolSchemaValidity.test.ts

@@ -217,6 +217,46 @@ function resolveReference(schema: unknown, reference: string): unknown {
   return value;
 }
 
+function localReferenceExpansions(schema: Record<string, unknown>): 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<string>): 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) => {
@@ -674,6 +714,14 @@ 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("rejects unresolved output-schema references even for uncalled tools", () => {
   const errors = schemaErrors([
     {