fix(request): copy thoughtSignature to all functionCall parts in the same turn to satisfy gemini-3.1-pro requirements

NoeFabris · NoeFabris · commit 3f9c14e36f12 · 2026-02-19T20:49:59.000Z
diff --git a/src/plugin/request.test.ts b/src/plugin/request.test.ts
@@ -1,6 +1,7 @@
 import { describe, it, expect } from "vitest";
 import {
   prepareAntigravityRequest,
+  transformAntigravityResponse,
   getPluginSessionId,
   isGenerativeLanguageRequest,
   __testExports,
@@ -517,6 +518,31 @@ describe("request.ts", () => {
   });
 
   describe("prepareAntigravityRequest", () => {
+    it("copies thoughtSignature to all functionCall parts in the same turn", () => {
+      const mockPayload = {
+        contents: [
+          {
+            role: "model",
+            parts: [
+              { functionCall: { name: "foo", args: {} }, thoughtSignature: "signature-123" },
+              { functionCall: { name: "bar", args: {} } }
+            ]
+          }
+        ]
+      };
+      
+      const result = prepareAntigravityRequest(
+        "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-pro:generateContent",
+        { method: "POST", body: JSON.stringify(mockPayload) },
+        "test-token",
+        "test-project"
+      );
+      
+      const parsedBody = JSON.parse(result.init.body as string);
+      expect(parsedBody.request.contents[0].parts[0].thoughtSignature).toBe("signature-123");
+      expect(parsedBody.request.contents[0].parts[1].thoughtSignature).toBe("signature-123");
+    });
+
     const mockAccessToken = "test-token";
     const mockProjectId = "test-project";
 
@@ -731,6 +757,61 @@ it("removes x-api-key header", () => {
       expect(result.streaming).toBe(false);
     });
 
+    it("removes contents entries with empty or invalid parts", () => {
+      const result = prepareAntigravityRequest(
+        "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent",
+        {
+          method: "POST",
+          body: JSON.stringify({
+            contents: [
+              { role: "user", parts: [] },
+              { role: "model", parts: [null, { text: "kept" }] },
+              { role: "user", parts: null },
+            ],
+            systemInstruction: {
+              role: "user",
+              parts: [null, { text: "system kept" }],
+            },
+          }),
+        },
+        mockAccessToken,
+        mockProjectId,
+        undefined,
+        "gemini-cli",
+      );
+
+      const wrapped = JSON.parse(result.init.body as string);
+      expect(wrapped.request.contents).toHaveLength(1);
+      expect(wrapped.request.contents[0]).toEqual({
+        role: "model",
+        parts: [{ text: "kept" }],
+      });
+      expect(wrapped.request.systemInstruction.parts).toEqual([{ text: "system kept" }]);
+    });
+
+    it("drops systemInstruction when all parts are invalid", () => {
+      const result = prepareAntigravityRequest(
+        "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent",
+        {
+          method: "POST",
+          body: JSON.stringify({
+            contents: [{ role: "user", parts: [{ text: "hi" }] }],
+            systemInstruction: {
+              role: "user",
+              parts: [null],
+            },
+          }),
+        },
+        mockAccessToken,
+        mockProjectId,
+        undefined,
+        "gemini-cli",
+      );
+
+      const wrapped = JSON.parse(result.init.body as string);
+      expect(wrapped.request.systemInstruction).toBeUndefined();
+    });
+
     it("preserves headerStyle in response", () => {
       const result = prepareAntigravityRequest(
         "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent",
@@ -768,6 +849,30 @@ it("removes x-api-key header", () => {
         expect(result.effectiveModel).toBe("gemini-3-pro-low");
       });
 
+      it("transforms gemini-3.1-pro-preview to gemini-3.1-pro-low for antigravity headerStyle", () => {
+        const result = prepareAntigravityRequest(
+          "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-pro-preview:generateContent",
+          { method: "POST", body: JSON.stringify({ contents: [] }) },
+          mockAccessToken,
+          mockProjectId,
+          undefined,
+          "antigravity"
+        );
+        expect(result.effectiveModel).toBe("gemini-3.1-pro-low");
+      });
+
+      it("transforms gemini-3.1-pro-preview-customtools to gemini-3.1-pro-low for antigravity headerStyle", () => {
+        const result = prepareAntigravityRequest(
+          "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-pro-preview-customtools:generateContent",
+          { method: "POST", body: JSON.stringify({ contents: [] }) },
+          mockAccessToken,
+          mockProjectId,
+          undefined,
+          "antigravity"
+        );
+        expect(result.effectiveModel).toBe("gemini-3.1-pro-low");
+      });
+
       it("transforms gemini-3-flash to gemini-3-flash-preview for gemini-cli headerStyle", () => {
         const result = prepareAntigravityRequest(
           "https://generativelanguage.googleapis.com/v1beta/models/gemini-3-flash:generateContent",
@@ -792,6 +897,30 @@ it("removes x-api-key header", () => {
         expect(result.effectiveModel).toBe("gemini-3-pro-preview");
       });
 
+      it("transforms gemini-3.1-pro-low to gemini-3.1-pro-preview for gemini-cli headerStyle", () => {
+        const result = prepareAntigravityRequest(
+          "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-pro-low:generateContent",
+          { method: "POST", body: JSON.stringify({ contents: [] }) },
+          mockAccessToken,
+          mockProjectId,
+          undefined,
+          "gemini-cli"
+        );
+        expect(result.effectiveModel).toBe("gemini-3.1-pro-preview");
+      });
+
+      it("keeps gemini-3.1-pro-preview-customtools unchanged for gemini-cli headerStyle", () => {
+        const result = prepareAntigravityRequest(
+          "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-pro-preview-customtools:generateContent",
+          { method: "POST", body: JSON.stringify({ contents: [] }) },
+          mockAccessToken,
+          mockProjectId,
+          undefined,
+          "gemini-cli"
+        );
+        expect(result.effectiveModel).toBe("gemini-3.1-pro-preview-customtools");
+      });
+
       it("keeps non-Gemini-3 models unchanged regardless of headerStyle", () => {
         const result = prepareAntigravityRequest(
           "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent",
@@ -805,4 +934,66 @@ it("removes x-api-key header", () => {
       });
     });
   });
+
+  describe("transformAntigravityResponse", () => {
+    it("does not misclassify generic INVALID_ARGUMENT as thinking recovery from debug metadata", async () => {
+      const response = new Response(
+        JSON.stringify({
+          error: {
+            code: 400,
+            message: "Request contains an invalid argument.",
+            status: "INVALID_ARGUMENT",
+          },
+        }),
+        {
+          status: 400,
+          headers: { "content-type": "application/json" },
+        },
+      );
+
+      const transformed = await transformAntigravityResponse(
+        response,
+        true,
+        undefined,
+        "antigravity-claude-opus-4-6-thinking",
+        "test-project",
+        "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse",
+        "claude-opus-4-6-thinking",
+        "session-1",
+        0,
+        "expected=1 found=0",
+      );
+
+      await expect(transformed.text()).resolves.toContain("Request contains an invalid argument.");
+    });
+
+    it("rethrows THINKING_RECOVERY_NEEDED for outer retry handling", async () => {
+      const response = new Response(
+        JSON.stringify({
+          error: {
+            code: 400,
+            message: "Thinking must start with a thinking block before tool use.",
+            status: "INVALID_ARGUMENT",
+          },
+        }),
+        {
+          status: 400,
+          headers: { "content-type": "application/json" },
+        },
+      );
+
+      await expect(
+        transformAntigravityResponse(
+          response,
+          true,
+          undefined,
+          "antigravity-claude-opus-4-6-thinking",
+          "test-project",
+          "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse",
+          "claude-opus-4-6-thinking",
+          "session-1",
+        ),
+      ).rejects.toMatchObject({ message: "THINKING_RECOVERY_NEEDED" });
+    });
+  });
 });
diff --git a/src/plugin/request.ts b/src/plugin/request.ts
@@ -20,6 +20,7 @@ import { defaultSignatureStore } from "./stores/signature-store";
 import {
   DEBUG_MESSAGE_PREFIX,
   isDebugEnabled,
+  isDebugTuiEnabled,
   logAntigravityDebugResponse,
   logCacheStats,
   type AntigravityDebugContext,
@@ -337,6 +338,65 @@ function stripInjectedDebugFromRequestPayload(payload: Record<string, unknown>):
   }
 }
 
+function isValidRequestPart(part: unknown): boolean {
+  if (!part || typeof part !== "object") {
+    return false;
+  }
+
+  const record = part as Record<string, unknown>;
+
+  return (
+    Object.prototype.hasOwnProperty.call(record, "text") ||
+    Object.prototype.hasOwnProperty.call(record, "functionCall") ||
+    Object.prototype.hasOwnProperty.call(record, "functionResponse") ||
+    Object.prototype.hasOwnProperty.call(record, "inlineData") ||
+    Object.prototype.hasOwnProperty.call(record, "fileData") ||
+    Object.prototype.hasOwnProperty.call(record, "executableCode") ||
+    Object.prototype.hasOwnProperty.call(record, "codeExecutionResult") ||
+    Object.prototype.hasOwnProperty.call(record, "thought")
+  );
+}
+
+function sanitizeRequestPayloadForAntigravity(payload: Record<string, unknown>): void {
+  const anyPayload = payload as any;
+
+  if (Array.isArray(anyPayload.contents)) {
+    anyPayload.contents = anyPayload.contents
+      .map((content: unknown) => {
+        if (!content || typeof content !== "object") {
+          return null;
+        }
+
+        const contentRecord = content as Record<string, unknown>;
+        const rawParts = Array.isArray(contentRecord.parts) ? contentRecord.parts : [];
+        const sanitizedParts = rawParts.filter(isValidRequestPart);
+
+        if (sanitizedParts.length === 0) {
+          return null;
+        }
+
+        return {
+          ...contentRecord,
+          parts: sanitizedParts,
+        };
+      })
+      .filter((content: unknown): content is Record<string, unknown> => content !== null);
+  }
+
+  const systemInstruction = anyPayload.systemInstruction;
+  if (systemInstruction && typeof systemInstruction === "object" && !Array.isArray(systemInstruction)) {
+    const sys = systemInstruction as Record<string, unknown>;
+    if (Array.isArray(sys.parts)) {
+      const sanitizedSystemParts = sys.parts.filter(isValidRequestPart);
+      if (sanitizedSystemParts.length > 0) {
+        sys.parts = sanitizedSystemParts;
+      } else {
+        delete anyPayload.systemInstruction;
+      }
+    }
+  }
+}
+
 function isGeminiToolUsePart(part: any): boolean {
   return !!(part && typeof part === "object" && (part.functionCall || part.tool_use || part.toolUse));
 }
@@ -1177,6 +1237,29 @@ export function prepareAntigravityRequest(
         const conversationKey = resolveConversationKey(requestPayload);
         signatureSessionKey = buildSignatureSessionKey(PLUGIN_SESSION_ID, effectiveModel, conversationKey, resolveProjectKey(projectId));
 
+        // Ensure thoughtSignature is present on all functionCall parts if at least one part has it
+        // This is required by Gemini 3.1 Pro which otherwise fails with 400 Bad Request
+        if (Array.isArray(requestPayload.contents)) {
+          requestPayload.contents = requestPayload.contents.map((content: any) => {
+            if (!content || !Array.isArray(content.parts)) return content;
+            
+            // Find if any part has a thoughtSignature
+            const signature = content.parts.find((p: any) => p && typeof p === "object" && typeof p.thoughtSignature === "string")?.thoughtSignature;
+            
+            if (signature) {
+              const newParts = content.parts.map((p: any) => {
+                if (p && typeof p === "object" && p.functionCall && !p.thoughtSignature) {
+                  return { ...p, thoughtSignature: signature };
+                }
+                return p;
+              });
+              return { ...content, parts: newParts };
+            }
+            
+            return content;
+          });
+        }
+
         // For Claude models, filter out unsigned thinking blocks (required by Claude API)
         // Attempts to restore signatures from cache for multi-turn conversations
         // Handle both Gemini-style contents[] and Anthropic-style messages[] payloads.
@@ -1316,6 +1399,7 @@ export function prepareAntigravityRequest(
         }
 
         stripInjectedDebugFromRequestPayload(requestPayload);
+        sanitizeRequestPayloadForAntigravity(requestPayload);
 
         const effectiveProjectId = projectId?.trim() || (headerStyle === "antigravity" ? generateSyntheticProjectId() : "");
         resolvedProjectId = effectiveProjectId;
@@ -1506,7 +1590,7 @@ export async function transformAntigravityResponse(
   // - If keep_thinking=true (but no debug): inject placeholder to trigger signature caching
   // Both use the same injection path (injectDebugThinking) for consistent behavior
   const debugText =
-    isDebugEnabled() && Array.isArray(debugLines) && debugLines.length > 0
+    isDebugTuiEnabled() && Array.isArray(debugLines) && debugLines.length > 0
       ? formatDebugLinesForThinking(debugLines)
       : getKeepThinking()
         ? SYNTHETIC_THINKING_PLACEHOLDER
@@ -1553,6 +1637,8 @@ export async function transformAntigravityResponse(
     });
   }
 
+  const responseFallback = response.clone();
+
   try {
     const headers = new Headers(response.headers);
     const text = await response.text();
@@ -1567,12 +1653,16 @@ export async function transformAntigravityResponse(
 
       // Inject Debug Info
       if (errorBody?.error) {
+        const rawErrorMessage =
+          typeof errorBody.error.message === "string" && errorBody.error.message.length > 0
+            ? errorBody.error.message
+            : "Unknown error";
+        const errorType = detectErrorType(rawErrorMessage);
         const debugInfo = `\n\n[Debug Info]\nRequested Model: ${requestedModel || "Unknown"}\nEffective Model: ${effectiveModel || "Unknown"}\nProject: ${projectId || "Unknown"}\nEndpoint: ${endpoint || "Unknown"}\nStatus: ${response.status}\nRequest ID: ${headers.get("x-request-id") || "N/A"}${toolDebugMissing !== undefined ? `\nTool Debug Missing: ${toolDebugMissing}` : ""}${toolDebugSummary ? `\nTool Debug Summary: ${toolDebugSummary}` : ""}${toolDebugPayload ? `\nTool Debug Payload: ${toolDebugPayload}` : ""}`;
         const injectedDebug = debugText ? `\n\n${debugText}` : "";
-        errorBody.error.message = (errorBody.error.message || "Unknown error") + debugInfo + injectedDebug;
+        errorBody.error.message = rawErrorMessage + debugInfo + injectedDebug;
 
         // Check if this is a recoverable thinking error - throw to trigger retry
-        const errorType = detectErrorType(errorBody.error.message || "");
         if (errorType === "thinking_block_order") {
           const recoveryError = new Error("THINKING_RECOVERY_NEEDED");
           (recoveryError as any).recoveryType = errorType;
@@ -1694,11 +1784,15 @@ export async function transformAntigravityResponse(
 
     return new Response(text, init);
   } catch (error) {
+    if (error instanceof Error && error.message === "THINKING_RECOVERY_NEEDED") {
+      throw error;
+    }
+
     logAntigravityDebugResponse(debugContext, response, {
       error,
       note: "Failed to transform Antigravity response",
     });
-    return response;
+    return responseFallback;
   }
 }
 
