fix(google): bound tool schema expansion

src/adapters/google-tool-schema.ts

@@ -7,6 +7,27 @@ type Schema = Record<string, unknown>;
 const ALLOWED_TYPES = new Set(["string", "integer", "number", "boolean", "array", "object"]);
 const MAX_SCHEMA_DEPTH = 24; // Google's documented nesting limit is 32; leave headroom for CCA.
 const MAX_DEREF_DEPTH = 16;
+const MAX_SCHEMA_NODES = 1_024;
+const BUDGET_EXHAUSTED = Symbol("schema-budget-exhausted");
+const MERGED_SCHEMA_KEYS = [
+  "type",
+  "nullable",
+  "description",
+  "format",
+  "enum",
+  "const",
+  "properties",
+  "items",
+  "required",
+  "anyOf",
+] as const;
+
+type SanitizeResult = Schema | typeof BUDGET_EXHAUSTED;
+
+interface SanitizeState {
+  activeRefs: Set<string>;
+  remainingNodes: number;
+}
 
 function isRecord(value: unknown): value is Schema {
   return !!value && typeof value === "object" && !Array.isArray(value);
@@ -34,6 +55,16 @@ function collectDefs(root: unknown, defs: Map<string, unknown>): void {
   }
 }
 
+function mergeRefTarget(target: Schema, overlay: Schema): Schema {
+  const merged: Schema = {};
+  if (Object.hasOwn(target, "$ref")) merged.$ref = target.$ref;
+  for (const key of MERGED_SCHEMA_KEYS) {
+    if (Object.hasOwn(overlay, key)) merged[key] = overlay[key];
+    else if (Object.hasOwn(target, key)) merged[key] = target[key];
+  }
+  return merged;
+}
+
 function normalizeType(value: unknown, out: Schema, preserveNullType: boolean): void {
   const candidates = Array.isArray(value) ? value : [value];
   let sawNull = false;
@@ -65,9 +96,16 @@ function normalizeAnyOf(
   defs: Map<string, unknown>,
   depth: number,
   refDepth: number,
+  state: SanitizeState,
 ): Schema {
   if (!Array.isArray(value) || value.length === 0) return {};
-  const schemas = value.map(item => sanitizeSchema(item, defs, depth + 1, refDepth, true));
+  const schemas: Schema[] = [];
+  for (let index = 0; index < value.length; index++) {
+    if (state.remainingNodes <= 0) return {};
+    const schema = sanitizeSchema(value[index], defs, depth + 1, refDepth, true, state);
+    if (schema === BUDGET_EXHAUSTED) return {};
+    schemas.push(schema);
+  }
 
   const nonNullSchemas = schemas.filter(schema => schema.type !== "null");
   const nullSchemas = schemas.filter(schema => schema.type === "null");
@@ -100,12 +138,17 @@ function sanitizeProperties(
   defs: Map<string, unknown>,
   depth: number,
   refDepth: number,
+  state: SanitizeState,
 ): Record<string, Schema> | undefined {
   if (!isRecord(value)) return undefined;
   const properties: Record<string, Schema> = Object.create(null) as Record<string, Schema>;
-  for (const [name, schema] of Object.entries(value)) {
+  for (const name in value) {
+    if (!Object.hasOwn(value, name)) continue;
+    if (state.remainingNodes <= 0) break;
     // Property names form a name bag and must never be interpreted as schema keywords.
-    properties[name] = sanitizeSchema(schema, defs, depth + 1, refDepth, false);
+    const schema = sanitizeSchema(value[name], defs, depth + 1, refDepth, false, state);
+    if (schema === BUDGET_EXHAUSTED) break;
+    properties[name] = schema;
   }
   return properties;
 }
@@ -116,17 +159,25 @@ function sanitizeSchema(
   depth: number,
   refDepth: number,
   preserveNullType: boolean,
-): Schema {
+  state: SanitizeState,
+): SanitizeResult {
+  if (state.remainingNodes <= 0) return BUDGET_EXHAUSTED;
+  state.remainingNodes -= 1;
   if (depth >= MAX_SCHEMA_DEPTH || !isRecord(node)) return {};
 
   if (typeof node.$ref === "string" && refDepth < MAX_DEREF_DEPTH) {
     const target = resolveRef(node.$ref, defs);
     if (isRecord(target)) {
-      const merged: Schema = { ...target };
-      for (const [key, value] of Object.entries(node)) {
-        if (key !== "$ref") merged[key] = value;
+      if (state.activeRefs.has(node.$ref)) return {};
+      state.activeRefs.add(node.$ref);
+      // Select only inputs the sanitizer can consume. Spreading an untrusted definition here would
+      // enumerate and allocate every unsupported annotation before the node budget can stop work.
+      const merged = mergeRefTarget(target, node);
+      try {
+        return sanitizeSchema(merged, defs, depth, refDepth + 1, preserveNullType, state);
+      } finally {
+        state.activeRefs.delete(node.$ref);
       }
-      return sanitizeSchema(merged, defs, depth, refDepth + 1, preserveNullType);
     }
   }
 
@@ -140,26 +191,40 @@ function sanitizeSchema(
   const enumValues = sanitizeEnum(node.enum ?? (typeof node.const === "string" ? [node.const] : undefined));
   if (enumValues) out.enum = enumValues;
 
-  const properties = sanitizeProperties(node.properties, defs, depth, refDepth);
+  const properties = sanitizeProperties(node.properties, defs, depth, refDepth, state);
   if (properties) out.properties = properties;
 
-  if (isRecord(node.items)) {
-    out.items = sanitizeSchema(node.items, defs, depth + 1, refDepth, false);
+  if (properties && Array.isArray(node.required)) {
+    const required = [...new Set(node.required.filter((item): item is string => (
+      typeof item === "string" && Object.hasOwn(properties, item)
+    )))];
+    if (required.length > 0) out.required = required;
   }
 
-  if (Array.isArray(node.required)) {
-    out.required = [...new Set(node.required.filter((item): item is string => typeof item === "string"))];
+  if (state.remainingNodes <= 0) return out;
+
+  if (isRecord(node.items)) {
+    const items = sanitizeSchema(node.items, defs, depth + 1, refDepth, false, state);
+    if (items !== BUDGET_EXHAUSTED) out.items = items;
   }
 
-  if (node.anyOf !== undefined) Object.assign(out, normalizeAnyOf(node.anyOf, defs, depth, refDepth));
+  if (state.remainingNodes <= 0) return out;
+  if (node.anyOf !== undefined) {
+    Object.assign(out, normalizeAnyOf(node.anyOf, defs, depth, refDepth, state));
+  }
   return out;
 }
 
 export function sanitizeGeminiToolParameters(parameters: unknown): Record<string, unknown> {
   try {
     const defs = new Map<string, unknown>();
     collectDefs(parameters, defs);
-    const root = sanitizeSchema(parameters, defs, 0, 0, false);
+    const state: SanitizeState = {
+      activeRefs: new Set(),
+      remainingNodes: MAX_SCHEMA_NODES,
+    };
+    const sanitized = sanitizeSchema(parameters, defs, 0, 0, false, state);
+    const root = sanitized === BUDGET_EXHAUSTED ? {} : sanitized;
 
     // Function arguments are always an object. Claude additionally rejects root composition and a
     // missing root type even when those forms are valid general-purpose JSON Schema.

tests/google-tool-schema.test.ts

@@ -1,6 +1,17 @@
 import { describe, expect, test } from "bun:test";
 import { sanitizeGeminiToolParameters } from "../src/adapters/google-tool-schema";
 
+function countSchemaNodes(value: unknown): number {
+  if (!value || typeof value !== "object" || Array.isArray(value)) return 0;
+  const schema = value as Record<string, unknown>;
+  let count = 1;
+  if (schema.properties && typeof schema.properties === "object" && !Array.isArray(schema.properties)) {
+    for (const child of Object.values(schema.properties)) count += countSchemaNodes(child);
+  }
+  if (schema.items !== undefined) count += countSchemaNodes(schema.items);
+  return count;
+}
+
 describe("sanitizeGeminiToolParameters", () => {
   test("drops JSON-Schema keywords outside Google's documented function-schema subset", () => {
     const out = sanitizeGeminiToolParameters({
@@ -332,15 +343,140 @@ describe("sanitizeGeminiToolParameters", () => {
     expect((node.properties as Record<string, Record<string, unknown>>).id.type).toBe("string");
   });
 
-  test("does not infinitely recurse on self-referential $defs", () => {
+  test("does not enumerate unsupported definition keys for repeated refs", () => {
+    let enumeratedDefinition = false;
+    const definition = new Proxy({
+      type: "object",
+      properties: { id: { type: "string" } },
+    }, {
+      ownKeys() {
+        enumeratedDefinition = true;
+        throw new Error("enumerated the full definition");
+      },
+    });
+    const out = sanitizeGeminiToolParameters({
+      type: "object",
+      properties: {
+        first: { $ref: "#/$defs/Wide" },
+        second: { $ref: "#/$defs/Wide" },
+      },
+      $defs: { Wide: definition },
+    });
+    const properties = out.properties as Record<string, Record<string, unknown>>;
+    expect(enumeratedDefinition).toBe(false);
+    expect(properties.first).toEqual({
+      type: "object",
+      properties: { id: { type: "string" } },
+    });
+    expect(properties.second).toEqual(properties.first);
+  });
+
+  test("widens recursive $refs without expanding sibling branches", () => {
     const out = sanitizeGeminiToolParameters({
       type: "object",
       properties: { tree: { $ref: "#/$defs/Tree" } },
-      $defs: { Tree: { type: "object", properties: { child: { $ref: "#/$defs/Tree" } } } },
+      $defs: {
+        Tree: {
+          type: "object",
+          properties: {
+            left: { $ref: "#/$defs/Tree" },
+            right: { $ref: "#/$defs/Tree" },
+          },
+        },
+      },
     });
-    expect(out.type).toBe("object");
     const tree = (out.properties as Record<string, Record<string, unknown>>).tree;
     expect(tree.type).toBe("object");
+    expect(tree.properties).toEqual({ left: {}, right: {} });
+  });
+
+  test("bounds acyclic shared-definition fan-out by node budget", () => {
+    const defs: Record<string, unknown> = {};
+    for (let index = 17; index >= 0; index--) {
+      defs[`Node${index}`] = index === 17
+        ? { type: "string" }
+        : {
+          type: "object",
+          properties: {
+            left: { $ref: `#/$defs/Node${index + 1}` },
+            right: { $ref: `#/$defs/Node${index + 1}` },
+          },
+        };
+    }
+
+    const out = sanitizeGeminiToolParameters({
+      type: "object",
+      properties: { tree: { $ref: "#/$defs/Node0" } },
+      $defs: defs,
+    });
+    expect(countSchemaNodes(out)).toBeLessThanOrEqual(1_024);
+  });
+
+  test("truncates wide properties without dangling required names", () => {
+    const names = Array.from({ length: 2_000 }, (_, index) => `field_${index}`);
+    const out = sanitizeGeminiToolParameters({
+      type: "object",
+      properties: Object.fromEntries(names.map(name => [name, { type: "string" }])),
+      required: names,
+    });
+    const properties = out.properties as Record<string, unknown>;
+    const retainedNames = Object.keys(properties);
+    expect(retainedNames).toHaveLength(1_023);
+    expect(out.required).toEqual(retainedNames);
+    expect(Object.hasOwn(properties, names[1_023])).toBe(false);
+    expect(countSchemaNodes(out)).toBe(1_024);
+  });
+
+  test("stops reading anyOf branches when the budget is exhausted", () => {
+    const branches = Array.from({ length: 2_000 }, () => ({ type: "string" }));
+    let readPastBudget = false;
+    Object.defineProperty(branches, 1_022, {
+      configurable: true,
+      get() {
+        readPastBudget = true;
+        throw new Error("read past schema budget");
+      },
+    });
+
+    const out = sanitizeGeminiToolParameters({
+      type: "object",
+      properties: {
+        choice: {
+          description: "kept",
+          anyOf: branches,
+        },
+      },
+    });
+    const choice = (out.properties as Record<string, Record<string, unknown>>).choice;
+    expect(readPastBudget).toBe(false);
+    expect(choice).toEqual({ description: "kept" });
+  });
+
+  test("does not read items after earlier traversal exhausts the budget", () => {
+    const container: Record<string, unknown> = {
+      type: "array",
+      properties: Object.fromEntries(Array.from(
+        { length: 1_022 },
+        (_, index) => [`field_${index}`, { type: "string" }],
+      )),
+    };
+    let readItems = false;
+    Object.defineProperty(container, "items", {
+      configurable: true,
+      get() {
+        readItems = true;
+        throw new Error("read items past schema budget");
+      },
+    });
+
+    const out = sanitizeGeminiToolParameters({
+      type: "object",
+      properties: { container },
+    });
+    const sanitized = (out.properties as Record<string, Record<string, unknown>>).container;
+    expect(readItems).toBe(false);
+    expect(sanitized.items).toBeUndefined();
+    expect(countSchemaNodes(out)).toBe(1_024);
   });
 
   test("falls back to an object schema for non-object input", () => {