feat(oauth): add account pool support and per-account quota probing for Google Antigravity (#1062) - #1084
Conversation
…or Google Antigravity (lidge-jun#1062)
github-actions
Bot
added
the
intake: hygiene-blocked
|
📝 WalkthroughWalkthroughGoogle Antigravity now supports provider-specific account-pool management and per-account quota probing. Quota results include ChangesGoogle Antigravity account pooling
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Review readiness checklistThis PR is kept in draft until every requirement below is fulfilled. The tickable checklist has been added to your PR description — tick all four boxes there.
3/4 boxes ticked. This PR stays in draft until every box above is ticked. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/providers/quota.ts`:
- Around line 676-724: Add focused mocked-fetch regression tests for
fetchAntigravityQuotaWithToken covering Gemini and Claude-family model rows,
missing quota metadata, reset-time conversion, and unavailable project IDs.
Place them beside the existing quota-provider tests, and assert the returned
custom windows and null outcomes for each case.
In `@src/server/management/oauth-account-routes.ts`:
- Around line 321-322: Update the provider pool update flow around poolKey so
every existing setting read at the referenced threshold, strategy, and
sticky-limit handling uses config[poolKey] rather than
config.anthropicAccountPool, preserving Google Antigravity values during partial
updates. Add a regression test that configures distinct Google Antigravity pool
values and verifies a partial PUT or PATCH retains the omitted settings.
In `@tests/oauth-accounts-api.test.ts`:
- Around line 183-192: Remove the duplicate const declarations for getJson and
putJson within the same OAuth account API test callback, retaining one
declaration of each and reusing those variables throughout the callback. Keep
declarations in separate test callbacks unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: b66d89d9-aefc-4ff6-8f26-5e9c614245ab
📒 Files selected for processing (4)
src/providers/quota.tssrc/server/management/oauth-account-routes.tssrc/types.tstests/oauth-accounts-api.test.ts
| 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(), | ||
| }; | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add focused quota-probe regression tests.
This new parser handles untyped upstream model metadata, quota units, family classification, and reset timestamps. The supplied test changes cover only pool configuration endpoints. They do not exercise this quota probe.
Add mocked-fetch tests for Gemini and Claude-family rows, missing quota metadata, reset-time conversion, and unavailable project IDs. As per path instructions, “A behavior change in src/ should come with a focused regression test near the existing tests for that subsystem.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/providers/quota.ts` around lines 676 - 724, Add focused mocked-fetch
regression tests for fetchAntigravityQuotaWithToken covering Gemini and
Claude-family model rows, missing quota metadata, reset-time conversion, and
unavailable project IDs. Place them beside the existing quota-provider tests,
and assert the returned custom windows and null outcomes for each case.
Source: Path instructions
| const poolKey = provider === "google-antigravity" ? "googleAntigravityAccountPool" : "anthropicAccountPool"; | ||
| let enabled = config[poolKey]?.enabled === true; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve the selected provider's existing pool settings.
poolKey selects Google Antigravity only for enabled. Lines 327, 339, and 347 still read config.anthropicAccountPool. A partial update for google-antigravity can reset its stored threshold, strategy, or sticky limit to Anthropic values or defaults.
Read all existing fields from config[poolKey].
Proposed fix
const poolKey = provider === "google-antigravity" ? "googleAntigravityAccountPool" : "anthropicAccountPool";
-let enabled = config[poolKey]?.enabled === true;
+const existingPool = config[poolKey] ?? {};
+let enabled = existingPool.enabled === true;
...
-let threshold = config.anthropicAccountPool?.autoSwitchThreshold ?? 80;
+let threshold = existingPool.autoSwitchThreshold ?? 80;
...
-let strategy = config.anthropicAccountPool?.strategy;
+let strategy = existingPool.strategy;
...
-let stickyLimit = config.anthropicAccountPool?.stickyLimit;
+let stickyLimit = existingPool.stickyLimit;Add a regression test that configures Google Antigravity values, then sends a partial PUT or PATCH request.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| const poolKey = provider === "google-antigravity" ? "googleAntigravityAccountPool" : "anthropicAccountPool"; | |
| let enabled = config[poolKey]?.enabled === true; | |
| const poolKey = provider === "google-antigravity" ? "googleAntigravityAccountPool" : "anthropicAccountPool"; | |
| const existingPool = config[poolKey] ?? {}; | |
| let enabled = existingPool.enabled === true; |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/server/management/oauth-account-routes.ts` around lines 321 - 322, Update
the provider pool update flow around poolKey so every existing setting read at
the referenced threshold, strategy, and sticky-limit handling uses
config[poolKey] rather than config.anthropicAccountPool, preserving Google
Antigravity values during partial updates. Add a regression test that configures
distinct Google Antigravity pool values and verifies a partial PUT or PATCH
retains the omitted settings.
| 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 }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Remove the duplicate const declarations.
getJson is declared twice in one test callback. putJson is declared three times in the same callback. Bun cannot parse this file, so the OAuth account API tests cannot run.
Keep one declaration for each variable. Based on learnings, repeated const declarations are valid only in separate test callbacks; these declarations share one callback scope.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/oauth-accounts-api.test.ts` around lines 183 - 192, Remove the
duplicate const declarations for getJson and putJson within the same OAuth
account API test callback, retaining one declaration of each and reusing those
variables throughout the callback. Keep declarations in separate test callbacks
unchanged.
Source: Learnings
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0a9ffac019
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (provider === "google-antigravity") config.googleAntigravityAccountPool = poolObj; | ||
| else config.anthropicAccountPool = poolObj; |
There was a problem hiding this comment.
Enabling this setting has no effect on Antigravity requests: src/server/responses/core.ts only invokes pool selection and 429 failover when route.providerName === "anthropic" (lines 1576 and 2903-2907), while Antigravity continues through getValidAccessTokenSnapshot, which always selects the single active account. Consequently, the advertised threshold, rotation strategy, affinity, and failover never run; implement the corresponding Antigravity request-routing path and cover actual account selection and 429 rotation rather than only testing config persistence.
AGENTS.md reference: AGENTS.md:L228-L230
Useful? React with 👍 / 👎.
| const poolKey = provider === "google-antigravity" ? "googleAntigravityAccountPool" : "anthropicAccountPool"; | ||
| let enabled = config[poolKey]?.enabled === true; |
There was a problem hiding this comment.
When a Google Antigravity caller sends a partial PATCH, only enabled is initialized through poolKey; threshold, strategy, and stickyLimit below are still read from config.anthropicAccountPool. For example, toggling an existing Antigravity pool off and back on silently replaces its saved threshold and strategy with Anthropic's values or defaults. Read every omitted field from config[poolKey] before rebuilding the provider-specific object.
Useful? React with 👍 / 👎.
| 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); |
There was a problem hiding this comment.
For provider: "google-antigravity", this newly accepted route still calls clearAnthropicAccountCooldown. It therefore always reports no Antigravity cooldown cleared, and if the Google and Anthropic accounts share an identity-derived eight-character ID (IDs are hashed from account ID/email in src/oauth/store.ts), it can instead clear the unrelated Anthropic account's cooldown. Dispatch to provider-owned Antigravity state, or reject this provider until such state exists.
Useful? React with 👍 / 👎.
| googleAntigravityAccountPool?: { | ||
| enabled?: boolean; | ||
| autoSwitchThreshold?: number; | ||
| strategy?: OcxAccountPoolRotationStrategy; | ||
| stickyLimit?: number; | ||
| }; |
There was a problem hiding this comment.
This adds a user-facing configuration object, but docs-site/src/content/docs/reference/configuration/providers.md and its translated counterparts still document only anthropicAccountPool, leaving users without the setting name, defaults, ranges, or experimental behavior needed to configure the feature. Add the Antigravity pool section and keep localized references consistent with the English source.
AGENTS.md reference: AGENTS.md:L231-L232
Useful? React with 👍 / 👎.
| async function fetchAntigravityQuotaWithToken( | ||
| accessToken: string, | ||
| projectId: string, | ||
| baseUrl = "https://daily-cloudcode-pa.googleapis.com", | ||
| ): Promise<ProviderQuota | null> { |
There was a problem hiding this comment.
The per-account probe duplicates the Antigravity service URL as a default argument even though the canonical endpoint already lives in src/providers/registry.ts. If that pinned endpoint changes, ordinary Antigravity traffic can follow the registry while every per-account quota probe continues calling the obsolete host and reports all rows unavailable. Resolve the fixed, allowlisted host from the registry or share the existing provider-level probe instead of maintaining a second provider fact.
AGENTS.md reference: src/AGENTS.md:L18-L18
Useful? React with 👍 / 👎.
| 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; |
There was a problem hiding this comment.
The only test added by this commit exercises pool-config GET/PUT; it never calls the new Antigravity quota branch. As a result, account-specific bearer selection, each credential's projectId, response parsing, and independent failure handling can regress while the submitted test remains green. Extend tests/provider-account-quota.test.ts with multiple Antigravity accounts and distinct tokens/projects, including a failed sibling probe.
AGENTS.md reference: src/AGENTS.md:L24-L27
Useful? React with 👍 / 👎.
| 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 ?? {}); |
There was a problem hiding this comment.
For a normal Google token exchange that omits id_token, this provider cannot retain multiple accounts even though the new route exposes a pool: credentialsFromPayload tries to derive email only by decoding the optional ID token or the access token, but the requested scopes omit openid and there is no userinfo request, so an opaque Google access token leaves both email and accountId unset. saveCredential treats such identityless credentials as replacement-style and overwrites the active slot, meaning addAccount: true can replace the first Antigravity login instead of appending a second one. Request and persist a stable Google identity, for example via OpenID claims or the userinfo endpoint, before enabling pooling.
Useful? React with 👍 / 👎.
|
Maintainer triage (code-level, against
Process notes: |
|
Closing this draft for now — the direction (Antigravity account pool) is wanted, but the current cut implements configuration without the runtime that would use it: (1) no pool-routing consumer reads the added config; (2) the cooldown endpoint accepts |
2 participants
Summary
Fixes #1062.
Adds Account Pooling, Auto-Failover, and Aggregated & Per-Account Quota Tracking support for Google Antigravity / Gemini accounts.
UI Screenshots
Key Changes
Per-Account Quota Probing (
src/providers/quota.ts):supportsPerAccountQuotaforgoogle-antigravity.fetchAntigravityQuotaWithTokento probe individual account quotas using account-specific access tokens and Cloud Code AssistprojectId.Management API Pool Routes (
src/server/management/oauth-account-routes.ts):/api/oauth/accounts/pool(GET/PUT/PATCH) and/api/oauth/accounts/clear-cooldownto supportgoogle-antigravityalongsideanthropic.Config Schema (
src/types.ts):googleAntigravityAccountPoolsettings object (enabled,autoSwitchThreshold,strategy,stickyLimit).Automated Tests (
tests/oauth-accounts-api.test.ts):google-antigravitypool configuration GET and PUT/PATCH endpoints.Readiness Checklist
Review readiness checklist
This PR stays in draft until every box below is ticked. Tick all four boxes once the requirements are met:
Verification
Ran
bun run typecheckandbun test tests/oauth-accounts-api.test.ts:6 pass, 0 failtsc --noEmitpassed with 0 errors.Summary by CodeRabbit