feat(plugin): add Gemini 3 Flash support and comprehensive tests

shekohex · shekohex · commit b360cc12f93d · 2025-12-17T20:31:54.000+02:00
diff --git a/src/plugin/request.test.ts b/src/plugin/request.test.ts
@@ -1,5 +1,7 @@
 import { expect, test, describe } from "bun:test";
-import { prepareAntigravityRequest } from "./request";
+import { prepareAntigravityRequest, isGenerativeLanguageRequest } from "./request";
+import { overrideEndpointForRequest } from "./fetch-wrapper";
+import { CODE_ASSIST_ENDPOINT, CODE_ASSIST_ENDPOINT_FALLBACKS } from "../constants";
 
 describe("Interleaved Thinking Headers", () => {
   test("adds interleaved thinking header for claude thinking models", async () => {
@@ -60,3 +62,110 @@ describe("Interleaved Thinking Headers", () => {
     expect(headers.get("anthropic-beta")).toBe("interleaved-thinking-2025-05-14");
   });
 });
+
+describe("URL Transformation", () => {
+  test("isGenerativeLanguageRequest detects Gemini API URLs", () => {
+    expect(isGenerativeLanguageRequest("https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent")).toBe(true);
+    expect(isGenerativeLanguageRequest("https://generativelanguage.googleapis.com/v1beta/models/gemini-3-flash:generateContent")).toBe(true);
+    expect(isGenerativeLanguageRequest("https://example.com/api")).toBe(false);
+    expect(isGenerativeLanguageRequest("https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent")).toBe(false);
+  });
+
+  test("transforms Gemini API URL to CODE_ASSIST_ENDPOINT", async () => {
+    const originalUrl = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent";
+    
+    const result = await prepareAntigravityRequest(
+      originalUrl,
+      { method: "POST", body: JSON.stringify({ contents: [] }) },
+      "dummy-token",
+      "dummy-project"
+    );
+    
+    expect(result.request).toBe(`${CODE_ASSIST_ENDPOINT}/v1internal:streamGenerateContent?alt=sse`);
+  });
+
+  test("transforms non-streaming URL correctly", async () => {
+    const originalUrl = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent";
+    
+    const result = await prepareAntigravityRequest(
+      originalUrl,
+      { method: "POST", body: JSON.stringify({ contents: [] }) },
+      "dummy-token",
+      "dummy-project"
+    );
+    
+    expect(result.request).toBe(`${CODE_ASSIST_ENDPOINT}/v1internal:generateContent`);
+    expect(result.streaming).toBe(false);
+  });
+
+  test("sets streaming flag for streamGenerateContent action", async () => {
+    const streamingUrl = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent";
+    const nonStreamingUrl = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent";
+    
+    const streamingResult = await prepareAntigravityRequest(
+      streamingUrl,
+      { method: "POST", body: JSON.stringify({ contents: [] }) },
+      "dummy-token",
+      "dummy-project"
+    );
+    
+    const nonStreamingResult = await prepareAntigravityRequest(
+      nonStreamingUrl,
+      { method: "POST", body: JSON.stringify({ contents: [] }) },
+      "dummy-token",
+      "dummy-project"
+    );
+    
+    expect(streamingResult.streaming).toBe(true);
+    expect(nonStreamingResult.streaming).toBe(false);
+  });
+
+  test("extracts and returns requested model from URL", async () => {
+    const url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3-flash:streamGenerateContent";
+    
+    const result = await prepareAntigravityRequest(
+      url,
+      { method: "POST", body: JSON.stringify({ contents: [] }) },
+      "dummy-token",
+      "dummy-project"
+    );
+    
+    expect(result.requestedModel).toBe("gemini-3-flash");
+  });
+});
+
+describe("Endpoint Fallback Override", () => {
+  test("overrideEndpointForRequest replaces base URL with string input", () => {
+    const preparedUrl = "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse";
+    const autopushEndpoint = CODE_ASSIST_ENDPOINT_FALLBACKS[1];
+    
+    const result = overrideEndpointForRequest(preparedUrl, autopushEndpoint!);
+    
+    expect(result).toBe("https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse");
+  });
+
+  test("overrideEndpointForRequest replaces base URL with URL input", () => {
+    const preparedUrl = new URL("https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse");
+    const prodEndpoint = CODE_ASSIST_ENDPOINT_FALLBACKS[2];
+    
+    const result = overrideEndpointForRequest(preparedUrl, prodEndpoint!);
+    
+    expect(result).toBe("https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse");
+  });
+
+  test("overrideEndpointForRequest works with Request object", () => {
+    const preparedUrl = "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse";
+    const request = new Request(preparedUrl, { method: "POST" });
+    const autopushEndpoint = CODE_ASSIST_ENDPOINT_FALLBACKS[1];
+    
+    const result = overrideEndpointForRequest(request, autopushEndpoint!) as Request;
+    
+    expect(result.url).toBe("https://autopush-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse");
+  });
+
+  test("fallback endpoints are in correct order", () => {
+    expect(CODE_ASSIST_ENDPOINT_FALLBACKS[0]).toContain("daily");
+    expect(CODE_ASSIST_ENDPOINT_FALLBACKS[1]).toContain("autopush");
+    expect(CODE_ASSIST_ENDPOINT_FALLBACKS[2]).toBe("https://cloudcode-pa.googleapis.com");
+  });
+});
diff --git a/src/plugin/request.ts b/src/plugin/request.ts
@@ -71,7 +71,7 @@ function transformSseLine(line: string, onError?: (body: GeminiApiBody) => Gemin
   }
   try {
     let parsed = JSON.parse(json) as unknown;
-    
+
     // Handle array-wrapped responses
     if (Array.isArray(parsed)) {
       parsed = parsed.find((item) => typeof item === "object" && item !== null);
@@ -127,7 +127,7 @@ export function createSseTransformStream(onError?: (body: GeminiApiBody) => Gemi
           if (!sessionId || !family) return;
           const response = body.response as any;
           if (!response?.candidates) return;
-          
+
           response.candidates.forEach((candidate: any, index: number) => {
             if (candidate.groundingMetadata) {
               log.debug("SSE Grounding metadata found", { groundingMetadata: candidate.groundingMetadata });
@@ -199,17 +199,17 @@ export async function prepareAntigravityRequest(
       new Headers(init.headers).forEach((value, key) => reqHeaders.set(key, value));
     }
     requestInit.headers = reqHeaders;
-    
+
     // If body isn't in init, try to get it from request
     if (!originalBody && input.body) {
       // We need to clone to avoid consuming the original request if possible, 
       // but standard Request cloning is sync. 
       // We'll try to read text if we can.
       try {
-         // Note: If input is a Request object that has been used, this might fail.
-         // But usually in this context it's fresh.
-         const cloned = input.clone();
-         originalBody = await cloned.text();
+        // Note: If input is a Request object that has been used, this might fail.
+        // But usually in this context it's fresh.
+        const cloned = input.clone();
+        originalBody = await cloned.text();
       } catch (e) {
         // If clone fails (e.g. body used), we might be in trouble or it's empty.
       }
@@ -256,33 +256,33 @@ export async function prepareAntigravityRequest(
 
       if (isWrapped) {
         if (isClaudeModel) {
-            const context: TransformContext = {
-                model: effectiveModel,
-                family: getModelFamily(effectiveModel),
-                projectId: (parsedBody.project as string) || projectId,
-                streaming,
-                requestId: generateRequestId(),
-                sessionId: getSessionId(),
-            };
-            const innerRequest = parsedBody.request as Record<string, unknown>;
-            const result = transformClaudeRequest(context, innerRequest);
-            body = result.body;
-            transformDebugInfo = result.debugInfo;
-
-            if (transformDebugInfo) {
-                log.debug("Using transformer (wrapped)", { transformer: transformDebugInfo.transformer, model: effectiveModel, family: context.family, toolCount: transformDebugInfo.toolCount });
-            }
+          const context: TransformContext = {
+            model: effectiveModel,
+            family: getModelFamily(effectiveModel),
+            projectId: (parsedBody.project as string) || projectId,
+            streaming,
+            requestId: generateRequestId(),
+            sessionId: getSessionId(),
+          };
+          const innerRequest = parsedBody.request as Record<string, unknown>;
+          const result = transformClaudeRequest(context, innerRequest);
+          body = result.body;
+          transformDebugInfo = result.debugInfo;
+
+          if (transformDebugInfo) {
+            log.debug("Using transformer (wrapped)", { transformer: transformDebugInfo.transformer, model: effectiveModel, family: context.family, toolCount: transformDebugInfo.toolCount });
+          }
         } else {
-            const wrappedBody = {
+          const wrappedBody = {
             ...parsedBody,
             model: effectiveModel,
             userAgent: "antigravity",
             requestId: generateRequestId(),
-            } as Record<string, unknown>;
-            if (wrappedBody.request && typeof wrappedBody.request === "object") {
+          } as Record<string, unknown>;
+          if (wrappedBody.request && typeof wrappedBody.request === "object") {
             (wrappedBody.request as Record<string, unknown>).sessionId = getSessionId();
-            }
-            body = JSON.stringify(wrappedBody);
+          }
+          body = JSON.stringify(wrappedBody);
         }
       } else {
         const context: TransformContext = {
@@ -369,22 +369,22 @@ export async function transformAntigravityResponse(
     return response;
   }
 
-    const errorHandler = (body: GeminiApiBody): GeminiApiBody | null => {
-        const previewErrorFixed = rewriteGeminiPreviewAccessError(body, response.status, requestedModel);
-        const rateLimitErrorFixed = rewriteGeminiRateLimitError(body);
+  const errorHandler = (body: GeminiApiBody): GeminiApiBody | null => {
+    const previewErrorFixed = rewriteGeminiPreviewAccessError(body, response.status, requestedModel);
+    const rateLimitErrorFixed = rewriteGeminiRateLimitError(body);
 
-        const patched = previewErrorFixed ?? rateLimitErrorFixed;
+    const patched = previewErrorFixed ?? rateLimitErrorFixed;
 
-        if (previewErrorFixed?.error) {
-             client.tui.showToast({
-                body: { message: previewErrorFixed.error.message ?? "You need access to gemini 3", title: "Gemini 3 Access Required", variant: "error" }
-            }).catch(() => {});
-        }
+    if (previewErrorFixed?.error) {
+      client.tui.showToast({
+        body: { message: previewErrorFixed.error.message ?? "You need access to gemini 3", title: "Gemini 3 Access Required", variant: "error" }
+      }).catch(() => { });
+    }
 
-        return patched;
-    };
+    return patched;
+  };
 
-    if (streaming && response.ok && isEventStreamResponse && response.body) {
+  if (streaming && response.ok && isEventStreamResponse && response.body) {
     logAntigravityDebugResponse(debugContext, response, {
       note: "Streaming SSE (passthrough mode)",
     });
@@ -419,25 +419,25 @@ export async function transformAntigravityResponse(
     if (sessionId && parsed) {
       const responseBody = parsed.response as any;
       if (responseBody?.candidates) {
-         responseBody.candidates.forEach((candidate: any) => {
-             if (candidate.groundingMetadata) {
-               log.debug("Grounding metadata found", { groundingMetadata: candidate.groundingMetadata });
-             }
-             let fullText = "";
-             let signature = "";
-             if (candidate.content?.parts) {
-                 candidate.content.parts.forEach((part: any) => {
-                     if (part.thought) {
-                          if (part.text) fullText += part.text;
-                          if (part.thoughtSignature) signature = part.thoughtSignature;
-                     }
-                 });
-             }
-             if (fullText && signature) {
-                  cacheSignature(family, sessionId, fullText, signature);
-                  log.debug("Cached signature", { family, sessionId, textLen: fullText.length });
+        responseBody.candidates.forEach((candidate: any) => {
+          if (candidate.groundingMetadata) {
+            log.debug("Grounding metadata found", { groundingMetadata: candidate.groundingMetadata });
+          }
+          let fullText = "";
+          let signature = "";
+          if (candidate.content?.parts) {
+            candidate.content.parts.forEach((part: any) => {
+              if (part.thought) {
+                if (part.text) fullText += part.text;
+                if (part.thoughtSignature) signature = part.thoughtSignature;
               }
-         });
+            });
+          }
+          if (fullText && signature) {
+            cacheSignature(family, sessionId, fullText, signature);
+            log.debug("Cached signature", { family, sessionId, textLen: fullText.length });
+          }
+        });
       }
     }
 
@@ -487,7 +487,7 @@ export async function transformAntigravityResponse(
             variant: "error",
           },
         });
-      } catch {}
+      } catch { }
     }
 
     if (streaming && response.ok && isEventStreamResponse) {
diff --git a/tests/integration/google-search.test.ts b/tests/integration/google-search.test.ts
@@ -116,6 +116,57 @@ describe("Google Search Tool", () => {
     });
   });
 
+  describe("with Gemini 3 Flash", () => {
+    let sessionId: string;
+
+    beforeEach(async () => {
+      if (process.env.CI) return;
+      sessionId = await createSession(ctx);
+    });
+
+    afterEach(async () => {
+      if (process.env.CI) return;
+      if (sessionId) {
+        try {
+          await deleteSession(ctx, sessionId);
+        } catch {}
+      }
+    });
+
+    itSkipInCI("should enable google_search tool and return grounded results", async () => {
+      const response = await sendPrompt(
+        ctx,
+        sessionId,
+        "Search the web for today's top technology news.",
+        {
+          model: TEST_MODELS.gemini3Flash,
+          tools: { google_search: true },
+        }
+      );
+
+      expect(response).not.toBeNull();
+      expect(response!.info.role).toBe("assistant");
+
+      const assistantText = extractTextFromParts(response!.parts);
+      expect(assistantText.length).toBeGreaterThan(0);
+    });
+
+    itSkipInCI("should work without google_search tool", async () => {
+      const response = await sendPrompt(
+        ctx,
+        sessionId,
+        "What is the capital of France?",
+        {
+          model: TEST_MODELS.gemini3Flash,
+        }
+      );
+
+      expect(response).not.toBeNull();
+      const assistantText = extractTextFromParts(response!.parts);
+      expect(assistantText.toLowerCase()).toContain("paris");
+    });
+  });
+
   describe("with Claude Sonnet via Antigravity", () => {
     let sessionId: string;
 
diff --git a/tests/integration/setup.ts b/tests/integration/setup.ts
@@ -11,6 +11,7 @@ export interface TestContext {
 export const TEST_MODELS = {
   gemini25Flash: { providerID: "google", modelID: "gemini-2.5-flash" },
   gemini3Pro: { providerID: "google", modelID: "gemini-3-pro-preview" },
+  gemini3Flash: { providerID: "google", modelID: "gemini-3-flash" },
   claudeSonnet: { providerID: "google", modelID: "gemini-claude-sonnet-4-5-thinking-medium" },
 } as const;
 
