From 0a9ffac0194820f45f7b51e91e026ee4d92317d6 Mon Sep 17 00:00:00 2001 From: agentHits Date: Thu, 6 Aug 2026 01:11:24 +0300 Subject: [PATCH] feat(oauth): add account pool support and per-account quota probing for Google Antigravity (#1062) --- src/providers/quota.ts | 60 ++++++++++++++++++- src/server/management/oauth-account-routes.ts | 19 ++++-- src/types.ts | 6 ++ tests/oauth-accounts-api.test.ts | 23 +++++++ 4 files changed, 100 insertions(+), 8 deletions(-) diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 1e5b46fb56..4fba584697 100644 --- a/src/providers/quota.ts +++ b/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 { + 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(); + 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 => { 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. diff --git a/src/server/management/oauth-account-routes.ts b/src/server/management/oauth-account-routes.ts index f3db21c6a6..4aed22c585 100644 --- a/src/server/management/oauth-account-routes.ts +++ b/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); diff --git a/src/types.ts b/src/types.ts index 18c8a63393..e213257853 100644 --- a/src/types.ts +++ b/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). */ diff --git a/tests/oauth-accounts-api.test.ts b/tests/oauth-accounts-api.test.ts index 7d50b81dc7..e5637928ca 100644 --- a/tests/oauth-accounts-api.test.ts +++ b/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 {