From 96f5b58bd34697687e9adc3e16432b5b0c8bc1de Mon Sep 17 00:00:00 2001 From: Mannan Kant <32628694+mannan24@users.noreply.github.com> Date: Mon, 16 Mar 2026 12:21:20 +0800 Subject: [PATCH] =?UTF-8?q?fix(pi-ai):=20address=20review=20comments=20on?= =?UTF-8?q?=20#504=20=E2=80=94=20exhaustive=20switch,=20tests,=20cleanup?= =?UTF-8?q?=20(#587)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Restore exhaustive never check in mapStopReason (throw on unhandled FinishReason) - Add 12 unit tests for sanitizeSchemaForGoogle covering patternProperties removal, const→enum conversion at various depths, arrays, deeply nested objects, pass-through - Simplify redundant recursion branches into single typeof object catch-all - Fix misleading comment ("only in anyOf/oneOf") — conversion happens everywhere - Drop unnecessary (p: Part) annotation; TypeScript infers it from @google/genai types Co-authored-by: Claude Sonnet 4.6 --- .../pi-ai/src/providers/google-shared.test.ts | 137 ++++++++++++++++++ packages/pi-ai/src/providers/google-shared.ts | 29 ++-- 2 files changed, 147 insertions(+), 19 deletions(-) create mode 100644 packages/pi-ai/src/providers/google-shared.test.ts diff --git a/packages/pi-ai/src/providers/google-shared.test.ts b/packages/pi-ai/src/providers/google-shared.test.ts new file mode 100644 index 000000000..4468ac231 --- /dev/null +++ b/packages/pi-ai/src/providers/google-shared.test.ts @@ -0,0 +1,137 @@ +import { describe, it } from "node:test"; +import assert from "node:assert/strict"; +import { sanitizeSchemaForGoogle } from "./google-shared.js"; + +// ═══════════════════════════════════════════════════════════════════════════ +// sanitizeSchemaForGoogle +// ═══════════════════════════════════════════════════════════════════════════ + +describe("sanitizeSchemaForGoogle", () => { + it("passes through primitives unchanged", () => { + assert.equal(sanitizeSchemaForGoogle(null), null); + assert.equal(sanitizeSchemaForGoogle(42), 42); + assert.equal(sanitizeSchemaForGoogle("hello"), "hello"); + assert.equal(sanitizeSchemaForGoogle(true), true); + }); + + it("passes through a valid schema with no banned fields", () => { + const schema = { + type: "object", + properties: { + name: { type: "string" }, + age: { type: "number" }, + }, + required: ["name"], + }; + assert.deepEqual(sanitizeSchemaForGoogle(schema), schema); + }); + + it("removes top-level patternProperties", () => { + const schema = { + type: "object", + patternProperties: { "^S_": { type: "string" } }, + properties: { foo: { type: "string" } }, + }; + const result = sanitizeSchemaForGoogle(schema) as Record; + assert.ok(!("patternProperties" in result)); + assert.deepEqual(result.properties, { foo: { type: "string" } }); + }); + + it("removes nested patternProperties", () => { + const schema = { + type: "object", + properties: { + nested: { + type: "object", + patternProperties: { ".*": { type: "string" } }, + }, + }, + }; + const result = sanitizeSchemaForGoogle(schema) as any; + assert.ok(!("patternProperties" in result.properties.nested)); + }); + + it("converts top-level const to enum", () => { + const schema = { const: "fixed-value" }; + const result = sanitizeSchemaForGoogle(schema) as Record; + assert.deepEqual(result.enum, ["fixed-value"]); + assert.ok(!("const" in result)); + }); + + it("converts const to enum inside anyOf", () => { + const schema = { + anyOf: [{ const: "a" }, { const: "b" }, { type: "string" }], + }; + const result = sanitizeSchemaForGoogle(schema) as any; + assert.deepEqual(result.anyOf[0], { enum: ["a"] }); + assert.deepEqual(result.anyOf[1], { enum: ["b"] }); + assert.deepEqual(result.anyOf[2], { type: "string" }); + }); + + it("converts const to enum inside oneOf", () => { + const schema = { + oneOf: [{ const: "x" }, { const: "y" }], + }; + const result = sanitizeSchemaForGoogle(schema) as any; + assert.deepEqual(result.oneOf[0], { enum: ["x"] }); + assert.deepEqual(result.oneOf[1], { enum: ["y"] }); + }); + + it("recursively sanitizes deeply nested schemas", () => { + const schema = { + type: "object", + properties: { + level1: { + type: "object", + properties: { + level2: { + anyOf: [{ const: "deep" }, { type: "null" }], + patternProperties: { ".*": { type: "string" } }, + }, + }, + }, + }, + }; + const result = sanitizeSchemaForGoogle(schema) as any; + const level2 = result.properties.level1.properties.level2; + assert.deepEqual(level2.anyOf[0], { enum: ["deep"] }); + assert.ok(!("patternProperties" in level2)); + }); + + it("sanitizes items in array schemas", () => { + const schema = { + type: "array", + items: { + anyOf: [{ const: "foo" }, { type: "string" }], + }, + }; + const result = sanitizeSchemaForGoogle(schema) as any; + assert.deepEqual(result.items.anyOf[0], { enum: ["foo"] }); + }); + + it("sanitizes arrays of schemas", () => { + const input = [{ const: "a" }, { const: "b" }]; + const result = sanitizeSchemaForGoogle(input) as any[]; + assert.deepEqual(result[0], { enum: ["a"] }); + assert.deepEqual(result[1], { enum: ["b"] }); + }); + + it("preserves non-string const values unchanged", () => { + // Only string const values are converted; number const is passed through + const schema = { const: 42 }; + const result = sanitizeSchemaForGoogle(schema) as Record; + assert.equal(result.const, 42); + assert.ok(!("enum" in result)); + }); + + it("sanitizes additionalProperties", () => { + const schema = { + type: "object", + additionalProperties: { + patternProperties: { "^x-": { type: "string" } }, + }, + }; + const result = sanitizeSchemaForGoogle(schema) as any; + assert.ok(!("patternProperties" in result.additionalProperties)); + }); +}); diff --git a/packages/pi-ai/src/providers/google-shared.ts b/packages/pi-ai/src/providers/google-shared.ts index 0ae58171b..255928c81 100644 --- a/packages/pi-ai/src/providers/google-shared.ts +++ b/packages/pi-ai/src/providers/google-shared.ts @@ -204,7 +204,7 @@ export function convertMessages(model: Model, contex // Cloud Code Assist API requires all function responses to be in a single user turn. // Check if the last content is already a user turn with function responses and merge. const lastContent = contents[contents.length - 1]; - if (lastContent?.role === "user" && lastContent.parts?.some((p: Part) => p.functionResponse)) { + if (lastContent?.role === "user" && lastContent.parts?.some((p) => p.functionResponse)) { lastContent.parts.push(functionResponsePart); } else { contents.push({ @@ -237,7 +237,7 @@ export function convertMessages(model: Model, contex * This is needed for providers like `google-antigravity` when proxying Claude models, * since Google Cloud Code Assist uses a restricted subset of JSON Schema. */ -function sanitizeSchemaForGoogle(schema: unknown): unknown { +export function sanitizeSchemaForGoogle(schema: unknown): unknown { if (!schema || typeof schema !== "object") { return schema; } @@ -250,29 +250,19 @@ function sanitizeSchemaForGoogle(schema: unknown): unknown { const sanitized: Record = {}; for (const [key, value] of Object.entries(obj)) { - // Skip patternProperties entirely + // Skip patternProperties entirely — not supported by Google's API if (key === "patternProperties") { continue; } - // Convert const to enum in anyOf/oneOf blocks + // Convert const to enum — Google's API rejects the const keyword if (key === "const" && typeof value === "string") { - // Only convert if we're inside anyOf/oneOf; otherwise leave as-is - // This will be handled by the anyOf/oneOf case below sanitized.enum = [value]; continue; } - // Recursively sanitize nested objects and arrays - if (key === "properties" && typeof value === "object") { - sanitized[key] = sanitizeSchemaForGoogle(value); - } else if (key === "items" && typeof value === "object") { - sanitized[key] = sanitizeSchemaForGoogle(value); - } else if (key === "anyOf" || key === "oneOf" || key === "allOf") { - sanitized[key] = sanitizeSchemaForGoogle(value); - } else if (key === "additionalProperties" && typeof value === "object") { - sanitized[key] = sanitizeSchemaForGoogle(value); - } else if (typeof value === "object" && !Array.isArray(value)) { + // Recursively sanitize all nested objects and arrays + if (typeof value === "object") { sanitized[key] = sanitizeSchemaForGoogle(value); } else { sanitized[key] = value; @@ -352,9 +342,10 @@ export function mapStopReason(reason: FinishReason): StopReason { case FinishReason.UNEXPECTED_TOOL_CALL: case FinishReason.NO_IMAGE: return "error"; - default: - // Fallback for new/unknown FinishReason values - return "error"; + default: { + const _exhaustive: never = reason; + throw new Error(`Unhandled stop reason: ${_exhaustive}`); + } } }