fix(antigravity): surface image quotas in Provider Limits (#12311)

Rebase the focused image-quota visibility fix onto the current release/v3.8.51 tip while preserving the newer shared Antigravity catalog and AGY live-discovery changes.

open-sse/config/antigravityModelAliases.ts

@@ -81,6 +81,13 @@ const ANTIGRAVITY_NON_CHAT_MODEL_IDS = new Set([
   "tab_jump_flash_lite_preview",
 ]);
 
+// Non-chat models that still expose user-facing quota buckets. Keep these out of
+// chat discovery while allowing Provider Limits to surface their live quota.
+const ANTIGRAVITY_QUOTA_VISIBLE_NON_CHAT_MODEL_IDS = new Set([
+  "gemini-3-pro-image-preview",
+  "gemini-3.1-flash-image",
+]);
+
 const ANTIGRAVITY_RETIRED_MODEL_IDS = new Set([
   "gemini-3-pro-preview",
   "gemini-3.1-pro",
@@ -161,3 +168,17 @@ export function isDiscoverableAntigravityModelId(modelId: string): boolean {
   }
   return !ANTIGRAVITY_NON_CHAT_MODEL_PATTERN.test(id);
 }
+
+/**
+ * Return whether an Antigravity model quota should be visible to users. Quota
+ * visibility is intentionally broader than chat discovery: image-only models
+ * are callable through /v1/images/generations and have their own live quota
+ * buckets, but must remain excluded from the chat model catalog.
+ */
+export function isUserVisibleAntigravityQuotaModelId(modelId: string): boolean {
+  const id = modelId.trim();
+  if (!id) return false;
+  return (
+    isDiscoverableAntigravityModelId(id) || ANTIGRAVITY_QUOTA_VISIBLE_NON_CHAT_MODEL_IDS.has(id)
+  );
+}

open-sse/services/usage/antigravity.ts

@@ -17,7 +17,7 @@ import {
   getAntigravityFetchAvailableModelsUrls,
 } from "../../config/antigravityUpstream.ts";
 import {
-  isDiscoverableAntigravityModelId,
+  isUserVisibleAntigravityQuotaModelId,
   toClientAntigravityQuotaModelId,
 } from "../../config/antigravityModelAliases.ts";
 import { isDiscoverableAgyModelId } from "../../config/agyModels.ts";
@@ -646,7 +646,7 @@ export async function getAntigravityUsage(
         info.isInternal === true ||
         !(provider === "agy"
           ? isDiscoverableAgyModelId(modelKey)
-          : isDiscoverableAntigravityModelId(modelKey)) ||
+          : isUserVisibleAntigravityQuotaModelId(modelKey)) ||
         Object.keys(quotaInfo).length === 0
       ) {
         continue;
@@ -703,7 +703,7 @@ export async function getAntigravityUsage(
         quotas[modelKey] ||
         !(provider === "agy"
           ? isDiscoverableAgyModelId(modelKey)
-          : isDiscoverableAntigravityModelId(modelKey))
+          : isUserVisibleAntigravityQuotaModelId(modelKey))
       ) {
         continue;
       }

tests/unit/antigravity-image-quota-visibility.test.ts

@@ -0,0 +1,123 @@
+/**
+ * Antigravity image generation uses a non-chat model, but Google reports a
+ * model-specific quota bucket for it. Provider Limits must keep that bucket
+ * without adding image/TTS/tab models to chat discovery.
+ */
+import test from "node:test";
+import assert from "node:assert/strict";
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+
+const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-ag-image-quota-"));
+process.env.DATA_DIR = TEST_DATA_DIR;
+process.env.API_KEY_SECRET = "test-ag-image-quota-secret";
+
+const core = await import("../../src/lib/db/core.ts");
+const { getUsageForProvider } = await import("../../open-sse/services/usage.ts");
+const {
+  isDiscoverableAntigravityModelId,
+  isUserVisibleAntigravityQuotaModelId,
+} = await import("../../open-sse/config/antigravityModelAliases.ts");
+const originalFetch = globalThis.fetch;
+
+const RESET_IN_2_HOURS = new Date(Date.now() + 2 * 60 * 60 * 1000).toISOString();
+
+test.after(() => {
+  globalThis.fetch = originalFetch;
+  core.resetDbInstance();
+  fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 });
+});
+
+test("image quota visibility stays separate from chat discovery", () => {
+  assert.equal(isDiscoverableAntigravityModelId("gemini-3.1-flash-image"), false);
+  assert.equal(isUserVisibleAntigravityQuotaModelId("gemini-3.1-flash-image"), true);
+  assert.equal(isUserVisibleAntigravityQuotaModelId("gemini-3-pro-image-preview"), true);
+
+  assert.equal(isUserVisibleAntigravityQuotaModelId("gemini-3.8-flash-high"), true);
+  assert.equal(isUserVisibleAntigravityQuotaModelId("gemini-3.1-flash-tts-preview"), false);
+  assert.equal(isUserVisibleAntigravityQuotaModelId("tab_flash_lite_preview"), false);
+  assert.equal(isUserVisibleAntigravityQuotaModelId("gemini-3.5-flash-preview"), false);
+  assert.equal(isUserVisibleAntigravityQuotaModelId(""), false);
+});
+
+test("Provider Limits surfaces gemini-3.1-flash-image quota but hides unrelated non-chat buckets", async () => {
+  core.resetDbInstance();
+
+  globalThis.fetch = (async (input: RequestInfo | URL) => {
+    const url =
+      typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
+
+    if (url.includes("retrieveUserQuotaSummary")) {
+      return { ok: false, status: 404, json: async () => ({}) } as Response;
+    }
+
+    if (url.includes("retrieveUserQuota")) {
+      return {
+        ok: true,
+        json: async () => ({
+          buckets: [
+            {
+              modelId: "gemini-3.1-flash-image",
+              remainingFraction: 0.25,
+              resetTime: RESET_IN_2_HOURS,
+            },
+            {
+              modelId: "gemini-3.1-flash-tts-preview",
+              remainingFraction: 0.5,
+              resetTime: RESET_IN_2_HOURS,
+            },
+          ],
+        }),
+      } as Response;
+    }
+
+    if (url.includes("fetchAvailableModels")) {
+      return {
+        ok: true,
+        json: async () => ({
+          models: {
+            "gemini-3.1-flash-image": {
+              quotaInfo: { remainingFraction: 1, resetTime: RESET_IN_2_HOURS },
+            },
+            "gemini-3.1-flash-tts-preview": {
+              quotaInfo: { remainingFraction: 1, resetTime: RESET_IN_2_HOURS },
+            },
+          },
+        }),
+      } as Response;
+    }
+
+    return {
+      ok: true,
+      json: async () => ({
+        cloudaicompanionProject: { id: "test-project" },
+        tierId: "FREE",
+        subscriptionType: "free",
+      }),
+    } as Response;
+  }) as typeof fetch;
+
+  const result = await getUsageForProvider(
+    {
+      id: "conn-image-quota",
+      provider: "antigravity",
+      accessToken: "fake-token-image-quota",
+      providerSpecificData: { clientProfile: "cli" },
+      projectId: "test-project",
+    },
+    { forceRefresh: true }
+  );
+
+  assert.ok(result && "quotas" in result, "should return quotas");
+  const quotas = (
+    result as {
+      quotas: Record<string, { remainingPercentage?: number; quotaSource?: string }>;
+    }
+  ).quotas;
+
+  assert.ok(quotas["gemini-3.1-flash-image"], "image quota should be visible");
+  assert.equal(quotas["gemini-3.1-flash-image"].remainingPercentage, 25);
+  assert.equal(quotas["gemini-3.1-flash-image"].quotaSource, "retrieveUserQuota");
+  assert.equal(quotas["gemini-3.1-flash-tts-preview"], undefined);
+});