diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 96e44b509c..6be47b88ad 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -968,6 +968,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", diff --git a/src/oauth/account-quota-rank.ts b/src/oauth/account-quota-rank.ts index 0cae484fa5..42590a8c2c 100644 --- a/src/oauth/account-quota-rank.ts +++ b/src/oauth/account-quota-rank.ts @@ -13,6 +13,38 @@ 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) + || /(?:^|[^a-z])gpt[-_]oss(?:[^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 +80,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 +115,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 +141,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 +182,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 +197,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)); } /** diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 6b444ea435..b4c912c70d 100644 --- a/src/oauth/generic-account-failover.ts +++ b/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(); /** Provider -> recent eligible-account count. TTL-bounded; never holds credential material. */ const presence = new Map(); -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)); diff --git a/src/server/responses/adapter-continuation.ts b/src/server/responses/adapter-continuation.ts index fcf2bce705..bc50107dc6 100644 --- a/src/server/responses/adapter-continuation.ts +++ b/src/server/responses/adapter-continuation.ts @@ -25,7 +25,6 @@ import { fetchWithTransientRetry, fetchWithResetRetry, applyUpstreamRecoveryInit, - isNonReplayableResponse, prepareSameTarget429Wait, } from "../../lib/upstream-retry"; import { redactSecretString } from "../../lib/redact"; @@ -265,9 +264,6 @@ export function createAdapterContinuations( // loop; only after the attempts are exhausted does the continuation fail over. while ( response.status === 429 - // A synthesized replay refusal is not a rate limit; replaying the continuation on - // it would re-send a turn whose first send may already have been processed. - && !isNonReplayableResponse(response) && rateLimitPolicy !== null && adapterExchange.rateLimitRetries < rateLimitPolicy.attempts // The main recovery loop and the passthrough ladder both consult the shared remainder @@ -315,7 +311,7 @@ export function createAdapterContinuations( } } - if (response.status === 429 && !isNonReplayableResponse(response) && hasKeyPoolFailover(route.provider)) { + if (response.status === 429 && hasKeyPoolFailover(route.provider)) { const rotated = rotateProviderTransportOn429(config, route.providerName, route.provider, { retryAfter: response.headers.get("retry-after"), now: Date.now(), @@ -350,7 +346,6 @@ export function createAdapterContinuations( } if ( response.status === 429 - && !isNonReplayableResponse(response) && transportState.anthropicPoolAccountId && transportState.anthropicPoolFailovers < ANTHROPIC_POOL_MAX_FAILOVERS_PER_REQUEST ) { @@ -392,7 +387,6 @@ export function createAdapterContinuations( // the per-request bound cannot be silently re-armed by reaching a different loop. if ( response.status === 429 - && !isNonReplayableResponse(response) && transportState.genericFailoverAccountId && transportState.genericFailovers < GENERIC_OAUTH_MAX_FAILOVERS_PER_REQUEST && isGenericOAuthFailoverEnabled(config, route.providerName) @@ -416,6 +410,8 @@ export function createAdapterContinuations( route.providerName, transportState.genericFailoverAccountId, response.headers.get("retry-after"), + Date.now(), + route.modelId, ) : null; if (!nextAccountId) hop.permit?.release(); diff --git a/src/server/responses/adapter-dispatch.ts b/src/server/responses/adapter-dispatch.ts index aed413118c..42e87efe54 100644 --- a/src/server/responses/adapter-dispatch.ts +++ b/src/server/responses/adapter-dispatch.ts @@ -1,4 +1,3 @@ -import { isNonReplayableResponse } from "../../lib/upstream-retry"; import type { ResponsesRequestContext, ResponsesAdmissionState } from "./core-options"; import type { PreparedResponsesRequest } from "./request-prepare"; import type { ResponsesTransport } from "./request-transport"; @@ -527,12 +526,6 @@ export async function prepareAdapterExchange( }; // Keep recovery kinds in sync with the native Responses `passthroughRecovery:` loop above. recovery: for (;;) { - // Preserve the terminal verdict through adapter and combo error formatting. - // This also covers a reset reached by a 401/429/413 recovery refetch. - if (isNonReplayableResponse(upstreamResponse)) { - cleanupUpstreamAbort(); - return upstreamResponse; - } if ( upstreamResponse.status === 401 && isOAuth401ReplayProvider @@ -621,11 +614,6 @@ export async function prepareAdapterExchange( const result = await rebuildAndRefetch("key-401"); if ("failed" in result) return result.failed; upstreamResponse = result; - // A recovery refetch can itself die on an ambiguous pre-header reset, and the refusal - // that answers it is a 429. Every arm below keys on 429, so letting it fall through - // hands the marked refusal to the next waiting arm and replays the send it exists to - // stop. Re-enter the loop guard instead, which returns it unchanged. - if (isNonReplayableResponse(upstreamResponse)) continue recovery; } // Same-target 429 wait-and-retry (opt-in `retryOn429`, issue #487). Codex never retries @@ -666,9 +654,6 @@ export async function prepareAdapterExchange( const result = await rebuildAndRefetch("rate-limit-429"); if ("failed" in result) return result.failed; upstreamResponse = result; - // The refusal is a 429 too: without this the while condition is still true and the - // next configured attempt replays it on the same target. - if (isNonReplayableResponse(upstreamResponse)) continue recovery; } // Multi-key 429 failover: rotate to the next pool key (cooldown-aware) and retry the @@ -700,9 +685,6 @@ export async function prepareAdapterExchange( const result = await rebuildAndRefetch("key-429"); if ("failed" in result) return result.failed; upstreamResponse = result; - // Rotating on the refusal would also write a cooldown against a key that rate-limited - // nothing, which outlives the request. - if (isNonReplayableResponse(upstreamResponse)) continue recovery; } // Opt-in Anthropic OAuth account pool (#294): cool the failed account and retry @@ -739,7 +721,6 @@ export async function prepareAdapterExchange( const result = await rebuildAndRefetch("anthropic-oauth-429"); if ("failed" in result) return result.failed; upstreamResponse = result; - if (isNonReplayableResponse(upstreamResponse)) continue recovery; } catch { break; } @@ -786,6 +767,8 @@ export async function prepareAdapterExchange( route.providerName, transportState.genericFailoverAccountId, upstreamResponse.headers.get("retry-after"), + Date.now(), + route.modelId, ); if (!nextAccountId) { hop.permit?.release(); @@ -833,9 +816,6 @@ export async function prepareAdapterExchange( return result.failed; } upstreamResponse = result; - // The hop's permit is already settled by the dispatch boundary above; continuing - // only skips the remaining arms, it does not abandon a reservation. - if (isNonReplayableResponse(upstreamResponse)) continue recovery; } catch { // A throw before the send — snapshot fetch, credential application, adapter // resolution — must hand the reservation back. Without this the ladder charges the diff --git a/src/server/responses/passthrough-dispatch.ts b/src/server/responses/passthrough-dispatch.ts index 19a4de0bac..cfff2a045d 100644 --- a/src/server/responses/passthrough-dispatch.ts +++ b/src/server/responses/passthrough-dispatch.ts @@ -102,7 +102,6 @@ import { fetchWithTransientRetry, applyUpstreamRecoveryInit, TRANSIENT_RETRY_MAX_ATTEMPTS, - isNonReplayableResponse, prepareSameTarget429Wait, sleepWithAbort, } from "../../lib/upstream-retry"; @@ -1118,10 +1117,6 @@ export async function preparePassthroughExchange( // the same quorum, cooldown and request budget here, before any client bytes flow. if ( upstreamResponse.status === 429 - // Not a provider rate limit when this proxy synthesized it for a refused reset - // replay; rotating accounts on it would re-send an inference that may already - // have run and would cool down an account that refused nothing. - && !isNonReplayableResponse(upstreamResponse) && transportState.genericFailoverAccountId && transportState.genericFailovers < GENERIC_OAUTH_MAX_FAILOVERS_PER_REQUEST && isGenericOAuthFailoverEnabled(config, route.providerName) @@ -1138,6 +1133,8 @@ export async function preparePassthroughExchange( const nextAccountId = rotateGenericOAuthAccountOn429( config, route.providerName, transportState.genericFailoverAccountId, upstreamResponse.headers.get("retry-after"), + Date.now(), + route.modelId, ); let snapshot: OAuthAccessSnapshot | undefined; if (nextAccountId) { @@ -1176,7 +1173,6 @@ export async function preparePassthroughExchange( // keep their pool logic below (rateLimitRetryPolicyFor returns null for them). while ( upstreamResponse.status === 429 - && !isNonReplayableResponse(upstreamResponse) && rateLimitPolicy !== null && rateLimitRetries < rateLimitPolicy.attempts // Checked here rather than inside the helper: prepareSameTarget429Wait releases the 429 diff --git a/src/server/responses/request-transport.ts b/src/server/responses/request-transport.ts index c5b92a177d..3e878959e2 100644 --- a/src/server/responses/request-transport.ts +++ b/src/server/responses/request-transport.ts @@ -487,7 +487,7 @@ export async function prepareResponsesTransport( // measured as spent. A null answer means "use the active account", so every provider // without quota evidence keeps the resolution it has today. const preferredAccountId = isGenericFailoverProvider(route.providerName, route.provider) - ? preferredInitialAccount(config, route.providerName) + ? preferredInitialAccount(config, route.providerName, Date.now(), route.modelId) : null; // Resolved account-scoped, NOT through failoverAccountSnapshot: that helper marks a // rotation site, and rotation sites must apply their credential through @@ -550,7 +550,7 @@ export async function prepareResponsesTransport( // Advance the pool cursor only now that this account is actually admitted. The // helper returns immediately unless the kernel is on AND the strategy is // round-robin, so quota and fill-first pools reach it without being touched. - noteGenericPoolSelection(config, route.providerName, resolved.accountId); + noteGenericPoolSelection(config, route.providerName, resolved.accountId, route.modelId); } // Anthropic is excluded from isGenericFailoverProvider -- its own pool owns affinity and // a fail-closed local-cli credential rule -- so without this stamp its identity is diff --git a/src/server/responses/run-turn-execution.ts b/src/server/responses/run-turn-execution.ts index f4f2e10228..501c13f702 100644 --- a/src/server/responses/run-turn-execution.ts +++ b/src/server/responses/run-turn-execution.ts @@ -237,6 +237,8 @@ export async function executeResponsesRunTurn( route.providerName, transportState.genericFailoverAccountId, null, + Date.now(), + route.modelId, ); if (!nextAccountId) { hop.permit?.release(); @@ -374,7 +376,6 @@ export async function executeResponsesRunTurn( stallTimeoutSec: config.stallTimeoutSec, hideThinkingSummary: parsed.options.hideThinkingSummary, declaredToolNames, - enforceDeclaredToolNames: inboundWire !== "chat" && inboundWire !== "anthropic", toolParameterSchemas, ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), ...(routedCompaction ? { compaction: true } : {}), @@ -445,7 +446,6 @@ export async function executeResponsesRunTurn( hideThinkingSummary: parsed.options.hideThinkingSummary, toolNsMap, declaredToolNames, - enforceDeclaredToolNames: inboundWire !== "chat" && inboundWire !== "anthropic", toolParameterSchemas, freeformToolNames, toolSearchToolNames, diff --git a/src/server/responses/sidecar-execution.ts b/src/server/responses/sidecar-execution.ts index 7aeb9d452d..de1156b222 100644 --- a/src/server/responses/sidecar-execution.ts +++ b/src/server/responses/sidecar-execution.ts @@ -186,6 +186,8 @@ export async function executeResponsesSidecars( route.providerName, transportState.genericFailoverAccountId, retryAfter, + Date.now(), + route.modelId, ); if (!nextAccountId) { hop.permit?.release(); diff --git a/structure/transports/inventory.md b/structure/transports/inventory.md index 6b55bbaf0d..c99e28f04d 100644 --- a/structure/transports/inventory.md +++ b/structure/transports/inventory.md @@ -152,3 +152,14 @@ Translated audio/file admission follows the [final-adapter input contract](../ad Canonical Responses identity sanitation and narrowly scoped pre-output combo recovery follow [request-local target compatibility](../runtime.md#request-local-target-compatibility); other adapter contracts remain unchanged. Shared response-log retention and native SSE inspection pacing follow the [bounded inspection contract](byte-accounting.md#response-log-inspection); other subsystem behavior remains unchanged. + +## Model-family-aware OAuth headroom + +`src/oauth/account-quota-rank.ts` ranks Antigravity custom windows for the requested +Gemini or Claude family, including GPT-OSS in the Claude family. An unknown model +retains all-window ranking; absent matching evidence retains the existing unranked behavior. +`src/server/responses/request-transport.ts` passes the routed model at initial selection. +The passthrough, adapter, continuation, sidecar and run-turn execution owners pass +the same routed model during account rotation, without bypassing their send-budget +admission or account-snapshot pairing. The forwarding contract is covered in +`tests/oauth/oauth-account-quota-rank.test.ts`; the core facade remains orchestration-only. diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index c1f826b816..d371be27a6 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -794,6 +794,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", diff --git a/tests/oauth/oauth-account-quota-rank.test.ts b/tests/oauth/oauth-account-quota-rank.test.ts new file mode 100644 index 0000000000..cd6a4b3f24 --- /dev/null +++ b/tests/oauth/oauth-account-quota-rank.test.ts @@ -0,0 +1,320 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, readFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + classifyModelFamilyForQuota, + hasHeadroomEvidence, + isAccountQuotaExhausted, + rankAccountsByHeadroom, +} from "../../src/oauth/account-quota-rank"; +import { + clearGenericFailoverHealth, + eligibleFailoverAccounts, + genericFailoverRetryAfterSeconds, + noteGenericPoolSelection, + preferredInitialAccount, + rotateGenericOAuthAccountOn429, +} from "../../src/oauth/generic-account-failover"; +import { + clearAccountQuotaCache, + setCachedProviderAccountQuotaForTests, +} from "../../src/providers/quota"; +import { + getAccountSet, + saveCredential, + setActiveAccount, +} from "../../src/oauth/store"; +import type { OcxConfig, OcxProviderConfig } from "../../src/types"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; +import { repoPath } from "../helpers/repo-root"; + +const originalHome = process.env.OPENCODEX_HOME; +let home: string; + +beforeEach(() => { + home = mkdtempSync(join(tmpdir(), "ocx-quota-rank-")); + process.env.OPENCODEX_HOME = home; + clearGenericFailoverHealth(); +}); + +afterEach(() => { + clearGenericFailoverHealth(); + clearAccountQuotaCache("google-antigravity"); + clearAccountQuotaCache("xai"); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + removeTreeWithRetry(home); +}); + +const PROVIDER = { + adapter: "google", + authMode: "oauth", +} as unknown as OcxProviderConfig; + +function config(strategy?: "quota" | "round-robin" | "fill-first"): OcxConfig { + return { + providers: { + "google-antigravity": { + ...PROVIDER, + oauthAccountFailover: { + enabled: true, + ...(strategy ? { strategy } : {}), + autoSwitchThreshold: 80, + }, + }, + }, + oauthAccountFailover: { enabled: true }, + ...(strategy ? { pool: { kernel: true } } : {}), + } as unknown as OcxConfig; +} + +function seedWindows(accountId: string, gem: number, cla: number): void { + setCachedProviderAccountQuotaForTests("google-antigravity", accountId, { + updatedAt: Date.now(), + customWindows: [ + { label: "Gem", percent: gem }, + { label: "Gem (Weekly)", percent: gem }, + { label: "Cla", percent: cla }, + { label: "Cla (Weekly)", percent: cla }, + ], + }); +} + +async function seedAntigravityAccounts(count = 2): Promise { + for (let i = 1; i <= count; i++) { + await saveCredential("google-antigravity", { + access: `tok-${i}`, + refresh: `ref-${i}`, + expires: Date.now() + 3600_000, + accountId: `acct-${i}`, + }); + } + const set = getAccountSet("google-antigravity")!; + return set.accounts.map(a => a.id); +} + +describe("classifyModelFamilyForQuota", () => { + test("maps Gemini, Claude, and GPT-OSS models, and ignores Gemma", () => { + expect(classifyModelFamilyForQuota("google-antigravity", "gemini-3.8-flash")).toBe("gem"); + expect(classifyModelFamilyForQuota("google-antigravity", "gemini-pro-agent")).toBe("gem"); + expect(classifyModelFamilyForQuota("google-antigravity", "gemini-3.1-flash-image")).toBe("gem"); + expect(classifyModelFamilyForQuota("google-antigravity", "claude-sonnet-4-5")).toBe("cla"); + expect(classifyModelFamilyForQuota("google-antigravity", "claude-opus-4-6-thinking")).toBe("cla"); + expect(classifyModelFamilyForQuota("google-antigravity", "claude-3-7-sonnet")).toBe("cla"); + expect(classifyModelFamilyForQuota("google-antigravity", "gpt-oss-120b")).toBe("cla"); + expect(classifyModelFamilyForQuota("google-antigravity", "gpt_oss_20b")).toBe("cla"); + expect(classifyModelFamilyForQuota("google-antigravity", "gemma-3-27b")).toBeUndefined(); + expect(classifyModelFamilyForQuota("google-antigravity", "gemma-2-9b-it")).toBeUndefined(); + expect(classifyModelFamilyForQuota("google-antigravity", "gem-experimental")).toBeUndefined(); + expect(classifyModelFamilyForQuota("xai", "gemini-3.8-flash")).toBeUndefined(); + expect(classifyModelFamilyForQuota("google-antigravity", undefined)).toBeUndefined(); + expect(classifyModelFamilyForQuota("google-antigravity", null)).toBeUndefined(); + }); +}); + +describe("model-family headroom filtering", () => { + test("Gemini request ignores spent Claude window on Antigravity", () => { + seedWindows("a", 20, 100); + seedWindows("b", 80, 5); + const ranked = rankAccountsByHeadroom( + "google-antigravity", + ["b", "a"], + "gemini-3.8-flash", + ); + expect(ranked[0]).toBe("a"); + }); + + test("Claude request ignores healthy Gemini window on Antigravity", () => { + seedWindows("a", 20, 100); + seedWindows("b", 80, 5); + const ranked = rankAccountsByHeadroom( + "google-antigravity", + ["a", "b"], + "claude-sonnet-4-5", + ); + expect(ranked[0]).toBe("b"); + }); + + test("GPT-OSS request uses Claude 3P window on Antigravity", () => { + seedWindows("a", 10, 95); + seedWindows("b", 90, 15); + const ranked = rankAccountsByHeadroom( + "google-antigravity", + ["a", "b"], + "gpt-oss-120b", + ); + expect(ranked[0]).toBe("b"); + }); + + test("non-Antigravity provider ignores modelId and uses all windows", () => { + setCachedProviderAccountQuotaForTests("xai", "a", { + fiveHourPercent: 80, + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("xai", "b", { + fiveHourPercent: 30, + updatedAt: Date.now(), + }); + const ranked = rankAccountsByHeadroom("xai", ["a", "b"], "gemini-3.8-flash"); + expect(ranked[0]).toBe("b"); + }); + + test("without modelId, Antigravity uses all windows (backward compatibility)", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "a", { + customWindows: [ + { label: "Gem", percent: 20 }, + { label: "Cla", percent: 99 }, + ], + updatedAt: Date.now(), + }); + const ranked = rankAccountsByHeadroom("google-antigravity", ["a"]); + expect(ranked).toEqual(["a"]); + }); + + test("falls back to unranked ring when family labels are missing or drifted", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "a", { + updatedAt: Date.now(), + customWindows: [{ label: "UnknownWindowA", percent: 1 }], + }); + setCachedProviderAccountQuotaForTests("google-antigravity", "b", { + updatedAt: Date.now(), + customWindows: [{ label: "UnknownWindowB", percent: 99 }], + }); + expect(hasHeadroomEvidence("google-antigravity", ["a", "b"], "gemini-3.8-flash")).toBe(false); + expect(rankAccountsByHeadroom("google-antigravity", ["b", "a"], "gemini-3.8-flash")).toEqual(["b", "a"]); + }); + + test("does not misclassify gemma as Gemini (no Gem prefix pollution)", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "a", { + customWindows: [ + { label: "Gem", percent: 10 }, + { label: "Cla", percent: 90 }, + ], + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("google-antigravity", "b", { + customWindows: [ + { label: "Gem", percent: 90 }, + { label: "Cla", percent: 10 }, + ], + updatedAt: Date.now(), + }); + const ranked = rankAccountsByHeadroom("google-antigravity", ["a", "b"], "gemma-3-27b"); + expect(ranked).toEqual(["a", "b"]); + }); +}); + +describe("model-family-aware exhaustion check", () => { + test("global exhaustion check without requestedModelId evaluates all windows", () => { + seedWindows("a", 20, 100); + expect(isAccountQuotaExhausted("google-antigravity", "a")).toBe(true); + }); + + test("model-filtered exhaustion check respects requested model family", () => { + seedWindows("a", 20, 100); + expect(isAccountQuotaExhausted("google-antigravity", "a", "gemini-3.8-flash")).toBe(false); + expect(isAccountQuotaExhausted("google-antigravity", "a", "claude-sonnet-4-5")).toBe(true); + }); + + test("hasHeadroomEvidence respects model family filter", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "a", { + customWindows: [{ label: "Cla", percent: 50 }], + updatedAt: Date.now(), + }); + expect(hasHeadroomEvidence("google-antigravity", ["a"], "claude-sonnet-4-5")).toBe(true); + expect(hasHeadroomEvidence("google-antigravity", ["a"], "gemini-3.8-flash")).toBe(false); + }); +}); + +describe("Antigravity family-scoped cooldown and routing", () => { + test("a Claude 429 still keeps the account for Gemini", async () => { + const [id1, id2] = await seedAntigravityAccounts(2); + await setActiveAccount("google-antigravity", id1); + seedWindows(id1, 10, 100); + seedWindows(id2, 80, 5); + + const cfg = config(); + // Rotating on Claude 429 chooses id2 + expect(rotateGenericOAuthAccountOn429(cfg, "google-antigravity", id1, null, Date.now(), "claude-sonnet-4-5")).toBe(id2); + // id1 is in Claude cooldown, but still eligible for Gemini + expect(eligibleFailoverAccounts("google-antigravity", Date.now(), "gem")).toContain(id1); + expect(eligibleFailoverAccounts("google-antigravity", Date.now(), "cla")).not.toContain(id1); + expect(genericFailoverRetryAfterSeconds("google-antigravity")).toBeGreaterThan(0); + + // preferredInitialAccount for Gemini stays on id1 + expect(preferredInitialAccount(cfg, "google-antigravity", Date.now(), "gemini-3.8-flash")).toBeNull(); + }); + + test("preferredInitialAccount switches away when active account is spent for requested family", async () => { + const [id1, id2] = await seedAntigravityAccounts(2); + await setActiveAccount("google-antigravity", id1); + seedWindows(id1, 20, 100); + seedWindows(id2, 50, 20); + + const cfg = config(); + // Gemini: active account has 80% Gem headroom -> keep active (null) + expect(preferredInitialAccount(cfg, "google-antigravity", Date.now(), "gemini-3.8-flash")).toBeNull(); + // Claude: active account is spent for Claude -> switch to id2 + expect(preferredInitialAccount(cfg, "google-antigravity", Date.now(), "claude-sonnet-4-5")).toBe(id2); + }); +}); + +describe("Antigravity family strategies behind pool.kernel", () => { + test("fill-first stays on Gemini headroom when only Claude is over threshold", async () => { + const [id1, id2] = await seedAntigravityAccounts(2); + await setActiveAccount("google-antigravity", id1); + seedWindows(id1, 40, 90); + seedWindows(id2, 10, 10); + + const cfg = config("fill-first"); + // Gemini usage is 40% (under 80% threshold) -> stays active + expect(preferredInitialAccount(cfg, "google-antigravity", Date.now(), "gemini-3.8-flash")).toBeNull(); + // Claude usage is 90% (over 80% threshold) -> advances to id2 + expect(preferredInitialAccount(cfg, "google-antigravity", Date.now(), "claude-sonnet-4-5")).toBe(id2); + }); + + test("a Claude 429 does not hide the account from Gemini round-robin", async () => { + const [id1, id2] = await seedAntigravityAccounts(2); + await setActiveAccount("google-antigravity", id1); + seedWindows(id1, 20, 20); + seedWindows(id2, 20, 20); + + const cfg = config("round-robin"); + expect(rotateGenericOAuthAccountOn429(cfg, "google-antigravity", id1, null, Date.now(), "claude-sonnet-4-5")).toBe(id2); + expect(eligibleFailoverAccounts("google-antigravity", Date.now(), "gem")).toContain(id1); + expect(eligibleFailoverAccounts("google-antigravity", Date.now(), "cla")).not.toContain(id1); + }); + + test("noteGenericPoolSelection accepts requestedModelId and advances cursor", async () => { + const [id1, id2] = await seedAntigravityAccounts(2); + const cfg = config("round-robin"); + expect(() => noteGenericPoolSelection(cfg, "google-antigravity", id1, "gemini-3.8-flash")).not.toThrow(); + }); +}); + +describe("modular Responses model-family forwarding", () => { + test("initial selection receives the routed model in the transport owner", () => { + const source = readFileSync(repoPath("src/server/responses/request-transport.ts"), "utf8"); + expect(source).toContain("preferredInitialAccount(config, route.providerName, Date.now(), route.modelId)"); + }); + + for (const [owner, retryAfter] of [ + ["passthrough-dispatch.ts", 'upstreamResponse.headers.get("retry-after")'], + ["adapter-dispatch.ts", 'upstreamResponse.headers.get("retry-after")'], + ["adapter-continuation.ts", 'response.headers.get("retry-after")'], + ["sidecar-execution.ts", "retryAfter"], + ["run-turn-execution.ts", "null"], + ] as const) { + test(owner + " forwards the routed model on account rotation", () => { + const source = readFileSync(repoPath("src/server/responses", owner), "utf8"); + expect(source.match(/\brotateGenericOAuthAccountOn429\s*\(/g)).toHaveLength(1); + expect(source.replace(/\s+/g, " ")).toContain( + "rotateGenericOAuthAccountOn429( config, route.providerName, " + + "transportState.genericFailoverAccountId, " + retryAfter + + ", Date.now(), route.modelId, )", + ); + }); + } +}); +