feat(oauth): rank Antigravity failover by Gemini vs Claude quota family

scripts/test-layout/layout.json

@@ -952,6 +952,7 @@
     "nvidia-nim-hardening.test.ts": "providers",
     "oauth-account-attribution.test.ts": "oauth",
     "oauth-account-id-collision.test.ts": "oauth",
+    "oauth-account-quota-rank.test.ts": "oauth",
     "oauth-accounts-api.test.ts": "oauth",
     "oauth-callback-binds.test.ts": "oauth",
     "oauth-callback-server.test.ts": "oauth",

src/oauth/account-quota-rank.ts

@@ -13,6 +13,37 @@
 import { getCachedProviderAccountQuota, hasPassiveAccountQuota } from "../providers/quota";
 import { getKiroAccountExhaustion } from "../providers/kiro-usage";
 
+/** Antigravity hosts Gemini and Claude windows on one account; ranking must not mix them. */
+export type QuotaModelFamily = "gem" | "cla";
+
+export function classifyModelFamilyForQuota(
+  provider: string,
+  modelId?: string | null,
+): QuotaModelFamily | undefined {
+  if (provider !== "google-antigravity" || typeof modelId !== "string" || !modelId.trim()) {
+    return undefined;
+  }
+  const id = modelId.toLowerCase();
+  // Gemma is not Gemini: a substring/prefix match would poison Gemini ranking.
+  if (/(?:^|[^a-z])gemma(?:[^a-z]|$)/.test(id)) return undefined;
+  // Catalog ids are gemini-*, never a bare gem- token. Window labels still match Gem via
+  // windowMatchesFamily; this classifier is only for request model ids.
+  if (/(?:^|[^a-z])gemini(?:[^a-z]|$)/.test(id)) return "gem";
+  if (
+    /(?:^|[^a-z])claude(?:[^a-z]|$)/.test(id)
+    || /(?:^|[^a-z])opus(?:[^a-z]|$)/.test(id)
+    || /(?:^|[^a-z])sonnet(?:[^a-z]|$)/.test(id)
+    || /(?:^|[^a-z])haiku(?:[^a-z]|$)/.test(id)
+  ) return "cla";
+  return undefined;
+}
+
+function windowMatchesFamily(label: string, family: QuotaModelFamily): boolean {
+  const token = label.trim().split(/[\s(/]+/)[0] ?? "";
+  if (family === "gem") return /^gem(?:ini)?$/i.test(token);
+  return /^cla(?:ude)?$/i.test(token);
+}
+
 /** Lower sorts earlier. Unknown sits between measured-healthy and measured-empty. */
 const RANK_HEALTHY = 0;
 const RANK_UNKNOWN = 1;
@@ -48,15 +79,24 @@ const PASSIVE_HEADROOM_MAX_AGE_MS = 60 * 60_000;
 /**
  * Remaining headroom across every window the provider reports.
  *
- * The minimum wins: an account at 5% of its five-hour window is unusable right now even if
- * its monthly allowance is barely touched.
- */
-function headroomOf(provider: string, accountId: string): number | null {
+* The minimum wins: an account at 5% of its five-hour window is unusable right now even if
+* its monthly allowance is barely touched.
+*/
+function headroomOf(provider: string, accountId: string, requestedModelId?: string | null): number | null {
   const quota = getCachedProviderAccountQuota(provider, accountId);
   if (!quota) return null;
   // Null, not a low rank: this must reproduce "no evidence" so a stale roster degrades to
   // the unranked ring rather than to a differently wrong answer.
   if (hasPassiveAccountQuota(provider) && Date.now() - quota.updatedAt > PASSIVE_HEADROOM_MAX_AGE_MS) return null;
+  const family = classifyModelFamilyForQuota(provider, requestedModelId);
+  if (family) {
+    const percents = (quota.customWindows ?? [])
+      .filter(window => windowMatchesFamily(window.label, family))
+      .map(window => window.percent)
+      .filter((value): value is number => typeof value === "number");
+    if (percents.length === 0) return null;
+    return 100 - Math.max(...percents);
+  }
   const percents = [
     quota.fiveHourPercent,
     quota.weeklyPercent,
@@ -74,15 +114,23 @@ function headroomOf(provider: string, accountId: string): number | null {
  * than an ordering. Null stays null all the way out: a caller must decide what "unmeasured"
  * means for its own rule instead of being handed a fabricated 0 or 100.
  */
-export function accountHeadroomPercent(provider: string, accountId: string): number | null {
-  return headroomOf(provider, accountId);
+export function accountHeadroomPercent(
+  provider: string,
+  accountId: string,
+  requestedModelId?: string | null,
+): number | null {
+  return headroomOf(provider, accountId, requestedModelId);
 }
 
 /** Unknown usage is not exhaustion; Kiro's explicit overage verdict is authoritative. */
-export function isAccountQuotaExhausted(provider: string, accountId: string): boolean {
+export function isAccountQuotaExhausted(
+  provider: string,
+  accountId: string,
+  requestedModelId?: string | null,
+): boolean {
   const exhaustion = provider === "kiro" ? getKiroAccountExhaustion(`${provider}\u0000${accountId}`) : null;
   if (exhaustion !== null) return exhaustion.exhausted;
-  const headroom = headroomOf(provider, accountId);
+  const headroom = headroomOf(provider, accountId, requestedModelId);
   return headroom !== null && headroom <= 0;
 }
 
@@ -92,24 +140,28 @@ export function isAccountQuotaExhausted(provider: string, accountId: string): bo
  * Returns the input untouched when no candidate has quota evidence, which keeps every
  * provider without per-account quota on exactly the behaviour it has today.
  */
-export function rankAccountsByHeadroom(provider: string, ring: readonly string[]): string[] {
+export function rankAccountsByHeadroom(
+  provider: string,
+  ring: readonly string[],
+  requestedModelId?: string | null,
+): string[] {
   if (ring.length < 2) return [...ring];
 
   let sawEvidence = false;
   // Same rule as hasHeadroomEvidence: a passive provider's partial roster must not rank
   // at all. The failover path calls this directly (selectFailoverAccount), so the guard
   // cannot live only in the pre-dispatch predicate.
-  if (hasPassiveAccountQuota(provider) && !ring.every(id => headroomOf(provider, id) !== null)) {
+  if (hasPassiveAccountQuota(provider) && !ring.every(id => headroomOf(provider, id, requestedModelId) !== null)) {
     return [...ring];
   }
   const ranked: Ranked[] = ring.map((id, index) => {
     // A provider-declared exhaustion verdict outranks the percentage: an account may sit at
     // 100% and still be servable when overage is enabled, and the verdict knows that.
     const exhaustion = provider === "kiro" ? getKiroAccountExhaustion(`${provider}\u0000${id}`) : null;
-    const headroom = headroomOf(provider, id);
+    const headroom = headroomOf(provider, id, requestedModelId);
     if (exhaustion !== null || headroom !== null) sawEvidence = true;
 
-    if (isAccountQuotaExhausted(provider, id)) return { id, bucket: RANK_EXHAUSTED, headroom: 0, index };
+    if (isAccountQuotaExhausted(provider, id, requestedModelId)) return { id, bucket: RANK_EXHAUSTED, headroom: 0, index };
     if (headroom === null) return { id, bucket: RANK_UNKNOWN, headroom: 0, index };
     return { id, bucket: RANK_HEALTHY, headroom, index };
   });
@@ -129,7 +181,11 @@ export function rankAccountsByHeadroom(provider: string, ring: readonly string[]
  * told "ranked" when nothing was measured. Pre-dispatch selection asks this first so it
  * can decline to act on a roster it knows nothing about.
  */
-export function hasHeadroomEvidence(provider: string, ids: readonly string[]): boolean {
+export function hasHeadroomEvidence(
+  provider: string,
+  ids: readonly string[],
+  requestedModelId?: string | null,
+): boolean {
   // A PASSIVE provider needs evidence for EVERY candidate, not any one of them.
   //
   // A probe fills the whole roster in one pass (fetchProviderAccountQuotas), so "any"
@@ -140,10 +196,10 @@ export function hasHeadroomEvidence(provider: string, ids: readonly string[]): b
   // AWAY from an unmeasured account and TOWARD the one account known to be spent, which
   // is the exact inversion of what ranking is for.
   if (hasPassiveAccountQuota(provider)) {
-    return ids.length > 0 && ids.every(id => headroomOf(provider, id) !== null);
+    return ids.length > 0 && ids.every(id => headroomOf(provider, id, requestedModelId) !== null);
   }
   return ids.some(id =>
-    headroomOf(provider, id) !== null
+    headroomOf(provider, id, requestedModelId) !== null
     || (provider === "kiro" && getKiroAccountExhaustion(`${provider}\u0000${id}`) !== null));
 }
 /**

src/oauth/generic-account-failover.ts

@@ -22,6 +22,8 @@ import {
   hasHeadroomEvidence,
   isAccountQuotaExhausted,
   rankAccountsByHeadroom,
+  classifyModelFamilyForQuota,
+  type QuotaModelFamily,
 } from "./account-quota-rank";
 import {
   genericPoolKey,
@@ -81,13 +83,14 @@ const health = new Map<string, AccountHealth>();
 /** Provider -> recent eligible-account count. TTL-bounded; never holds credential material. */
 const presence = new Map<string, PresenceEntry>();
 
-const healthKey = (provider: string, accountId: string) => `${provider}\u0000${accountId}`;
+const healthKey = (provider: string, accountId: string, family?: QuotaModelFamily) =>
+  family ? `${provider}\u0000${accountId}\u0000${family}` : `${provider}\u0000${accountId}`;
 
-function isCooled(provider: string, accountId: string, now: number): boolean {
-  const entry = health.get(healthKey(provider, accountId));
+function isCooled(provider: string, accountId: string, now: number, family?: QuotaModelFamily): boolean {
+  const entry = health.get(healthKey(provider, accountId, family));
   if (!entry) return false;
   if (entry.cooldownUntil <= now) {
-    health.delete(healthKey(provider, accountId));
+    health.delete(healthKey(provider, accountId, family));
     return false;
   }
   return true;
@@ -175,11 +178,11 @@ function isProactivePreferenceEnabled(config: OcxConfig, providerName: string, n
 }
 
 /** Accounts that may serve traffic right now: not cooled, not flagged for reauth. */
-export function eligibleFailoverAccounts(providerName: string, now = Date.now()): string[] {
+export function eligibleFailoverAccounts(providerName: string, now = Date.now(), family?: QuotaModelFamily): string[] {
   const set = getAccountSet(providerName);
   if (!set) return [];
   return set.accounts
-    .filter(account => account.needsReauth !== true && !isCooled(providerName, account.id, now))
+    .filter(account => account.needsReauth !== true && !isCooled(providerName, account.id, now, family))
     .map(account => account.id);
 }
 
@@ -230,8 +233,8 @@ function stableGenericRoster(providerName: string): string[] {
  * statement about observed usage, and treating "no observation" as "spent" would evacuate every
  * quota-less provider off its active account on the very first request.
  */
-function isOverAutoSwitchThreshold(providerName: string, accountId: string, threshold: number): boolean {
-  const headroom = accountHeadroomPercent(providerName, accountId);
+function isOverAutoSwitchThreshold(providerName: string, accountId: string, threshold: number, requestedModelId?: string | null): boolean {
+  const headroom = accountHeadroomPercent(providerName, accountId, requestedModelId);
   if (headroom === null) return false;
   return 100 - headroom >= threshold;
 }
@@ -245,15 +248,17 @@ function pickFillFirstGenericAccount(
   providerName: string,
   activeId: string | undefined,
   now: number,
+  requestedModelId?: string | null,
 ): string | null {
   const stableAll = stableGenericRoster(providerName);
   if (stableAll.length < 2) return null;
-  const eligible = new Set(eligibleFailoverAccounts(providerName, now));
+  const family = classifyModelFamilyForQuota(providerName, requestedModelId);
+  const eligible = new Set(eligibleFailoverAccounts(providerName, now, family));
   const stored = config.providers?.[providerName]?.oauthAccountFailover?.autoSwitchThreshold;
   const threshold = typeof stored === "number" && Number.isInteger(stored) && stored >= 0 && stored <= 100
     ? stored
     : DEFAULT_GENERIC_AUTO_SWITCH_THRESHOLD;
-  if (activeId && eligible.has(activeId) && !isOverAutoSwitchThreshold(providerName, activeId, threshold)) {
+  if (activeId && eligible.has(activeId) && !isOverAutoSwitchThreshold(providerName, activeId, threshold, requestedModelId)) {
     return null;
   }
   const start = activeId ? stableAll.indexOf(activeId) : -1;
@@ -277,11 +282,17 @@ function pickFillFirstGenericAccount(
  * advance and round-robin would propose the same account forever. This is the same shape
  * `commitAnthropicSelectionRouting` already commits with.
  */
-export function noteGenericPoolSelection(config: OcxConfig, providerName: string, accountId: string): void {
+export function noteGenericPoolSelection(
+  config: OcxConfig,
+  providerName: string,
+  accountId: string,
+  requestedModelId?: string | null,
+): void {
   if (activeGenericStrategy(config, providerName) !== "round-robin") return;
   const poolKey = genericPoolKey(providerName);
   const limit = genericStickyLimit(config, providerName);
-  const picked = pickRoundRobinAccount(poolKey, eligibleFailoverAccounts(providerName), limit);
+  const family = classifyModelFamilyForQuota(providerName, requestedModelId);
+  const picked = pickRoundRobinAccount(poolKey, eligibleFailoverAccounts(providerName, Date.now(), family), limit);
   // The resolver may have admitted a different account than the ring proposed: a removal, a
   // reauth verdict or a manual selection can land during credential resolution. Realign the
   // cursor onto what actually served rather than leaving it on a road not taken.
@@ -302,6 +313,7 @@ export function rotateGenericOAuthAccountOn429(
   failedAccountId: string,
   retryAfterHeader: string | null | undefined,
   now = Date.now(),
+  requestedModelId?: string | null,
 ): string | null {
   if (!isGenericOAuthFailoverEnabled(config, providerName)) return null;
   const set = getAccountSet(providerName);
@@ -314,13 +326,14 @@ export function rotateGenericOAuthAccountOn429(
   // A Retry-After from upstream still wins — it is the server's own instruction.
   const exhausted = parsed === undefined ? exhaustedCooldownMs(providerName, failedAccountId, now) : null;
   const cooldownMs = exhausted ?? Math.min(parsed ?? DEFAULT_COOLDOWN_MS, MAX_COOLDOWN_MS);
-  health.set(healthKey(providerName, failedAccountId), {
+  const family = classifyModelFamilyForQuota(providerName, requestedModelId);
+  health.set(healthKey(providerName, failedAccountId, family), {
     cooldownUntil: now + cooldownMs,
     cooldownSource: parsed ? "retry-after" : "default",
   });
   sweepExpiredOnWrite(now);
 
-  const eligible = eligibleFailoverAccounts(providerName, now).filter(id => id !== failedAccountId);
+  const eligible = eligibleFailoverAccounts(providerName, now, family).filter(id => id !== failedAccountId);
   if (eligible.length === 0) return null;
   // A rotation means the roster in use just changed; do not answer the next activation question
   // from a count read before the failure.
@@ -359,7 +372,7 @@ export function rotateGenericOAuthAccountOn429(
   }
   // With no quota evidence this returns the ring untouched, so providers without
   // per-account quota keep exactly the traversal they have today.
-  return rankAccountsByHeadroom(providerName, candidates)[0] ?? null;
+  return rankAccountsByHeadroom(providerName, candidates, requestedModelId)[0] ?? null;
 }
 
 /**
@@ -392,6 +405,7 @@ export function preferredInitialAccount(
   config: OcxConfig,
   providerName: string,
   now = Date.now(),
+  requestedModelId?: string | null,
 ): string | null {
   // The PROACTIVE predicate, not the reactive one: this steers a request upstream has not
   // refused, so `oauthAccountFailover.enabled: false` must still be able to refuse it.
@@ -411,7 +425,8 @@ export function preferredInitialAccount(
   // would never reach its own test. Cooldowns and reauth are still honoured inside each pick.
   const strategy = activeGenericStrategy(config, providerName);
   if (strategy === "round-robin") {
-    const eligibleNow = eligibleFailoverAccounts(providerName, now);
+    const family = classifyModelFamilyForQuota(providerName, requestedModelId);
+    const eligibleNow = eligibleFailoverAccounts(providerName, now, family);
     if (eligibleNow.length === 0) return null;
     // PEEK, not pick: this proposal is discardable, and advancing the ring for an account the
     // resolver then rejects would skip a turn for nothing. noteGenericPoolSelection commits.
@@ -423,26 +438,26 @@ export function preferredInitialAccount(
     return picked && picked !== active ? picked : null;
   }
   if (strategy === "fill-first") {
-    const picked = pickFillFirstGenericAccount(config, providerName, active, now);
+    const picked = pickFillFirstGenericAccount(config, providerName, active, now, requestedModelId);
     return picked && picked !== active ? picked : null;
   }
 
   const activeRow = selected.accounts.find(account => account.id === active);
   if (activeRow && activeRow.needsReauth !== true
-    && !isCooled(providerName, activeRow.id, now)
-    && !isAccountQuotaExhausted(providerName, activeRow.id)) return null;
+    && !isCooled(providerName, activeRow.id, now, classifyModelFamilyForQuota(providerName, requestedModelId))
+    && !isAccountQuotaExhausted(providerName, activeRow.id, requestedModelId)) return null;
 
   // Evidence is required BEFORE eligibility narrows the field. Without this, a provider
   // with no quota data at all could still be redirected: cool the active account with a
   // 429 and the eligible list collapses to one candidate, which any ranking returns
   // unchanged — an answer that looks ranked but was never measured. The no-op guarantee
   // for quota-less providers has to be checked on the full roster.
-  if (!hasHeadroomEvidence(providerName, order)) return null;
+  if (!hasHeadroomEvidence(providerName, order, requestedModelId)) return null;
 
   // Cooldowns are respected here, unlike in the presence count: this picks the account to
   // send to right now, and one inside its 429 window is the single candidate we hold
   // positive evidence against.
-  const eligible = order.filter(id => !isCooled(providerName, id, now));
+  const eligible = order.filter(id => !isCooled(providerName, id, now, classifyModelFamilyForQuota(providerName, requestedModelId)));
   if (eligible.length === 0) return null;
 
   // Start the ring at the active account so an unranked outcome reproduces today's choice.
@@ -451,7 +466,7 @@ export function preferredInitialAccount(
   const candidates = ring.filter(id => eligible.includes(id));
   if (candidates.length === 0) return null;
 
-  const best = rankAccountsByHeadroom(providerName, candidates)[0] ?? null;
+  const best = rankAccountsByHeadroom(providerName, candidates, requestedModelId)[0] ?? null;
   // Nothing to do when the ranking agrees with the account we would have used anyway.
   //
   // A proposal still needs guarded selection commit after credential resolution: a
@@ -461,12 +476,10 @@ export function preferredInitialAccount(
 
 /** Earliest remaining cooldown, for a client-facing Retry-After when every account is cooled. */
 export function genericFailoverRetryAfterSeconds(providerName: string, now = Date.now()): number | null {
-  const set = getAccountSet(providerName);
-  if (!set) return null;
+  const prefix = `${providerName}\u0000`;
   let earliest: number | null = null;
-  for (const account of set.accounts) {
-    const entry = health.get(healthKey(providerName, account.id));
-    if (!entry || entry.cooldownUntil <= now) continue;
+  for (const [key, entry] of health) {
+    if (!key.startsWith(prefix) || entry.cooldownUntil <= now) continue;
     if (earliest === null || entry.cooldownUntil < earliest) earliest = entry.cooldownUntil;
   }
   return earliest === null ? null : Math.max(1, Math.ceil((earliest - now) / 1000));

src/server/responses/adapter-continuation.ts

@@ -395,6 +395,8 @@ export function createAdapterContinuations(
             route.providerName,
             transportState.genericFailoverAccountId,
             response.headers.get("retry-after"),
+            Date.now(),
+            route.modelId,
           )
           : null;
         if (!nextAccountId) hop.permit?.release();

src/server/responses/adapter-dispatch.ts

@@ -732,6 +732,8 @@ export async function prepareAdapterExchange(
           route.providerName,
           transportState.genericFailoverAccountId,
           upstreamResponse.headers.get("retry-after"),
+          Date.now(),
+          route.modelId,
         );
         if (!nextAccountId) {
           hop.permit?.release();