diff --git a/src/routing/quota.ts b/src/routing/quota.ts index 06cc8eecc..c239980e2 100644 --- a/src/routing/quota.ts +++ b/src/routing/quota.ts @@ -12,63 +12,100 @@ * carries an account reference (dry-run/evaluate), never invented. */ -import { codexQuotaWindowForPlan, getAccountQuota, isCodexQuotaExhausted } from "../codex/quota"; +import { + codexQuotaWindowForPlan, + getAccountQuota, + isCodexQuotaExhausted, + listAccountQuotas, +} from "../codex/quota"; import { getCachedProviderAccountQuota } from "../providers/quota"; import type { RouteQuotaEvidence } from "./trace"; export interface QuotaEvidenceInput { provider: string; model: string; - /** Opaque account reference for per-account quota sources. */ accountRef?: string; - /** Codex pool account id (provider "openai"). */ codexAccountId?: string; - /** The account's plan (provider "openai"); selects the governing quota window. */ codexAccountPlan?: string; } -/** - * Assemble quota evidence from canonical local caches only (no network). - * Unknown dimensions stay unknown - never zero. - */ +export interface CodexPoolQuotaAccount { + accountId: string; + plan?: string; +} + +function codexAccountQuotaEvidence(accountId: string, plan?: string): RouteQuotaEvidence { + const quota = getAccountQuota(accountId); + if (!quota) return { known: false }; + const monthly = codexQuotaWindowForPlan(plan) === "monthly"; + const percents = [ + ...(monthly ? [] : [quota.weeklyPercent]), + quota.monthlyPercent, + ].filter((value): value is number => typeof value === "number" && Number.isFinite(value)); + const maxPercent = percents.length > 0 ? Math.max(...percents) : undefined; + // Credits-only snapshots prove neither usage nor exhaustion. Unknown must not + // become healthy capacity merely because a cache row exists. + if (maxPercent === undefined) return { known: false }; + const resets = [ + ...(monthly ? [] : [quota.weeklyResetAt]), + quota.monthlyResetAt, + ].filter((value): value is number => typeof value === "number" && Number.isFinite(value)) + .filter(value => value > Date.now()); + return { + known: true, + headroom: Math.max(0, Math.min(1, 1 - maxPercent / 100)), + exhausted: isCodexQuotaExhausted(quota, plan), + ...(resets.length > 0 ? { resetAtMs: Math.min(...resets) } : {}), + source: "codex-pool", + }; +} + +export function codexPoolQuotaEvidence(accounts: readonly CodexPoolQuotaAccount[]): RouteQuotaEvidence { + if (accounts.length === 0) return { known: false }; + const evidence = accounts.map(account => codexAccountQuotaEvidence(account.accountId, account.plan)); + const known = evidence.filter(item => item.known); + if (known.length === 0) return { known: false }; + const usable = known.filter(item => item.exhausted !== true); + if (usable.length > 0) { + const headrooms = usable.map(item => item.headroom) + .filter((value): value is number => typeof value === "number" && Number.isFinite(value)); + return { + known: true, + exhausted: false, + ...(headrooms.length > 0 ? { headroom: Math.max(...headrooms) } : {}), + source: "codex-pool", + }; + } + if (known.length < evidence.length) return { known: false }; + const resets = known.map(item => item.resetAtMs) + .filter((value): value is number => typeof value === "number" && Number.isFinite(value)); + return { + known: true, + exhausted: true, + headroom: 0, + ...(resets.length > 0 ? { resetAtMs: Math.min(...resets) } : {}), + source: "codex-pool", + }; +} + export function quotaEvidenceForCandidate(input: QuotaEvidenceInput): RouteQuotaEvidence { if (input.provider === "openai" && input.codexAccountId) { - const quota = getAccountQuota(input.codexAccountId); - if (quota) { - // Go/Free accounts report a 30-day window only; weekly windows gate - // everything else. `codexQuotaWindowForPlan` is the single shared rule - // (parser, exhaustion, recovery), so select the plan-specific bars here - // too instead of always combining weekly + monthly. - const monthly = codexQuotaWindowForPlan(input.codexAccountPlan) === "monthly"; - const percents = [ - ...(monthly ? [] : [quota.weeklyPercent]), - quota.monthlyPercent, - ].filter((value): value is number => typeof value === "number" && Number.isFinite(value)); - const maxPercent = percents.length > 0 ? Math.max(...percents) : undefined; - const resets = [ - ...(monthly ? [] : [quota.weeklyResetAt]), - quota.monthlyResetAt, - ].filter((value): value is number => typeof value === "number" && Number.isFinite(value)) - .filter(value => value > Date.now()); - return { - known: true, - ...(maxPercent !== undefined - ? { headroom: Math.max(0, Math.min(1, 1 - maxPercent / 100)) } - : {}), - exhausted: isCodexQuotaExhausted(quota, input.codexAccountPlan), - ...(resets.length > 0 ? { resetAtMs: Math.min(...resets) } : {}), - source: "codex-pool", - }; + // listAccountQuotas() is the reconciled quota snapshot: config-generation + // reconciliation prunes removed accounts. Other eligibility dimensions + // (pause/reauth/cooldown/soft-avoid) remain health/pool-selector concerns. + if (input.codexAccountPlan !== undefined) { + const pool = [...listAccountQuotas()].map(([accountId]) => ({ + accountId, + ...(accountId === input.codexAccountId ? { plan: input.codexAccountPlan } : {}), + })); + if (pool.length > 1) return codexPoolQuotaEvidence(pool); } + return codexAccountQuotaEvidence(input.codexAccountId, input.codexAccountPlan); } if (input.provider === "anthropic" && input.accountRef) { const quota = getCachedProviderAccountQuota("anthropic", input.accountRef); if (quota) { - // Anthropic per-family buckets (e.g. Opus / Sonnet) are stricter than the - // broad account windows: fold the candidate model's matching bucket into - // headroom and exhaustion so a model-specific overage is not hidden by - // a healthy aggregate window. const family = anthropicFamilyWindow(input.model, quota.customWindows ?? []); const percents = [quota.fiveHourPercent, quota.weeklyPercent, quota.monthlyPercent, family?.percent] .filter((value): value is number => typeof value === "number" && Number.isFinite(value)); @@ -78,25 +115,16 @@ export function quotaEvidenceForCandidate(input: QuotaEvidenceInput): RouteQuota .filter(value => value > Date.now()); return { known: true, - ...(maxPercent !== undefined - ? { headroom: Math.max(0, Math.min(1, 1 - maxPercent / 100)) } - : {}), + ...(maxPercent !== undefined ? { headroom: Math.max(0, Math.min(1, 1 - maxPercent / 100)) } : {}), exhausted: maxPercent !== undefined && maxPercent >= 100, ...(resets.length > 0 ? { resetAtMs: Math.min(...resets) } : {}), source: "provider-report", }; } } - return { known: false }; } -/** - * Match the candidate model to an Anthropic per-family quota window. Window - * labels from the provider probe are "Opus" / "Sonnet"; model ids carry the - * family as a segment (e.g. `claude-opus-...`, `claude-sonnet-...`). Returns - * undefined when no family window is cached or no label matches. - */ function anthropicFamilyWindow( model: string, windows: Array<{ label: string; percent?: number; resetAt?: number }>, @@ -104,18 +132,11 @@ function anthropicFamilyWindow( const normalized = model.toLowerCase(); for (const window of windows) { const family = window.label.trim().toLowerCase(); - if (family && normalized.includes(family)) { - return window; - } + if (family && normalized.includes(family)) return window; } return undefined; } -/** - * Deterministic quota score in [0,1]: larger available headroom scores - * higher; exhausted evidence scores 0. Unknown evidence returns null so the - * caller can apply the profile's unknownEvidence policy. - */ export function quotaScore(evidence: RouteQuotaEvidence | undefined): number | null { if (!evidence || !evidence.known) return null; if (evidence.exhausted === true) return 0; diff --git a/tests/routing-policy-pool-quota.test.ts b/tests/routing-policy-pool-quota.test.ts new file mode 100644 index 000000000..89f3ce1d7 --- /dev/null +++ b/tests/routing-policy-pool-quota.test.ts @@ -0,0 +1,58 @@ +import { afterEach, describe, expect, test } from "bun:test"; + +import { clearAccountQuota, setAccountQuotaFromParsed } from "../src/codex/quota"; +import { codexPoolQuotaEvidence, quotaEvidenceForCandidate } from "../src/routing/quota"; + +afterEach(() => clearAccountQuota()); + +describe("Codex pool quota evidence for routing policies", () => { + test("uses the best known usable headroom instead of only one active account", () => { + setAccountQuotaFromParsed("low", { weeklyPercent: 95 }); + setAccountQuotaFromParsed("healthy", { weeklyPercent: 20 }); + expect(codexPoolQuotaEvidence([ + { accountId: "low", plan: "plus" }, + { accountId: "healthy", plan: "plus" }, + ])).toMatchObject({ known: true, exhausted: false, headroom: 0.8, source: "codex-pool" }); + }); + + test("the live policy evidence path aggregates the reconciled pool", () => { + setAccountQuotaFromParsed("active", { weeklyPercent: 96 }); + setAccountQuotaFromParsed("alternate", { weeklyPercent: 25 }); + expect(quotaEvidenceForCandidate({ + provider: "openai", + model: "gpt-5.6", + codexAccountId: "active", + codexAccountPlan: "plus", + })).toMatchObject({ known: true, exhausted: false, headroom: 0.75, source: "codex-pool" }); + }); + + test("reports exhausted only when every pool account is known exhausted", () => { + setAccountQuotaFromParsed("a", { weeklyPercent: 100, weeklyResetAt: Date.now() + 60_000 }); + setAccountQuotaFromParsed("b", { weeklyPercent: 100, weeklyResetAt: Date.now() + 120_000 }); + const evidence = codexPoolQuotaEvidence([ + { accountId: "a", plan: "plus" }, + { accountId: "b", plan: "plus" }, + ]); + expect(evidence.known).toBe(true); + expect(evidence.exhausted).toBe(true); + expect(evidence.headroom).toBe(0); + expect(evidence.resetAtMs).toBeDefined(); + }); + + test("does not call a partially unknown pool exhausted", () => { + setAccountQuotaFromParsed("known-exhausted", { weeklyPercent: 100 }); + expect(codexPoolQuotaEvidence([ + { accountId: "known-exhausted", plan: "plus" }, + { accountId: "unknown", plan: "plus" }, + ])).toEqual({ known: false }); + }); + + test("credits-only cached evidence stays unknown", () => { + setAccountQuotaFromParsed("known-exhausted", { weeklyPercent: 100 }); + setAccountQuotaFromParsed("credits-only", { resetCredits: 7 }); + expect(codexPoolQuotaEvidence([ + { accountId: "known-exhausted", plan: "plus" }, + { accountId: "credits-only", plan: "plus" }, + ])).toEqual({ known: false }); + }); +});