From 408c825e568e558d7d803ca2542d692acb35e6d1 Mon Sep 17 00:00:00 2001 From: luvs01 <27862058+luvs01@users.noreply.github.com> Date: Sun, 9 Aug 2026 20:36:07 +0900 Subject: [PATCH] fix(google): bound tool schema expansion --- src/adapters/google-tool-schema.ts | 95 ++++++++++++++++--- tests/google-tool-schema.test.ts | 142 ++++++++++++++++++++++++++++- 2 files changed, 219 insertions(+), 18 deletions(-) diff --git a/src/adapters/google-tool-schema.ts b/src/adapters/google-tool-schema.ts index 36273ff0d8c..aad7d96bb4f 100644 --- a/src/adapters/google-tool-schema.ts +++ b/src/adapters/google-tool-schema.ts @@ -7,6 +7,27 @@ type Schema = Record; 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; + 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): 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, 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, depth: number, refDepth: number, + state: SanitizeState, ): Record | undefined { if (!isRecord(value)) return undefined; const properties: Record = Object.create(null) as Record; - 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,18 +191,27 @@ 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; } @@ -159,7 +219,12 @@ export function sanitizeGeminiToolParameters(parameters: unknown): Record(); 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. diff --git a/tests/google-tool-schema.test.ts b/tests/google-tool-schema.test.ts index 745f5d5a50b..79d7b19a2b0 100644 --- a/tests/google-tool-schema.test.ts +++ b/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; + 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>).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>; + 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>).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 = {}; + 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; + 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>).choice; + expect(readPastBudget).toBe(false); + expect(choice).toEqual({ description: "kept" }); + }); + + test("does not read items after earlier traversal exhausts the budget", () => { + const container: Record = { + 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>).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", () => {