feat: display Antigravity credit balance in dashboard Limits & Quotas (#1338)

Integrated into release/v3.6.7

open-sse/executors/antigravity.ts

@@ -7,8 +7,11 @@ import { classify429, decide429, type Decision } from "../services/antigravity42
 import {
   injectCreditsField,
   shouldRetryWithCredits,
+  shouldUseCreditsFirst,
+  getCreditsMode,
   handleCreditsFailure,
 } from "../services/antigravityCredits.ts";
+import { persistCreditBalance, getAllPersistedCreditBalances } from "@/lib/db/creditBalance";
 import { obfuscateSensitiveWords } from "../services/antigravityObfuscation.ts";
 
 const MAX_RETRY_AFTER_MS = 60_000;
@@ -28,11 +31,30 @@ const creditsExhaustedUntil = new Map<string, number>();
  * Per-account GOOGLE_ONE_AI remaining credit balance cache.
  * Populated from the final SSE chunk's `remainingCredits` field after every
  * successful credit-injected request. Keyed by accountId.
+ * On first access, hydrated from the DB-persisted balances so values survive restarts.
  */
 const creditBalanceCache = new Map<string, number>();
+let creditCacheHydrated = false;
+
+function hydrateCreditCacheFromDb(): void {
+  if (creditCacheHydrated) return;
+  creditCacheHydrated = true;
+  try {
+    const persisted = getAllPersistedCreditBalances();
+    for (const [accountId, balance] of persisted) {
+      // Only fill in accounts not already populated by a live SSE response
+      if (!creditBalanceCache.has(accountId)) {
+        creditBalanceCache.set(accountId, balance);
+      }
+    }
+  } catch {
+    // DB not ready yet (build phase, etc.) — ignore silently
+  }
+}
 
 /** Read the last-known GOOGLE_ONE_AI credit balance for a given account. */
 export function getAntigravityRemainingCredits(accountId: string): number | null {
+  hydrateCreditCacheFromDb();
   const balance = creditBalanceCache.get(accountId);
   return balance !== undefined ? balance : null;
 }
@@ -40,6 +62,12 @@ export function getAntigravityRemainingCredits(accountId: string): number | null
 /** Update the balance cache — called when we parse `remainingCredits` from an SSE stream. */
 export function updateAntigravityRemainingCredits(accountId: string, balance: number): void {
   creditBalanceCache.set(accountId, balance);
+  // Persist to DB so the value survives server restarts
+  try {
+    persistCreditBalance(accountId, balance);
+  } catch {
+    // Non-critical — in-memory cache is the primary source
+  }
 }
 
 function isCreditsExhausted(accountId: string): boolean {
@@ -426,17 +454,31 @@ export class AntigravityExecutor extends BaseExecutor {
     // non-streaming Response so chatCore's non-streaming path stays unchanged.
     const upstreamStream = true;
 
-    // Account ID for credits-exhausted tracking.
-    // Key must match getAntigravityUsage() in fetcher.ts (providerSpecificData?.email || sub).
-    // credentials.email and credentials.sub are populated from the same OAuth token store,
-    // so the cache keys written here and read in the fetcher will always match.
-    const accountId: string = credentials?.email || credentials?.sub || "unknown";
+    // Account ID for credits tracking.
+    // Use connectionId as the stable cache key — it's available in both the executor
+    // (via credentials.connectionId) and the usage fetcher (via connection.id).
+    // The email-based key was unreliable because email isn't always on the credentials object.
+    const accountId: string = credentials?.connectionId || "unknown";
+
+    // Resolve credits mode once per execute() call. "always" injects
+    // enabledCreditTypes: ["GOOGLE_ONE_AI"] on the first request so the
+    // preflight normal call is skipped entirely.
+    const creditsMode = getCreditsMode();
+    const useCreditsFirst = shouldUseCreditsFirst(credentials?.accessToken || "", creditsMode);
 
     for (let urlIndex = 0; urlIndex < fallbackCount; urlIndex++) {
       const url = this.buildUrl(model, upstreamStream, urlIndex);
       const headers = this.buildHeaders(credentials, upstreamStream);
       mergeUpstreamExtraHeaders(headers, upstreamExtraHeaders);
-      const transformedBody = await this.transformRequest(model, body, upstreamStream, credentials);
+      let transformedBody = await this.transformRequest(model, body, upstreamStream, credentials);
+
+      // Credits-first: inject GOOGLE_ONE_AI upfront so we never try the normal
+      // quota path. If credits are exhausted / disabled shouldUseCreditsFirst()
+      // returns false and we fall back to the legacy retry-on-429 flow.
+      if (useCreditsFirst) {
+        transformedBody = injectCreditsField(transformedBody);
+        log?.debug?.("AG_CREDITS", "Credits-first enabled (ANTIGRAVITY_CREDITS=always)");
+      }
 
       // Initialize retry counter for this URL
       if (!retryAttemptsByUrl[urlIndex]) {
@@ -471,17 +513,29 @@ export class AntigravityExecutor extends BaseExecutor {
               // 1. Try to parse explicit retry time from message
               const parsedRetryMs = this.parseRetryFromErrorMessage(errorMessage);
 
-              // 2. Classify 429
-              const category = classify429(errorMessage);
+              // 2. Classify 429 (pass header-parsed retry hint as fallback
+              //    signal — multi-hour Retry-After upgrades rate_limited to
+              //    quota_exhausted so the GOOGLE_ONE_AI credits retry fires).
+              const effectiveRetryHintMs = retryMs ?? parsedRetryMs ?? null;
+              const category = classify429(errorMessage, effectiveRetryHintMs);
 
               // 3. For quota_exhausted, attempt Google One AI credits retry FIRST!
+              //    Skip if credits were already injected on the first call
+              //    (creditsMode === "always") — no point re-running with the
+              //    same body. Record the failure so the 5h breaker kicks in.
+              const creditsAlreadyInjected =
+                (transformedBody as { enabledCreditTypes?: unknown }).enabledCreditTypes != null;
+
+              if (category === "quota_exhausted" && creditsAlreadyInjected) {
+                handleCreditsFailure(credentials?.accessToken || "");
+                log?.warn?.("AG_CREDITS", "Credits-first request 429'd — credits likely exhausted");
+                markCreditsExhausted(accountId);
+              }
+
               if (
                 category === "quota_exhausted" &&
-                shouldRetryWithCredits(
-                  credentials?.accessToken || "",
-                  process.env.ANTIGRAVITY_CREDITS === "1" ||
-                    process.env.ANTIGRAVITY_CREDITS === "true"
-                )
+                !creditsAlreadyInjected &&
+                shouldRetryWithCredits(credentials?.accessToken || "", creditsMode !== "off")
               ) {
                 log?.info?.("AG_CREDITS", "Retrying with Google One AI credits");
                 const creditsBody = injectCreditsField(transformedBody);
@@ -627,7 +681,7 @@ export class AntigravityExecutor extends BaseExecutor {
         // For non-streaming clients, collect the SSE stream and return a synthetic
         // non-streaming Response so chatCore doesn't need to handle SSE conversion.
         if (!stream) {
-          return this.collectStreamToResponse(
+          const collected = await this.collectStreamToResponse(
             response,
             model,
             url,
@@ -636,6 +690,82 @@ export class AntigravityExecutor extends BaseExecutor {
             log,
             signal
           );
+          // When credits were injected (credits-first or credits-retry), the
+          // synthetic body contains _remainingCredits — mirror it into the
+          // balance cache so the dashboard stays fresh.
+          try {
+            const syntheticJson = await collected.response.clone().json();
+            const rc = syntheticJson?._remainingCredits;
+            if (Array.isArray(rc)) {
+              const googleCredit = rc.find(
+                (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI"
+              );
+              if (googleCredit) {
+                const balance = parseInt(googleCredit.creditAmount, 10);
+                if (!isNaN(balance)) updateAntigravityRemainingCredits(accountId, balance);
+              }
+            }
+          } catch {
+            /* balance cache is best-effort */
+          }
+          return collected;
+        }
+
+        // Streaming path: wrap the response body in a pass-through TransformStream
+        // that extracts remainingCredits from the final SSE chunk(s) without
+        // consuming the stream. The client receives the unmodified SSE data.
+        if (response.body) {
+          let sseBuffer = "";
+          const passThrough = new TransformStream({
+            transform(chunk, controller) {
+              controller.enqueue(chunk);
+              // Accumulate text to scan for remainingCredits
+              try {
+                const text = new TextDecoder().decode(chunk, { stream: true });
+                sseBuffer += text;
+              } catch {
+                /* decoding best-effort */
+              }
+            },
+            flush() {
+              // Parse the accumulated SSE data for remainingCredits
+              try {
+                const lines = sseBuffer.split("\n");
+                for (const line of lines) {
+                  const trimmed = line.trim();
+                  if (!trimmed.startsWith("data:")) continue;
+                  const payload = trimmed.slice(5).trim();
+                  if (!payload || payload === "[DONE]") continue;
+                  try {
+                    const parsed = JSON.parse(payload);
+                    if (Array.isArray(parsed?.remainingCredits)) {
+                      const googleCredit = parsed.remainingCredits.find(
+                        (c) => c?.creditType === "GOOGLE_ONE_AI"
+                      );
+                      if (googleCredit) {
+                        const balance = parseInt(googleCredit.creditAmount, 10);
+                        if (!isNaN(balance)) {
+                          updateAntigravityRemainingCredits(accountId, balance);
+                        }
+                      }
+                    }
+                  } catch {
+                    /* skip malformed lines */
+                  }
+                }
+              } catch {
+                /* credits extraction is best-effort */
+              }
+              sseBuffer = "";
+            },
+          });
+          const tappedBody = response.body.pipeThrough(passThrough);
+          const tappedResponse = new Response(tappedBody, {
+            status: response.status,
+            statusText: response.statusText,
+            headers: response.headers,
+          });
+          return { response: tappedResponse, url, headers, transformedBody };
         }
 
         return { response, url, headers, transformedBody };

open-sse/services/usage.ts

@@ -3,15 +3,24 @@
  */
 
 import { PROVIDERS } from "../config/constants.ts";
-import { getAntigravityFetchAvailableModelsUrls } from "../config/antigravityUpstream.ts";
+import {
+  getAntigravityFetchAvailableModelsUrls,
+  ANTIGRAVITY_BASE_URLS,
+} from "../config/antigravityUpstream.ts";
 import { getGlmQuotaUrl } from "../config/glmProvider.ts";
 import { safePercentage } from "@/shared/utils/formatting";
 import { fetchBailianQuota, type BailianTripleWindowQuota } from "./bailianQuotaFetcher.ts";
 import {
   antigravityUserAgent,
+  googApiClientHeader,
   getAntigravityHeaders,
   getAntigravityLoadCodeAssistMetadata,
 } from "./antigravityHeaders.ts";
+import {
+  getAntigravityRemainingCredits,
+  updateAntigravityRemainingCredits,
+} from "../executors/antigravity.ts";
+import { getCreditsMode } from "./antigravityCredits.ts";
 
 // GitHub API config
 const GITHUB_CONFIG = {
@@ -208,7 +217,7 @@ async function getBailianCodingPlanUsage(
  * @returns {Promise<unknown>} Usage data with quotas
  */
 export async function getUsageForProvider(connection) {
-  const { id, provider, accessToken, apiKey, providerSpecificData, projectId } = connection;
+  const { id, provider, accessToken, apiKey, providerSpecificData, projectId, email } = connection;
 
   switch (provider) {
     case "github":
@@ -216,7 +225,7 @@ export async function getUsageForProvider(connection) {
     case "gemini-cli":
       return await getGeminiUsage(accessToken, providerSpecificData, projectId);
     case "antigravity":
-      return await getAntigravityUsage(accessToken, undefined);
+      return await getAntigravityUsage(accessToken, providerSpecificData, projectId, id);
     case "claude":
       return await getClaudeUsage(accessToken);
     case "codex":
@@ -881,16 +890,129 @@ function getAntigravityPlanLabel(subscriptionInfo) {
   return "Free";
 }
 
+/**
+ * Proactive credit balance probe for Antigravity.
+ *
+ * Fires a minimal streamGenerateContent request with GOOGLE_ONE_AI credits enabled
+ * and maxOutputTokens=1 to extract the `remainingCredits` field from the SSE stream.
+ * This uses ~1 credit but lets us show the balance on the dashboard without waiting
+ * for a real user request.
+ *
+ * Returns the credit balance, or null if the probe failed.
+ */
+async function probeAntigravityCreditBalance(
+  accessToken: string,
+  accountId: string,
+  projectId?: string | null
+): Promise<number | null> {
+  try {
+    if (!projectId) return null;
+
+    // Try all base URLs (some accounts only work with specific endpoints)
+    for (const baseUrl of ANTIGRAVITY_BASE_URLS) {
+      const url = `${baseUrl}/v1internal:streamGenerateContent?alt=sse`;
+
+      const sessionId = `-${Math.floor(Math.random() * 9_000_000_000_000_000_000)}`;
+      const body = {
+        project: projectId,
+        model: "gemini-2-flash",
+        userAgent: "antigravity",
+        requestType: "agent",
+        requestId: `credits-probe-${Date.now()}`,
+        enabledCreditTypes: ["GOOGLE_ONE_AI"],
+        request: {
+          model: "gemini-2-flash",
+          contents: [{ role: "user", parts: [{ text: "hi" }] }],
+          generationConfig: { maxOutputTokens: 1 },
+          sessionId,
+        },
+      };
+
+      const headers = {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${accessToken}`,
+        "User-Agent": antigravityUserAgent(),
+        "X-Goog-Api-Client": googApiClientHeader(),
+        Accept: "text/event-stream",
+      };
+
+      try {
+        const res = await fetch(url, {
+          method: "POST",
+          headers,
+          body: JSON.stringify(body),
+          signal: AbortSignal.timeout(10_000),
+        });
+
+        if (!res.ok) continue;
+
+        // Read the full SSE response and scan for remainingCredits
+        const rawSSE = await res.text();
+        const lines = rawSSE.split("\n");
+
+        for (const line of lines) {
+          const trimmed = line.trim();
+          if (!trimmed.startsWith("data:")) continue;
+          const payload = trimmed.slice(5).trim();
+          if (payload === "[DONE]") break;
+          try {
+            const parsed = JSON.parse(payload);
+            if (Array.isArray(parsed?.remainingCredits)) {
+              const googleCredit = parsed.remainingCredits.find(
+                (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI"
+              );
+              if (googleCredit) {
+                const balance = parseInt(googleCredit.creditAmount, 10);
+                if (!isNaN(balance)) {
+                  updateAntigravityRemainingCredits(accountId, balance);
+                  return balance;
+                }
+              }
+            }
+          } catch {
+            // Skip malformed SSE lines
+          }
+        }
+      } catch {
+        // Individual endpoint failure; try next
+      }
+    }
+
+    return null;
+  } catch {
+    // Probe is best-effort — don't let it break the usage fetch
+    return null;
+  }
+}
+
 /**
  * Antigravity Usage - Fetch quota from Google Cloud Code API
  * Uses fetchAvailableModels API which returns ALL models (including Claude)
  * with per-model quotaInfo (remainingFraction, resetTime).
  * retrieveUserQuota only returns Gemini models — not suitable for Antigravity.
  */
-async function getAntigravityUsage(accessToken, providerSpecificData) {
+async function getAntigravityUsage(
+  accessToken,
+  providerSpecificData,
+  connectionProjectId?,
+  connectionId?
+) {
   try {
     const subscriptionInfo = await getAntigravitySubscriptionInfoCached(accessToken);
-    const projectId = subscriptionInfo?.cloudaicompanionProject || null;
+    const projectId = connectionProjectId || subscriptionInfo?.cloudaicompanionProject || null;
+
+    // Derive accountId for credit balance cache.
+    // Must match executor key: credentials.connectionId
+    const accountId: string = connectionId || "unknown";
+
+    // Read cached credit balance (hydrated from DB on first access)
+    let creditBalance = getAntigravityRemainingCredits(accountId);
+
+    // 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);
+    }
 
     // Fetch model list with quota info from fetchAvailableModels
     let response: Response | null = null;
@@ -987,7 +1109,18 @@ async function getAntigravityUsage(accessToken, providerSpecificData) {
 
     return {
       plan: getAntigravityPlanLabel(subscriptionInfo),
-      quotas,
+      quotas: {
+        ...quotas,
+        ...(creditBalance !== null && {
+          credits: {
+            used: 0,
+            total: 0,
+            remaining: creditBalance,
+            unlimited: false,
+            resetAt: null,
+          },
+        }),
+      },
       subscriptionInfo,
     };
   } catch (error) {

src/lib/db/creditBalance.ts

@@ -0,0 +1,92 @@
+/**
+ * Antigravity GOOGLE_ONE_AI credit balance persistence.
+ *
+ * Stores last-known credit balances in the `key_value` table so they survive
+ * server restarts. The in-memory `creditBalanceCache` in the executor is the
+ * primary source at runtime; this module provides the fallback read/write layer.
+ */
+
+import { getDbInstance, isBuildPhase, isCloud } from "./core";
+
+interface StatementLike<TRow = unknown> {
+  get: (...params: unknown[]) => TRow | undefined;
+  run: (...params: unknown[]) => { changes?: number };
+  all: (...params: unknown[]) => TRow[];
+}
+
+interface DbLike {
+  prepare: <TRow = unknown>(sql: string) => StatementLike<TRow>;
+}
+
+interface KeyValueRow {
+  key: string;
+  value: string;
+}
+
+interface CreditBalanceEntry {
+  balance: number;
+  updatedAt: string;
+}
+
+const NAMESPACE = "antigravityCreditBalance";
+
+function parseJson(raw: string): unknown {
+  try {
+    return JSON.parse(raw);
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Read the persisted credit balance for an accountId.
+ * Returns the balance number, or null if not found.
+ */
+export function getPersistedCreditBalance(accountId: string): number | null {
+  if (isBuildPhase || isCloud) return null;
+  const db = getDbInstance() as unknown as DbLike;
+  const row = db
+    .prepare("SELECT value FROM key_value WHERE namespace = ? AND key = ?")
+    .get(NAMESPACE, accountId) as KeyValueRow | undefined;
+  if (!row?.value) return null;
+  const parsed = parseJson(row.value) as CreditBalanceEntry | null;
+  if (!parsed || typeof parsed.balance !== "number") return null;
+  return parsed.balance;
+}
+
+/**
+ * Read all persisted credit balances.
+ * Returns a Map of accountId → balance.
+ */
+export function getAllPersistedCreditBalances(): Map<string, number> {
+  const result = new Map<string, number>();
+  if (isBuildPhase || isCloud) return result;
+  const db = getDbInstance() as unknown as DbLike;
+  const rows = db
+    .prepare("SELECT key, value FROM key_value WHERE namespace = ?")
+    .all(NAMESPACE) as KeyValueRow[];
+  for (const row of rows) {
+    const parsed = parseJson(row.value) as CreditBalanceEntry | null;
+    if (parsed && typeof parsed.balance === "number") {
+      result.set(row.key, parsed.balance);
+    }
+  }
+  return result;
+}
+
+/**
+ * Persist a credit balance for an accountId.
+ */
+export function persistCreditBalance(accountId: string, balance: number): void {
+  if (isBuildPhase || isCloud) return;
+  const db = getDbInstance() as unknown as DbLike;
+  const entry: CreditBalanceEntry = {
+    balance,
+    updatedAt: new Date().toISOString(),
+  };
+  db.prepare("INSERT OR REPLACE INTO key_value (namespace, key, value) VALUES (?, ?, ?)").run(
+    NAMESPACE,
+    accountId,
+    JSON.stringify(entry)
+  );
+}

src/lib/localDb.ts

@@ -267,3 +267,9 @@ export {
 } from "./db/providerLimits";
 
 export type { ProviderLimitsCacheEntry } from "./db/providerLimits";
+
+export {
+  getPersistedCreditBalance,
+  getAllPersistedCreditBalances,
+  persistCreditBalance,
+} from "./db/creditBalance";

src/lib/usage/fetcher.ts

@@ -3,9 +3,20 @@
  */
 
 import { GITHUB_CONFIG, GEMINI_CONFIG } from "@/lib/oauth/constants/oauth";
-import { getAntigravityHeaders } from "@omniroute/open-sse/services/antigravityHeaders.ts";
-import { getAntigravityFetchAvailableModelsUrls } from "@omniroute/open-sse/config/antigravityUpstream.ts";
-import { getAntigravityRemainingCredits } from "@omniroute/open-sse/executors/antigravity.ts";
+import {
+  getAntigravityHeaders,
+  antigravityUserAgent,
+  googApiClientHeader,
+} from "@omniroute/open-sse/services/antigravityHeaders.ts";
+import {
+  getAntigravityFetchAvailableModelsUrls,
+  ANTIGRAVITY_BASE_URLS,
+} from "@omniroute/open-sse/config/antigravityUpstream.ts";
+import {
+  getAntigravityRemainingCredits,
+  updateAntigravityRemainingCredits,
+} from "@omniroute/open-sse/executors/antigravity.ts";
+import { getCreditsMode } from "@omniroute/open-sse/services/antigravityCredits.ts";
 
 /**
  * Get usage data for a provider connection
@@ -21,7 +32,12 @@ export async function getUsageForProvider(connection) {
     case "gemini-cli":
       return await getGeminiUsage(accessToken);
     case "antigravity":
-      return await getAntigravityUsage(accessToken, providerSpecificData);
+      return await getAntigravityUsage(
+        accessToken,
+        providerSpecificData,
+        connection.projectId,
+        connection.id
+      );
     case "claude":
       return await getClaudeUsage(accessToken);
     case "codex":
@@ -147,23 +163,120 @@ async function getGeminiUsage(accessToken) {
   }
 }
 
+/**
+ * Proactive credit balance probe for Antigravity.
+ *
+ * Fires a minimal streamGenerateContent request with GOOGLE_ONE_AI credits enabled
+ * and maxOutputTokens=1 to extract the `remainingCredits` field from the SSE stream.
+ * This uses ~1 credit but lets us show the balance on the dashboard without waiting
+ * for a real user request.
+ *
+ * Returns the credit balance, or null if the probe failed.
+ */
+async function probeAntigravityCreditBalance(
+  accessToken: string,
+  accountId: string,
+  projectId?: string | null
+): Promise<number | null> {
+  try {
+    if (!projectId) return null; // Can't call streamGenerateContent without a projectId
+
+    const baseUrl = ANTIGRAVITY_BASE_URLS[0];
+    const url = `${baseUrl}/v1internal:streamGenerateContent?alt=sse`;
+
+    const body = {
+      project: projectId,
+      model: "gemini-2-flash",
+      userAgent: "antigravity",
+      requestType: "agent",
+      requestId: `credits-probe-${Date.now()}`,
+      enabledCreditTypes: ["GOOGLE_ONE_AI"],
+      request: {
+        model: "gemini-2-flash",
+        contents: [{ role: "user", parts: [{ text: "hi" }] }],
+        generationConfig: { maxOutputTokens: 1 },
+      },
+    };
+
+    const headers = {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${accessToken}`,
+      "User-Agent": antigravityUserAgent(),
+      "X-Goog-Api-Client": googApiClientHeader(),
+      Accept: "text/event-stream",
+    };
+
+    const res = await fetch(url, {
+      method: "POST",
+      headers,
+      body: JSON.stringify(body),
+      signal: AbortSignal.timeout(10_000),
+    });
+
+    if (!res.ok) return null;
+
+    // Read the full SSE response and scan for remainingCredits
+    const rawSSE = await res.text();
+    const lines = rawSSE.split("\n");
+
+    for (const line of lines) {
+      const trimmed = line.trim();
+      if (!trimmed.startsWith("data:")) continue;
+      const payload = trimmed.slice(5).trim();
+      if (payload === "[DONE]") break;
+      try {
+        const parsed = JSON.parse(payload);
+        if (Array.isArray(parsed?.remainingCredits)) {
+          const googleCredit = parsed.remainingCredits.find(
+            (c: { creditType?: string }) => c?.creditType === "GOOGLE_ONE_AI"
+          );
+          if (googleCredit) {
+            const balance = parseInt(googleCredit.creditAmount, 10);
+            if (!isNaN(balance)) {
+              // Cache the balance for future reads (also persists to DB)
+              updateAntigravityRemainingCredits(accountId, balance);
+              return balance;
+            }
+          }
+        }
+      } catch {
+        // Skip malformed SSE lines
+      }
+    }
+
+    return null;
+  } catch {
+    // Probe is best-effort — don't let it break the usage fetch
+    return null;
+  }
+}
+
 /**
  * Antigravity Usage
  * Calls fetchAvailableModels to get per-model quota fractions.
  * Credit balance (GOOGLE_ONE_AI) is read from the executor's in-memory cache,
  * which is populated automatically after each successful credit-injected SSE call.
+ * If the cache is empty and credits mode is enabled, fires a minimal probe request
+ * to fetch the balance proactively.
  */
 async function getAntigravityUsage(
   accessToken: string,
-  providerSpecificData: Record<string, unknown> = {}
+  providerSpecificData: Record<string, unknown> = {},
+  projectId?: string | null,
+  connectionId?: string | null
 ) {
   try {
-    // Derive accountId (same key used in AntigravityExecutor.execute)
-    const accountId: string =
-      (providerSpecificData?.email as string) || (providerSpecificData?.sub as string) || "unknown";
+    // Use connectionId as the cache key — matches executor's credentials.connectionId
+    const accountId: string = connectionId || "unknown";
 
     // Read cached credit balance from executor module (populated from SSE remainingCredits)
-    const creditBalance = getAntigravityRemainingCredits(accountId);
+    let creditBalance = getAntigravityRemainingCredits(accountId);
+
+    // 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);
+    }
 
     // fetchAvailableModels — resolves project from token, no projectId needed
     let res: Response | null = null;