feat(quota): add per-account Gem/Cla quota probing for Google Antigravity (#1082)

src/lib/state-store-registrations.ts

@@ -37,7 +37,7 @@ import { sweepExpiredApiKeyCooldowns } from "../providers/key-failover";
 import { reconcileProviderRequestPacing } from "../providers/request-pacing";
 import { sweepAbandonedResponseStateTemps, sweepExpiredResponseStates } from "../responses/state";
 import { sweepExpiredAntigravityReplay } from "../adapters/google-antigravity-replay";
-import { reconcileProviderAccountQuotaRows } from "../providers/quota";
+import { listLiveProviderAccountQuotaKeys, reconcileProviderAccountQuotaRows } from "../providers/quota";
 import { reconcileRouterWarningMemos } from "../router";
 import type { OcxConfig } from "../types";
 import {
@@ -62,13 +62,15 @@ export function reconcileLiveStateStores() {
 export function buildGenerationContext(): GenerationContext {
   if (!liveServerConfig) throw new Error("live server config is not installed");
   const providerNames = new Set(Object.keys(liveServerConfig.providers));
+  const oauthAccountKeys = listLiveOAuthAccountKeys(providerNames);
   return {
     generation: 0,
     providerNames,
     comboIds: new Set(Object.keys(liveServerConfig.combos ?? {})),
     comboTargets: listLiveComboTargetKeys(liveServerConfig),
     codexAccountIds: listLiveCodexAccountIds(liveServerConfig),
-    oauthAccountKeys: listLiveOAuthAccountKeys(providerNames),
+    oauthAccountKeys,
+    providerAccountQuotaKeys: listLiveProviderAccountQuotaKeys(liveServerConfig.providers, oauthAccountKeys),
     configRoots: listLiveConfigOwnershipRoots(getConfigDir()),
   };
 }

src/lib/state-store-sweeper.ts

@@ -7,6 +7,7 @@ export interface GenerationContext {
   comboTargets: ReadonlySet<string>;
   codexAccountIds: ReadonlySet<string>;
   oauthAccountKeys: ReadonlySet<string>;
+  providerAccountQuotaKeys: ReadonlySet<string>;
   configRoots: ReadonlySet<string>;
 }
 

src/providers/quota.ts

@@ -1,3 +1,7 @@
+const originalFetch = globalThis.fetch;
+import { providerDestinationConfigError } from "../lib/destination-policy";
+import { providerOutboundPost, providerRedirectError } from "../lib/provider-outbound";
+
 import { createHash } from "node:crypto";
 import {
   effectiveCodexAuthAccountId,
@@ -1432,11 +1436,13 @@ type AccountQuotaCacheEntry = {
 const accountQuotaCache = new Map<string, AccountQuotaCacheEntry>();
 const accountQuotaInflight = new Map<string, Promise<AccountQuotaCacheEntry>>();
 let lastReconciledGeneration = 0;
-let liveAccountQuotaKeys = new Set<string>();
+let liveAccountQuotaKeys = new Map<string, number>();
 let liveProviderQuotaKeys = new Set<string>();
 
 function mayCommitAccountQuotaKey(key: string, writerGeneration: number): boolean {
-  return writerGeneration >= lastReconciledGeneration || liveAccountQuotaKeys.has(key);
+  const liveSinceGeneration = liveAccountQuotaKeys.get(key);
+  return writerGeneration >= lastReconciledGeneration
+    || (liveSinceGeneration !== undefined && writerGeneration >= liveSinceGeneration);
 }
 
 function mayCommitProviderQuotaKey(key: string, writerGeneration: number): boolean {
@@ -1452,19 +1458,51 @@ export interface ProviderAccountQuota {
 
 /** Providers whose per-account quota can be probed. Extend as other OAuth APIs are covered. */
 export function supportsPerAccountQuota(provider: string): boolean {
-  return provider === "anthropic";
+  return provider === "anthropic" || provider === "google-antigravity";
+}
+
+function normalizeAntigravityDestination(baseUrl?: string): string {
+  const raw = baseUrl?.trim() || "https://daily-cloudcode-pa.googleapis.com";
+  try {
+    const u = new URL(raw);
+    return u.origin.toLowerCase() + u.pathname.replace(/\/+$/, "");
+  } catch {
+    return raw.replace(/\/+$/, "").toLowerCase();
+  }
+}
+
+function accountCacheKey(provider: string, accountId: string, destination = ""): string {
+  return destination ? provider + "\u0000" + accountId + "\u0000" + destination : provider + "\u0000" + accountId;
 }
 
-function accountCacheKey(provider: string, accountId: string): string {
-  return `${provider}\u0000${accountId}`;
+export function listLiveProviderAccountQuotaKeys(
+  providers: OcxConfig["providers"],
+  oauthAccountKeys: ReadonlySet<string>,
+): Set<string> {
+  const keys = new Set<string>();
+  for (const canonical of oauthAccountKeys) {
+    const separator = canonical.indexOf("\0");
+    const provider = canonical.slice(0, separator);
+    const accountId = canonical.slice(separator + 1);
+    const destination = provider === "google-antigravity"
+      ? normalizeAntigravityDestination(providers[provider]?.baseUrl)
+      : "";
+    keys.add(accountCacheKey(provider, accountId, destination));
+  }
+  return keys;
 }
 
 /**
  * Synchronous last-good per-account quota read for routing. Never probes the network.
  * Returns null when nothing is cached (or the cached row has no bars).
  */
-export function getCachedProviderAccountQuota(provider: string, accountId: string): ProviderQuota | null {
-  const entry = accountQuotaCache.get(accountCacheKey(provider, accountId));
+export function getCachedProviderAccountQuota(
+  provider: string,
+  accountId: string,
+  baseUrl?: string,
+): ProviderQuota | null {
+  const destination = provider === "google-antigravity" ? normalizeAntigravityDestination(baseUrl) : "";
+  const entry = accountQuotaCache.get(accountCacheKey(provider, accountId, destination));
   return entry?.quota ?? null;
 }
 
@@ -1494,9 +1532,16 @@ export function sweepExpiredProviderAccountQuotaRows(now = Date.now()): number {
 
 export function reconcileProviderAccountQuotaRows(context: GenerationContext): number {
   if (context.generation <= lastReconciledGeneration) return 0;
+  const nextLiveAccountQuotaKeys = new Map<string, number>();
+  for (const key of context.providerAccountQuotaKeys) {
+    nextLiveAccountQuotaKeys.set(
+      key,
+      liveAccountQuotaKeys.get(key) ?? (lastReconciledGeneration === 0 ? 0 : context.generation),
+    );
+  }
   let removed = 0;
   for (const key of accountQuotaCache.keys()) {
-    if (context.oauthAccountKeys.has(key)) continue;
+    if (context.providerAccountQuotaKeys.has(key)) continue;
     accountQuotaCache.delete(key);
     removed += 1;
   }
@@ -1505,7 +1550,7 @@ export function reconcileProviderAccountQuotaRows(context: GenerationContext): n
     removed += cache.response.reports.length - reports.length;
     cache = { ...cache, response: { ...cache.response, reports } };
   }
-  liveAccountQuotaKeys = new Set(context.oauthAccountKeys);
+  liveAccountQuotaKeys = nextLiveAccountQuotaKeys;
   liveProviderQuotaKeys = new Set(context.providerNames);
   lastReconciledGeneration = context.generation;
   return removed;
@@ -1514,7 +1559,7 @@ export function reconcileProviderAccountQuotaRows(context: GenerationContext): n
 /** Test-only reset so a direct reconcile call in one file cannot leak across files. */
 export function resetProviderQuotaReconcileStateForTests(): void {
   lastReconciledGeneration = 0;
-  liveAccountQuotaKeys = new Set();
+  liveAccountQuotaKeys = new Map();
   liveProviderQuotaKeys = new Set();
 }
 
@@ -1560,9 +1605,12 @@ async function getTokenForAccountQuotaProbe(provider: string, accountId: string)
 async function fetchAccountQuota(
   provider: string,
   accountId: string,
-  forceRefresh: boolean,
+  forceRefresh = false,
+  baseUrl?: string,
+  allowPrivateNetwork?: boolean,
 ): Promise<AccountQuotaCacheEntry> {
-  const key = accountCacheKey(provider, accountId);
+  const normalizedDest = provider === "google-antigravity" ? normalizeAntigravityDestination(baseUrl) : "";
+  const key = accountCacheKey(provider, accountId, normalizedDest);
   const writerGeneration = captureConfigGeneration();
   const cached = accountQuotaCache.get(key);
   if (!forceRefresh && cached && Date.now() - cached.ts < ACCOUNT_QUOTA_TTL_MS) return cached;
@@ -1571,8 +1619,39 @@ async function fetchAccountQuota(
 
   const probe = (async (): Promise<AccountQuotaCacheEntry> => {
     try {
-      const token = await getTokenForAccountQuotaProbe(provider, accountId);
-      const quota = await fetchAnthropicUsageQuota(token);
+      let quota: ProviderQuota | null = null;
+      if (provider === "anthropic") {
+        const token = await getTokenForAccountQuotaProbe(provider, accountId);
+        quota = await fetchAnthropicUsageQuota(token);
+      } else if (provider === "google-antigravity") {
+        const stored = getAccountCredential(provider, accountId);
+        if (!stored?.projectId) {
+          // Permanent configuration gap (projectId only appears after a fresh login):
+          // skip silently without marking the row unavailable, so the GUI does not
+          // surface a spurious quota error on every poll.
+          const entry: AccountQuotaCacheEntry = { ts: Date.now(), quota: null };
+          if (mayCommitAccountQuotaKey(key, writerGeneration)) {
+            accountQuotaCache.set(key, entry);
+            sweepExpiredOnWrite(entry.ts);
+          }
+          return entry;
+        }
+        // Pre-flight destination policy gate before acquiring or refreshing account token
+        const destError = providerDestinationConfigError("google-antigravity", {
+          baseUrl: baseUrl || "https://daily-cloudcode-pa.googleapis.com",
+          allowPrivateNetwork: allowPrivateNetwork ?? false,
+        });
+        if (destError) {
+          const entry: AccountQuotaCacheEntry = { ts: Date.now(), quota: null, unavailable: true };
+          if (mayCommitAccountQuotaKey(key, writerGeneration)) {
+            accountQuotaCache.set(key, entry);
+            sweepExpiredOnWrite(entry.ts);
+          }
+          return entry;
+        }
+        const token = await getTokenForAccountQuotaProbe(provider, accountId);
+        quota = await fetchAntigravityUsageQuota(token, stored.projectId, baseUrl, { allowPrivateNetwork });
+      }
       if (!quota) {
         // Preserve last-good bars and mark unavailable; advance TTL so failures
         // negative-cache instead of re-probing on every GUI poll.
@@ -1619,12 +1698,14 @@ async function fetchAccountQuota(
 export async function fetchProviderAccountQuotas(
   provider: string,
   forceRefresh = false,
+  baseUrl?: string,
+  allowPrivateNetwork?: boolean,
 ): Promise<ProviderAccountQuota[]> {
   if (!supportsPerAccountQuota(provider)) return [];
   const set = getAccountSet(provider);
   if (!set) return [];
   return await Promise.all(set.accounts.map(async account => {
-    const entry = await fetchAccountQuota(provider, account.id, forceRefresh);
+    const entry = await fetchAccountQuota(provider, account.id, forceRefresh, baseUrl, allowPrivateNetwork);
     return {
       accountId: account.id,
       quota: entry.quota,
@@ -2126,31 +2207,49 @@ function antigravityUsedPercent(quotaInfo: Record<string, unknown>): number | un
       ? toFiniteNumber(quotaInfo.remainingPercentage)! * 100
       : undefined);
   if (remaining === undefined) return undefined;
-  return normalizePercent(100 - remaining);
-}
-
-async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise<ProviderQuotaReport | null> {
-  const credential = getCredential("google-antigravity");
-  if (!credential?.projectId) return null;
-  let accessToken: string;
+  return Math.min(100, Math.max(0, normalizePercent(100 - remaining) ?? 0));
+}
+
+export async function fetchAntigravityUsageQuota(
+  accessToken: string,
+  projectId: string,
+  baseUrl = "https://daily-cloudcode-pa.googleapis.com",
+  options?: {
+    allowPrivateNetwork?: boolean;
+    fetch?: typeof globalThis.fetch;
+    outboundPost?: typeof providerOutboundPost;
+  },
+): Promise<ProviderQuota | null> {
+  const outboundPost = options?.outboundPost ?? providerOutboundPost;
+  const normalizedUrl = (baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, "");
+  const targetUrl = normalizedUrl + "/v1internal:fetchAvailableModels";
+  let response: Response;
   try {
-    accessToken = await getValidAccessToken("google-antigravity");
+    const activeFetch = options?.fetch ?? (globalThis.fetch !== originalFetch ? globalThis.fetch : undefined);
+    response = await outboundPost(
+      "google-antigravity",
+      {
+        baseUrl: normalizedUrl,
+        allowPrivateNetwork: options?.allowPrivateNetwork ?? false,
+        ...(activeFetch ? { fetch: activeFetch } : {}),
+      },
+      targetUrl,
+      {
+        headers: {
+          Accept: "application/json",
+          "Content-Type": "application/json",
+          "User-Agent": antigravityUserAgent(),
+          Authorization: "Bearer " + accessToken,
+        },
+        body: JSON.stringify({ project: projectId }),
+        signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
+      },
+    );
   } catch {
     return null;
   }
-  const baseUrl = (config.baseUrl || "https://daily-cloudcode-pa.googleapis.com").replace(/\/+$/, "");
-  const response = await fetch(`${baseUrl}/v1internal:fetchAvailableModels`, {
-    method: "POST",
-    headers: {
-      Accept: "application/json",
-      "Content-Type": "application/json",
-      "User-Agent": antigravityUserAgent(),
-      Authorization: `Bearer ${accessToken}`,
-    },
-    body: JSON.stringify({ project: credential.projectId }),
-    signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
-  });
-  if (!response.ok) return null;
+  const redirectError = await providerRedirectError(response, targetUrl);
+  if (redirectError || !response.ok) return null;
   const body = asRecord(await readQuotaJson(response));
   const models = asRecord(body?.models);
   if (!models) return null;
@@ -2177,8 +2276,34 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig
     return window ? [window] : [];
   });
   if (customWindows.length === 0) return null;
+  return { customWindows, updatedAt: Date.now() };
+}
+
+async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig): Promise<ProviderQuotaReport | null> {
+  const destError = providerDestinationConfigError("google-antigravity", config);
+  if (destError) return null;
+  const probedAccountId = getAccountSet("google-antigravity")?.activeAccountId;
+  const normalizedDest = normalizeAntigravityDestination(config.baseUrl);
+  const probedAccountKey = probedAccountId ? accountCacheKey("google-antigravity", probedAccountId, normalizedDest) : null;
+  const writerGeneration = captureConfigGeneration();
+  const credential = getCredential("google-antigravity");
+  if (!credential?.projectId) return null;
+  let accessToken: string;
+  try {
+    accessToken = await getValidAccessToken("google-antigravity");
+  } catch {
+    return null;
+  }
+  const quota = await fetchAntigravityUsageQuota(accessToken, credential.projectId, config.baseUrl, { allowPrivateNetwork: config.allowPrivateNetwork });
+  if (!quota) return null;
+  if (probedAccountId && probedAccountKey) {
+    const stillOwnsToken = getAccountCredential("google-antigravity", probedAccountId)?.access === accessToken;
+    if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) {
+      accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota });
+    }
+  }
   return report(provider, "google-antigravity:fetchAvailableModels", {
-    customWindows,
+    ...quota,
     updatedAt: Date.now(),
   });
 }

src/server/management/oauth-account-routes.ts

@@ -279,7 +279,7 @@ export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise<
     const forceRefresh = url.searchParams.get("refresh") === "1";
     // Probing may refresh the active credential and mark needsReauth — project health
     // from the post-probe store so the response is not stale.
-    const rows = await fetchProviderAccountQuotas(provider, forceRefresh);
+    const rows = await fetchProviderAccountQuotas(provider, forceRefresh, ctx.config.providers[provider]?.baseUrl, ctx.config.providers[provider]?.allowPrivateNetwork);
     const byId = new Map(rows.map(row => [row.accountId, row]));
     const projected = projectAccounts();
     return jsonResponse({

tests/combos.test.ts

@@ -809,6 +809,7 @@ describe("combo generation reconciliation", () => {
       comboTargets: new Set(["free::a/m1", "free::c/m3"]),
       codexAccountIds: new Set(),
       oauthAccountKeys: new Set(),
+      providerAccountQuotaKeys: new Set(),
       configRoots: new Set(),
     });
     expect(removed).toBeGreaterThan(0);
@@ -841,6 +842,7 @@ describe("combo generation reconciliation", () => {
       comboTargets: new Set(["free::a/m1", "free::c/m3"]),
       codexAccountIds: new Set(),
       oauthAccountKeys: new Set(),
+      providerAccountQuotaKeys: new Set(),
       configRoots: new Set(),
     });
 

tests/oauth-store-multi.test.ts

@@ -327,6 +327,7 @@ describe("multi-account auth store", () => {
       comboTargets: new Set(),
       codexAccountIds: new Set(),
       oauthAccountKeys: new Set(),
+      providerAccountQuotaKeys: new Set(),
       configRoots: new Set(),
     });
     release();