From f0d210d1c5806b5eb199289c1448570af2ee2774 Mon Sep 17 00:00:00 2001 From: testikun <320479488+testikun@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:04:19 +0800 Subject: [PATCH 1/2] fix(antigravity): preserve local schema refs --- docs/design/ANTIGRAVITY_LOCAL_SCHEMA_REFS.md | 20 +++++ .../antigravity/google-conversion.ts | 84 ++++++++++++++++++- .../ai-providers/antigravity.test.ts | 72 ++++++++++++++++ 3 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 docs/design/ANTIGRAVITY_LOCAL_SCHEMA_REFS.md diff --git a/docs/design/ANTIGRAVITY_LOCAL_SCHEMA_REFS.md b/docs/design/ANTIGRAVITY_LOCAL_SCHEMA_REFS.md new file mode 100644 index 00000000..d3ae9cc5 --- /dev/null +++ b/docs/design/ANTIGRAVITY_LOCAL_SCHEMA_REFS.md @@ -0,0 +1,20 @@ +# Antigravity Local Schema References + +Status: validated +Created: 2026-09-08 +Verified: 2026-09-08 +Issue: [#465](https://github.com/openpi-dev/openpi/issues/465) + +The Antigravity boundary expands only same-document JSON Pointer references +before applying Cloud Code Assist's unsupported-keyword sanitizer. External, +unresolved, recursive, or over-limit references fail before a model request. +Expansion is bounded by depth, nodes, and serialized bytes. The original Pi +schema remains the authority for local validation; this only preserves its +meaning in the provider declaration. + +The ablation is explicit: removing expansion reproduces the original empty +`{}` property for a `$ref` result contract, while allowing references without +bounds could make provider preparation unbounded. Both the expansion and +limits are retained. + +Validation: `node --test --experimental-strip-types tests/extensions/ai-providers/antigravity.test.ts` (39/39) and `bun run check` passed. diff --git a/extensions/ai-providers/antigravity/google-conversion.ts b/extensions/ai-providers/antigravity/google-conversion.ts index 1c1165e3..581bab80 100644 --- a/extensions/ai-providers/antigravity/google-conversion.ts +++ b/extensions/ai-providers/antigravity/google-conversion.ts @@ -34,6 +34,10 @@ const JSON_SCHEMA_META_DECLARATIONS = new Set([ "definitions", ]); +const MAX_SCHEMA_REF_DEPTH = 16; +const MAX_SCHEMA_REF_NODES = 512; +const MAX_SCHEMA_REF_BYTES = 256 * 1024; + interface GoogleFunctionCall { id?: string; name: string; @@ -394,6 +398,80 @@ export function convertMessages( return contents; } +function expandLocalSchemaRefs(schema: unknown): unknown { + if (typeof schema !== "object" || schema === null || Array.isArray(schema)) { + return schema; + } + const root = schema as Record; + let nodes = 0; + let bytes = 0; + const active = new Set(); + + const pointer = (ref: string) => { + if (!ref.startsWith("#/") && ref !== "#") { + throw new Error( + `Antigravity tool schema has unsupported external $ref: ${ref}`, + ); + } + let value: unknown = root; + if (ref !== "#") { + for (const segment of ref.slice(2).split("/")) { + if (value === null || typeof value !== "object") value = undefined; + else + value = (value as Record)[ + segment.replaceAll("~1", "/").replaceAll("~0", "~") + ]; + } + } + if (value === undefined) + throw new Error(`Antigravity tool schema has unresolved $ref: ${ref}`); + return value; + }; + + const visit = (value: unknown, depth: number, path: string): unknown => { + if (++nodes > MAX_SCHEMA_REF_NODES || depth > MAX_SCHEMA_REF_DEPTH) { + throw new Error( + "Antigravity tool schema exceeds local $ref expansion limits", + ); + } + bytes += Buffer.byteLength(JSON.stringify(value) ?? ""); + if (bytes > MAX_SCHEMA_REF_BYTES) { + throw new Error( + "Antigravity tool schema exceeds local $ref expansion byte limit", + ); + } + if (Array.isArray(value)) + return value.map((entry, index) => + visit(entry, depth + 1, `${path}/${index}`), + ); + if (value === null || typeof value !== "object") return value; + const object = value as Record; + if (typeof object.$ref === "string") { + const ref = object.$ref; + if (active.has(ref)) + throw new Error(`Antigravity tool schema has recursive $ref: ${ref}`); + active.add(ref); + const target = visit(pointer(ref), depth + 1, ref) as Record< + string, + unknown + >; + active.delete(ref); + const siblings = Object.fromEntries( + Object.entries(object).filter(([key]) => key !== "$ref"), + ); + return visit({ ...target, ...siblings }, depth + 1, path); + } + return Object.fromEntries( + Object.entries(object).map(([key, entry]) => [ + key, + visit(entry, depth + 1, `${path}/${key}`), + ]), + ); + }; + + return visit(schema, 0, "#"); +} + function sanitizeForOpenApi( schema: unknown, insidePropertiesMap = false, @@ -425,7 +503,11 @@ export function convertTools( name: tool.name, description: tool.description, ...(useParameters - ? { parameters: sanitizeForOpenApi(tool.parameters) } + ? { + parameters: sanitizeForOpenApi( + expandLocalSchemaRefs(tool.parameters), + ), + } : { parametersJsonSchema: tool.parameters }), })), }, diff --git a/tests/extensions/ai-providers/antigravity.test.ts b/tests/extensions/ai-providers/antigravity.test.ts index ea83c49c..9f8b2d51 100644 --- a/tests/extensions/ai-providers/antigravity.test.ts +++ b/tests/extensions/ai-providers/antigravity.test.ts @@ -21,6 +21,7 @@ import { import { fetchAntigravityModels } from "../../../extensions/ai-providers/antigravity/discovery.ts"; import { convertMessages, + convertTools, isThinkingPart, mapStopReasonString, retainThoughtSignature, @@ -442,6 +443,77 @@ test("sanitizeSchemaForCca preserves property names that match schema keywords", ); }); +test("convertTools expands bounded local refs before CCA sanitization", () => { + const declarations = convertTools( + [ + { + name: "structured_output", + description: "Return the result", + parameters: { + type: "object", + properties: { answer: { $ref: "#/$defs/Answer" } }, + required: ["answer"], + $defs: { + Answer: { + type: "object", + properties: { + verdict: { type: "string", enum: ["pass", "fail"] }, + }, + required: ["verdict"], + }, + }, + }, + } as never, + ], + true, + ); + assert.deepEqual(declarations?.[0]?.functionDeclarations[0]?.parameters, { + type: "object", + properties: { + answer: { + type: "object", + properties: { verdict: { type: "string", enum: ["pass", "fail"] } }, + required: ["verdict"], + }, + }, + required: ["answer"], + }); +}); + +test("convertTools rejects unsafe local refs before the request", () => { + assert.throws( + () => + convertTools( + [ + { + name: "bad", + description: "bad", + parameters: { $ref: "https://example.test/schema" }, + } as never, + ], + true, + ), + /unsupported external \$ref/, + ); + assert.throws( + () => + convertTools( + [ + { + name: "loop", + description: "loop", + parameters: { + $defs: { Node: { $ref: "#/$defs/Node" } }, + $ref: "#/$defs/Node", + }, + } as never, + ], + true, + ), + /recursive \$ref/, + ); +}); + test("buildRequestBody strips CCA-rejected keywords and spills constraints into description", () => { const contextWithTools = { ...SIMPLE_CONTEXT, From 1904fe3a6258801a59b402311c86d635b955d1b6 Mon Sep 17 00:00:00 2001 From: testikun <320479488+testikun@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:39:07 +0800 Subject: [PATCH 2/2] fix(antigravity): avoid repeated schema ref traversal --- .../antigravity/google-conversion.ts | 42 ++++++------- .../ai-providers/antigravity.test.ts | 59 +++++++++++++++++++ 2 files changed, 81 insertions(+), 20 deletions(-) diff --git a/extensions/ai-providers/antigravity/google-conversion.ts b/extensions/ai-providers/antigravity/google-conversion.ts index 581bab80..1281982d 100644 --- a/extensions/ai-providers/antigravity/google-conversion.ts +++ b/extensions/ai-providers/antigravity/google-conversion.ts @@ -404,7 +404,6 @@ function expandLocalSchemaRefs(schema: unknown): unknown { } const root = schema as Record; let nodes = 0; - let bytes = 0; const active = new Set(); const pointer = (ref: string) => { @@ -428,22 +427,14 @@ function expandLocalSchemaRefs(schema: unknown): unknown { return value; }; - const visit = (value: unknown, depth: number, path: string): unknown => { + const visit = (value: unknown, depth: number): unknown => { if (++nodes > MAX_SCHEMA_REF_NODES || depth > MAX_SCHEMA_REF_DEPTH) { throw new Error( "Antigravity tool schema exceeds local $ref expansion limits", ); } - bytes += Buffer.byteLength(JSON.stringify(value) ?? ""); - if (bytes > MAX_SCHEMA_REF_BYTES) { - throw new Error( - "Antigravity tool schema exceeds local $ref expansion byte limit", - ); - } if (Array.isArray(value)) - return value.map((entry, index) => - visit(entry, depth + 1, `${path}/${index}`), - ); + return value.map((entry) => visit(entry, depth + 1)); if (value === null || typeof value !== "object") return value; const object = value as Record; if (typeof object.$ref === "string") { @@ -451,25 +442,36 @@ function expandLocalSchemaRefs(schema: unknown): unknown { if (active.has(ref)) throw new Error(`Antigravity tool schema has recursive $ref: ${ref}`); active.add(ref); - const target = visit(pointer(ref), depth + 1, ref) as Record< - string, - unknown - >; - active.delete(ref); + let target: unknown; + try { + target = visit(pointer(ref), depth + 1); + } finally { + active.delete(ref); + } const siblings = Object.fromEntries( - Object.entries(object).filter(([key]) => key !== "$ref"), + Object.entries(object) + .filter(([key]) => key !== "$ref") + .map(([key, entry]) => [key, visit(entry, depth + 2)]), ); - return visit({ ...target, ...siblings }, depth + 1, path); + return { ...(target as Record), ...siblings }; } return Object.fromEntries( Object.entries(object).map(([key, entry]) => [ key, - visit(entry, depth + 1, `${path}/${key}`), + visit(entry, depth + 1), ]), ); }; - return visit(schema, 0, "#"); + const expanded = visit(schema, 0); + if ( + Buffer.byteLength(JSON.stringify(expanded) ?? "") > MAX_SCHEMA_REF_BYTES + ) { + throw new Error( + "Antigravity tool schema exceeds local $ref expansion byte limit", + ); + } + return expanded; } function sanitizeForOpenApi( diff --git a/tests/extensions/ai-providers/antigravity.test.ts b/tests/extensions/ai-providers/antigravity.test.ts index 9f8b2d51..655c59a5 100644 --- a/tests/extensions/ai-providers/antigravity.test.ts +++ b/tests/extensions/ai-providers/antigravity.test.ts @@ -480,6 +480,65 @@ test("convertTools expands bounded local refs before CCA sanitization", () => { }); }); +test("convertTools counts each expanded ref subtree once", () => { + const fields = Object.fromEntries( + Array.from({ length: 60 }, (_, index) => [ + `field${index}`, + { type: "string" }, + ]), + ); + const declarations = convertTools( + [ + { + name: "structured_output", + description: "Return the result", + parameters: { + type: "object", + properties: { + first: { $ref: "#/$defs/Result" }, + second: { $ref: "#/$defs/Result" }, + }, + $defs: { + Result: { type: "object", properties: fields }, + }, + }, + } as never, + ], + true, + ); + const parameters = declarations?.[0]?.functionDeclarations[0]?.parameters as { + properties?: Record }>; + }; + assert.deepEqual(parameters.properties?.first, parameters.properties?.second); + assert.equal( + parameters.properties?.first?.properties?.field59 !== undefined, + true, + ); +}); + +test("convertTools bounds the serialized expanded schema size", () => { + assert.throws( + () => + convertTools( + [ + { + name: "oversized", + description: "oversized", + parameters: { + type: "object", + properties: { value: { $ref: "#/$defs/Value" } }, + $defs: { + Value: { type: "string", description: "x".repeat(256 * 1024) }, + }, + }, + } as never, + ], + true, + ), + /expansion byte limit/, + ); +}); + test("convertTools rejects unsafe local refs before the request", () => { assert.throws( () =>