fix(antigravity): stabilize streaming and usage refresh (#1748)

Integrated into release/v3.7.5 (PR #1748)

open-sse/config/antigravityModelAliases.ts

@@ -39,6 +39,11 @@ const CLIENT_VISIBLE_MODEL_NAMES = Object.freeze(
   }, {})
 );
 
+const PUBLIC_MODEL_IDS = new Set(ANTIGRAVITY_PUBLIC_MODELS.map((model) => model.id));
+const UPSTREAM_PUBLIC_MODEL_IDS = new Set(
+  ANTIGRAVITY_PUBLIC_MODELS.map((model) => resolveAntigravityModelId(model.id))
+);
+
 export function resolveAntigravityModelId(modelId: string): string {
   if (!modelId) return modelId;
   return (ANTIGRAVITY_MODEL_ALIASES as AntigravityModelAliasMap)[modelId] || modelId;
@@ -55,3 +60,10 @@ export function getClientVisibleAntigravityModelName(
 ): string {
   return CLIENT_VISIBLE_MODEL_NAMES[modelId] || fallbackName || modelId;
 }
+
+export function isUserCallableAntigravityModelId(modelId: string): boolean {
+  if (!modelId) return false;
+  const clientId = toClientAntigravityModelId(modelId);
+  const upstreamId = resolveAntigravityModelId(modelId);
+  return PUBLIC_MODEL_IDS.has(clientId) || UPSTREAM_PUBLIC_MODEL_IDS.has(upstreamId);
+}

open-sse/services/usage.ts

@@ -8,6 +8,7 @@ import {
   getAntigravityFetchAvailableModelsUrls,
   ANTIGRAVITY_BASE_URLS,
 } from "../config/antigravityUpstream.ts";
+import { isUserCallableAntigravityModelId } from "../config/antigravityModelAliases.ts";
 import { getGlmQuotaUrl } from "../config/glmProvider.ts";
 import {
   CURSOR_REGISTRY_VERSION,
@@ -599,7 +600,7 @@ async function getBailianCodingPlanUsage(
  * @param {Object} connection - Provider connection with accessToken
  * @returns {Promise<unknown>} Usage data with quotas
  */
-export async function getUsageForProvider(connection) {
+export async function getUsageForProvider(connection, options: { forceRefresh?: boolean } = {}) {
   const { id, provider, accessToken, apiKey, providerSpecificData, projectId, email } = connection;
 
   switch (provider) {
@@ -608,7 +609,7 @@ export async function getUsageForProvider(connection) {
     case "gemini-cli":
       return await getGeminiUsage(accessToken, providerSpecificData, projectId);
     case "antigravity":
-      return await getAntigravityUsage(accessToken, providerSpecificData, projectId, id);
+      return await getAntigravityUsage(accessToken, providerSpecificData, projectId, id, options);
     case "claude":
       return await getClaudeUsage(accessToken);
     case "codex":
@@ -1205,6 +1206,80 @@ function getGeminiCliPlanLabel(subscriptionInfo) {
 // Key: truncated accessToken → { data, fetchedAt }
 const _antigravitySubCache = new Map();
 const ANTIGRAVITY_CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes
+const ANTIGRAVITY_MODELS_CACHE_TTL_MS = 60 * 1000;
+const ANTIGRAVITY_CREDIT_PROBE_TTL_MS = 5 * 60 * 1000;
+const _antigravityAvailableModelsCache = new Map<string, { data: unknown; fetchedAt: number }>();
+const _antigravityAvailableModelsInflight = new Map<string, Promise<unknown>>();
+const _antigravityCreditProbeCache = new Map<string, { data: number | null; fetchedAt: number }>();
+const _antigravityCreditProbeInflight = new Map<string, Promise<number | null>>();
+
+interface AntigravityUsageOptions {
+  forceRefresh?: boolean;
+}
+
+function buildAntigravityUsageCacheKey(accessToken: string, projectId?: string | null): string {
+  return `${accessToken.substring(0, 16)}:${projectId || "default"}`;
+}
+
+async function fetchAntigravityAvailableModelsCached(
+  accessToken: string,
+  projectId?: string | null,
+  options: AntigravityUsageOptions = {}
+): Promise<unknown> {
+  if (!accessToken) throw new Error("Access token is required");
+
+  const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId);
+  const cached = _antigravityAvailableModelsCache.get(cacheKey);
+  if (!options.forceRefresh && cached && Date.now() - cached.fetchedAt < ANTIGRAVITY_MODELS_CACHE_TTL_MS) {
+    return cached.data;
+  }
+
+  const inflight = _antigravityAvailableModelsInflight.get(cacheKey);
+  if (inflight) return inflight;
+
+  const promise = (async () => {
+    let response: Response | null = null;
+    let lastError: Error | null = null;
+
+    for (const quotaApiUrl of ANTIGRAVITY_CONFIG.quotaApiUrls) {
+      try {
+        response = await fetch(quotaApiUrl, {
+          method: "POST",
+          headers: getAntigravityHeaders("fetchAvailableModels", accessToken),
+          body: JSON.stringify(projectId ? { project: projectId } : {}),
+          signal: AbortSignal.timeout(10000),
+        });
+
+        if (response.ok || response.status === 401 || response.status === 403) {
+          break;
+        }
+      } catch (error) {
+        lastError = error as Error;
+      }
+    }
+
+    if (!response) {
+      throw lastError || new Error("Antigravity API unavailable");
+    }
+
+    if (response.status === 403) {
+      return { __antigravityForbidden: true };
+    }
+
+    if (!response.ok) {
+      throw new Error(`Antigravity API error: ${response.status}`);
+    }
+
+    const data = await response.json();
+    _antigravityAvailableModelsCache.set(cacheKey, { data, fetchedAt: Date.now() });
+    return data;
+  })().finally(() => {
+    _antigravityAvailableModelsInflight.delete(cacheKey);
+  });
+
+  _antigravityAvailableModelsInflight.set(cacheKey, promise);
+  return promise;
+}
 
 /**
  * Map raw loadCodeAssist tier data to short display labels.
@@ -1277,6 +1352,40 @@ function getAntigravityPlanLabel(subscriptionInfo) {
  * Returns the credit balance, or null if the probe failed.
  */
 async function probeAntigravityCreditBalance(
+  accessToken: string,
+  accountId: string,
+  projectId?: string | null,
+  options: AntigravityUsageOptions = {}
+): Promise<number | null> {
+  if (!accessToken) return null;
+
+  const cacheKey = buildAntigravityUsageCacheKey(accessToken, projectId || accountId);
+  const cached = _antigravityCreditProbeCache.get(cacheKey);
+  if (!options.forceRefresh && cached && Date.now() - cached.fetchedAt < ANTIGRAVITY_CREDIT_PROBE_TTL_MS) {
+    return cached.data;
+  }
+
+  const inflight = _antigravityCreditProbeInflight.get(cacheKey);
+  if (inflight) return inflight;
+
+  const promise = probeAntigravityCreditBalanceUncached(accessToken, accountId, projectId).then(
+    (data) => {
+      _antigravityCreditProbeCache.set(cacheKey, { data, fetchedAt: Date.now() });
+      return data;
+    },
+    (error) => {
+      _antigravityCreditProbeCache.set(cacheKey, { data: null, fetchedAt: Date.now() });
+      throw error;
+    }
+  ).finally(() => {
+    _antigravityCreditProbeInflight.delete(cacheKey);
+  });
+
+  _antigravityCreditProbeInflight.set(cacheKey, promise);
+  return promise;
+}
+
+async function probeAntigravityCreditBalanceUncached(
   accessToken: string,
   accountId: string,
   projectId?: string | null
@@ -1371,8 +1480,13 @@ async function getAntigravityUsage(
   accessToken,
   providerSpecificData,
   connectionProjectId?,
-  connectionId?
+  connectionId?,
+  options: AntigravityUsageOptions = {}
 ) {
+  if (!accessToken) {
+    return { plan: "Free", message: "Antigravity access token not available." };
+  }
+
   try {
     const subscriptionInfo = await getAntigravitySubscriptionInfoCached(accessToken);
     const projectId = connectionProjectId || subscriptionInfo?.cloudaicompanionProject || null;
@@ -1386,68 +1500,18 @@ async function getAntigravityUsage(
 
     // If no cached balance and credits mode is enabled, fire a minimal probe
     const creditsMode = getCreditsMode();
-    if (creditBalance === null && creditsMode !== "off") {
-      creditBalance = await probeAntigravityCreditBalance(accessToken, accountId, projectId);
+    if ((options.forceRefresh || creditBalance === null) && creditsMode !== "off") {
+      creditBalance = await probeAntigravityCreditBalance(accessToken, accountId, projectId, options);
     }
 
-    // Fetch model list with quota info from fetchAvailableModels
-    let response: Response | null = null;
-    let lastError: Error | null = null;
-
-    for (const quotaApiUrl of ANTIGRAVITY_CONFIG.quotaApiUrls) {
-      try {
-        response = await fetch(quotaApiUrl, {
-          method: "POST",
-          headers: getAntigravityHeaders("fetchAvailableModels", accessToken),
-          body: JSON.stringify(projectId ? { project: projectId } : {}),
-          signal: AbortSignal.timeout(10000),
-        });
-
-        if (response.ok || response.status === 401 || response.status === 403) {
-          break;
-        }
-      } catch (error) {
-        lastError = error as Error;
-      }
-    }
-
-    if (!response) {
-      throw lastError || new Error("Antigravity API unavailable");
-    }
-
-    if (response.status === 403) {
+    const data = await fetchAntigravityAvailableModelsCached(accessToken, projectId, options);
+    const dataObj = toRecord(data);
+    if (dataObj.__antigravityForbidden === true) {
       return { message: "Antigravity access forbidden. Check subscription." };
     }
-
-    if (!response.ok) {
-      throw new Error(`Antigravity API error: ${response.status}`);
-    }
-
-    const data = await response.json();
-    const dataObj = toRecord(data);
     const modelEntries = toRecord(dataObj.models);
     const quotas: Record<string, UsageQuota> = {};
 
-    // Models excluded from quota display — internal/special-purpose models that
-    // the Antigravity API returns quota for but are not user-callable via
-    // generateContent.  Matches CLIProxyAPI's hardcoded exclusion list.
-    const ANTIGRAVITY_EXCLUDED_MODELS = new Set([
-      "chat_20706",
-      "chat_23310",
-      "tab_flash_lite_preview",
-      "tab_jump_flash_lite_preview",
-      "gemini-2.5-flash-thinking",
-      "gemini-2.5-pro", // browser subagent model — not user-callable
-      "gemini-2.5-flash", // internal — quota always exhausted on free tier
-      "gemini-2.5-flash-lite", // internal — quota always exhausted on free tier
-      "gemini-2.5-flash-preview-image-generation", // image-gen only, not usable for chat
-      "gemini-3.1-flash-image-preview", // image-gen preview, not usable for chat
-      "gemini-3-flash-agent", // internal agent model — not user-callable
-      "gemini-3.1-flash-lite", // not usable for chat
-      "gemini-3-pro-low", // not usable for chat
-      "gemini-3-pro-high", // not usable for chat
-    ]);
-
     // Parse per-model quota info from fetchAvailableModels response.
     for (const [modelKey, infoValue] of Object.entries(modelEntries)) {
       const info = toRecord(infoValue);
@@ -1456,7 +1520,7 @@ async function getAntigravityUsage(
       // Skip internal, excluded, and models without quota info
       if (
         info.isInternal === true ||
-        ANTIGRAVITY_EXCLUDED_MODELS.has(modelKey) ||
+        !isUserCallableAntigravityModelId(modelKey) ||
         Object.keys(quotaInfo).length === 0
       ) {
         continue;

open-sse/utils/streamReadiness.ts

@@ -50,6 +50,9 @@ function hasUsefulValue(value: unknown): boolean {
     "function_call_output",
     "output",
     "content_block",
+    "response",
+    "choices",
+    "candidates",
     "parts",
   ]) {
     if (hasUsefulValue(value[key])) return true;
@@ -60,21 +63,7 @@ function hasUsefulValue(value: unknown): boolean {
 
 function hasUsefulJsonPayload(payload: unknown): boolean {
   if (!isRecord(payload)) return false;
-
-  const type = typeof payload.type === "string" ? payload.type : "";
-  if (
-    type.includes("delta") ||
-    type.includes("tool") ||
-    type.includes("function") ||
-    type.includes("content_block") ||
-    type.includes("output_item")
-  ) {
-    if (hasUsefulValue(payload)) return true;
-  }
-
-  return (
-    hasUsefulValue(payload.choices) || hasUsefulValue(payload.candidates) || hasUsefulValue(payload)
-  );
+  return hasUsefulValue(payload);
 }
 
 export function hasUsefulStreamContent(text: string): boolean {

src/app/api/providers/[id]/models/route.ts

@@ -48,6 +48,7 @@ import {
 import {
   ANTIGRAVITY_PUBLIC_MODELS,
   getClientVisibleAntigravityModelName,
+  isUserCallableAntigravityModelId,
   toClientAntigravityModelId,
 } from "@omniroute/open-sse/config/antigravityModelAliases.ts";
 import { getEmbeddingProvider } from "@omniroute/open-sse/config/embeddingRegistry.ts";
@@ -64,6 +65,8 @@ import {
 
 type JsonRecord = Record<string, unknown>;
 
+const antigravityDiscoveryInflight = new Map<string, Promise<Array<{ id: string; name: string }>>>();
+
 function asRecord(value: unknown): JsonRecord {
   return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {};
 }
@@ -173,6 +176,10 @@ function normalizeAntigravityModelsResponse(data: unknown): Array<{ id: string;
     .filter((value): value is { id: string; name: string } => Boolean(value));
 }
 
+function filterUserCallableAntigravityModels(models: Array<{ id: string; name: string }>) {
+  return models.filter((model) => isUserCallableAntigravityModelId(model.id));
+}
+
 function mapAntigravityModelForClient(model: { id: string; name: string }): {
   id: string;
   name: string;
@@ -184,6 +191,58 @@ function mapAntigravityModelForClient(model: { id: string; name: string }): {
   };
 }
 
+async function fetchAntigravityDiscoveryModelsCached(
+  accessToken: string,
+  connectionId: string,
+  proxy: unknown
+): Promise<Array<{ id: string; name: string }>> {
+  const cacheKey = `${connectionId}:${accessToken.substring(0, 16)}`;
+  const inflight = antigravityDiscoveryInflight.get(cacheKey);
+  if (inflight) return inflight;
+
+  const promise = (async () => {
+    await resolveAntigravityVersion();
+
+    for (const discoveryUrl of getAntigravityModelsDiscoveryUrls()) {
+      try {
+        const response = await safeOutboundFetch(discoveryUrl, {
+          ...SAFE_OUTBOUND_FETCH_PRESETS.modelsDiscovery,
+          guard: getProviderOutboundGuard(),
+          proxyConfig: proxy,
+          method: "POST",
+          headers: getAntigravityHeaders("models", accessToken),
+          body: JSON.stringify({}),
+        });
+
+        if (!response.ok) {
+          const errorText = await response.text();
+          console.warn(
+            `[models] antigravity discovery failed at ${discoveryUrl} (${response.status}): ${errorText}`
+          );
+          continue;
+        }
+
+        const models = filterUserCallableAntigravityModels(
+          normalizeAntigravityModelsResponse(await response.json())
+        ).map(mapAntigravityModelForClient);
+        if (models.length > 0) {
+          return models;
+        }
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        console.warn(`[models] antigravity discovery threw for ${discoveryUrl}: ${message}`);
+      }
+    }
+
+    return [];
+  })().finally(() => {
+    antigravityDiscoveryInflight.delete(cacheKey);
+  });
+
+  antigravityDiscoveryInflight.set(cacheKey, promise);
+  return promise;
+}
+
 function normalizeDataRobotCatalogResponse(data: unknown): Array<{ id: string; name: string }> {
   const items = Array.isArray(asRecord(data).data) ? (asRecord(data).data as unknown[]) : [];
 
@@ -1561,7 +1620,6 @@ export async function GET(
       if (autoFetchDisabledResponse) return autoFetchDisabledResponse;
 
       const staticModels = STATIC_MODEL_PROVIDERS.antigravity();
-      const discoveryUrls = getAntigravityModelsDiscoveryUrls();
 
       if (!accessToken) {
         const fallback = buildDiscoveryFallbackResponse({
@@ -1578,37 +1636,9 @@ export async function GET(
         });
       }
 
-      await resolveAntigravityVersion();
-
-      for (const discoveryUrl of discoveryUrls) {
-        try {
-          const response = await safeOutboundFetch(discoveryUrl, {
-            ...SAFE_OUTBOUND_FETCH_PRESETS.modelsDiscovery,
-            guard: getProviderOutboundGuard(),
-            proxyConfig: proxy,
-            method: "POST",
-            headers: getAntigravityHeaders("models", accessToken),
-            body: JSON.stringify({}),
-          });
-
-          if (!response.ok) {
-            const errorText = await response.text();
-            console.warn(
-              `[models] antigravity discovery failed at ${discoveryUrl} (${response.status}): ${errorText}`
-            );
-            continue;
-          }
-
-          const remoteModels = normalizeAntigravityModelsResponse(await response.json()).map(
-            mapAntigravityModelForClient
-          );
-          if (remoteModels.length > 0) {
-            return buildApiDiscoveryResponse(remoteModels);
-          }
-        } catch (error) {
-          const message = error instanceof Error ? error.message : String(error);
-          console.warn(`[models] antigravity discovery threw for ${discoveryUrl}: ${message}`);
-        }
+      const remoteModels = await fetchAntigravityDiscoveryModelsCached(accessToken, connectionId, proxy);
+      if (remoteModels.length > 0) {
+        return buildApiDiscoveryResponse(remoteModels);
       }
 
       const fallback = buildDiscoveryFallbackResponse();

src/lib/usage/providerLimits.ts

@@ -239,6 +239,16 @@ export function getCachedProviderLimitsMap(): Record<string, ProviderLimitsCache
 export async function fetchLiveProviderLimits(connectionId: string): Promise<{
   connection: ProviderConnectionLike;
   usage: JsonRecord;
+}> {
+  return fetchLiveProviderLimitsWithOptions(connectionId, { forceRefresh: false });
+}
+
+async function fetchLiveProviderLimitsWithOptions(
+  connectionId: string,
+  options: { forceRefresh?: boolean } = {}
+): Promise<{
+  connection: ProviderConnectionLike;
+  usage: JsonRecord;
 }> {
   let connection = (await getProviderConnectionById(connectionId)) as ProviderConnectionLike | null;
   if (!connection) {
@@ -250,7 +260,7 @@ export async function fetchLiveProviderLimits(connectionId: string): Promise<{
   }
 
   if (connection.authType !== "oauth") {
-    const usage = (await getUsageForProvider(connection)) as JsonRecord;
+    const usage = (await getUsageForProvider(connection, options)) as JsonRecord;
     if (isRecord(usage.quotas)) {
       setQuotaCache(connectionId, connection.provider, usage.quotas);
     }
@@ -274,7 +284,7 @@ export async function fetchLiveProviderLimits(connectionId: string): Promise<{
         await syncToCloudIfEnabled();
       }
 
-      const usageData = (await getUsageForProvider(conn)) as JsonRecord;
+      const usageData = (await getUsageForProvider(conn, options)) as JsonRecord;
       connection = conn;
       return { usage: usageData };
     });
@@ -330,7 +340,9 @@ export async function fetchAndPersistProviderLimits(
   usage: JsonRecord;
   cache: ProviderLimitsCacheEntry;
 }> {
-  const { connection, usage } = await fetchLiveProviderLimits(connectionId);
+  const { connection, usage } = await fetchLiveProviderLimitsWithOptions(connectionId, {
+    forceRefresh: source === "manual",
+  });
   const cache = toProviderLimitsCacheEntry(usage, source);
   setProviderLimitsCache(connectionId, cache);
   return { connection, usage, cache };
@@ -360,7 +372,9 @@ export async function syncAllProviderLimits(
     const chunk = connections.slice(i, i + concurrency);
     const results = await Promise.allSettled(
       chunk.map(async (connection) => {
-        const { usage } = await fetchLiveProviderLimits(connection.id);
+        const { usage } = await fetchLiveProviderLimitsWithOptions(connection.id, {
+          forceRefresh: source === "manual",
+        });
         const cache = toProviderLimitsCacheEntry(usage, source);
         return { connectionId: connection.id, cache };
       })

tests/unit/antigravity-model-aliases.test.ts

@@ -2,6 +2,7 @@ import test from "node:test";
 import assert from "node:assert/strict";
 
 import {
+  isUserCallableAntigravityModelId,
   resolveAntigravityModelId,
   toClientAntigravityModelId,
 } from "../../open-sse/config/antigravityModelAliases.ts";
@@ -32,6 +33,15 @@ test("toClientAntigravityModelId exposes client-visible aliases for known upstre
   assert.equal(toClientAntigravityModelId("claude-opus-4-6-thinking"), "claude-opus-4-6-thinking");
 });
 
+test("isUserCallableAntigravityModelId only allows public chat-capable model IDs", () => {
+  assert.equal(isUserCallableAntigravityModelId("gemini-3-pro-preview"), true);
+  assert.equal(isUserCallableAntigravityModelId("gemini-3.1-pro-high"), true);
+  assert.equal(isUserCallableAntigravityModelId("claude-sonnet-4-6"), true);
+  assert.equal(isUserCallableAntigravityModelId("gemini-3-flash-agent"), false);
+  assert.equal(isUserCallableAntigravityModelId("tab_flash_lite_preview"), false);
+  assert.equal(isUserCallableAntigravityModelId("unknown-model"), false);
+});
+
 test("AntigravityExecutor.transformRequest resolves alias models before dispatching upstream", async () => {
   const executor = new AntigravityExecutor();
   const result = await executor.transformRequest(

tests/unit/stream-readiness.test.ts

@@ -66,6 +66,18 @@ test("hasUsefulStreamContent detects text, reasoning, and tool deltas", () => {
     ),
     true
   );
+  assert.equal(
+    hasUsefulStreamContent(
+      `data: ${JSON.stringify({ type: "response.output_text.delta", delta: "hello" })}\n\n`
+    ),
+    true
+  );
+  assert.equal(
+    hasUsefulStreamContent(
+      `data: ${JSON.stringify({ response: { candidates: [{ content: { parts: [{ text: "hello" }] } }] } })}\n\n`
+    ),
+    true
+  );
 });
 
 test("ensureStreamReadiness preserves buffered chunks when stream starts", async () => {

tests/unit/usage-fetcher-antigravity.test.ts

@@ -1,59 +0,0 @@
-import test from "node:test";
-import assert from "node:assert/strict";
-
-const usageFetcher = await import("../../src/lib/usage/fetcher.ts");
-
-const originalFetch = globalThis.fetch;
-
-test.afterEach(() => {
-  globalThis.fetch = originalFetch;
-});
-
-test("usage fetcher retries Antigravity quota discovery across shared fallback URLs", async () => {
-  const calls = [];
-
-  globalThis.fetch = async (url, init = {}) => {
-    calls.push({ url: String(url), init });
-
-    const urlStr = String(url);
-    if (
-      urlStr === "https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels" ||
-      urlStr === "https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels"
-    ) {
-      return new Response("unavailable", { status: 503 });
-    }
-
-    return new Response(
-      JSON.stringify({
-        models: {
-          "claude-sonnet-4-6": {
-            quotaInfo: {
-              remainingFraction: 0.4,
-              resetTime: new Date(Date.now() + 60_000).toISOString(),
-            },
-          },
-        },
-      }),
-      { status: 200, headers: { "Content-Type": "application/json" } }
-    );
-  };
-
-  const usage: any = await usageFetcher.getUsageForProvider({
-    provider: "antigravity",
-    accessToken: "ag-token",
-    providerSpecificData: { email: "coder@example.com" },
-  });
-
-  assert.deepEqual(
-    calls.map((call) => call.url),
-    [
-      "https://cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels",
-      "https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels",
-      "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:fetchAvailableModels",
-    ]
-  );
-  assert.match(calls[1].init.headers["User-Agent"], /^antigravity\//);
-  assert.equal(usage.plan, "Antigravity");
-  assert.equal(usage.quotas.models.total, 1);
-  assert.equal(usage.modelQuotas["claude-sonnet-4-6"].remaining, 40);
-});

tests/unit/usage-service-hardening.test.ts

@@ -5,9 +5,15 @@ const usageService = await import("../../open-sse/services/usage.ts");
 const { __testing } = usageService;
 
 const originalFetch = globalThis.fetch;
+const originalCreditsMode = process.env.ANTIGRAVITY_CREDITS;
 
 test.afterEach(() => {
   globalThis.fetch = originalFetch;
+  if (originalCreditsMode === undefined) {
+    delete process.env.ANTIGRAVITY_CREDITS;
+  } else {
+    process.env.ANTIGRAVITY_CREDITS = originalCreditsMode;
+  }
 });
 
 test("usage service covers GitHub free-plan parsing, auth denial and unsupported providers", async () => {
@@ -290,7 +296,7 @@ test("usage service covers Antigravity quota parsing, exclusions and forbidden a
             "gemini-unlimited": {
               quotaInfo: {},
             },
-            "gemini-open": {
+            "gemini-3.1-pro-high": {
               quotaInfo: { remainingFraction: 1 },
             },
             "internal-model": {
@@ -312,10 +318,10 @@ test("usage service covers Antigravity quota parsing, exclusions and forbidden a
   });
 
   assert.equal(usage.plan, "Ultra");
-  assert.deepEqual(Object.keys(usage.quotas).sort(), ["claude-sonnet-4-6", "gemini-open"]);
+  assert.deepEqual(Object.keys(usage.quotas).sort(), ["claude-sonnet-4-6", "gemini-3.1-pro-high"]);
   assert.equal(usage.quotas["claude-sonnet-4-6"].used, 600);
-  assert.equal(usage.quotas["gemini-open"].total, 0);
-  assert.equal(usage.quotas["gemini-open"].remainingPercentage, 100);
+  assert.equal(usage.quotas["gemini-3.1-pro-high"].total, 0);
+  assert.equal(usage.quotas["gemini-3.1-pro-high"].remainingPercentage, 100);
   const loadCodeAssistCall = calls.find((call) => call.url.includes("loadCodeAssist"));
   assert.equal(loadCodeAssistCall?.init.headers["User-Agent"], "google-api-nodejs-client/10.3.0");
   assert.equal(
@@ -407,6 +413,74 @@ test("usage service retries Antigravity fetchAvailableModels across the shared f
   assert.equal(usage.quotas["claude-sonnet-4-6"].used, 500);
 });
 
+test("usage service manual Antigravity refresh bypasses usage TTL caches", async () => {
+  process.env.ANTIGRAVITY_CREDITS = "retry";
+  let probeCalls = 0;
+  let modelCalls = 0;
+
+  globalThis.fetch = async (url) => {
+    const urlStr = String(url);
+    if (urlStr.includes("loadCodeAssist")) {
+      return new Response(JSON.stringify({ cloudaicompanionProject: "ag-project" }), { status: 200 });
+    }
+
+    if (urlStr.includes("streamGenerateContent")) {
+      probeCalls++;
+      return new Response(
+        `data: ${JSON.stringify({ remainingCredits: [{ creditType: "GOOGLE_ONE_AI", creditAmount: String(100 - probeCalls) }] })}\n\n`,
+        { status: 200, headers: { "Content-Type": "text/event-stream" } }
+      );
+    }
+
+    if (urlStr.includes("fetchAvailableModels")) {
+      modelCalls++;
+      return new Response(
+        JSON.stringify({
+          models: {
+            "claude-sonnet-4-6": {
+              quotaInfo: { remainingFraction: 1 },
+            },
+          },
+        }),
+        { status: 200, headers: { "Content-Type": "application/json" } }
+      );
+    }
+
+    throw new Error(`unexpected fetch: ${url}`);
+  };
+
+  const connection = {
+    id: "ag-manual-refresh-service-test",
+    provider: "antigravity",
+    accessToken: "ag-manual-service-token",
+    projectId: "ag-project",
+  };
+
+  await usageService.getUsageForProvider(connection, { forceRefresh: true });
+  await usageService.getUsageForProvider(connection, { forceRefresh: true });
+
+  assert.equal(probeCalls, 2);
+  assert.equal(modelCalls, 2);
+});
+
+test("usage service handles missing Antigravity access tokens without probing upstream", async () => {
+  let fetchCalls = 0;
+
+  globalThis.fetch = async () => {
+    fetchCalls++;
+    return new Response("unexpected", { status: 500 });
+  };
+
+  const usage: any = await usageService.getUsageForProvider({
+    provider: "antigravity",
+    accessToken: undefined,
+  });
+
+  assert.equal(fetchCalls, 0);
+  assert.equal(usage.plan, "Free");
+  assert.match(usage.message, /Antigravity access token not available/i);
+});
+
 test("usage service covers Antigravity tier fallbacks and non-403 upstream failures", async () => {
   globalThis.fetch = async (url) => {
     if (String(url).includes("loadCodeAssist")) {