From 91573181fa76febdf0e684205d13e90a5fb6f580 Mon Sep 17 00:00:00 2001 From: chilung Date: Fri, 11 Sep 2026 18:28:26 +0000 Subject: [PATCH 1/4] feat(oauth): model-family aware quota headroom ranking and failover - Filter quota custom windows by requested model family (Gemini vs Claude) for multi-window providers like google-antigravity. - Pass route.modelId through generic OAuth pre-dispatch and 429 rotation hooks in src/server/responses/core.ts. - Prevent cross-family quota starvation where an exhausted Claude window would wrongfully deprioritize Gemini requests. - Safely degenerate to unranked ring when customWindows labels drift or have no matching family prefix. --- scripts/test-layout/layout.json | 1 + src/oauth/account-quota-rank.ts | 45 ++- src/oauth/generic-account-failover.ts | 10 +- src/server/responses/adapter-continuation.ts | 2 + src/server/responses/adapter-dispatch.ts | 2 + src/server/responses/passthrough-dispatch.ts | 2 + src/server/responses/request-transport.ts | 2 +- src/server/responses/run-turn-execution.ts | 2 + src/server/responses/sidecar-execution.ts | 2 + tests/fixtures/test-layout-expected.json | 1 + tests/oauth/oauth-account-quota-rank.test.ts | 358 +++++++++++++++++++ 11 files changed, 411 insertions(+), 16 deletions(-) create mode 100644 tests/oauth/oauth-account-quota-rank.test.ts diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 885d2abcd4..cd21f3c691 100644 --- a/scripts/test-layout/layout.json +++ b/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", diff --git a/src/oauth/account-quota-rank.ts b/src/oauth/account-quota-rank.ts index 0cae484fa5..b5c8b947cd 100644 --- a/src/oauth/account-quota-rank.ts +++ b/src/oauth/account-quota-rank.ts @@ -45,23 +45,42 @@ interface Ranked { */ const PASSIVE_HEADROOM_MAX_AGE_MS = 60 * 60_000; +/** + * Map a requested model ID to the Antigravity quota-window family prefix. + * Returns "Gem" for Gemini models, "Cla" for Claude/Opus/Sonnet, or undefined + * for unknown models (which falls back to all-window ranking). + */ +function classifyModelFamilyForQuota(modelId: string): string | undefined { + const lower = modelId.toLowerCase(); + if (lower.includes("gemini") || lower.startsWith("gem")) return "Gem"; + if (lower.includes("claude") || lower.includes("opus") || lower.includes("sonnet") || lower.includes("gpt-oss") || lower.includes("gpt_oss")) return "Cla"; + return undefined; +} + /** * 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 { +function headroomOf(provider: string, accountId: string, requestedModelId?: string): 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 familyPrefix = provider === "google-antigravity" && requestedModelId + ? classifyModelFamilyForQuota(requestedModelId) + : undefined; + const percents = [ quota.fiveHourPercent, quota.weeklyPercent, quota.monthlyPercent, - ...(quota.customWindows ?? []).map(window => window.percent), + ...(quota.customWindows ?? []) + .filter(window => !familyPrefix || window.label.startsWith(familyPrefix)) + .map(window => window.percent), ].filter((value): value is number => typeof value === "number"); if (percents.length === 0) return null; return 100 - Math.max(...percents); @@ -79,10 +98,10 @@ export function accountHeadroomPercent(provider: string, accountId: string): num } /** 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): 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 +111,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, +): 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 +152,7 @@ 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): 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 +163,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..3d8da93e89 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -302,6 +302,7 @@ export function rotateGenericOAuthAccountOn429( failedAccountId: string, retryAfterHeader: string | null | undefined, now = Date.now(), + requestedModelId?: string, ): string | null { if (!isGenericOAuthFailoverEnabled(config, providerName)) return null; const set = getAccountSet(providerName); @@ -359,7 +360,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 +393,7 @@ export function preferredInitialAccount( config: OcxConfig, providerName: string, now = Date.now(), + requestedModelId?: string, ): 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. @@ -430,14 +432,14 @@ export function preferredInitialAccount( 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; + && !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 @@ -451,7 +453,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 diff --git a/src/server/responses/adapter-continuation.ts b/src/server/responses/adapter-continuation.ts index a1db9398d5..ba1533ae75 100644 --- a/src/server/responses/adapter-continuation.ts +++ b/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(); diff --git a/src/server/responses/adapter-dispatch.ts b/src/server/responses/adapter-dispatch.ts index 57eadf6a28..28a7c9f8c6 100644 --- a/src/server/responses/adapter-dispatch.ts +++ b/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(); diff --git a/src/server/responses/passthrough-dispatch.ts b/src/server/responses/passthrough-dispatch.ts index f5c5b94694..1d5557d4d1 100644 --- a/src/server/responses/passthrough-dispatch.ts +++ b/src/server/responses/passthrough-dispatch.ts @@ -1121,6 +1121,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) { diff --git a/src/server/responses/request-transport.ts b/src/server/responses/request-transport.ts index c80f242121..ba913bb2c7 100644 --- a/src/server/responses/request-transport.ts +++ b/src/server/responses/request-transport.ts @@ -431,7 +431,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 diff --git a/src/server/responses/run-turn-execution.ts b/src/server/responses/run-turn-execution.ts index 24e96802fb..9e2de70158 100644 --- a/src/server/responses/run-turn-execution.ts +++ b/src/server/responses/run-turn-execution.ts @@ -214,6 +214,8 @@ export async function executeResponsesRunTurn( route.providerName, transportState.genericFailoverAccountId, null, + Date.now(), + route.modelId, ); if (!nextAccountId) { hop.permit?.release(); diff --git a/src/server/responses/sidecar-execution.ts b/src/server/responses/sidecar-execution.ts index 7987d5ca93..9075d85a4f 100644 --- a/src/server/responses/sidecar-execution.ts +++ b/src/server/responses/sidecar-execution.ts @@ -184,6 +184,8 @@ export async function executeResponsesSidecars( route.providerName, transportState.genericFailoverAccountId, retryAfter, + Date.now(), + route.modelId, ); if (!nextAccountId) { hop.permit?.release(); diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index d2eb5d244b..e647184832 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -780,6 +780,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..4e32b78a85 --- /dev/null +++ b/tests/oauth/oauth-account-quota-rank.test.ts @@ -0,0 +1,358 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + rankAccountsByHeadroom, + isAccountQuotaExhausted, + hasHeadroomEvidence, +} from "../../src/oauth/account-quota-rank"; +import { + clearGenericFailoverHealth, + 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"; + +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); +}); + +describe("model-family headroom filtering", () => { + test("Gemini request ignores spent Claude window on Antigravity", () => { + // acct-1 has Claude at 99% (spent) but Gemini at 20% (healthy) + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-1", { + customWindows: [ + { label: "Gem", percent: 20 }, + { label: "Gem (Weekly)", percent: 10 }, + { label: "Cla", percent: 99 }, + { label: "Cla (Weekly)", percent: 85 }, + ], + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-2", { + customWindows: [ + { label: "Gem", percent: 50 }, + { label: "Gem (Weekly)", percent: 30 }, + { label: "Cla", percent: 10 }, + { label: "Cla (Weekly)", percent: 5 }, + ], + updatedAt: Date.now(), + }); + // For a Gemini model, acct-1 (80% Gem headroom) ranks above acct-2 (50% Gem headroom) + const ranked = rankAccountsByHeadroom( + "google-antigravity", + ["acct-2", "acct-1"], + "gemini-3.8-flash", + ); + expect(ranked[0]).toBe("acct-1"); + }); + + test("Claude request ignores healthy Gemini window on Antigravity", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-1", { + customWindows: [ + { label: "Gem", percent: 5 }, + { label: "Cla", percent: 95 }, + ], + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-2", { + customWindows: [ + { label: "Gem", percent: 90 }, + { label: "Cla", percent: 30 }, + ], + updatedAt: Date.now(), + }); + // For a Claude model, acct-2 (70% Cla headroom) ranks above acct-1 (5% Cla headroom) + const ranked = rankAccountsByHeadroom( + "google-antigravity", + ["acct-1", "acct-2"], + "claude-opus-4-6-thinking", + ); + expect(ranked[0]).toBe("acct-2"); + }); + + test("non-Antigravity provider ignores modelId and uses all windows", () => { + setCachedProviderAccountQuotaForTests("xai", "acct-1", { + fiveHourPercent: 80, + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("xai", "acct-2", { + fiveHourPercent: 30, + updatedAt: Date.now(), + }); + const ranked = rankAccountsByHeadroom("xai", ["acct-1", "acct-2"], "grok-3"); + expect(ranked[0]).toBe("acct-2"); + }); + + test("without modelId, Antigravity uses all windows (backward compatibility)", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-1", { + customWindows: [ + { label: "Gem", percent: 20 }, + { label: "Cla", percent: 99 }, + ], + updatedAt: Date.now(), + }); + const ranked = rankAccountsByHeadroom("google-antigravity", ["acct-1"]); + expect(ranked).toEqual(["acct-1"]); + }); +}); + +describe("model-family-aware exhaustion check", () => { + test("global exhaustion check without requestedModelId", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-1", { + customWindows: [ + { label: "Gem", percent: 20 }, + { label: "Cla", percent: 100 }, + ], + updatedAt: Date.now(), + }); + // Global check without requestedModelId evaluates all windows (Cla at 100% -> exhausted) + expect(isAccountQuotaExhausted("google-antigravity", "acct-1")).toBe(true); + }); + + test("model-filtered exhaustion check respects requested model family", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-1", { + customWindows: [ + { label: "Gem", percent: 20 }, + { label: "Cla", percent: 100 }, + ], + updatedAt: Date.now(), + }); + + // For Gemini, acct-1 has 80% Gem headroom -> not exhausted + expect(isAccountQuotaExhausted("google-antigravity", "acct-1", "gemini-3.8-flash")).toBe(false); + + // For Claude, acct-1 has 0% Cla headroom -> exhausted + expect(isAccountQuotaExhausted("google-antigravity", "acct-1", "claude-3-7-sonnet")).toBe(true); + }); + + test("ranking does not treat account as exhausted when unrelated family window is spent", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-1", { + customWindows: [ + { label: "Gem", percent: 20 }, + { label: "Cla", percent: 100 }, + ], + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-2", { + customWindows: [ + { label: "Gem", percent: 60 }, + { label: "Cla", percent: 40 }, + ], + updatedAt: Date.now(), + }); + + // For Gemini: acct-1 has 80% Gem headroom, acct-2 has 40% Gem headroom. + // acct-1 must not be marked exhausted by the 100% Cla window. + const rankedGemini = rankAccountsByHeadroom( + "google-antigravity", + ["acct-2", "acct-1"], + "gemini-3.8-flash", + ); + expect(rankedGemini[0]).toBe("acct-1"); + + // For Claude: acct-1 is exhausted (Cla 100%), acct-2 has 60% Cla headroom. + const rankedClaude = rankAccountsByHeadroom( + "google-antigravity", + ["acct-1", "acct-2"], + "claude-opus-4-6-thinking", + ); + expect(rankedClaude[0]).toBe("acct-2"); + }); + + test("hasHeadroomEvidence respects model family filter", () => { + setCachedProviderAccountQuotaForTests("google-antigravity", "acct-1", { + customWindows: [{ label: "Cla", percent: 50 }], + updatedAt: Date.now(), + }); + + // Only Claude window exists: has evidence for Claude, but not for Gemini + expect(hasHeadroomEvidence("google-antigravity", ["acct-1"], "claude-3-7-sonnet")).toBe(true); + expect(hasHeadroomEvidence("google-antigravity", ["acct-1"], "gemini-3.8-flash")).toBe(false); + }); +}); + +describe("model-aware failover and pre-dispatch integration", () => { + const dummyConfig = { + providers: { + "google-antigravity": { + authMode: "oauth", + destination: "https://example.com", + oauthAccountFailover: { + enabled: true, + preferHealthyAccounts: true, + }, + }, + }, + } as unknown as OcxConfig; + + test("preferredInitialAccount keeps active account if it has headroom for requested model family", async () => { + await saveCredential("google-antigravity", { + access: "tok-1", + refresh: "ref-1", + expires: Date.now() + 3600_000, + accountId: "acct-1", + }); + await saveCredential("google-antigravity", { + access: "tok-2", + refresh: "ref-2", + expires: Date.now() + 3600_000, + accountId: "acct-2", + }); + + const set = getAccountSet("google-antigravity")!; + const id1 = set.accounts.find(a => a.credential.accountId === "acct-1")!.id; + const id2 = set.accounts.find(a => a.credential.accountId === "acct-2")!.id; + await setActiveAccount("google-antigravity", id1); + + // id1 is active. Cla is spent (100%), but Gem is healthy (20%). + setCachedProviderAccountQuotaForTests("google-antigravity", id1, { + customWindows: [ + { label: "Gem", percent: 20 }, + { label: "Cla", percent: 100 }, + ], + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("google-antigravity", id2, { + customWindows: [ + { label: "Gem", percent: 50 }, + { label: "Cla", percent: 20 }, + ], + updatedAt: Date.now(), + }); + + // For a Gemini request: active account (id1) is healthy for Gemini, so returns null + const preferredGemini = preferredInitialAccount( + dummyConfig, + "google-antigravity", + Date.now(), + "gemini-3.8-flash", + ); + expect(preferredGemini).toBeNull(); + + // For a Claude request: active account (id1) is spent for Claude, so switches to id2 + const preferredClaude = preferredInitialAccount( + dummyConfig, + "google-antigravity", + Date.now(), + "claude-opus-4-6-thinking", + ); + expect(preferredClaude).toBe(id2); + }); + + test("rotateGenericOAuthAccountOn429 uses model-filtered ranking", async () => { + await saveCredential("google-antigravity", { + access: "tok-fail", + refresh: "ref-fail", + expires: Date.now() + 3600_000, + accountId: "acct-fail", + }); + await saveCredential("google-antigravity", { + access: "tok-gem", + refresh: "ref-gem", + expires: Date.now() + 3600_000, + accountId: "acct-gem-healthy", + }); + await saveCredential("google-antigravity", { + access: "tok-cla", + refresh: "ref-cla", + expires: Date.now() + 3600_000, + accountId: "acct-cla-healthy", + }); + + const set = getAccountSet("google-antigravity")!; + const idFail = set.accounts.find(a => a.credential.accountId === "acct-fail")!.id; + const idGem = set.accounts.find(a => a.credential.accountId === "acct-gem-healthy")!.id; + const idCla = set.accounts.find(a => a.credential.accountId === "acct-cla-healthy")!.id; + + // idGem has Gem 10%, Cla 90% + setCachedProviderAccountQuotaForTests("google-antigravity", idGem, { + customWindows: [ + { label: "Gem", percent: 10 }, + { label: "Cla", percent: 90 }, + ], + updatedAt: Date.now(), + }); + // idCla has Gem 90%, Cla 10% + setCachedProviderAccountQuotaForTests("google-antigravity", idCla, { + customWindows: [ + { label: "Gem", percent: 90 }, + { label: "Cla", percent: 10 }, + ], + updatedAt: Date.now(), + }); + + // Rotate on 429 for Gemini request -> should pick idGem (90% Gem headroom) + const nextGemini = rotateGenericOAuthAccountOn429( + dummyConfig, + "google-antigravity", + idFail, + null, + Date.now(), + "gemini-3.8-flash", + ); + expect(nextGemini).toBe(idGem); + + // Rotate on 429 for Claude request -> should pick idCla (90% Cla headroom) + const nextClaude = rotateGenericOAuthAccountOn429( + dummyConfig, + "google-antigravity", + idFail, + null, + Date.now(), + "claude-opus-4-6-thinking", + ); + expect(nextClaude).toBe(idCla); + }); + + test("degenerates safely to unranked ring when customWindows labels drift or have no matching prefix", () => { + const idDrift1 = "antigravity-drift-1"; + const idDrift2 = "antigravity-drift-2"; + const ring = [idDrift1, idDrift2]; + + // Both accounts have labels that drift from standard "Gem" / "Cla" prefix (e.g. upstream renamed window labels) + setCachedProviderAccountQuotaForTests("google-antigravity", idDrift1, { + customWindows: [ + { label: "UnknownModelWindowA", percent: 95 }, + ], + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("google-antigravity", idDrift2, { + customWindows: [ + { label: "UnknownModelWindowB", percent: 50 }, + ], + updatedAt: Date.now(), + }); + + // When requestedModelId is specified (gemini), but no windows match "Gem" prefix: + // headroomOf returns null for both accounts, ranking preserves unranked ring fallback + const ranked = rankAccountsByHeadroom("google-antigravity", ring, "gemini-3.8-flash"); + expect(ranked).toEqual(ring); + }); +}); + From f08f531d6931432d0842d6f5498a06f0be4df966 Mon Sep 17 00:00:00 2001 From: chilung Date: Fri, 11 Sep 2026 18:40:49 +0000 Subject: [PATCH 2/4] feat(oauth): harden model family classification against gemma prefix collision - Tighten classifyModelFamilyForQuota to require explicit gemini keyword, preventing gemma from matching Gem family. - Add regression test in tests/oauth/oauth-account-quota-rank.test.ts. --- src/oauth/account-quota-rank.ts | 4 ++-- tests/oauth/oauth-account-quota-rank.test.ts | 25 ++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/src/oauth/account-quota-rank.ts b/src/oauth/account-quota-rank.ts index b5c8b947cd..409b7297f4 100644 --- a/src/oauth/account-quota-rank.ts +++ b/src/oauth/account-quota-rank.ts @@ -50,9 +50,9 @@ const PASSIVE_HEADROOM_MAX_AGE_MS = 60 * 60_000; * Returns "Gem" for Gemini models, "Cla" for Claude/Opus/Sonnet, or undefined * for unknown models (which falls back to all-window ranking). */ -function classifyModelFamilyForQuota(modelId: string): string | undefined { +function classifyModelFamilyForQuota(modelId: string): "Gem" | "Cla" | undefined { const lower = modelId.toLowerCase(); - if (lower.includes("gemini") || lower.startsWith("gem")) return "Gem"; + if (lower.includes("gemini") || lower === "gemini" || lower.startsWith("gemini-")) return "Gem"; if (lower.includes("claude") || lower.includes("opus") || lower.includes("sonnet") || lower.includes("gpt-oss") || lower.includes("gpt_oss")) return "Cla"; return undefined; } diff --git a/tests/oauth/oauth-account-quota-rank.test.ts b/tests/oauth/oauth-account-quota-rank.test.ts index 4e32b78a85..443ead15cc 100644 --- a/tests/oauth/oauth-account-quota-rank.test.ts +++ b/tests/oauth/oauth-account-quota-rank.test.ts @@ -354,5 +354,30 @@ describe("model-aware failover and pre-dispatch integration", () => { const ranked = rankAccountsByHeadroom("google-antigravity", ring, "gemini-3.8-flash"); expect(ranked).toEqual(ring); }); + + test("does not misclassify gemma as Gemini (no Gem prefix pollution)", () => { + const idGem = "antigravity-gem"; + const idCla = "antigravity-cla"; + const ring = [idGem, idCla]; + + setCachedProviderAccountQuotaForTests("google-antigravity", idGem, { + customWindows: [ + { label: "Gem", percent: 10 }, + { label: "Cla", percent: 90 }, + ], + updatedAt: Date.now(), + }); + setCachedProviderAccountQuotaForTests("google-antigravity", idCla, { + customWindows: [ + { label: "Gem", percent: 90 }, + { label: "Cla", percent: 10 }, + ], + updatedAt: Date.now(), + }); + + // When requestedModelId is gemma, it should return undefined family and compare all windows (Math.max of 10 and 90 = 90 for both) + const ranked = rankAccountsByHeadroom("google-antigravity", ring, "gemma-2-9b-it"); + expect(ranked).toEqual(ring); + }); }); From 40b9cacd0996305a43018344c44124337ca805a0 Mon Sep 17 00:00:00 2001 From: chilung Date: Sun, 13 Sep 2026 18:47:57 +0000 Subject: [PATCH 3/4] refactor(oauth): remove redundant gemini substring checks in model family classifier --- src/oauth/account-quota-rank.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/oauth/account-quota-rank.ts b/src/oauth/account-quota-rank.ts index 409b7297f4..d786f5a57d 100644 --- a/src/oauth/account-quota-rank.ts +++ b/src/oauth/account-quota-rank.ts @@ -52,7 +52,7 @@ const PASSIVE_HEADROOM_MAX_AGE_MS = 60 * 60_000; */ function classifyModelFamilyForQuota(modelId: string): "Gem" | "Cla" | undefined { const lower = modelId.toLowerCase(); - if (lower.includes("gemini") || lower === "gemini" || lower.startsWith("gemini-")) return "Gem"; + if (lower.includes("gemini")) return "Gem"; if (lower.includes("claude") || lower.includes("opus") || lower.includes("sonnet") || lower.includes("gpt-oss") || lower.includes("gpt_oss")) return "Cla"; return undefined; } From 2a1b32a915d2548695b4dcc53404d279cd70b919 Mon Sep 17 00:00:00 2001 From: JUN Date: Tue, 15 Sep 2026 21:37:39 +0900 Subject: [PATCH 4/4] test(oauth): guard model-family forwarding in split Responses owners Ported the six pre-modularization call sites to their current owners while retaining send-budget and snapshot pairing. Static inspection only; runtime checks were not run on the connected host. Co-authored-by: chilung --- structure/transports/inventory.md | 11 ++++++++ tests/oauth/oauth-account-quota-rank.test.ts | 29 ++++++++++++++++++-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/structure/transports/inventory.md b/structure/transports/inventory.md index ca80373a8e..bf17edb755 100644 --- a/structure/transports/inventory.md +++ b/structure/transports/inventory.md @@ -148,3 +148,14 @@ Renamed fixed-key providers receive [missing reasoning metadata](../catalog.md#r Translated audio/file admission follows the [final-adapter input contract](../adapters/registry.md#untranslated-input-media); native raw passthrough remains separate. 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. + +## 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/oauth/oauth-account-quota-rank.test.ts b/tests/oauth/oauth-account-quota-rank.test.ts index 443ead15cc..c504f0afbe 100644 --- a/tests/oauth/oauth-account-quota-rank.test.ts +++ b/tests/oauth/oauth-account-quota-rank.test.ts @@ -1,5 +1,5 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; -import { mkdtempSync } from "node:fs"; +import { mkdtempSync, readFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { @@ -21,8 +21,9 @@ import { saveCredential, setActiveAccount, } from "../../src/oauth/store"; -import type { OcxConfig, OcxProviderConfig } from "../../src/types"; +import type { OcxConfig } from "../../src/types"; import { removeTreeWithRetry } from "../helpers/remove-tree"; +import { repoPath } from "../helpers/repo-root"; const originalHome = process.env.OPENCODEX_HOME; let home: string; @@ -381,3 +382,27 @@ describe("model-aware failover and pre-dispatch integration", () => { }); }); +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, )", + ); + }); + } +});