feat(oauth): add account pool support and per-account quota probing for Google Antigravity (#1062)

src/providers/quota.ts

@@ -568,7 +568,7 @@ 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 accountCacheKey(provider: string, accountId: string): string {
@@ -673,6 +673,56 @@ async function getTokenForAccountQuotaProbe(provider: string, accountId: string)
   return getValidAccessTokenForAccount(provider, accountId);
 }
 
+async function fetchAntigravityQuotaWithToken(
+  accessToken: string,
+  projectId: string,
+  baseUrl = "https://daily-cloudcode-pa.googleapis.com",
+): Promise<ProviderQuota | null> {
+  const response = await fetch(`${baseUrl.replace(/\/+$/, "")}/v1internal:fetchAvailableModels`, {
+    method: "POST",
+    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),
+  });
+  if (!response.ok) return null;
+  const body = asRecord(await response.json().catch(() => null));
+  const models = asRecord(body?.models);
+  if (!models) return null;
+
+  const windows = new Map<string, ProviderQuotaWindow>();
+  for (const [modelId, rawModelInfo] of Object.entries(models)) {
+    const modelInfo = asRecord(rawModelInfo);
+    if (!modelInfo) continue;
+    for (const quotaInfo of quotaInfoEntries(modelInfo)) {
+      const label = classifyAntigravityFamily(modelId, modelInfo, quotaInfo);
+      if (!label || windows.has(label)) continue;
+      const percent = antigravityUsedPercent(quotaInfo);
+      if (percent === undefined) continue;
+      windows.set(label, {
+        label,
+        percent,
+        ...(normalizeResetAt(quotaInfo.resetTime) !== undefined ? { resetAt: normalizeResetAt(quotaInfo.resetTime) } : {}),
+      });
+    }
+  }
+
+  const customWindows = ["Gem", "Cla"].flatMap(label => {
+    const window = windows.get(label);
+    return window ? [window] : [];
+  });
+  if (customWindows.length === 0) return null;
+
+  return {
+    customWindows,
+    updatedAt: Date.now(),
+  };
+}
+
 async function fetchAccountQuota(
   provider: string,
   accountId: string,
@@ -688,7 +738,13 @@ 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") {
+        quota = await fetchAnthropicUsageQuota(token);
+      } else if (provider === "google-antigravity") {
+        const cred = getAccountCredential(provider, accountId);
+        quota = cred?.projectId ? await fetchAntigravityQuotaWithToken(token, cred.projectId) : null;
+      }
       if (!quota) {
         // Preserve last-good bars and mark unavailable; advance TTL so failures
         // negative-cache instead of re-probing on every GUI poll.

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

@@ -289,8 +289,10 @@ export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise<
   // Opt-in Anthropic OAuth account pool (#294): enable/threshold/strategy + clear cooldown.
   if (url.pathname === "/api/oauth/accounts/pool" && req.method === "GET") {
     const provider = (url.searchParams.get("provider") ?? "").trim().toLowerCase();
-    if (provider !== "anthropic") return jsonResponse({ error: "pool config is only supported for anthropic" }, 400);
-    const pool = config.anthropicAccountPool ?? {};
+    if (provider !== "anthropic" && provider !== "google-antigravity") {
+      return jsonResponse({ error: "pool config is only supported for anthropic and google-antigravity" }, 400);
+    }
+    const pool = provider === "google-antigravity" ? (config.googleAntigravityAccountPool ?? {}) : (config.anthropicAccountPool ?? {});
     return jsonResponse({
       provider,
       enabled: pool.enabled === true,
@@ -313,8 +315,11 @@ export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise<
       stickyLimit?: unknown;
     };
     const provider = typeof body.provider === "string" ? body.provider.trim().toLowerCase() : "";
-    if (provider !== "anthropic") return jsonResponse({ error: "pool config is only supported for anthropic" }, 400);
-    let enabled = config.anthropicAccountPool?.enabled === true;
+    if (provider !== "anthropic" && provider !== "google-antigravity") {
+      return jsonResponse({ error: "pool config is only supported for anthropic and google-antigravity" }, 400);
+    }
+    const poolKey = provider === "google-antigravity" ? "googleAntigravityAccountPool" : "anthropicAccountPool";
+    let enabled = config[poolKey]?.enabled === true;
     if (body.enabled !== undefined) {
       if (typeof body.enabled !== "boolean") return jsonResponse({ error: "enabled must be a boolean" }, 400);
       enabled = body.enabled;
@@ -347,12 +352,14 @@ export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise<
       }
       stickyLimit = parsed;
     }
-    config.anthropicAccountPool = {
+    const poolObj = {
       enabled,
       autoSwitchThreshold: threshold,
       ...(strategy !== undefined ? { strategy } : {}),
       ...(stickyLimit !== undefined ? { stickyLimit } : {}),
     };
+    if (provider === "google-antigravity") config.googleAntigravityAccountPool = poolObj;
+    else config.anthropicAccountPool = poolObj;
     saveConfigPreservingClaudeCode(config);
     reconcileLiveStateStores();
     return jsonResponse({
@@ -369,7 +376,7 @@ export async function handleOauthAccountRoutes(ctx: ManagementContext): Promise<
     const body = await readManagementJsonBodyOr(req, {}) as { provider?: unknown; accountId?: unknown };
     const provider = typeof body.provider === "string" ? body.provider.trim().toLowerCase() : "";
     const accountId = typeof body.accountId === "string" ? body.accountId.trim() : "";
-    if (provider !== "anthropic") return jsonResponse({ error: "clear-cooldown is only supported for anthropic" }, 400);
+    if (provider !== "anthropic" && provider !== "google-antigravity") return jsonResponse({ error: "clear-cooldown is supported for anthropic and google-antigravity" }, 400);
     if (!accountId) return jsonResponse({ error: "missing accountId" }, 400);
     const { clearAnthropicAccountCooldown } = await import("../../oauth/anthropic-routing");
     const cleared = clearAnthropicAccountCooldown(accountId);

src/types.ts

@@ -783,6 +783,12 @@ export interface OcxConfig {
    * Failover on 429 + sticky affinity; new sessions may pick lowest known 5h usage.
    * Experimental — see docs and GUI warning before enabling.
    */
+  googleAntigravityAccountPool?: {
+    enabled?: boolean;
+    autoSwitchThreshold?: number;
+    strategy?: OcxAccountPoolRotationStrategy;
+    stickyLimit?: number;
+  };
   anthropicAccountPool?: {
     enabled?: boolean;
     /** Usage % threshold for new-session auto-pick. Default 80. 0 = disabled (affinity/active only). */

tests/oauth-accounts-api.test.ts

@@ -175,6 +175,29 @@ describe("multiauth accounts API", () => {
     }
   });
 
+  test("pool config GET and PUT work for google-antigravity", async () => {
+    const server = startServer(0);
+    try {
+      const getRes = await fetch(new URL("/api/oauth/accounts/pool?provider=google-antigravity", server.url));
+      expect(getRes.status).toBe(200);
+      const getJson = await getRes.json() as { provider: string; enabled: boolean };
+      expect(getJson.provider).toBe("google-antigravity");
+      expect(getJson.enabled).toBe(false);
+
+      const putRes = await fetch(new URL("/api/oauth/accounts/pool", server.url), {
+        method: "PUT", headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({ provider: "google-antigravity", enabled: true, autoSwitchThreshold: 85 }),
+      });
+      expect(putRes.status).toBe(200);
+      const putJson = await putRes.json() as { ok: boolean; enabled: boolean; autoSwitchThreshold: number };
+      expect(putJson.ok).toBe(true);
+      expect(putJson.enabled).toBe(true);
+      expect(putJson.autoSwitchThreshold).toBe(85);
+    } finally {
+      await server.stop(true);
+    }
+  });
+
   test("DELETE removes one account; active removal promotes the other", async () => {
     const server = startServer(0);
     try {