From 4dfcd56893be15c6b1357f848c9d13db30f4f301 Mon Sep 17 00:00:00 2001
From: Abhi <40645221+AbhisekBasu1@users.noreply.github.com>
Date: Mon, 12 Jan 2026 18:46:58 +0530
Subject: [PATCH] Fix pi-tools test ordering and clean-for-gemini handling -
 which fixes the 400 error people are experiencing trying to use antigravity
 on opus

---
 src/agents/pi-tools.test.ts           | 26 +++++++
 src/agents/schema/clean-for-gemini.ts | 98 +++++++++++++++++++++++++--
 2 files changed, 118 insertions(+), 6 deletions(-)

diff --git a/src/agents/pi-tools.test.ts b/src/agents/pi-tools.test.ts
index e3797b51e..21e5e293f 100644
--- a/src/agents/pi-tools.test.ts
+++ b/src/agents/pi-tools.test.ts
@@ -103,6 +103,32 @@ describe("createClawdbotCodingTools", () => {
     });
   });
 
+  it("flattens simple anyOf/oneOf unions into single types", () => {
+    const cleaned = __testing.cleanToolSchemaForGemini({
+      type: "object",
+      properties: {
+        parentId: { anyOf: [{ type: "string" }, { type: "null" }] },
+        count: { oneOf: [{ type: "string" }, { type: "number" }] },
+      },
+    }) as {
+      properties?: Record<string, unknown>;
+    };
+
+    const parentId = cleaned.properties?.parentId as
+      | { type?: unknown; anyOf?: unknown; oneOf?: unknown }
+      | undefined;
+    expect(parentId?.anyOf).toBeUndefined();
+    expect(parentId?.oneOf).toBeUndefined();
+    expect(parentId?.type).toBe("string");
+
+    const count = cleaned.properties?.count as
+      | { type?: unknown; anyOf?: unknown; oneOf?: unknown }
+      | undefined;
+    expect(count?.anyOf).toBeUndefined();
+    expect(count?.oneOf).toBeUndefined();
+    expect(count?.type).toBe("string");
+  });
+
   it("preserves action enums in normalized schemas", () => {
     const tools = createClawdbotCodingTools();
     const toolNames = [
diff --git a/src/agents/schema/clean-for-gemini.ts b/src/agents/schema/clean-for-gemini.ts
index 219319fb8..2d738bda6 100644
--- a/src/agents/schema/clean-for-gemini.ts
+++ b/src/agents/schema/clean-for-gemini.ts
@@ -67,6 +67,60 @@ function tryFlattenLiteralAnyOf(
   return null;
 }
 
+const TYPE_UNION_IGNORED_KEYS = new Set([
+  ...UNSUPPORTED_SCHEMA_KEYWORDS,
+  "description",
+  "title",
+  "default",
+]);
+
+function tryFlattenTypeUnion(
+  variants: unknown[],
+): { type: string } | null {
+  if (variants.length === 0) return null;
+
+  const types = new Set<string>();
+  for (const variant of variants) {
+    if (!variant || typeof variant !== "object" || Array.isArray(variant)) {
+      return null;
+    }
+    const record = variant as Record<string, unknown>;
+    const keys = Object.keys(record).filter(
+      (key) => !TYPE_UNION_IGNORED_KEYS.has(key),
+    );
+    if (keys.length !== 1 || keys[0] !== "type") return null;
+
+    const typeValue = record.type;
+    if (typeof typeValue === "string") {
+      types.add(typeValue);
+      continue;
+    }
+    if (
+      Array.isArray(typeValue) &&
+      typeValue.every((entry) => typeof entry === "string")
+    ) {
+      for (const entry of typeValue) types.add(entry);
+      continue;
+    }
+    return null;
+  }
+
+  if (types.size === 0) return null;
+
+  const pickType = () => {
+    if (types.has("string")) return "string";
+    if (types.has("number")) return "number";
+    if (types.has("integer")) return "number";
+    if (types.has("boolean")) return "boolean";
+    if (types.has("object")) return "object";
+    if (types.has("array")) return "array";
+    const nonNull = Array.from(types).find((value) => value !== "null");
+    return nonNull ?? "string";
+  };
+
+  return { type: pickType() };
+}
+
 type SchemaDefs = Map<string, unknown>;
 
 function extendSchemaDefs(
@@ -166,6 +220,16 @@ function cleanSchemaForGeminiWithDefs(
 
   const hasAnyOf = "anyOf" in obj && Array.isArray(obj.anyOf);
   const hasOneOf = "oneOf" in obj && Array.isArray(obj.oneOf);
+  const cleanedAnyOf = hasAnyOf
+    ? (obj.anyOf as unknown[]).map((variant) =>
+        cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
+      )
+    : undefined;
+  const cleanedOneOf = hasOneOf
+    ? (obj.oneOf as unknown[]).map((variant) =>
+        cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
+      )
+    : undefined;
 
   if (hasAnyOf) {
     const flattened = tryFlattenLiteralAnyOf(obj.anyOf as unknown[]);
@@ -179,6 +243,15 @@ function cleanSchemaForGeminiWithDefs(
       }
       return result;
     }
+
+    const flattenedTypes = tryFlattenTypeUnion(cleanedAnyOf ?? []);
+    if (flattenedTypes) {
+      const result: Record<string, unknown> = { ...flattenedTypes };
+      for (const key of ["description", "title", "default"]) {
+        if (key in obj && obj[key] !== undefined) result[key] = obj[key];
+      }
+      return result;
+    }
   }
 
   if (hasOneOf) {
@@ -193,6 +266,15 @@ function cleanSchemaForGeminiWithDefs(
       }
       return result;
     }
+
+    const flattenedTypes = tryFlattenTypeUnion(cleanedOneOf ?? []);
+    if (flattenedTypes) {
+      const result: Record<string, unknown> = { ...flattenedTypes };
+      for (const key of ["description", "title", "default"]) {
+        if (key in obj && obj[key] !== undefined) result[key] = obj[key];
+      }
+      return result;
+    }
   }
 
   const cleaned: Record<string, unknown> = {};
@@ -218,13 +300,17 @@ function cleanSchemaForGeminiWithDefs(
     } else if (key === "items" && value && typeof value === "object") {
       cleaned[key] = cleanSchemaForGeminiWithDefs(value, nextDefs, refStack);
     } else if (key === "anyOf" && Array.isArray(value)) {
-      cleaned[key] = value.map((variant) =>
-        cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
-      );
+      cleaned[key] =
+        cleanedAnyOf ??
+        value.map((variant) =>
+          cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
+        );
     } else if (key === "oneOf" && Array.isArray(value)) {
-      cleaned[key] = value.map((variant) =>
-        cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
-      );
+      cleaned[key] =
+        cleanedOneOf ??
+        value.map((variant) =>
+          cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),
+        );
     } else if (key === "allOf" && Array.isArray(value)) {
       cleaned[key] = value.map((variant) =>
         cleanSchemaForGeminiWithDefs(variant, nextDefs, refStack),