diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 51779de4c1..71901089c3 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -1,409 +1,191 @@ -import { randomUUID } from "node:crypto"; import { saveConfigPreservingClaudeCode } from "../config"; -import { isCodexAccountGenerationLive, readCodexAccountRecord, type CodexRefreshProvenance } from "./account-store"; +import { isCodexAccountGenerationLive } from "./account-store"; import { codexAccountLogLabel } from "./account-label"; -import { NATIVE_RESERVE_MODEL } from "./catalog/native-models"; import { isCodexAccountPaused } from "./account-pause"; -import { clearCodexAccountPin, codexAccountPriorityLookup, pinnedCodexAccountId } from "./account-priority"; +import { clearCodexAccountPin, pinnedCodexAccountId } from "./account-priority"; import { isCodexAccountUsable, type CodexAccountUsabilityOptions } from "./account-usability"; -import { clearAccountNeedsReauth, isAccountNeedsReauth, markAccountNeedsReauth } from "./account-runtime-state"; -import { - POOL_KEY_CODEX, - normalizeAccountPoolStickyLimit, - normalizeCodexAccountPoolStrategy, - notePoolRotationFailure, - notePoolRotationSuccess, - peekRoundRobinAccount, - pickRoundRobinAccount, - seedPoolRotationAccount, - selectPriorityTier, -} from "./pool-rotation"; -import { - CODEX_EXHAUSTED_USAGE_PERCENT, - CODEX_UNKNOWN_USAGE_SCORE, - getAccountQuota, - isRetiredCodexSparkModel, - resetAtToMs, -} from "./quota"; -import { codexPlanKey, isThirtyDayOnlyCodexPlan } from "./plan"; -import { - MAIN_CODEX_ACCOUNT_ID, - getMainAccountPlan, - hasMainAccountRefreshGrant, -} from "./main-account"; +import { isAccountNeedsReauth, markAccountNeedsReauth } from "./account-runtime-state"; +import { POOL_KEY_CODEX, notePoolRotationFailure } from "./pool-rotation"; +import { getAccountQuota, isRetiredCodexSparkModel } from "./quota"; +import { MAIN_CODEX_ACCOUNT_ID } from "./main-account"; import { isSelectableCodexPoolAccount } from "./account-id"; import type { OcxConfig } from "../types"; import { captureConfigGeneration, type GenerationContext } from "../lib/state-store-sweeper"; -import { isCanonicalOpenAiForwardProvider } from "../providers/openai-tiers"; -import { retainedUtf8Bytes } from "../lib/admission"; import { recordUpstreamHostFailure } from "./upstream-host-health"; - -type ThreadAffinityEntry = { - accountId: string; - generation: number; - createdAt: number; - lastUsedAt: number; - // Last time the bound account's quota threshold was re-evaluated for this - // thread (interval-gated to avoid per-request flapping). See REEVAL_INTERVAL_MS. - lastReevalAt: number; - // When a transient failure streak first forced this thread onto another account - // while the binding was HELD (#4546). Cleared the moment the bound account serves - // again; once it ages past CODEX_TRANSIENT_AFFINITY_HOLD_MS the binding is - // released through the ordinary path instead of detouring forever. - transientHoldSince?: number; - // Which account is serving this thread while its own is held under a transient hold. - // Remembered rather than re-picked per request: under round-robin a fresh pick each turn - // would walk the ring and start cold on every hop, which is the behaviour the hold exists - // to prevent. Cleared with transientHoldSince when the bound account serves again. - transientDetourAccountId?: string; -}; - -export type CodexThreadResolution = - | { status: "selected"; accountId: string; affinity?: CodexAffinityDecision } - | { status: "none"; affinity?: CodexAffinityDecision } - | { status: "expired"; accountId: string; affinity?: CodexAffinityDecision }; - -/** What happened to this thread's binding on this request (#4546). */ -export type CodexAffinityMove = - /** Served by its own bound account, which was healthy. */ - | "reused" - /** Served by its own bound account while something transient was wrong with it. */ - | "held" - /** Served by another account while the binding stayed put. */ - | "detour" - /** The binding was released and a different account took the thread. */ - | "rebound" - /** There was no live binding; this request established one. */ - | "new_bind" - /** The binding was released without a replacement on this request. */ - | "cleared"; - -/** - * Why. A move is the expensive event -- it discards the prompt-cache prefix warmed on the old - * account -- so the operator should not have to infer it from account labels across log lines, - * which is how #4546 had to be diagnosed. - */ -export type CodexAffinityReason = - | "healthy" - | "quota_headroom" - | "quota_refusal" - | "transient" - | "transient_hold_expired" - | "unusable" - | "paused" - | "plan_excluded" - | "cooldown" - | "quota_avoided" - | "generation" - | "expired" - | "model_lane"; - -export interface CodexAffinityDecision { - move: CodexAffinityMove; - reason: CodexAffinityReason; -} - -/** The decision to report once a binding has been released and selection starts over. */ -function affinityAfterRelease( - threadId: string | null, - releaseReason: CodexAffinityReason | undefined, -): CodexAffinityDecision { - // Reported now, so it must not be reported again by the next request. - clearPendingReleaseReason(threadId); - return releaseReason === undefined - ? { move: "new_bind", reason: "healthy" } - : { move: "rebound", reason: releaseReason }; -} - -/** - * What to report when selection produced no account at all. The binding is gone and nothing took - * it, which is a `cleared`, and the pending reason is deliberately NOT consumed: a no-account - * result reaches no auth context and therefore no usage entry, so the next resolve that does - * produce one is the first place this release can actually be seen. - */ -function affinityOnNoAccount( - threadId: string | null, - releaseReason: CodexAffinityReason | undefined, -): CodexAffinityDecision | undefined { - if (releaseReason === undefined) return undefined; - // Hand it forward as well as reporting it. A reason derived from the entry this request just - // released lives only in a local, so without this the next resolve finds no entry and no - // pending reason and calls the rebind a fresh healthy bind. - notePendingReleaseReason(threadId, releaseReason); - return { move: "cleared", reason: releaseReason }; -} - -/** - * Process-local cursor for automatic RR/fill-first (and quota-429 when not - * sync-writing) picks. Keeps unrelated `saveConfig` from persisting transient - * rotation as the operator's `activeCodexAccountId`. Manual selection clears it - * so disk/`config.activeCodexAccountId` remains authoritative. - */ -let runtimeActiveCodexAccountId: string | undefined; - -type CodexUpstreamHealth = { - consecutiveFailures: number; - /** Consecutive healthy terminals observed while recovering from escalation level 2+. */ - consecutiveSuccesses?: number; - lastFailureStatus?: number; - lastFailureAt?: number; - /** Hard cooldown (quota 429). Survives a later 2xx; blocks auth + selection. */ - cooldownUntil?: number; - /** - * How long a quota refusal keeps selection away from this account (or this native quota - * group), as opposed to how long it is hard-blocked. - * - * The two are deliberately different lengths. {@link CODEX_MAX_RESET_DERIVED_COOLDOWN_MS} - * caps the hard cooldown at 15 minutes because a reset announcement is advisory and plan - * quota usually frees up before it — an account must stay reachable so the pool can find - * that out (#433). The window the refusal announced is not 15 minutes, though, so once the - * cooldown lapses the account is selectable again while its burst window is still spent, - * and the strategy picks it straight back: this proxy reads a weekly bar a burst limit never - * touches, so a refused account still scores as the coolest in the pool. Every request then - * earns the same 429 until the process restarts, which is the only thing that drops this map. - * - * So the announcement governs avoidance and the cap still governs blocking. Avoidance is soft - * in the {@link softAvoidUntil} sense: it reorders the pool and releases a bound thread, and - * the last-resort paths still reach the account when nothing else can serve, so one pessimistic - * announcement cannot stall routing. - */ - quotaAvoidUntil?: number; - /** When the current cooldown was recorded; origin of the probe interval clock. */ - cooldownSince?: number; - /** - * What produced the cooldown. An explicit Retry-After is a literal retry - * directive and is never probed; a quota resetAt only announces a window - * refresh, so it may be probed early (#433). - */ - cooldownSource?: CodexCooldownSource; - /** - * Bumped on every cooldown write. A probe lease records the generation it was - * issued for so a lease cannot clear a cooldown that a later 429 replaced. - */ - cooldownGeneration?: number; - /** - * Identity of the in-flight probe. A cooled-down account sends no traffic, so - * no organic 2xx can prove recovery; only the outcome carrying this id may - * clear the cooldown. - */ - probeLeaseId?: string; - /** Cooldown generation at the moment the lease was granted. */ - probeLeaseGeneration?: number; - /** Last probe grant or conclusion; paces the probe interval. */ - lastProbeAt?: number; - /** - * Soft avoid after connect_error / timeout / transient 5xx. Cleared on 2xx. - * Blocks pool selection + thread affinity reuse so a sticky session can leave a - * flaky account without throwing CodexAccountCooldownError (hard-only). - */ - softAvoidUntil?: number; - /** - * Credential generation a 401/403 quarantine was derived from (#2892 gap 4). - * - * Provenance lives ON the entry rather than in a side map keyed by account id. A side map spends - * "whatever health is current when the old credential is found dead", which deletes a later - * unrelated entry: a G1 401, then a G2 save, then a genuine G2 503 would lose the 503. Only the - * entry that carries this field can be spent, and any later write simply replaces it. - */ - credentialFailureGeneration?: number; -}; - -const CODEX_DEFAULT_QUOTA_COOLDOWN_MS = 60_000; -const CODEX_MAX_QUOTA_COOLDOWN_MS = 24 * 60 * 60_000; -/** - * A weekly/monthly quota `resetAt` announces when the window refreshes; it is not - * a "come back after this" directive like Retry-After. Plan quota routinely frees - * up long before the advertised reset, so cap reset-derived cooldowns far below - * the Retry-After ceiling (#433). - */ -const CODEX_MAX_RESET_DERIVED_COOLDOWN_MS = 15 * 60_000; -/** - * Ceiling on quota-refusal avoidance. Generous enough to cover a full five-hour burst window, - * tight enough that a weekly or monthly reset four days out cannot take an account out of - * rotation for the {@link CODEX_MAX_QUOTA_COOLDOWN_MS} day the Retry-After ceiling allows. - */ -const CODEX_MAX_QUOTA_AVOID_MS = 6 * 60 * 60_000; -/** Minimum gap between probe leases for one cooled-down account. */ -export const CODEX_QUOTA_PROBE_INTERVAL_MS = 5 * 60_000; -export const CODEX_FAILURE_WINDOW_MS = 5 * 60_000; -/** - * How recently a 100% burst reading must have been OBSERVED to exclude an account when it - * carries no reset timestamp (#3425). Deliberately far tighter than the 6h disk-hydration - * horizon in `quota.ts`: shorter than any plausible five-hour burst window, so a persisted - * reading can never strand a recovered account, and long enough that a snapshot taken at - * admission is still fresh when selection reads it. - */ -export const TERMINAL_SHORT_WINDOW_FRESHNESS_MS = 5 * 60_000; -/** How long a transient failure keeps the account out of pool selection. */ -export const CODEX_TRANSIENT_SOFT_AVOID_MS = 30_000; -const CODEX_TRANSIENT_SOFT_AVOID_ESCALATION_MS = [ +import { + classifyCodexUpstreamOutcome, + computeCodexUsageScore, + computeQuotaCooldown, + quotaAvoidUntilFor, + CODEX_FAILURE_WINDOW_MS, + CODEX_TRANSIENT_SOFT_AVOID_ESCALATION_MS, + type CodexUpstreamOutcome, + type CodexUpstreamOutcomeMeta, +} from "./routing/cooldown-math"; +import { + codexPoolKeyForScope, + codexQuotaScopeForModel, + deleteAccountHealth, + deleteAllScopedHealth, + deleteScopedHealth, + dropSpentCredentialFailure, + getAccountHealth, + getCodexAccountCooldownUntil, + getCodexAccountSoftAvoidUntil, + getCodexQuotaHealthSnapshot, + isCodexAccountSoftAvoided, + isCodexQuotaAvoided, + isHealthAccountAdmissible, + isHealthGenerationReconciled, + isIndependentCodexQuotaScope, + preservedCooldownFields, + pruneHealthAccountsForContext, + commitHealthReconcile, + clearUpstreamHealthState, + resetHealthReconcileState, + deleteAllHealthForAccount, + scopedHealthFor, + setAccountHealth, + setScopedHealth, + type CodexQuotaScope, + type CodexUpstreamHealth, +} from "./routing/health-store"; +import { ownsProbeLease, probeMayClearCooldown, withProbeLeaseReleased } from "./routing/probe-lease"; +import { + affinityAfterRelease, + affinityOnNoAccount, + bindModelDetourAffinity, + bindThreadAffinity, + clearThreadAccountMapForAccount, + deleteModelDetourAffinity, + deleteThreadAffinity, + deleteThreadAffinitiesForAccount, + getThreadAffinity, + getThreadAffinityScopes, + getModelDetourAffinity, + isThreadAffinityExpired, + isThreadAffinityGenerationLive, + peekPendingReleaseReason, + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS, + CODEX_TRANSIENT_AFFINITY_HOLD_MS, + type CodexAffinityReason, + type CodexThreadResolution, + type ThreadAffinityEntry, +} from "./routing/thread-affinity"; +import { + accountPoolStrategyForScope, + applyFailureFailover, + applyQuotaAutoSwitch, + codexAccountBlockReason, + getEligiblePoolAccounts, + getPoolAccountPlanForSelection, + hasCodexQuotaHeadroom, + isCodexAccountPlanExcluded, + isCacheAffinityEnabled, + isCodexAccountSelectable, + isHealthySharedCodexSelection, + isUnknownUsage, + pickAlternateCodexAccount, + pickLowerUsageAccount, + pickLowestUsageAmong, + pickLowestUsageCodexAccount, + pickPriorityPreemption, + pickResetFirstCodexAccount, + pickUnboundStrategyAccount, + sharedStateSelectionOptions, + strategySelectionOptionsForModelDetour, + shouldFailover, + peekAlternateCodexAccount, +} from "./routing/selection"; +import { + clearAllManualPreferences, + consumeManualPreference, + forgetManualPreference, + forgetRoutingPreferencesOutside, + forgetRuntimeActiveCodexAccount, + getEffectiveActiveCodexAccountId, + manualPreferenceBlocks, + promoteActiveCodexAccount, + rememberActiveCodexAccount, + setActiveCodexAccount, +} from "./routing/active-account"; + +export { + CODEX_QUOTA_PROBE_INTERVAL_MS, + CODEX_FAILURE_WINDOW_MS, + TERMINAL_SHORT_WINDOW_FRESHNESS_MS, CODEX_TRANSIENT_SOFT_AVOID_MS, - 2 * 60_000, - 10 * 60_000, - 30 * 60_000, -] as const; -export const CODEX_THREAD_AFFINITY_IDLE_TTL_MS = 24 * 60 * 60_000; -export const CODEX_THREAD_AFFINITY_MAX_ENTRIES = 2048; -const MAX_AFFINITY_COMPONENT_BYTES = 512; -// Min interval between quota threshold re-evaluations for a single bound thread. -// Well under the 5h/weekly quota windows, but enough to stop per-request flapping. -export const CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS = 60_000; - -/** - * How long a live binding outlives a TRANSIENT failure streak on its own account (#4546). - * - * Being unable to send right now is not the same as losing ownership of the conversation. - * A 5xx streak is frequently provider-wide rather than account-specific, and deleting the - * binding for it discards a prompt-cache prefix that the next turn then pays for again -- - * the same cost the quota threshold used to impose, arriving through a different door. - * So the request detours to another account while the binding is held here. - * - * Bounded, because an unbounded hold is its own defect: an account that never recovers - * would keep a thread detouring indefinitely while the conversation's real warm prefix - * accumulates somewhere else. Ten minutes is longer than the whole soft-avoid escalation - * ladder up to its final step, so an ordinary outage resolves inside the hold and a - * genuine one converts to a real rebind instead of a permanent detour. - */ -export const CODEX_TRANSIENT_AFFINITY_HOLD_MS = 10 * 60_000; - -const upstreamHealth = new Map(); -/** - * Reset-derived 429s can describe a quota owned by one native model family, - * rather than the whole ChatGPT account. Keep those advisory cooldowns apart - * from account-wide Retry-After/default throttles and transient health. - */ -const quotaScopedHealth = new Map>(); -/** - * Spend a credential-failure health entry whose credential no longer exists (#2892 gap 4). - * - * A 401/403 describes one CREDENTIAL, not an account, and a replacement can land at any point after - * the outcome is recorded — so re-reading the store inside `recordCodexUpstreamOutcome` narrows the - * window without closing it. The reader decides instead, and it may only spend an entry that - * actually carries credential provenance: a later transient or quota write replaces the entry and - * with it the tag, so this can never delete evidence that belongs to a different failure. - */ -function dropSpentCredentialFailure(accountId: string): void { - const health = upstreamHealth.get(accountId); - const generation = health?.credentialFailureGeneration; - if (health === undefined || generation === undefined) return; - if (isCodexAccountGenerationLive(accountId, generation)) return; - upstreamHealth.delete(accountId); -} -let lastReconciledGeneration = 0; -let liveHealthAccountIds = new Set(); - -export type CodexUpstreamOutcome = number | "connect_error" | "timeout" | "connect_neutral"; -export type CodexUpstreamOutcomeClass = "success" | "credential" - | "workspace" | "quota" | "transient" | "caller" | "neutral" | "unknown"; -export type CodexCooldownSource = "retry-after" | "reset-derived" | "default"; -/** - * Native Codex quota groups known to be independent upstream. Keep the mapping - * deliberately conservative: unlisted models share the normal native group. - * Add a new explicit group here only when its independent upstream quota is - * confirmed, so shared limits never receive cross-model bypasses. - */ -export type CodexQuotaScope = "shared" | "reserve"; - -export type CodexQuotaRecoveryProbeClaim = { - accountId: string; - scope?: CodexQuotaScope; - leaseId: string; - cooldownGeneration: number; - credentialGeneration: number; - /** Claim-time `replacedAt`; unchanged after a probe-owned refresh, stamped on external replacement. */ - credentialReplacedAt?: number; -}; - -export type CodexQuotaRecoveryProbeProof = { - credentialGeneration?: number; -}; - -/** - * Requests without a resolved native model retain the historic one-account-per- - * thread behavior. Requests with a known quota scope get an independent - * affinity so a Reserve failover cannot displace the same thread's Terra/Luna - * account (and vice versa). - */ -type BaseThreadAffinityScope = CodexQuotaScope | "legacy"; -type ModelDetourAffinityScope = `model-detour:${BaseThreadAffinityScope}:${string}`; -type ThreadAffinityScope = BaseThreadAffinityScope | ModelDetourAffinityScope; -const LEGACY_THREAD_AFFINITY_SCOPE = "legacy" as const; -const threadAccountMap = new Map>(); -let threadAffinityEntryTotal = 0; - -function isModelDetourAffinityScope(scope: ThreadAffinityScope): scope is ModelDetourAffinityScope { - return scope.startsWith("model-detour:"); -} - -const NATIVE_MODEL_QUOTA_SCOPES: Readonly> = { - [NATIVE_RESERVE_MODEL]: "reserve", -}; - -export function codexQuotaScopeForModel(modelId: string | undefined): CodexQuotaScope | undefined { - if (!modelId?.trim()) return undefined; - return NATIVE_MODEL_QUOTA_SCOPES[modelId.trim().toLowerCase()] ?? "shared"; -} - -/** Independent quota groups must not mutate the shared active-account cursor. */ -function isIndependentCodexQuotaScope(quotaScope?: CodexQuotaScope): boolean { - return quotaScope !== undefined && quotaScope !== "shared"; -} - -function codexPoolKeyForScope(quotaScope?: CodexQuotaScope): string { - return isIndependentCodexQuotaScope(quotaScope) ? `${POOL_KEY_CODEX}:${quotaScope}` : POOL_KEY_CODEX; -} - -export type CodexUpstreamOutcomeMeta = { - retryAfter?: string | null; - resetAt?: unknown | unknown[]; - now?: number; - /** (provider, host) ledger key for account-neutral reachability failures (#914). */ - hostKey?: string; - /** - * Upstream denial evidence for a 403. A workspace/entitlement denial means the CREDENTIAL - * is fine and the account simply cannot reach this workspace, so it must not be quarantined - * for reauthentication (#1789). Absent evidence keeps the historical credential handling. - */ - denial?: "workspace" | "entitlement"; - /** Stable transport code recorded alongside a neutral host failure. */ - lastFailureCode?: string; - /** Native model selected for this request; used only for confirmed scoped quotas. */ - modelId?: string; - /** When set, clears affinity for this thread immediately on transient failure. */ - threadId?: string | null; - /** - * Suppress Pool rotation and quota/transient affinity mutations for an account-qualified - * request. Credential failures still sweep stale affinities because reauthentication is - * account-wide. - */ - fixedAccount?: boolean; - /** - * Probe lease held by this request, when it was admitted through an active - * quota cooldown. Only the outcome carrying the current lease may clear the - * cooldown (#433). - */ - probeLeaseId?: string; - /** Scope of `probeLeaseId` when it was granted against a model-scoped cooldown. */ - probeQuotaScope?: CodexQuotaScope; - /** - * Already-chosen alternate for same-request 429 retry. When set, promotion - * reuses this account instead of calling {@link pickAlternateCodexAccount} - * again (which would advance a round-robin ring twice). - */ - promoteAccountId?: string; - /** Generation captured when this routed account was selected. */ - writerGeneration?: number; - /** - * Credential generation this request's bearer was read at. Distinct from - * `writerGeneration`, which tracks the config store. - * - * A 401 that arrives after the credential was already replaced is evidence about a - * token nobody is using any more, so it must not quarantine the replacement. Absent - * means the caller cannot supply lineage and the historical unfenced handling stands. - */ - credentialGeneration?: number; -}; - + classifyCodexUpstreamOutcome, + computeCodexUsageScore, + computeQuotaCooldown, + computeQuotaCooldownUntil, + parseRetryAfterMs, + parseResetCooldownMs, +} from "./routing/cooldown-math"; +export type { + CodexUpstreamOutcome, + CodexUpstreamOutcomeClass, + CodexCooldownSource, + CodexUpstreamOutcomeMeta, +} from "./routing/cooldown-math"; +export { + codexQuotaScopeForModel, + listLiveCodexAccountIds, + getCodexUpstreamHealth, + getCodexAccountCooldownUntil, + getCodexAccountHealthSnapshot, + getCodexQuotaHealthSnapshot, + isCodexAccountInCooldown, + clearCodexAccountCooldown, + getCodexAccountSoftAvoidUntil, + isCodexAccountSoftAvoided, +} from "./routing/health-store"; +export type { CodexQuotaScope } from "./routing/health-store"; +export { + tryAcquireCodexQuotaProbeLease, + canAcquireCodexQuotaProbeLease, + claimDueCodexQuotaRecoveryProbes, + claimManualResetCooldowns, + settleManualResetCooldown, + settleCodexQuotaRecoveryProbe, + tryAcquireCodexQuotaScopeProbeLease, + canAcquireCodexQuotaScopeProbeLease, + releaseCodexQuotaProbeLease, + releaseCodexQuotaScopeProbeLease, +} from "./routing/probe-lease"; +export type { + CodexQuotaRecoveryProbeClaim, + CodexQuotaRecoveryProbeProof, + ManualResetCooldownClaim, + ManualResetRefreshLineage, +} from "./routing/probe-lease"; +export { + CODEX_THREAD_AFFINITY_IDLE_TTL_MS, + CODEX_THREAD_AFFINITY_MAX_ENTRIES, + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS, + CODEX_TRANSIENT_AFFINITY_HOLD_MS, + clearThreadAccountMap, + clearThreadAccountMapForAccount, + debugCodexAffinityGenerations, + handOffThreadAffinityGeneration, +} from "./routing/thread-affinity"; +export type { + CodexThreadResolution, + CodexAffinityMove, + CodexAffinityReason, + CodexAffinityDecision, +} from "./routing/thread-affinity"; +export { + isCodexAccountPlanExcluded, + getPoolAccountPlan, + pickLowestUsageCodexAccount, + pickAlternateCodexAccount, +} from "./routing/selection"; +export { + resetCodexRoutingForManualSelection, + getEffectiveActiveCodexAccountId, + isEffectiveCodexAccountPinned, +} from "./routing/active-account"; function hasConfiguredPoolAccount( config: OcxConfig, accountId: string, @@ -416,1284 +198,41 @@ function hasConfiguredPoolAccount( .some(account => isSelectableCodexPoolAccount(account) && account.id === accountId); } -export function listLiveCodexAccountIds(config: OcxConfig): ReadonlySet { - const ids = new Set((config.codexAccounts ?? []).map(account => account.id)); - const openai = config.providers.openai; - if (openai && openai.disabled !== true && isCanonicalOpenAiForwardProvider(openai)) { - ids.add(MAIN_CODEX_ACCOUNT_ID); - } - return ids; -} - -export function clearThreadAccountMap(): void { - threadAccountMap.clear(); - threadAffinityEntryTotal = 0; -} - -export function clearThreadAccountMapForAccount( - accountId: string, - reason: CodexAffinityReason = "unusable", -): void { - for (const [threadId, affinities] of threadAccountMap) { - for (const [scope, entry] of affinities) { - if (entry.accountId === accountId && affinities.delete(scope)) { - threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); - notePendingReleaseReason(threadId, reason); - } - } - if (affinities.size === 0) threadAccountMap.delete(threadId); - } -} - -/** - * Why a binding was released, held until that thread's next resolve can report it (#4546). - * - * A release and the request that pays for it are two different moments: a 429 clears the pin - * inside the outcome recorder, and the next request arrives with nothing left to explain why it - * is starting cold. Bounded, because it is a diagnostic and must not become a leak. - */ -const pendingReleaseReasons = new Map(); -const MAX_PENDING_RELEASE_REASONS = 4096; - -function notePendingReleaseReason(threadId: string | null, reason: CodexAffinityReason): void { - if (threadId === null) return; - if (!pendingReleaseReasons.has(threadId) && pendingReleaseReasons.size >= MAX_PENDING_RELEASE_REASONS) { - const oldest = pendingReleaseReasons.keys().next(); - if (!oldest.done) pendingReleaseReasons.delete(oldest.value); - } - pendingReleaseReasons.set(threadId, reason); -} - -function peekPendingReleaseReason(threadId: string | null): CodexAffinityReason | undefined { - if (threadId === null) return undefined; - return pendingReleaseReasons.get(threadId); -} - -/** - * Forget a release only once it has actually been reported. - * - * Consuming it at derivation time lost it whenever selection then failed to produce an account: - * a no-account return carries no payload, so the release went unrecorded and the next successful - * resolve claimed a fresh healthy bind (#4598). A release survives until some resolve reports it. - */ -function clearPendingReleaseReason(threadId: string | null): void { - if (threadId !== null) pendingReleaseReasons.delete(threadId); -} - export function clearCodexUpstreamHealth(): void { // Operator preferences are routing state, not health, but they live and die with the same - // reset points. Leaving them behind lets a selection from one context suppress the - // automatic cursor in the next one. - manualPreference.clear(); - upstreamHealth.clear(); - quotaScopedHealth.clear(); - runtimeActiveCodexAccountId = undefined; - // The reconcile watermark is part of this state, not something that outlives it. Keeping - // it across a full reset is incoherent: there is no health left to protect, yet - // recordCodexUpstreamOutcome would still drop a writer whose generation predates the - // watermark for any account missing from the equally stale live set. Left behind, it also - // leaks between test files, which is how it was found. - lastReconciledGeneration = 0; - liveHealthAccountIds = new Set(); -} - -export function clearCodexUpstreamHealthForAccount(accountId: string): void { - upstreamHealth.delete(accountId); - quotaScopedHealth.delete(accountId); - // Deletion is the third operator exit, next to pause and exclusion, and it is the one - // with no reconcile path behind it: once the account is gone nothing can succeed on it, - // so an unspent preference naming it would suppress the automatic cursor for every other - // account until the process restarts. - forgetManualPreference(accountId); -} - -export function reconcileCodexRoutingHealth(context: GenerationContext): number { - if (context.generation <= lastReconciledGeneration) return 0; - let removed = 0; - for (const accountId of upstreamHealth.keys()) { - if (context.codexAccountIds.has(accountId)) continue; - upstreamHealth.delete(accountId); - removed += 1; - } - for (const accountId of quotaScopedHealth.keys()) { - if (context.codexAccountIds.has(accountId)) continue; - quotaScopedHealth.delete(accountId); - removed += 1; - } - // Sweep preferences the same way, for the account set this generation actually has. The - // delete path above is the direct route; this is the one that catches an account removed - // by an edit the runtime never saw. Deliberately not counted in `removed`, which reports - // health rows. - for (const [poolKey, preferred] of manualPreference) { - if (context.codexAccountIds.has(preferred)) continue; - manualPreference.delete(poolKey); - } - liveHealthAccountIds = new Set(context.codexAccountIds); - lastReconciledGeneration = context.generation; - return removed; -} - -export function getCodexUpstreamHealth( - accountId: string, -): CodexUpstreamHealth | null { - dropSpentCredentialFailure(accountId); - return upstreamHealth.get(accountId) ?? null; -} - -function scopedHealthFor(accountId: string, scope: CodexQuotaScope): CodexUpstreamHealth | undefined { - return quotaScopedHealth.get(accountId)?.get(scope); -} - -function setScopedHealth(accountId: string, scope: CodexQuotaScope, health: CodexUpstreamHealth): void { - let scopes = quotaScopedHealth.get(accountId); - if (!scopes) { - scopes = new Map(); - quotaScopedHealth.set(accountId, scopes); - } - scopes.set(scope, health); -} - -function deleteScopedHealth(accountId: string, scope: CodexQuotaScope): void { - const scopes = quotaScopedHealth.get(accountId); - if (!scopes) return; - scopes.delete(scope); - if (scopes.size === 0) quotaScopedHealth.delete(accountId); -} - -export function computeCodexUsageScore(quota: { - weeklyPercent?: number; - monthlyPercent?: number; - shortPercent?: number; - shortResetAt?: number; - shortObservedAt?: number; -} | null, plan?: unknown, now: number = Date.now()): number { - if (!quota) return CODEX_UNKNOWN_USAGE_SCORE; - const finite = (value: unknown): value is number => typeof value === "number" && Number.isFinite(value); - const longWindows = isThirtyDayOnlyCodexPlan(plan) - ? [quota.monthlyPercent] - : [quota.weeklyPercent, quota.monthlyPercent]; - const knownLong = longWindows.filter(finite); - // The short burst window only REFINES a known long-window position; it cannot stand in for - // one. A snapshot carrying just `shortPercent: 0` would otherwise score a flat 0 and make an - // account whose weekly/monthly usage is entirely unverified look like the emptiest in the - // pool, so `pickLowestUsageAmong` would send every request to it. Unknown has to stay - // unknown until a governing window is actually observed. - // - // A FULL burst window is the exception (#3029). It is not an optimistic guess about an - // unobserved window — it is a direct observation that the account cannot serve a request - // right now, whatever its monthly position turns out to be. Unknown-means-selectable is - // correct for uncertainty and wrong for a measured refusal: the account stays selected, - // `applyQuotaAutoSwitch` never fires, and the pool wedges on an exhausted credential. - if (knownLong.length === 0) { - return isTerminalShortWindow(quota, now) ? CODEX_EXHAUSTED_USAGE_PERCENT : CODEX_UNKNOWN_USAGE_SCORE; - } - const values = finite(quota.shortPercent) ? [...knownLong, quota.shortPercent] : knownLong; - return Math.max(...values); -} - -/** - * A short-only reading that proves the account is blocked NOW. - * - * Freshness is not optional. `getAccountQuota` performs no expiry check, partial updates - * carry a still-open short tuple forward, and disk hydration accepts a persisted reading for - * hours — so scoring 100 from `shortPercent` alone would keep excluding an account whose - * five-hour window has since reset. Merge no longer carries an elapsed shortResetAt, but an - * explicit incoming elapsed tuple is still stored, and a missing reset cannot be aged there. - * That is #3029 pointed the other way: the issue is that - * an exhausted account stays selected, and "a recovered account stays excluded" trades one - * unusable pool for another. - * - * A reading with no `shortResetAt` cannot be aged, so it stays unknown. The conservative - * direction here is the one that keeps an account selectable: a wrongly-selected account - * fails one request, while a wrongly-excluded one is invisible until someone reads the pool - * by hand. - * - * A missing reset can instead be aged by shortObservedAt (#3425). General updatedAt is not - * sufficient: credit-only updates preserve the old short tuple but advance that timestamp. - * Old disk snapshots without short-window provenance remain unknown. - */ -function isTerminalShortWindow( - quota: { shortPercent?: number; shortResetAt?: number; shortObservedAt?: number }, - now: number, -): boolean { - if (typeof quota.shortPercent !== "number" || !Number.isFinite(quota.shortPercent)) return false; - if (quota.shortPercent < CODEX_EXHAUSTED_USAGE_PERCENT) return false; - const resetAt = quota.shortResetAt; - if (typeof resetAt !== "number" || !Number.isFinite(resetAt) || resetAt <= 0) { - const observedAt = quota.shortObservedAt; - if (typeof observedAt !== "number" || !Number.isFinite(observedAt)) return false; - const age = now - observedAt; - return age >= 0 && age <= TERMINAL_SHORT_WINDOW_FRESHNESS_MS; - } - // Seconds and milliseconds both reach storage, so the split lives in one place next to the - // merge that also ages a stored reset instant (`resetAtToMs`, src/codex/quota.ts). - return resetAtToMs(resetAt) > now; -} - -export function classifyCodexUpstreamOutcome( - outcome: CodexUpstreamOutcome, - denial?: "workspace" | "entitlement", -): CodexUpstreamOutcomeClass { - if (outcome === "connect_neutral") return "neutral"; - if (outcome === "connect_error" || outcome === "timeout") return "transient"; - if (!Number.isFinite(outcome)) return "unknown"; - if (outcome >= 200 && outcome < 300) return "success"; - // Explicit 3xx policy (#914): a redirect response is relayed as-is and is - // never account or host health evidence — it proves the host is reachable - // and says nothing about the credential. Relayed as the neutral class so a - // stray 3xx cannot increment an account's transient streak. - if (outcome >= 300 && outcome < 400) return "neutral"; - // 401 is always a credential problem. A 403 is only a credential problem when nothing - // tells us otherwise: a workspace/entitlement denial (#1789) means the credential is valid - // and the account simply lacks access here, so quarantining it for reauth is wrong advice. - // Absent denial evidence the historical mapping stands, so the change fails safe. - if (outcome === 403 && denial !== undefined) return "workspace"; - if (outcome === 401 || outcome === 403) return "credential"; - // 402 Payment Required is treated as quota exhaustion for pool cooldown/failover - // (same-request alternate retry records this outcome for the depleted account). - if (outcome === 429 || outcome === 402) return "quota"; - if (outcome >= 400 && outcome < 500) return "caller"; - if (outcome >= 500 && outcome < 600) return "transient"; - return "unknown"; -} - -function clampCooldownMs(ms: number): number { - return Math.min(Math.max(ms, 1), CODEX_MAX_QUOTA_COOLDOWN_MS); -} - -export function parseRetryAfterMs(value: string | null | undefined, now = Date.now()): number | undefined { - const text = value?.trim(); - if (!text) return undefined; - if (/^\d+(?:\.\d+)?$/.test(text)) { - const seconds = Number(text); - if (Number.isFinite(seconds) && seconds > 0) return clampCooldownMs(Math.ceil(seconds * 1000)); - } - const timestamp = Date.parse(text); - if (!Number.isFinite(timestamp)) return undefined; - const delay = timestamp - now; - return delay > 0 ? clampCooldownMs(delay) : undefined; -} - -function resetTimestampMs(value: unknown): number | undefined { - const numeric = typeof value === "number" - ? value - : typeof value === "string" && value.trim() !== "" - ? Number(value) - : undefined; - if (typeof numeric !== "number" || !Number.isFinite(numeric) || numeric <= 0) return undefined; - return numeric < 1_000_000_000_000 ? numeric * 1000 : numeric; -} - -export function parseResetCooldownMs(resetAt: unknown | unknown[] | undefined, now = Date.now()): number | undefined { - const values = Array.isArray(resetAt) ? resetAt : [resetAt]; - let best: number | undefined; - for (const value of values) { - const timestamp = resetTimestampMs(value); - if (timestamp === undefined) continue; - const delay = timestamp - now; - if (delay <= 0) continue; - // A far-future reset must not pin the account for the full Retry-After - // ceiling: quota usually frees up well before the advertised window (#433). - const clamped = Math.min(clampCooldownMs(delay), CODEX_MAX_RESET_DERIVED_COOLDOWN_MS); - if (best === undefined || clamped < best) best = clamped; - } - return best; -} - -export function computeQuotaCooldown(meta: CodexUpstreamOutcomeMeta = {}): { - until: number; - source: CodexCooldownSource; -} { - const now = meta.now ?? Date.now(); - const retryAfterMs = parseRetryAfterMs(meta.retryAfter, now); - if (retryAfterMs !== undefined) return { until: now + retryAfterMs, source: "retry-after" }; - const resetCooldownMs = parseResetCooldownMs(meta.resetAt, now); - if (resetCooldownMs !== undefined) return { until: now + resetCooldownMs, source: "reset-derived" }; - return { until: now + CODEX_DEFAULT_QUOTA_COOLDOWN_MS, source: "default" }; -} - -/** - * When the pool should stop preferring an account after it refused on quota. - * - * The earliest window the refusal actually announced, bounded by {@link CODEX_MAX_QUOTA_AVOID_MS}, - * and never shorter than the cooldown the same refusal produced — a Retry-After directive that - * outlasts every announcement still governs. - */ -function quotaAvoidUntilFor(meta: CodexUpstreamOutcomeMeta, now: number, cooldownUntil: number): number { - const values = Array.isArray(meta.resetAt) ? meta.resetAt : [meta.resetAt]; - let announced: number | undefined; - for (const value of values) { - const timestamp = resetTimestampMs(value); - if (timestamp === undefined) continue; - const delay = timestamp - now; - if (delay <= 0) continue; - const until = now + Math.min(delay, CODEX_MAX_QUOTA_AVOID_MS); - if (announced === undefined || until < announced) announced = until; - } - return Math.max(cooldownUntil, announced ?? 0); -} - -/** Live quota-refusal avoidance for an account, including the lane the request belongs to. */ -function codexQuotaAvoidUntil( - accountId: string, - quotaScope: CodexQuotaScope | undefined, - now: number, -): number | null { - const live = (value: number | undefined): number | null => - typeof value === "number" && Number.isFinite(value) && value > now ? value : null; - const account = live(upstreamHealth.get(accountId)?.quotaAvoidUntil); - const scoped = quotaScope === undefined - ? null - : live(scopedHealthFor(accountId, quotaScope)?.quotaAvoidUntil); - if (account === null) return scoped; - return scoped === null ? account : Math.max(account, scoped); -} - -function isCodexQuotaAvoided( - accountId: string, - quotaScope: CodexQuotaScope | undefined, - now: number, -): boolean { - return codexQuotaAvoidUntil(accountId, quotaScope, now) !== null; -} - -export function computeQuotaCooldownUntil(meta: CodexUpstreamOutcomeMeta = {}): number { - return computeQuotaCooldown(meta).until; -} - -/** - * Grant at most one probe lease per interval for a cooled-down account. - * - * A cooled-down account is short-circuited locally, so it never sends traffic and - * no organic 2xx can prove that upstream quota recovered — the cooldown can only - * end by expiry or a proxy restart (#433). Releasing a single probe breaks that - * deadlock. Explicit Retry-After cooldowns are excluded: those are literal retry - * directives, not window announcements. - * - * Returns the lease id, or null when no probe may go out right now. - */ -export function tryAcquireCodexQuotaProbeLease(accountId: string, now = Date.now()): string | null { - if (!canAcquireCodexQuotaProbeLease(accountId, now)) return null; - const health = upstreamHealth.get(accountId)!; - const probeLeaseId = randomUUID(); - upstreamHealth.set(accountId, { - ...health, - probeLeaseId, - probeLeaseGeneration: health.cooldownGeneration ?? 0, - lastProbeAt: now, - }); - return probeLeaseId; -} - -/** Side-effect-free check mirroring {@link tryAcquireCodexQuotaProbeLease} eligibility. */ -export function canAcquireCodexQuotaProbeLease(accountId: string, now = Date.now()): boolean { - return canAcquireQuotaProbeLease(upstreamHealth.get(accountId), now); -} - -function canAcquireQuotaProbeLease(health: CodexUpstreamHealth | undefined, now: number): boolean { - if (!health) return false; - const cooldownUntil = health.cooldownUntil; - if (typeof cooldownUntil !== "number" || !Number.isFinite(cooldownUntil) || cooldownUntil <= now) return false; - if (health.cooldownSource === "retry-after") return false; - if (health.probeLeaseId !== undefined) return false; - const origin = health.lastProbeAt ?? health.cooldownSince ?? cooldownUntil; - return now - origin >= CODEX_QUOTA_PROBE_INTERVAL_MS; -} - -/** - * Claim due reset-derived cooldown probes without consulting account selection. - * Added Pool credentials only; owned main usage recovery is handled separately. - */ -export function claimDueCodexQuotaRecoveryProbes( - config: OcxConfig, - limit: number, - now = Date.now(), -): CodexQuotaRecoveryProbeClaim[] { - const boundedLimit = Math.max(0, Math.floor(limit)); - if (boundedLimit === 0) return []; - const candidates: Array<{ - accountId: string; - scope?: CodexQuotaScope; - health: CodexUpstreamHealth; - credentialGeneration: number; - credentialReplacedAt?: number; - order: number; - }> = []; - for (const [order, account] of (config.codexAccounts ?? []).entries()) { - if (!isSelectableCodexPoolAccount(account) - || isCodexAccountPaused(config, account.id) - || isAccountNeedsReauth(account.id)) continue; - const record = readCodexAccountRecord(account.id); - if (!record?.credential || record.deletedAt != null) continue; - const due = [ - { scope: undefined, health: upstreamHealth.get(account.id) }, - ...[...(quotaScopedHealth.get(account.id) ?? [])].map(([scope, health]) => ({ scope, health })), - ].filter((entry): entry is { scope?: CodexQuotaScope; health: CodexUpstreamHealth } => - // Generic WHAM evidence can recover only ordinary quota, never Reserve. - // Do not spend this account's one claim per pass on an independent scope and - // delay the shared scope that the response can actually recover. - (entry.scope === undefined || entry.scope === "shared") - && entry.health?.cooldownSource === "reset-derived" - && canAcquireQuotaProbeLease(entry.health, now)) - .sort((a, b) => - (a.health.lastProbeAt ?? a.health.cooldownSince ?? 0) - - (b.health.lastProbeAt ?? b.health.cooldownSince ?? 0)); - const candidate = due[0]; - if (candidate) candidates.push({ - accountId: account.id, - ...(candidate.scope ? { scope: candidate.scope } : {}), - health: candidate.health, - credentialGeneration: record.generation, - ...(record.replacedAt !== undefined ? { credentialReplacedAt: record.replacedAt } : {}), - order, - }); - } - candidates.sort((a, b) => { - const age = (a.health.lastProbeAt ?? a.health.cooldownSince ?? 0) - - (b.health.lastProbeAt ?? b.health.cooldownSince ?? 0); - return age || a.order - b.order; - }); - return candidates.slice(0, boundedLimit).map(candidate => { - const leaseId = randomUUID(); - const next = { - ...candidate.health, - probeLeaseId: leaseId, - probeLeaseGeneration: candidate.health.cooldownGeneration ?? 0, - lastProbeAt: now, - }; - if (candidate.scope) setScopedHealth(candidate.accountId, candidate.scope, next); - else upstreamHealth.set(candidate.accountId, next); - return { - accountId: candidate.accountId, - ...(candidate.scope ? { scope: candidate.scope } : {}), - leaseId, - cooldownGeneration: candidate.health.cooldownGeneration ?? 0, - credentialGeneration: candidate.credentialGeneration, - ...(candidate.credentialReplacedAt !== undefined - ? { credentialReplacedAt: candidate.credentialReplacedAt } - : {}), - }; - }); -} - -type CooldownRecoveryLease = Pick; - -export type ManualResetCooldownClaim = - | { kind: "pool"; probe: CodexQuotaRecoveryProbeClaim } - | { kind: "main"; probe: CooldownRecoveryLease }; - -function manualResetAccountEligible(config: OcxConfig, accountId: string): boolean { - return !isCodexAccountPaused(config, accountId) && !isAccountNeedsReauth(accountId) - && (accountId === MAIN_CODEX_ACCOUNT_ID - || (config.codexAccounts ?? []).some(account => account.id === accountId && isSelectableCodexPoolAccount(account))); -} - -/** Explicit reset bypasses probe pacing, never another owner's lease or quota scope. */ -export function claimManualResetCooldowns( - config: OcxConfig, - accountId: string, - now = Date.now(), - expectedPoolGeneration?: number, -): ManualResetCooldownClaim[] { - if (!manualResetAccountEligible(config, accountId)) return []; - const record = accountId === MAIN_CODEX_ACCOUNT_ID ? undefined : readCodexAccountRecord(accountId); - if (accountId !== MAIN_CODEX_ACCOUNT_ID && (!record?.credential || record.deletedAt != null)) return []; - if (record && expectedPoolGeneration !== undefined && record.generation !== expectedPoolGeneration) return []; - const claims: ManualResetCooldownClaim[] = []; - for (const scope of [undefined, "shared"] as const) { - const health = scope ? scopedHealthFor(accountId, scope) : upstreamHealth.get(accountId); - if (!health || health.cooldownSource !== "reset-derived" || health.probeLeaseId !== undefined - || !Number.isFinite(health.cooldownUntil) || !(health.cooldownUntil! > now)) continue; - const leaseId = randomUUID(); - const cooldownGeneration = health.cooldownGeneration ?? 0; - const next = { ...health, probeLeaseId: leaseId, probeLeaseGeneration: cooldownGeneration, lastProbeAt: now }; - if (scope) setScopedHealth(accountId, scope, next); - else upstreamHealth.set(accountId, next); - const probe = { accountId, scope, leaseId, cooldownGeneration }; - claims.push(record ? { kind: "pool", probe: { - ...probe, credentialGeneration: record.generation, credentialReplacedAt: record.replacedAt, - } } : { kind: "main", probe }); - } - return claims; -} - -export type ManualResetRefreshLineage = Readonly<{ - fromGeneration: number; - toGeneration: number; - provenance: CodexRefreshProvenance; -}>; - -type ManualResetQuotaProof = CodexQuotaRecoveryProbeProof & { - refreshLineage?: ManualResetRefreshLineage; -}; - -/** Main proof is checked by the already-owned auth operation, never by a Pool record. */ -export function settleManualResetCooldown( - config: OcxConfig, - claim: ManualResetCooldownClaim, - recovered: boolean, - proof: ManualResetQuotaProof = {}, - now = Date.now(), -): boolean { - if (!recovered) return settleCooldownRecoveryLease(claim.probe, false, now); - const eligible = manualResetAccountEligible(config, claim.probe.accountId); - if (claim.kind === "main") return settleCooldownRecoveryLease(claim.probe, eligible, now); - const lineage = proof.refreshLineage; - // Equal wall-clock replacement stamps do not establish ancestry. Manual +1 - // recovery additionally needs the actual forced-refresh result for this edge. - const ownedGeneration = proof.credentialGeneration === claim.probe.credentialGeneration - || (proof.credentialGeneration === claim.probe.credentialGeneration + 1 - && lineage?.fromGeneration === claim.probe.credentialGeneration - && lineage.toGeneration === proof.credentialGeneration - && (lineage.provenance === "self-refresh" || lineage.provenance === "joined-lineage")); - return settleCodexQuotaRecoveryProbe(claim.probe, eligible && ownedGeneration, proof, now); -} - -/** Settle one background recovery claim without mutating account-wide outcome state. */ -export function settleCodexQuotaRecoveryProbe( - claim: CodexQuotaRecoveryProbeClaim, - recovered: boolean, - proof: CodexQuotaRecoveryProbeProof, - now = Date.now(), -): boolean { - const health = claim.scope - ? scopedHealthFor(claim.accountId, claim.scope) - : upstreamHealth.get(claim.accountId); - if (!health || health.probeLeaseId !== claim.leaseId) return false; - const currentRecord = readCodexAccountRecord(claim.accountId); - const proofGeneration = proof.credentialGeneration; - // A probe-owned token refresh (getValidCodexToken) advances the credential generation by - // exactly one while preserving `replacedAt`; an external credential replacement bumps the - // generation too but stamps a fresh `replacedAt`. Accept the +1 transition only when the - // claim-time lineage is intact AND the generation the fresh quota was proven under is live. - const generationFenced = proofGeneration !== undefined - && (proofGeneration === claim.credentialGeneration - ? isCodexAccountGenerationLive(claim.accountId, proofGeneration) - : proofGeneration === claim.credentialGeneration + 1 - && currentRecord?.replacedAt === claim.credentialReplacedAt - && isCodexAccountGenerationLive(claim.accountId, proofGeneration)); - return settleCooldownRecoveryLease(claim, recovered && generationFenced, now); -} - -function settleCooldownRecoveryLease(claim: CooldownRecoveryLease, recovered: boolean, now: number): boolean { - const health = claim.scope ? scopedHealthFor(claim.accountId, claim.scope) : upstreamHealth.get(claim.accountId); - if (!health || health.probeLeaseId !== claim.leaseId) return false; - const fenced = (claim.scope === undefined || claim.scope === "shared") - && health.cooldownSource === "reset-derived" - && (health.cooldownGeneration ?? 0) === claim.cooldownGeneration - && (health.probeLeaseGeneration ?? 0) === claim.cooldownGeneration; - if (!recovered || !fenced) { - const released = withProbeLeaseReleased(health, now); - if (claim.scope) setScopedHealth(claim.accountId, claim.scope, released); - else upstreamHealth.set(claim.accountId, released); - return false; - } - if (claim.scope) { - deleteScopedHealth(claim.accountId, claim.scope); - } else { - const { - cooldownUntil: _until, - cooldownSince: _since, - cooldownSource: _source, - probeLeaseId: _leaseId, - probeLeaseGeneration: _leaseGeneration, - // "The quota window moved" is a statement about the whole refusal, so the avoidance it - // announced goes with the block it produced. Leaving it would make this escape hatch stop - // escaping: the account would still be passed over by every selection it is meant to win. - quotaAvoidUntil: _avoid, - ...rest - } = health; - upstreamHealth.set(claim.accountId, { - ...rest, - cooldownGeneration: claim.cooldownGeneration + 1, - lastProbeAt: now, - }); - } - return true; -} - -/** Acquire the recovery probe for one confirmed model-specific quota group. */ -export function tryAcquireCodexQuotaScopeProbeLease( - accountId: string, - scope: CodexQuotaScope, - now = Date.now(), -): string | null { - const health = scopedHealthFor(accountId, scope); - if (!canAcquireQuotaProbeLease(health, now)) return null; - const probeLeaseId = randomUUID(); - setScopedHealth(accountId, scope, { - ...health!, - probeLeaseId, - probeLeaseGeneration: health!.cooldownGeneration ?? 0, - lastProbeAt: now, - }); - return probeLeaseId; -} - -/** Side-effect-free check for a confirmed model-specific quota probe. */ -export function canAcquireCodexQuotaScopeProbeLease( - accountId: string, - scope: CodexQuotaScope, - now = Date.now(), -): boolean { - return canAcquireQuotaProbeLease(scopedHealthFor(accountId, scope), now); -} - -/** - * Hand a probe lease back without recording an upstream outcome. Used by paths - * that take a lease and then fail before any request reaches upstream. - */ -export function releaseCodexQuotaProbeLease(accountId: string, leaseId: string, now = Date.now()): void { - const health = upstreamHealth.get(accountId); - if (!health || health.probeLeaseId !== leaseId) return; - upstreamHealth.set(accountId, withProbeLeaseReleased(health, now)); -} - -/** Release a model-specific quota probe when the request never reaches upstream. */ -export function releaseCodexQuotaScopeProbeLease( - accountId: string, - scope: CodexQuotaScope, - leaseId: string, - now = Date.now(), -): void { - const health = scopedHealthFor(accountId, scope); - if (!health || health.probeLeaseId !== leaseId) return; - setScopedHealth(accountId, scope, withProbeLeaseReleased(health, now)); -} - -/** - * True when this outcome belongs to the account's in-flight probe. The - * undefined-id guard matters: without it an outcome carrying no lease would match - * an account holding no lease and be mistaken for the probe owner. - */ -function ownsProbeLease(health: CodexUpstreamHealth | undefined, meta: CodexUpstreamOutcomeMeta): boolean { - return meta.probeLeaseId !== undefined && meta.probeLeaseId === health?.probeLeaseId; -} - -/** - * True when the owning probe may still clear the cooldown. A later 429 bumps the - * generation, so a probe that started under an older cooldown must not erase the - * newer restriction (which may carry an explicit Retry-After). - */ -function probeMayClearCooldown(health: CodexUpstreamHealth | undefined, meta: CodexUpstreamOutcomeMeta): boolean { - return ownsProbeLease(health, meta) - && (health!.probeLeaseGeneration ?? 0) === (health!.cooldownGeneration ?? 0); -} - -/** Strip the in-flight lease while preserving every hard-cooldown field. */ -function withProbeLeaseReleased(health: CodexUpstreamHealth, now: number): CodexUpstreamHealth { - const { probeLeaseId: _id, probeLeaseGeneration: _gen, ...rest } = health; - return { ...rest, lastProbeAt: now }; -} - -/** - * Hard-cooldown bookkeeping that ordinary success/transient transitions rebuild - * their health object from. Dropping these would let one late unrelated response - * erase a Retry-After source, a cooldown generation, or someone else's live probe. - */ -function preservedCooldownFields(health: CodexUpstreamHealth | undefined): Partial { - if (!health) return {}; - // `credentialFailureGeneration` is provenance for ONE credential failure, so it must not survive - // into a later transient or quota entry — otherwise that entry inherits the tag and gets spent - // when the old credential dies, deleting evidence that was never about it (#2892 gap 4 review). - const { - consecutiveFailures: _f, consecutiveSuccesses: _s, lastFailureStatus: _st, lastFailureAt: _at, - softAvoidUntil: _sa, credentialFailureGeneration: _cg, ...cooldownFields - } = health; - return cooldownFields; -} - -/** Manual selection resets transient routing evidence without bypassing a real 429 cooldown. */ -export function resetCodexRoutingForManualSelection(accountId: string): void { - clearThreadAccountMap(); - // Manual selection is the operator source of truth — drop any automatic runtime cursor. - runtimeActiveCodexAccountId = undefined; - // Record the pick as an unspent one-shot on the SHARED scope only. An independent scope - // gets no entry on purpose: every write site the guard protects is already skipped for - // independent scopes, so an entry there would be state nothing reads — and state nothing - // reads is what the next reader mistakes for a rule. - // - // Seeding happens ONLY here. A pool-driven promote must never create or move a preference, - // or the pool would manufacture an operator intent nobody expressed. - manualPreference.set(POOL_KEY_CODEX, accountId); - // Seed the RR ring so the next unbound new session honors the manually selected account - // under round-robin (affinity-cleared threads / null threadId). Fill-first already follows - // config.activeCodexAccountId, which the caller persists before invoking this. - seedPoolRotationAccount(POOL_KEY_CODEX, accountId); - for (const scope of new Set(Object.values(NATIVE_MODEL_QUOTA_SCOPES))) { - if (isIndependentCodexQuotaScope(scope)) { - seedPoolRotationAccount(codexPoolKeyForScope(scope), accountId); - } - } - // Quota avoidance is a preference, like the soft avoid dropped above, and an operator naming - // this account has overruled it. The hard cooldown is the part that survives. - const overrule = (health: CodexUpstreamHealth) => { - const { quotaAvoidUntil: _avoid, ...retained } = preservedCooldownFields(health); - return retained; - }; - const current = upstreamHealth.get(accountId); - if (current) { - const retained = overrule(current); - if (Object.keys(retained).length === 0) upstreamHealth.delete(accountId); - else upstreamHealth.set(accountId, { consecutiveFailures: 0, ...retained }); - } - // A reset-derived refusal records its avoidance on the SCOPED map and returns before the - // account-wide entry is written, so naming the account has to reach that map too. Stopping - // at `upstreamHealth` — and returning early when it holds nothing — overruled nothing in - // the case that produces the avoidance this function exists to overrule. - for (const [scope, health] of [...(quotaScopedHealth.get(accountId) ?? [])]) { - const retained = overrule(health); - if (Object.keys(retained).length === 0) deleteScopedHealth(accountId, scope); - else setScopedHealth(accountId, scope, { consecutiveFailures: 0, ...retained }); - } -} - -export function getCodexAccountCooldownUntil(accountId: string, now = Date.now()): number | null { - const cooldownUntil = upstreamHealth.get(accountId)?.cooldownUntil; - return typeof cooldownUntil === "number" && Number.isFinite(cooldownUntil) && cooldownUntil > now ? cooldownUntil : null; -} - -/** Read-only cooldown snapshot for shared OAuth health projection (no write side effects). */ -export function getCodexAccountHealthSnapshot(accountId: string, now = Date.now()): { - cooldownUntil?: number; - cooldownSource?: CodexCooldownSource; -} | null { - const cooldownUntil = getCodexAccountCooldownUntil(accountId, now); - if (cooldownUntil === null) return null; - const source = upstreamHealth.get(accountId)?.cooldownSource; - return { - cooldownUntil, - ...(source ? { cooldownSource: source } : {}), - }; -} - -/** - * Read the cooldown relevant to a routed native model. Account-wide cooldowns - * (Retry-After/default) always win; reset-derived scoped state applies only to - * its confirmed quota group. - */ -export function getCodexQuotaHealthSnapshot( - accountId: string, - quotaScope: CodexQuotaScope | undefined, - now = Date.now(), -): { - cooldownUntil?: number; - cooldownSource?: CodexCooldownSource; - quotaScope?: CodexQuotaScope; -} | null { - const account = getCodexAccountHealthSnapshot(accountId, now); - if (account) return account; - if (!quotaScope) return null; - const scoped = scopedHealthFor(accountId, quotaScope); - const cooldownUntil = scoped?.cooldownUntil; - if (typeof cooldownUntil !== "number" || !Number.isFinite(cooldownUntil) || cooldownUntil <= now) return null; - return { - cooldownUntil, - ...(scoped?.cooldownSource ? { cooldownSource: scoped.cooldownSource } : {}), - quotaScope, - }; -} - -export function isCodexAccountInCooldown(accountId: string, now = Date.now()): boolean { - return getCodexAccountCooldownUntil(accountId, now) !== null; -} - -/** - * Manually lift a hard quota cooldown without touching failure history. - * - * Injected Codex routing makes this proxy the ONLY model path for Codex Desktop, so a - * cooldown that outlives the real upstream limit reads to the user as "the whole app is - * broken" with no escape but editing config.toml. This is that escape hatch. - * - * Deliberately narrow: - * - Failure counters and softAvoid survive. Clearing a cooldown says "the quota window - * moved", not "this account is healthy"; failover must keep its knowledge. - * - Dropping `probeLeaseId` is what stops a stale in-flight probe from later "proving" - * recovery against a NEWER cooldown: {@link ownsProbeLease} needs the id to match. - * `cooldownGeneration` is preserved and bumped as redundancy only — a fresh 429 already - * bumps it in {@link recordCodexUpstreamOutcome}, so the bump here is not load-bearing - * today and is kept so the invariant survives a future change that retains the lease. - * - * Returns false when the account carried neither a live cooldown nor a live avoidance window. - * The window outlives the cooldown by design — the cooldown caps at fifteen minutes and the - * window runs up to six hours — so the moment an operator actually reaches for this escape - * hatch is usually after the cooldown lapsed and only the window is still keeping the account - * out of rotation. Refusing to look at the window then would leave the hatch shut in the one - * case it exists for. - */ -export function clearCodexAccountCooldown(accountId: string, now = Date.now()): boolean { - const clear = (health: CodexUpstreamHealth): CodexUpstreamHealth | null => { - const cooldownUntil = health.cooldownUntil; - const liveCooldown = typeof cooldownUntil === "number" && Number.isFinite(cooldownUntil) && cooldownUntil > now; - const avoidUntil = health.quotaAvoidUntil; - const liveAvoidance = typeof avoidUntil === "number" && Number.isFinite(avoidUntil) && avoidUntil > now; - if (!liveCooldown && !liveAvoidance) return null; - const { - cooldownUntil: _until, - cooldownSince: _since, - cooldownSource: _source, - probeLeaseId: _leaseId, - probeLeaseGeneration: _leaseGeneration, - // Same reasoning as the probe recovery above: "the quota window moved" is a statement - // about the whole refusal, so the avoidance it announced goes with the block it - // produced. Keeping it would leave this escape hatch not escaping, because selection - // would still pass over the account for as long as the announced window runs. - quotaAvoidUntil: _avoid, - ...rest - } = health; - return { - ...rest, - cooldownGeneration: (health.cooldownGeneration ?? 0) + 1, - lastProbeAt: now, - }; - }; - - let cleared = false; - const accountHealth = upstreamHealth.get(accountId); - if (accountHealth) { - const next = clear(accountHealth); - if (next) { - upstreamHealth.set(accountId, next); - cleared = true; - } - } - for (const [scope, health] of quotaScopedHealth.get(accountId) ?? []) { - const next = clear(health); - if (next) { - setScopedHealth(accountId, scope, next); - cleared = true; - } - } - return cleared; -} - -export function getCodexAccountSoftAvoidUntil(accountId: string, now = Date.now()): number | null { - const softAvoidUntil = upstreamHealth.get(accountId)?.softAvoidUntil; - return typeof softAvoidUntil === "number" && Number.isFinite(softAvoidUntil) && softAvoidUntil > now - ? softAvoidUntil - : null; -} - -export function isCodexAccountSoftAvoided(accountId: string, now = Date.now()): boolean { - return getCodexAccountSoftAvoidUntil(accountId, now) !== null; -} - -/** - * Plan keys the operator excluded from automatic rotation. Absent or empty means no policy, so an - * existing install rotates exactly as before. Compared with `codexPlanKey` because the stored plan - * is an unrestricted provider string whose casing this repository does not control. - */ -function excludedCodexPoolPlanKeys(config: OcxConfig): ReadonlySet | undefined { - const configured = config.codexPool?.excludedPlans; - if (!configured?.length) return undefined; - const keys = configured - .map(plan => codexPlanKey(plan)) - .filter((key): key is string => key !== undefined); - return keys.length > 0 ? new Set(keys) : undefined; -} - -/** - * Whether the operator's plan policy removes this account from automatic selection. - * - * Modelled on pause rather than usability: an excluded account keeps its credential, quota history, - * and affinity, stays visible on the account surface, and is still reachable by explicit account - * selection. Only automatic rotation skips it, which is the distinction #4211 asked for. - * - * It is checked in the same two places pause is checked, and that is not redundancy. The eligible - * list is consulted only when routing picks a NEW account; an already-active or already-affined - * account is served straight from {@link isCodexAccountSelectable}. A lapsed subscription leaves - * behind exactly that account, so a policy that filtered only the eligible list would miss the case - * it exists for. - * - * `__main__` is exempt. {@link getPoolAccountPlanForSelection} withholds the main plan during a - * selection-only drain so routing never reads the fenced native credential for it, so a rule that - * covered main would disagree with itself between drain and ordinary routing. - */ -export function isCodexAccountPlanExcluded( - config: OcxConfig, - accountId: string, - precomputed?: ReadonlySet, -): boolean { - if (accountId === MAIN_CODEX_ACCOUNT_ID) return false; - // Callers that test a whole list pass the set once rather than rebuilding it per row. - const excluded = precomputed ?? excludedCodexPoolPlanKeys(config); - if (!excluded) return false; - const plan = codexPlanKey(getPoolAccountPlan(config, accountId)); - return plan !== undefined && excluded.has(plan); -} - -function isCodexAccountSelectable( - config: OcxConfig, - accountId: string, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): boolean { - return !isCodexAccountPaused(config, accountId) - && !isCodexAccountPlanExcluded(config, accountId) - && getCodexQuotaHealthSnapshot(accountId, quotaScope, now) === null - && !isCodexQuotaAvoided(accountId, quotaScope, now) - && !isCodexAccountSoftAvoided(accountId, now) - && isCodexAccountUsable(config, accountId, selectionOptions); -} - -/** - * Which guard in {@link isCodexAccountSelectable} refused this account, if any. - * - * Deliberately the same predicates in the same order as that function, because the point is to - * REPORT the guard that actually fired rather than to re-derive a plausible-looking cause. An - * earlier version of the release reason checked only a subset and let a paused, plan-excluded, - * cooled-down or quota-avoided release fall through to a quota fallback, which named something - * routing never used -- a diagnostic that is confidently wrong in exactly the cases an operator - * would consult it for (#4598). - */ -function codexAccountBlockReason( - config: OcxConfig, - accountId: string, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): CodexAffinityReason | undefined { - if (isCodexAccountPaused(config, accountId)) return "paused"; - if (isCodexAccountPlanExcluded(config, accountId)) return "plan_excluded"; - if (getCodexQuotaHealthSnapshot(accountId, quotaScope, now) !== null) return "cooldown"; - if (isCodexQuotaAvoided(accountId, quotaScope, now)) return "quota_avoided"; - if (isCodexAccountSoftAvoided(accountId, now)) return "transient"; - if (!isCodexAccountUsable(config, accountId, selectionOptions)) return "unusable"; - return undefined; -} - -function threadAffinityScope(quotaScope?: CodexQuotaScope): BaseThreadAffinityScope { - return quotaScope ?? LEGACY_THREAD_AFFINITY_SCOPE; -} - -function admissibleAffinityComponent(value: string): boolean { - return retainedUtf8Bytes(value) <= MAX_AFFINITY_COMPONENT_BYTES; -} - -function modelDetourAffinityScope( - modelId: string | undefined, - quotaScope?: CodexQuotaScope, -): ModelDetourAffinityScope | undefined { - const canonicalModelId = modelId?.trim().toLowerCase(); - if (!canonicalModelId || !admissibleAffinityComponent(canonicalModelId)) return undefined; - return `model-detour:${threadAffinityScope(quotaScope)}:${canonicalModelId}`; -} - -function getThreadAffinityForScope( - threadId: string, - scope: ThreadAffinityScope, -): ThreadAffinityEntry | undefined { - if (!admissibleAffinityComponent(threadId)) return undefined; - return threadAccountMap.get(threadId)?.get(scope); -} - -function getThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): ThreadAffinityEntry | undefined { - return getThreadAffinityForScope(threadId, threadAffinityScope(quotaScope)); -} - -function getModelDetourAffinity( - threadId: string, - modelId: string | undefined, - quotaScope?: CodexQuotaScope, -): ThreadAffinityEntry | undefined { - const scope = modelDetourAffinityScope(modelId, quotaScope); - return scope ? getThreadAffinityForScope(threadId, scope) : undefined; -} - -function deleteThreadAffinityForScope(threadId: string, scope: ThreadAffinityScope): void { - if (!admissibleAffinityComponent(threadId)) return; - const affinities = threadAccountMap.get(threadId); - if (!affinities) return; - if (affinities.delete(scope)) { - threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); - } - if (affinities.size === 0) threadAccountMap.delete(threadId); -} - -function deleteThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): void { - deleteThreadAffinityForScope(threadId, threadAffinityScope(quotaScope)); -} - -function deleteModelDetourAffinity( - threadId: string, - modelId: string | undefined, - quotaScope?: CodexQuotaScope, -): void { - const scope = modelDetourAffinityScope(modelId, quotaScope); - if (scope) deleteThreadAffinityForScope(threadId, scope); -} - -/** Remove only the matching failed account's affinities for one thread. */ -function deleteThreadAffinitiesForAccount(threadId: string, accountId: string): void { - if (!admissibleAffinityComponent(threadId) || !admissibleAffinityComponent(accountId)) return; - const affinities = threadAccountMap.get(threadId); - if (!affinities) return; - for (const [scope, entry] of affinities) { - if (entry.accountId === accountId && affinities.delete(scope)) { - threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); - } - } - if (affinities.size === 0) threadAccountMap.delete(threadId); -} - -function threadAffinityEntryCount(): number { - return threadAffinityEntryTotal; -} - -function isThreadAffinityExpired(entry: ThreadAffinityEntry, now: number): boolean { - return now - entry.lastUsedAt > CODEX_THREAD_AFFINITY_IDLE_TTL_MS; -} - -function isThreadAffinityGenerationLive(entry: ThreadAffinityEntry): boolean { - if (entry.accountId === MAIN_CODEX_ACCOUNT_ID) return entry.generation === 0; - return isCodexAccountGenerationLive(entry.accountId, entry.generation); -} - -/** Generations this account's affinity entries are bound at. Test observability only. */ -export function debugCodexAffinityGenerations(accountId: string): number[] { - const generations: number[] = []; - for (const affinities of threadAccountMap.values()) { - for (const entry of affinities.values()) { - if (entry.accountId === accountId) generations.push(entry.generation); - } - } - return generations; -} - -/** - * Advance this account's affinity entries from the generation a rejected credential - * was bound under to the generation its own refresh produced. - * - * A 401 refresh-and-replay keeps the request on the same account, but the CAS write - * moves the credential from G to G+1, and {@link isThreadAffinityGenerationLive} - * demands exact equality — so without this the entry the replay just preserved is - * dead on the next request. Not quarantining an account is not the same as keeping - * its affinity. - * - * Lineage is proven by the CALLER, which must pass only a generation its own refresh - * produced. Re-deriving it here from `replacedAt` cannot work: the caller reads that - * field after the refresh and this function would re-read the same record, so the - * comparison is tautological and an external replacement passes it. An external - * replacement must retire the affinity, because that credential may belong to a - * different upstream identity. - */ -export function handOffThreadAffinityGeneration( - accountId: string, - fromGeneration: number, - toGeneration: number, -): boolean { - if (accountId === MAIN_CODEX_ACCOUNT_ID) return false; - if (toGeneration !== fromGeneration + 1) return false; - const record = readCodexAccountRecord(accountId); - if (!record?.credential || record.deletedAt != null) return false; - if (record.generation !== toGeneration) return false; - let handedOff = false; - for (const affinities of threadAccountMap.values()) { - for (const entry of affinities.values()) { - if (entry.accountId !== accountId || entry.generation !== fromGeneration) continue; - entry.generation = toGeneration; - handedOff = true; - } - } - return handedOff; -} - -function pruneExpiredThreadAffinities(now: number): void { - for (const [threadId, affinities] of threadAccountMap) { - for (const [scope, entry] of affinities) { - if (isThreadAffinityExpired(entry, now) && affinities.delete(scope)) { - threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); - } - } - if (affinities.size === 0) threadAccountMap.delete(threadId); - } -} - -function pruneLruThreadAffinities(): void { - if (threadAffinityEntryCount() <= CODEX_THREAD_AFFINITY_MAX_ENTRIES) return; - while (threadAffinityEntryCount() > CODEX_THREAD_AFFINITY_MAX_ENTRIES) { - let oldestThreadId: string | null = null; - let oldestScope: ThreadAffinityScope | null = null; - let oldestLastUsedAt = Number.POSITIVE_INFINITY; - let oldestIsDetour = false; - for (const [threadId, affinities] of threadAccountMap) { - for (const [scope, entry] of affinities) { - const candidateIsDetour = isModelDetourAffinityScope(scope); - if ( - (candidateIsDetour && !oldestIsDetour) - || (candidateIsDetour === oldestIsDetour && entry.lastUsedAt < oldestLastUsedAt) - ) { - oldestThreadId = threadId; - oldestScope = scope; - oldestLastUsedAt = entry.lastUsedAt; - oldestIsDetour = candidateIsDetour; - } - } - } - if (!oldestThreadId || !oldestScope) return; - deleteThreadAffinityForScope(oldestThreadId, oldestScope); - } -} - -function bindThreadAffinityForScope( - threadId: string, - accountId: string, - now: number, - scope: ThreadAffinityScope, -): void { - if (!admissibleAffinityComponent(threadId) || !admissibleAffinityComponent(accountId)) return; - const record = accountId === MAIN_CODEX_ACCOUNT_ID ? undefined : readCodexAccountRecord(accountId); - if (accountId !== MAIN_CODEX_ACCOUNT_ID && (!record?.credential || record.deletedAt != null)) return; - pruneExpiredThreadAffinities(now); - const affinities = threadAccountMap.get(threadId) ?? new Map(); - const previous = affinities.get(scope); - affinities.set(scope, { - accountId, - generation: accountId === MAIN_CODEX_ACCOUNT_ID ? 0 : record!.generation, - createdAt: previous?.createdAt ?? now, - lastUsedAt: now, - lastReevalAt: now, - }); - if (!previous) threadAffinityEntryTotal += 1; - threadAccountMap.set(threadId, affinities); - pruneLruThreadAffinities(); -} - -function bindThreadAffinity( - threadId: string, - accountId: string, - now: number, - quotaScope?: CodexQuotaScope, -): void { - bindThreadAffinityForScope(threadId, accountId, now, threadAffinityScope(quotaScope)); -} - -function bindModelDetourAffinity( - threadId: string, - accountId: string, - now: number, - modelId: string | undefined, - quotaScope?: CodexQuotaScope, -): void { - const scope = modelDetourAffinityScope(modelId, quotaScope); - if (scope) bindThreadAffinityForScope(threadId, accountId, now, scope); -} - -function getEligiblePoolAccounts( - config: OcxConfig, - excludeId?: string, - now = Date.now(), - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, - skipFailoverReadyCandidates = false, -): readonly string[] { - const excludedPlans = excludedCodexPoolPlanKeys(config); - const ids = (config.codexAccounts ?? []) - .filter(account => isSelectableCodexPoolAccount(account) - && account.id !== excludeId - && !isCodexAccountPaused(config, account.id) - && !isCodexAccountPlanExcluded(config, account.id, excludedPlans) - && !isAccountNeedsReauth(account.id) - && (!skipFailoverReadyCandidates || !shouldFailover(config, account.id, now))) - .filter(account => getCodexQuotaHealthSnapshot(account.id, quotaScope, now) === null) - .filter(account => !isCodexAccountSoftAvoided(account.id, now)) - .filter(account => !isCodexQuotaAvoided(account.id, quotaScope, now)) - .filter(account => isCodexAccountUsable(config, account.id, selectionOptions)) - .map(account => account.id); - // The main Codex account is not stored in config.codexAccounts; include it as a - // first-class rotation candidate when its read-only token is usable (Option A). - if ( - excludeId !== MAIN_CODEX_ACCOUNT_ID - && !isCodexAccountPaused(config, MAIN_CODEX_ACCOUNT_ID) - && (!isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID) || hasMainAccountRefreshGrant()) - && getCodexQuotaHealthSnapshot(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) === null - && !isCodexAccountSoftAvoided(MAIN_CODEX_ACCOUNT_ID, now) - // The main login is not in `config.codexAccounts`, so it never passes through the - // filters above and this is the only place an avoidance window can exclude it. Without - // this the window a refusal announced applies to the pool but not to the account that - // earned it: the cooldown caps at fifteen minutes, the window runs up to six hours, and - // in between the main account returns as a first-class candidate. - && !isCodexQuotaAvoided(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) - && (!skipFailoverReadyCandidates || !shouldFailover(config, MAIN_CODEX_ACCOUNT_ID, now)) - && isCodexAccountUsable(config, MAIN_CODEX_ACCOUNT_ID, selectionOptions) - ) { - ids.unshift(MAIN_CODEX_ACCOUNT_ID); - } - // Single choke point for selection order: every strategy, failover, and preview - // reaches the pool through here, so tiering applies once rather than per picker. - // Eligibility above is unchanged — this only narrows an already-eligible list. - return selectPriorityTier( - ids, - codexAccountPriorityLookup(config), - id => hasCodexQuotaHeadroom(config, id, selectionOptions, now), - pinnedCodexAccountId(config), - ); -} - -function listEligibleCodexAccountIds( - config: OcxConfig, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): readonly string[] { - return getEligiblePoolAccounts(config, undefined, now, quotaScope, selectionOptions); -} - -/** Shared reset timestamps are not evidence for independent model-quota groups. */ -function accountPoolStrategyForScope(config: OcxConfig, quotaScope?: CodexQuotaScope) { - const strategy = normalizeCodexAccountPoolStrategy(config.accountPoolStrategy); - return strategy === "reset-first" && isIndependentCodexQuotaScope(quotaScope) ? "quota" : strategy; -} - -function stickyLimitForConfig(config: OcxConfig): number { - return normalizeAccountPoolStickyLimit(config.accountPoolStickyLimit); -} - -/** - * Whether an account still has quota to give under the auto-switch threshold. - * - * Fill-first and the priority tier filter share this predicate, and share both of - * its escape hatches. A disabled threshold means only health, pause, and reauth - * may drain an account; unknown usage is a guess, so it must neither force - * fill-first off the active account nor drain a tier that was simply never - * primed. A genuinely exhausted account 429s into cooldown and leaves - * eligibility on its own. - */ -function hasCodexQuotaHeadroom( - config: OcxConfig, - accountId: string, - selectionOptions?: CodexAccountUsabilityOptions, - now: number = Date.now(), -): boolean { - const threshold = config.autoSwitchThreshold ?? 80; - if (threshold <= 0) return true; - const usage = computeCodexUsageScore( - getAccountQuota(accountId), - getPoolAccountPlanForSelection(config, accountId, selectionOptions), - now, - ); - if (isUnknownUsage(usage)) return true; - return usage < threshold; + // reset points. Leaving them behind lets a selection from one context suppress the + // automatic cursor in the next one. + clearAllManualPreferences(); + clearUpstreamHealthState(); + forgetRuntimeActiveCodexAccount(); + // The reconcile watermark is part of this state, not something that outlives it. Keeping + // it across a full reset is incoherent: there is no health left to protect, yet + // recordCodexUpstreamOutcome would still drop a writer whose generation predates the + // watermark for any account missing from the equally stale live set. Left behind, it also + // leaks between test files, which is how it was found. + resetHealthReconcileState(); } -/** - * Is a live binding held for its prompt cache? - * - * Unset means yes. Cache affinity shipped as an opt-in flag (#4292) and then #4546 measured - * what the default costs: a pool whose accounts all sit in the 80-99% band hands a bound - * conversation from account to account, and because provider prompt caches are account-isolated - * every hop re-sends the entire prefix. An install that has never heard of this flag is exactly - * the install that gets hurt by it, so the protection cannot be something you have to find. - * - * `false` restores capacity-first routing byte-for-byte. It is a real choice -- a pinned thread - * on a busy account pays latency -- and it stays available; it is just no longer the default. - */ -function isCacheAffinityEnabled(config: OcxConfig): boolean { - return config.pool?.cacheAffinity !== false; +export function clearCodexUpstreamHealthForAccount(accountId: string): void { + deleteAllHealthForAccount(accountId); + // Deletion is the third operator exit, next to pause and exclusion, and it is the one + // with no reconcile path behind it: once the account is gone nothing can succeed on it, + // so an unspent preference naming it would suppress the automatic cursor for every other + // account until the process restarts. + forgetManualPreference(accountId); } +export function reconcileCodexRoutingHealth(context: GenerationContext): number { + if (isHealthGenerationReconciled(context.generation)) return 0; + const removed = pruneHealthAccountsForContext(context.codexAccountIds); + // Sweep preferences the same way, for the account set this generation actually has. The + // delete path above is the direct route; this is the one that catches an account removed + // by an edit the runtime never saw. Deliberately not counted in `removed`, which reports + // health rows. + forgetRoutingPreferencesOutside(context.codexAccountIds); + commitHealthReconcile(context.generation, context.codexAccountIds); + return removed; +} /** * Is a transient failure streak the ONLY thing standing between this thread and its account? * @@ -1742,7 +281,7 @@ function isTransientHoldExpired(entry: ThreadAffinityEntry, now: number): boolea * that chance away. */ function isTransientHoldSpentForAccount(threadId: string, accountId: string, now: number): boolean { - const affinities = threadAccountMap.get(threadId); + const affinities = getThreadAffinityScopes(threadId); if (!affinities) return false; let matched = false; for (const entry of affinities.values()) { @@ -1784,431 +323,6 @@ function transientDetourAccount( : pickAlternateCodexAccount(config, entry.accountId, now, quotaScope, selectionOptions); } -/** Earliest future shared short/weekly reset; missing evidence and ties use usage order. */ -function pickResetFirstCodexAccount( - config: OcxConfig, - ids: readonly string[], - now: number, - selectionOptions?: CodexAccountUsabilityOptions, -): string | null { - const available = ids.filter(id => hasCodexQuotaHeadroom(config, id, selectionOptions, now)); - if (available.length === 0) return pickLowestUsageAmong(config, ids, selectionOptions, now); - let earliest = Number.POSITIVE_INFINITY; - let candidates: string[] = []; - for (const id of available) { - const quota = getAccountQuota(id); - const resets = [quota?.shortResetAt, quota?.weeklyResetAt] - .filter((reset): reset is number => typeof reset === "number" && Number.isFinite(reset)) - .map(resetAtToMs) - .filter(reset => reset > now); - const next = Math.min(...resets); - if (next < earliest) { - earliest = next; - candidates = [id]; - } else if (next === earliest) candidates.push(id); - } - return pickLowestUsageAmong(config, candidates, selectionOptions, now); -} - -/** - * Fill-first: keep selectable active under threshold; otherwise advance to the next - * eligible id in stable sorted order after the current active (wrapping). - */ -function pickFillFirstCodexAccount( - config: OcxConfig, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): string | null { - const eligible = listEligibleCodexAccountIds(config, now, quotaScope, selectionOptions); - if (eligible.length === 0) return null; - - const active = getEffectiveActiveCodexAccountId(config); - if (active && eligible.includes(active) && hasCodexQuotaHeadroom(config, active, selectionOptions, now)) { - return active; - } - - return pickNextFillFirstCodexAccount(config, active ?? null, eligible, now, selectionOptions); -} - -/** Next eligible account in stable order after `afterId` (wrapping). */ -function pickNextFillFirstCodexAccount( - config: OcxConfig, - afterId: string | null, - eligible: readonly string[] = listEligibleCodexAccountIds(config, Date.now()), - now = Date.now(), - selectionOptions?: CodexAccountUsabilityOptions, -): string | null { - if (eligible.length === 0) return null; - const ordered = [...eligible].sort((a, b) => a.localeCompare(b)); - if (!afterId) { - // Prefer an under-threshold account when starting with no active cursor. - for (const id of ordered) { - if (hasCodexQuotaHeadroom(config, id, selectionOptions, now)) return id; - } - return ordered[0] ?? null; - } - - const allConfigured = [ - ...(isCodexAccountUsable(config, MAIN_CODEX_ACCOUNT_ID, selectionOptions) || afterId === MAIN_CODEX_ACCOUNT_ID - ? [MAIN_CODEX_ACCOUNT_ID] - : []), - ...(config.codexAccounts ?? []).filter(account => !account.isMain).map(account => account.id), - ]; - const stableAll = [...new Set(allConfigured)].sort((a, b) => a.localeCompare(b)); - const startIdx = stableAll.indexOf(afterId); - if (startIdx < 0) { - for (const id of ordered) { - if (hasCodexQuotaHeadroom(config, id, selectionOptions, now)) return id; - } - return ordered[0] ?? null; - } - - // Skip successors that are also at/above threshold (known drained usage). - let fallback: string | null = null; - for (let step = 1; step <= stableAll.length; step++) { - const candidate = stableAll[(startIdx + step) % stableAll.length]!; - if (!eligible.includes(candidate)) continue; - if (!fallback) fallback = candidate; - if (hasCodexQuotaHeadroom(config, candidate, selectionOptions, now)) return candidate; - } - return fallback ?? ordered[0] ?? null; -} - -/** - * Unbound new-session pick for round-robin / fill-first. Returns null to fall through - * to the legacy quota path (or when the strategy is quota). - * - * When `commit` is true (resolve path), advances RR state. `commitSharedActive` - * and `commitAffinity` independently control the two cross-request side effects: - * model-scoped entitlement selection can bind a new task without replacing an - * existing task binding or global active choice. Preview remains a dry-run peek. - * - * Automatic strategy picks never sync-write config; only manual selection persists active. - * - * Known limitation (follow-up): when a subagent preview peeks an RR account and the request - * then falls back to a non-Codex provider, the ring is not reserved/committed. Prefer seeding - * the peeked account if that path becomes load-bearing. - */ -function pickUnboundStrategyAccount( - config: OcxConfig, - threadId: string | null, - now: number, - commit: boolean, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, - commitSharedActive = commit, - commitAffinity = commit, -): string | null { - const strategy = accountPoolStrategyForScope(config, quotaScope); - if (strategy === "quota") return null; - const poolKey = codexPoolKeyForScope(quotaScope); - - let picked: string | null = null; - if (strategy === "round-robin") { - const eligible = listEligibleCodexAccountIds(config, now, quotaScope, selectionOptions); - const limit = stickyLimitForConfig(config); - if (!commit) { - return peekRoundRobinAccount(poolKey, eligible, limit); - } - picked = pickRoundRobinAccount(poolKey, eligible, limit); - if (!picked) return null; - if (commitSharedActive) { - if (!isIndependentCodexQuotaScope(quotaScope) - && !manualPreferenceBlocks(codexPoolKeyForScope(quotaScope), picked)) { - rememberActiveCodexAccount(config, picked); - } - } - if (commitAffinity && threadId) bindThreadAffinity(threadId, picked, now, quotaScope); - notePoolRotationSuccess(poolKey, picked, limit); - return picked; - } - - if (strategy === "fill-first" || strategy === "reset-first") { - picked = strategy === "reset-first" - ? pickResetFirstCodexAccount(config, listEligibleCodexAccountIds(config, now, quotaScope, selectionOptions), now, selectionOptions) - : pickFillFirstCodexAccount(config, now, quotaScope, selectionOptions); - if (!picked) return null; - if (commitSharedActive) { - if (!isIndependentCodexQuotaScope(quotaScope) - && !manualPreferenceBlocks(codexPoolKeyForScope(quotaScope), picked)) { - rememberActiveCodexAccount(config, picked); - } - } - if (commitAffinity && threadId) bindThreadAffinity(threadId, picked, now, quotaScope); - return picked; - } - - return null; -} - -export function getPoolAccountPlan(config: OcxConfig, accountId: string): string | undefined { - if (accountId === MAIN_CODEX_ACCOUNT_ID) return getMainAccountPlan(); - return (config.codexAccounts ?? []) - .find(account => isSelectableCodexPoolAccount(account) && account.id === accountId)?.plan; -} - -/** Selection-only main routing must not lazily read the fenced native credential for its plan. */ -function getPoolAccountPlanForSelection( - config: OcxConfig, - accountId: string, - selectionOptions?: CodexAccountUsabilityOptions, -): string | undefined { - if (accountId === MAIN_CODEX_ACCOUNT_ID && selectionOptions?.nativeMainSelectionOnly === true) { - return undefined; - } - return getPoolAccountPlan(config, accountId); -} - -/** Shared routing state must ignore a request-scoped entitlement roster. */ -function sharedStateSelectionOptions( - selectionOptions?: CodexAccountUsabilityOptions, -): Pick< - CodexAccountUsabilityOptions, - "nativeMainSelectionOnly" | "isMainAccountTokenLive" -> | undefined { - if (!selectionOptions) return undefined; - return { - ...(selectionOptions.nativeMainSelectionOnly !== undefined - ? { nativeMainSelectionOnly: selectionOptions.nativeMainSelectionOnly } - : {}), - ...(selectionOptions.isMainAccountTokenLive - ? { isMainAccountTokenLive: selectionOptions.isMainAccountTokenLive } - : {}), - }; -} - -function pickLowerUsageAccount( - config: OcxConfig, - active: string, - activeUsage: number, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, - skipFailoverReadyCandidates = false, -): string { - let best = active; - let bestUsage = activeUsage; - for (const id of getEligiblePoolAccounts( - config, - active, - now, - quotaScope, - selectionOptions, - skipFailoverReadyCandidates, - )) { - const usage = computeCodexUsageScore( - getAccountQuota(id), - getPoolAccountPlanForSelection(config, id, selectionOptions), - now, - ); - if (usage < bestUsage) { - best = id; - bestUsage = usage; - } - } - return best; -} - -/** Coolest account in an already-selected candidate list; first index wins ties. */ -function pickLowestUsageAmong( - config: OcxConfig, - ids: readonly string[], - selectionOptions?: CodexAccountUsabilityOptions, - now: number = Date.now(), -): string | null { - let best: string | null = null; - let bestUsage = Number.POSITIVE_INFINITY; - for (const id of ids) { - const usage = computeCodexUsageScore( - getAccountQuota(id), - getPoolAccountPlanForSelection(config, id, selectionOptions), - now, - ); - if (usage < bestUsage) { - best = id; - bestUsage = usage; - } - } - return best; -} - -export function pickLowestUsageCodexAccount( - config: OcxConfig, - excludeId?: string, - now = Date.now(), - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): string | null { - return pickLowestUsageAmong( - config, - getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions), - selectionOptions, - now, - ); -} - -/** - * Strategy-aware alternate after a cooled/excluded account (same-request 429 retry - * and active promotion). Quota keeps lowest-usage; fill-first advances stable order; - * round-robin takes the next ring pick (caller should have noted the failure). - */ -export function pickAlternateCodexAccount( - config: OcxConfig, - excludeId: string, - now = Date.now(), - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): string | null { - const strategy = accountPoolStrategyForScope(config, quotaScope); - // The exclusion is passed into eligibility rather than post-filtered off its - // result: when the excluded account is the only healthy member of the top - // tier, the tier walk must be free to descend instead of selecting that tier - // and then handing back an empty list. - if (strategy === "round-robin") { - const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); - return pickRoundRobinAccount(codexPoolKeyForScope(quotaScope), eligible, stickyLimitForConfig(config)); - } - if (strategy === "fill-first") { - const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); - return pickNextFillFirstCodexAccount(config, excludeId, eligible, now, selectionOptions); - } - if (strategy === "reset-first") { - return pickResetFirstCodexAccount(config, getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions), now, selectionOptions); - } - return pickLowestUsageCodexAccount(config, excludeId, now, quotaScope, selectionOptions); -} - -/** - * The account {@link pickAlternateCodexAccount} WOULD return, without returning it. - * - * Only the round-robin branch has a side effect -- `pickRoundRobinAccount` commits the pick and - * advances the ring -- so every other strategy delegates rather than growing a second copy of - * the selection rule that could drift from it. - * - * This exists because preview and resolve have to agree on the FIRST transient detour, not just - * on later ones. Preview feeds subagent model-availability scoring, so a preview that reported - * the bound account while resolve was about to serve from a cool sibling could retire a model - * over usage the request would never have touched. - */ -function peekAlternateCodexAccount( - config: OcxConfig, - excludeId: string, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): string | null { - if (accountPoolStrategyForScope(config, quotaScope) === "round-robin") { - const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); - return peekRoundRobinAccount(codexPoolKeyForScope(quotaScope), eligible, stickyLimitForConfig(config)); - } - return pickAlternateCodexAccount(config, excludeId, now, quotaScope, selectionOptions); -} - -/** Effective active: automatic runtime cursor, else operator/persisted selection. */ -/** - * Unspent operator selections, keyed by pool scope. - * - * Codex has no account-side equivalent of the Anthropic `selectionRevision`, so staleness - * cannot be detected by comparing values: a pool-driven promote legitimately moves the - * persisted active account, and reading that as staleness would silently spend the - * operator's one-shot. Invalidation is keyed to the OPERATOR path instead — another manual - * selection, the account leaving the pool, or a successful dispatch on it. - */ -const manualPreference = new Map(); - -/** - * Spend the one-shot for a pool scope once a dispatch on that account actually succeeded. - * This is the Codex analogue of `commitAnthropicSelectionRouting`, which Codex lacks. - * - * Wiring this BEFORE the guard below is not a style choice. Measured: with the guard in - * place and no consume site, the first manual selection freezes the automatic cursor - * permanently and 15 of 69 rotation tests fail. - */ -function consumeManualPreference(accountId: string, poolKey: string): void { - if (manualPreference.get(poolKey) === accountId) manualPreference.delete(poolKey); -} - -/** - * Drop an account's preference in every scope. Pause and exclusion do not route through - * `resetCodexRoutingForManualSelection`, so without this a preference could outlive the - * account it names and keep suppressing the automatic cursor. - */ -function forgetManualPreference(accountId: string): void { - for (const [poolKey, preferred] of manualPreference) { - if (preferred === accountId) manualPreference.delete(poolKey); - } -} - -/** - * True while an unspent operator selection for this scope names a DIFFERENT account than - * the automatic pick about to be recorded. - * - * Callers pass their own scope: an independent quota scope keeps its own entry and must - * never read the shared one. The failover promote does NOT consult this — see its call - * site for why. - */ -function manualPreferenceBlocks(poolKey: string, accountId: string): boolean { - const preferred = manualPreference.get(poolKey); - return preferred !== undefined && preferred !== accountId; -} - -export function getEffectiveActiveCodexAccountId(config: OcxConfig): string | undefined { - return runtimeActiveCodexAccountId ?? config.activeCodexAccountId; -} - -/** - * Whether the account routing is currently on is there because an operator asked - * for it, rather than because a strategy landed on it. Surfaces read this instead - * of comparing the stored pin themselves, which would report a pin that a later - * automatic pick has already moved past. - */ -export function isEffectiveCodexAccountPinned(config: OcxConfig): boolean { - const pinned = pinnedCodexAccountId(config); - return pinned !== undefined && pinned === getEffectiveActiveCodexAccountId(config); -} - -/** - * Automatic strategy / failover cursor only — never mutates `config.activeCodexAccountId` - * so an unrelated `saveConfig` cannot persist transient rotation as operator selection. - */ -function rememberActiveCodexAccount(_config: OcxConfig, accountId: string): void { - runtimeActiveCodexAccountId = accountId; -} - -/** - * End the manual pin when routing moves to a different account. Returns whether - * the pin changed so the caller can fold it into a write it was already making. - */ -function releaseCodexAccountPinFor(config: OcxConfig, accountId: string): boolean { - const pinned = pinnedCodexAccountId(config); - if (pinned === undefined || pinned === accountId) return false; - clearCodexAccountPin(config); - return true; -} - -/** Persist operator (or quota-strategy) active selection to config + disk. */ -function setActiveCodexAccount(config: OcxConfig, accountId: string): void { - runtimeActiveCodexAccountId = undefined; - const releasedPin = releaseCodexAccountPinFor(config, accountId); - if (config.activeCodexAccountId === accountId && !releasedPin) return; - config.activeCodexAccountId = accountId; - saveConfigPreservingClaudeCode(config); -} - -/** Quota strategy persists; RR/fill-first keep a process-local cursor only. */ -function promoteActiveCodexAccount(config: OcxConfig, accountId: string): void { - if (normalizeCodexAccountPoolStrategy(config.accountPoolStrategy) === "quota") { - setActiveCodexAccount(config, accountId); - return; - } - // Runtime-only, like the cursor itself: a caller that persists (pause, delete) - // saves this release with its own write; a transient failover does not, so the - // pin survives a restart that also clears the failure history behind it. - releaseCodexAccountPinFor(config, accountId); - rememberActiveCodexAccount(config, accountId); -} - /** * Reconcile the effective active account after an administrative exclusion such as pause. * The operator's persisted selection is cleared when it names the excluded account; quota @@ -2234,54 +348,12 @@ export function reconcileCodexActiveAfterExclusion( clearCodexAccountPin(config, excludedAccountId); if (!wasEffective) return getEffectiveActiveCodexAccountId(config) ?? null; - runtimeActiveCodexAccountId = undefined; + forgetRuntimeActiveCodexAccount(); const fallback = pickAlternateCodexAccount(config, excludedAccountId, now); if (fallback) promoteActiveCodexAccount(config, fallback); return fallback; } -function isUnknownUsage(usage: number): boolean { - return usage >= CODEX_UNKNOWN_USAGE_SCORE; -} - -/** - * Move an unbound request back up when a higher tier regains headroom — the - * weekly-reset case. Returns null when nothing should change. - * - * Downward moves are deliberately left to {@link applyQuotaAutoSwitch}: this only - * fires when the tier filter has already excluded `active`, and only toward a - * tier that strictly outranks it. Threads bound by affinity never reach here. - */ -function pickPriorityPreemption( - config: OcxConfig, - active: string, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, -): string | null { - const eligible = getEligiblePoolAccounts(config, undefined, now, quotaScope, selectionOptions); - if (eligible.length === 0 || eligible.includes(active)) return null; - const pinned = pinnedCodexAccountId(config); - // A live pin already lowered the tier ceiling; never preempt past an explicit - // operator choice. Same liveness test the tier filter applies, so preview and - // resolve agree even before the pin is garbage-collected. - if ( - pinned !== undefined - && eligible.includes(pinned) - && hasCodexQuotaHeadroom(config, pinned, selectionOptions, now) - ) return null; - const priorityOf = codexAccountPriorityLookup(config); - if (priorityOf(eligible[0]!) <= priorityOf(active)) return null; - // Members without headroom are in the tier only because a sibling has some; - // picking one would hand the request straight back to a drained account. - return pickLowestUsageAmong( - config, - eligible.filter(id => hasCodexQuotaHeadroom(config, id, selectionOptions, now)), - selectionOptions, - now, - ); -} - /** * Release a pin whose account is durably drained. "Use this account now" ends * when the account crosses the auto-switch threshold or stops being selectable @@ -2316,107 +388,6 @@ function releaseDrainedCodexAccountPin( saveConfigPreservingClaudeCode(config); } -function applyQuotaAutoSwitch( - config: OcxConfig, - active: string, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, - commitSharedSelection = true, -): string { - const threshold = config.autoSwitchThreshold ?? 80; - if (threshold <= 0) return active; - const quota = getAccountQuota(active); - const activeUsage = computeCodexUsageScore( - quota, - getPoolAccountPlanForSelection(config, active, selectionOptions), - now, - ); - // Unknown usage is not evidence that a user's explicit selection crossed the - // threshold. Wait for quota priming instead of rotating among guesses. - if (isUnknownUsage(activeUsage)) return active; - if (activeUsage < threshold) return active; - const best = pickLowerUsageAccount(config, active, activeUsage, now, quotaScope, selectionOptions); - if (best !== active) { - if (commitSharedSelection && !isIndependentCodexQuotaScope(quotaScope)) { - setActiveCodexAccount(config, best); - } - return best; - } - - return active; -} - -function shouldFailover(config: OcxConfig, accountId: string, now: number): boolean { - const threshold = config.upstreamFailoverThreshold ?? 3; - if (threshold <= 0) return false; - dropSpentCredentialFailure(accountId); - const health = upstreamHealth.get(accountId); - if (health?.lastFailureAt && now - health.lastFailureAt > CODEX_FAILURE_WINDOW_MS) return false; - return !!health && health.consecutiveFailures >= threshold; -} - -function isHealthySharedCodexSelection( - config: OcxConfig, - accountId: string, - now: number, - quotaScope: CodexQuotaScope | undefined, - selectionOptions: CodexAccountUsabilityOptions | undefined, -): boolean { - return isCodexAccountSelectable(config, accountId, now, quotaScope, selectionOptions) - && hasCodexQuotaHeadroom(config, accountId, selectionOptions, now) - && !shouldFailover(config, accountId, now); -} - -function strategySelectionOptionsForModelDetour( - config: OcxConfig, - now: number, - quotaScope: CodexQuotaScope | undefined, - selectionOptions: CodexAccountUsabilityOptions | undefined, -): CodexAccountUsabilityOptions | undefined { - if (selectionOptions?.modelEligibleAccountIds === undefined) return selectionOptions; - const sharedSelectionOptions = sharedStateSelectionOptions(selectionOptions) ?? {}; - return { - ...selectionOptions, - modelEligibleAccountIds: new Set( - [...selectionOptions.modelEligibleAccountIds].filter(accountId => - isHealthySharedCodexSelection( - config, - accountId, - now, - quotaScope, - sharedSelectionOptions, - ) - ), - ), - }; -} - -function applyFailureFailover( - config: OcxConfig, - active: string, - now: number, - quotaScope?: CodexQuotaScope, - selectionOptions?: CodexAccountUsabilityOptions, - commitSharedSelection = true, -): string { - if (!shouldFailover(config, active, now)) return active; - const best = pickAlternateCodexAccount(config, active, now, quotaScope, selectionOptions); - if (best) { - // The scope still routes away from the failing account — that is this request's - // own decision — but an independent one must not persist a new shared active - // account. recordCodexUpstreamOutcome only suppresses the promotion it makes at - // the moment of the failure; the streak outlives the soft avoid, so a later - // scoped resolve reaches here with the streak still tripped and would otherwise - // move the shared cursor after all. - if (commitSharedSelection && !isIndependentCodexQuotaScope(quotaScope)) { - promoteActiveCodexAccount(config, best); - } - return best; - } - return active; -} - export function resolveCodexAccountForThread( threadId: string | null, config: OcxConfig, @@ -2462,7 +433,7 @@ function carriesQuotaRefusal(health: CodexUpstreamHealth | undefined): boolean { * quota group, so a spent Spark window still cannot displace the same thread's Terra binding. */ function hasUnrecoveredCodexQuotaRefusal(accountId: string, quotaScope?: CodexQuotaScope): boolean { - if (carriesQuotaRefusal(upstreamHealth.get(accountId))) return true; + if (carriesQuotaRefusal(getAccountHealth(accountId))) return true; return quotaScope !== undefined && carriesQuotaRefusal(scopedHealthFor(accountId, quotaScope)); } @@ -3144,6 +1115,7 @@ export function resolveCodexAccountForThreadDetailed( return { status: "selected", accountId: active, affinity: affinityAfterRelease(threadId, releaseReason) }; } + export function recordCodexUpstreamOutcome( config: OcxConfig, accountId: string | null, @@ -3159,7 +1131,7 @@ export function recordCodexUpstreamOutcome( } if (!accountId) return; const writerGeneration = meta.writerGeneration ?? captureConfigGeneration(); - if (writerGeneration < lastReconciledGeneration && !liveHealthAccountIds.has(accountId)) return; + if (!isHealthAccountAdmissible(accountId, writerGeneration)) return; const now = meta.now ?? Date.now(); const outcomeClass = classifyCodexUpstreamOutcome(outcome, meta.denial); // Reject retired quota evidence before stale-credential cleanup or any shared mutation. @@ -3204,12 +1176,12 @@ export function recordCodexUpstreamOutcome( if (Object.keys(retained).length > 1) setScopedHealth(accountId, quotaScope, retained); else deleteScopedHealth(accountId, quotaScope); } - const current = upstreamHealth.get(accountId); + const current = getAccountHealth(accountId); const cooldownUntil = getCodexAccountCooldownUntil(accountId, now); // A leased probe that is still on its own cooldown generation proves the // account recovered: clear the hard cooldown outright (#433). if (cooldownUntil && probeMayClearCooldown(current, meta)) { - upstreamHealth.delete(accountId); + deleteAccountHealth(accountId); return; } // Owning probe on a stale generation: the lease is done, but a newer 429 @@ -3221,7 +1193,7 @@ export function recordCodexUpstreamOutcome( if (failoverEnabled && current && current.consecutiveFailures >= 2) { const consecutiveSuccesses = (current.consecutiveSuccesses ?? 0) + 1; if (consecutiveSuccesses < 2) { - upstreamHealth.set(accountId, { + setAccountHealth(accountId, { ...base!, ...preserved, consecutiveSuccesses, @@ -3231,14 +1203,14 @@ export function recordCodexUpstreamOutcome( } // Level 1 clears immediately; escalated accounts need two consecutive healthy terminals. // Hard quota cooldown intentionally survives either recovery path. - if (cooldownUntil) upstreamHealth.set(accountId, { consecutiveFailures: 0, ...preserved }); - else upstreamHealth.delete(accountId); + if (cooldownUntil) setAccountHealth(accountId, { consecutiveFailures: 0, ...preserved }); + else deleteAccountHealth(accountId); return; } if (outcomeClass === "caller") { // A 4xx does not change account health, but it does conclude an in-flight // probe — otherwise the lease would never be handed back. - const current = upstreamHealth.get(accountId); + const current = getAccountHealth(accountId); const scopedProbe = meta.probeQuotaScope ? scopedHealthFor(accountId, meta.probeQuotaScope) : undefined; @@ -3246,7 +1218,7 @@ export function recordCodexUpstreamOutcome( setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); } if (ownsProbeLease(current, meta)) { - upstreamHealth.set(accountId, withProbeLeaseReleased(current!, now)); + setAccountHealth(accountId, withProbeLeaseReleased(current!, now)); } return; } @@ -3257,7 +1229,7 @@ export function recordCodexUpstreamOutcome( // it and must not happen (#914). Conclude any owned probe lease, record the // failure under the (provider, host) ledger when one is named, and leave // account health, thread affinity, and the active account untouched. - const current = upstreamHealth.get(accountId); + const current = getAccountHealth(accountId); const scopedProbe = meta.probeQuotaScope ? scopedHealthFor(accountId, meta.probeQuotaScope) : undefined; @@ -3265,7 +1237,7 @@ export function recordCodexUpstreamOutcome( setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); } if (ownsProbeLease(current, meta)) { - upstreamHealth.set(accountId, withProbeLeaseReleased(current!, now)); + setAccountHealth(accountId, withProbeLeaseReleased(current!, now)); } return; } @@ -3276,8 +1248,8 @@ export function recordCodexUpstreamOutcome( // Record the failure so routing stops preferring it, but do not mark it for // reauthentication and do not sweep its thread affinities: telling the user to // re-login is wrong advice that cannot fix a workspace grant. - upstreamHealth.set(accountId, { - consecutiveFailures: (upstreamHealth.get(accountId)?.consecutiveFailures ?? 0) + 1, + setAccountHealth(accountId, { + consecutiveFailures: (getAccountHealth(accountId)?.consecutiveFailures ?? 0) + 1, lastFailureStatus, lastFailureAt: now, }); @@ -3315,7 +1287,7 @@ export function recordCodexUpstreamOutcome( * Affinity sweeping needs no tag: an affinity entry already carries a credential generation and * self-invalidates on the next check, and re-adding swept entries would be a worse bug. */ - upstreamHealth.set(accountId, { + setAccountHealth(accountId, { consecutiveFailures: 1, lastFailureStatus, lastFailureAt: now, @@ -3324,7 +1296,7 @@ export function recordCodexUpstreamOutcome( ? { credentialFailureGeneration: meta.credentialGeneration } : {}), }); - quotaScopedHealth.delete(accountId); + deleteAllScopedHealth(accountId); // The reauth flag carries the same provenance, so a replacement landing after this call cannot // inherit a quarantine that was never about it. markAccountNeedsReauth(accountId, writerGeneration, meta.credentialGeneration); @@ -3385,13 +1357,13 @@ export function recordCodexUpstreamOutcome( if (scopedProbe && meta.probeQuotaScope && ownsProbeLease(scopedProbe, meta)) { setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); } - const prior = upstreamHealth.get(accountId); + const prior = getAccountHealth(accountId); // Every cooldown write bumps the generation so a probe issued against the // previous cooldown can no longer clear this one (#433). const cooldownGeneration = (prior?.cooldownGeneration ?? 0) + 1; // A failed probe concludes its lease; an unrelated 429 leaves the live probe alone. const ownsLease = ownsProbeLease(prior, meta); - upstreamHealth.set(accountId, { + setAccountHealth(accountId, { consecutiveFailures: 0, lastFailureStatus, lastFailureAt: now, @@ -3431,7 +1403,7 @@ export function recordCodexUpstreamOutcome( } // transient (connect_error / timeout / 5xx) - const current = upstreamHealth.get(accountId); + const current = getAccountHealth(accountId); const scopedProbe = meta.probeQuotaScope ? scopedHealthFor(accountId, meta.probeQuotaScope) : undefined; @@ -3457,7 +1429,7 @@ export function recordCodexUpstreamOutcome( now + escalationMs, ) : undefined; - upstreamHealth.set(accountId, { + setAccountHealth(accountId, { ...preservedCooldownFields(transientBase), consecutiveFailures, lastFailureStatus, diff --git a/src/codex/routing/active-account.ts b/src/codex/routing/active-account.ts new file mode 100644 index 0000000000..e53ef06b26 --- /dev/null +++ b/src/codex/routing/active-account.ts @@ -0,0 +1,194 @@ +import { saveConfigPreservingClaudeCode } from "../../config"; +import { clearCodexAccountPin, pinnedCodexAccountId } from "../account-priority"; +import { + POOL_KEY_CODEX, + normalizeCodexAccountPoolStrategy, + seedPoolRotationAccount, +} from "../pool-rotation"; +import type { OcxConfig } from "../../types"; +import { clearThreadAccountMap } from "./thread-affinity"; +import { + NATIVE_MODEL_QUOTA_SCOPES, + codexPoolKeyForScope, + deleteAccountHealth, + deleteScopedHealth, + getAccountHealth, + isIndependentCodexQuotaScope, + listScopedHealthEntries, + preservedCooldownFields, + setAccountHealth, + setScopedHealth, + type CodexUpstreamHealth, +} from "./health-store"; + +/** + * Process-local cursor for automatic RR/fill-first (and quota-429 when not + * sync-writing) picks. Keeps unrelated `saveConfig` from persisting transient + * rotation as the operator's `activeCodexAccountId`. Manual selection clears it + * so disk/`config.activeCodexAccountId` remains authoritative. + */ +let runtimeActiveCodexAccountId: string | undefined; + +/** Manual selection resets transient routing evidence without bypassing a real 429 cooldown. */ +export function resetCodexRoutingForManualSelection(accountId: string): void { + clearThreadAccountMap(); + // Manual selection is the operator source of truth — drop any automatic runtime cursor. + runtimeActiveCodexAccountId = undefined; + // Record the pick as an unspent one-shot on the SHARED scope only. An independent scope + // gets no entry on purpose: every write site the guard protects is already skipped for + // independent scopes, so an entry there would be state nothing reads — and state nothing + // reads is what the next reader mistakes for a rule. + // + // Seeding happens ONLY here. A pool-driven promote must never create or move a preference, + // or the pool would manufacture an operator intent nobody expressed. + manualPreference.set(POOL_KEY_CODEX, accountId); + // Seed the RR ring so the next unbound new session honors the manually selected account + // under round-robin (affinity-cleared threads / null threadId). Fill-first already follows + // config.activeCodexAccountId, which the caller persists before invoking this. + seedPoolRotationAccount(POOL_KEY_CODEX, accountId); + for (const scope of new Set(Object.values(NATIVE_MODEL_QUOTA_SCOPES))) { + if (isIndependentCodexQuotaScope(scope)) { + seedPoolRotationAccount(codexPoolKeyForScope(scope), accountId); + } + } + // Quota avoidance is a preference, like the soft avoid dropped above, and an operator naming + // this account has overruled it. The hard cooldown is the part that survives. + const overrule = (health: CodexUpstreamHealth) => { + const { quotaAvoidUntil: _avoid, ...retained } = preservedCooldownFields(health); + return retained; + }; + const current = getAccountHealth(accountId); + if (current) { + const retained = overrule(current); + if (Object.keys(retained).length === 0) deleteAccountHealth(accountId); + else setAccountHealth(accountId, { consecutiveFailures: 0, ...retained }); + } + // A reset-derived refusal records its avoidance on the SCOPED map and returns before the + // account-wide entry is written, so naming the account has to reach that map too. Stopping + // at `upstreamHealth` — and returning early when it holds nothing — overruled nothing in + // the case that produces the avoidance this function exists to overrule. + for (const [scope, health] of [...(listScopedHealthEntries(accountId))]) { + const retained = overrule(health); + if (Object.keys(retained).length === 0) deleteScopedHealth(accountId, scope); + else setScopedHealth(accountId, scope, { consecutiveFailures: 0, ...retained }); + } +} + +/** Effective active: automatic runtime cursor, else operator/persisted selection. */ +/** + * Unspent operator selections, keyed by pool scope. + * + * Codex has no account-side equivalent of the Anthropic `selectionRevision`, so staleness + * cannot be detected by comparing values: a pool-driven promote legitimately moves the + * persisted active account, and reading that as staleness would silently spend the + * operator's one-shot. Invalidation is keyed to the OPERATOR path instead — another manual + * selection, the account leaving the pool, or a successful dispatch on it. + */ +const manualPreference = new Map(); + +/** + * Spend the one-shot for a pool scope once a dispatch on that account actually succeeded. + * This is the Codex analogue of `commitAnthropicSelectionRouting`, which Codex lacks. + * + * Wiring this BEFORE the guard below is not a style choice. Measured: with the guard in + * place and no consume site, the first manual selection freezes the automatic cursor + * permanently and 15 of 69 rotation tests fail. + */ +export function consumeManualPreference(accountId: string, poolKey: string): void { + if (manualPreference.get(poolKey) === accountId) manualPreference.delete(poolKey); +} + +/** + * Drop an account's preference in every scope. Pause and exclusion do not route through + * `resetCodexRoutingForManualSelection`, so without this a preference could outlive the + * account it names and keep suppressing the automatic cursor. + */ +export function forgetManualPreference(accountId: string): void { + for (const [poolKey, preferred] of manualPreference) { + if (preferred === accountId) manualPreference.delete(poolKey); + } +} + +/** + * True while an unspent operator selection for this scope names a DIFFERENT account than + * the automatic pick about to be recorded. + * + * Callers pass their own scope: an independent quota scope keeps its own entry and must + * never read the shared one. The failover promote does NOT consult this — see its call + * site for why. + */ +export function manualPreferenceBlocks(poolKey: string, accountId: string): boolean { + const preferred = manualPreference.get(poolKey); + return preferred !== undefined && preferred !== accountId; +} + +export function getEffectiveActiveCodexAccountId(config: OcxConfig): string | undefined { + return runtimeActiveCodexAccountId ?? config.activeCodexAccountId; +} + +/** + * Whether the account routing is currently on is there because an operator asked + * for it, rather than because a strategy landed on it. Surfaces read this instead + * of comparing the stored pin themselves, which would report a pin that a later + * automatic pick has already moved past. + */ +export function isEffectiveCodexAccountPinned(config: OcxConfig): boolean { + const pinned = pinnedCodexAccountId(config); + return pinned !== undefined && pinned === getEffectiveActiveCodexAccountId(config); +} + +/** + * Automatic strategy / failover cursor only — never mutates `config.activeCodexAccountId` + * so an unrelated `saveConfig` cannot persist transient rotation as operator selection. + */ +export function rememberActiveCodexAccount(_config: OcxConfig, accountId: string): void { + runtimeActiveCodexAccountId = accountId; +} + +/** + * End the manual pin when routing moves to a different account. Returns whether + * the pin changed so the caller can fold it into a write it was already making. + */ +function releaseCodexAccountPinFor(config: OcxConfig, accountId: string): boolean { + const pinned = pinnedCodexAccountId(config); + if (pinned === undefined || pinned === accountId) return false; + clearCodexAccountPin(config); + return true; +} + +/** Persist operator (or quota-strategy) active selection to config + disk. */ +export function setActiveCodexAccount(config: OcxConfig, accountId: string): void { + runtimeActiveCodexAccountId = undefined; + const releasedPin = releaseCodexAccountPinFor(config, accountId); + if (config.activeCodexAccountId === accountId && !releasedPin) return; + config.activeCodexAccountId = accountId; + saveConfigPreservingClaudeCode(config); +} + +/** Quota strategy persists; RR/fill-first keep a process-local cursor only. */ +export function promoteActiveCodexAccount(config: OcxConfig, accountId: string): void { + if (normalizeCodexAccountPoolStrategy(config.accountPoolStrategy) === "quota") { + setActiveCodexAccount(config, accountId); + return; + } + // Runtime-only, like the cursor itself: a caller that persists (pause, delete) + // saves this release with its own write; a transient failover does not, so the + // pin survives a restart that also clears the failure history behind it. + releaseCodexAccountPinFor(config, accountId); + rememberActiveCodexAccount(config, accountId); +} + +export function clearAllManualPreferences(): void { + manualPreference.clear(); +} + +export function forgetRuntimeActiveCodexAccount(): void { + runtimeActiveCodexAccountId = undefined; +} + +export function forgetRoutingPreferencesOutside(codexAccountIds: ReadonlySet): void { + for (const [poolKey, preferred] of manualPreference) { + if (codexAccountIds.has(preferred)) continue; + manualPreference.delete(poolKey); + } +} diff --git a/src/codex/routing/cooldown-math.ts b/src/codex/routing/cooldown-math.ts new file mode 100644 index 0000000000..123da5e3b1 --- /dev/null +++ b/src/codex/routing/cooldown-math.ts @@ -0,0 +1,275 @@ +import { + CODEX_EXHAUSTED_USAGE_PERCENT, + CODEX_UNKNOWN_USAGE_SCORE, + resetAtToMs, +} from "../quota"; +import { isThirtyDayOnlyCodexPlan } from "../plan"; +import type { CodexQuotaScope } from "./health-store"; + +export const CODEX_DEFAULT_QUOTA_COOLDOWN_MS = 60_000; +export const CODEX_MAX_QUOTA_COOLDOWN_MS = 24 * 60 * 60_000; +/** + * A weekly/monthly quota `resetAt` announces when the window refreshes; it is not + * a "come back after this" directive like Retry-After. Plan quota routinely frees + * up long before the advertised reset, so cap reset-derived cooldowns far below + * the Retry-After ceiling (#433). + */ +export const CODEX_MAX_RESET_DERIVED_COOLDOWN_MS = 15 * 60_000; +/** + * Ceiling on quota-refusal avoidance. Generous enough to cover a full five-hour burst window, + * tight enough that a weekly or monthly reset four days out cannot take an account out of + * rotation for the {@link CODEX_MAX_QUOTA_COOLDOWN_MS} day the Retry-After ceiling allows. + */ +export const CODEX_MAX_QUOTA_AVOID_MS = 6 * 60 * 60_000; +/** Minimum gap between probe leases for one cooled-down account. */ +export const CODEX_QUOTA_PROBE_INTERVAL_MS = 5 * 60_000; +export const CODEX_FAILURE_WINDOW_MS = 5 * 60_000; +/** + * How recently a 100% burst reading must have been OBSERVED to exclude an account when it + * carries no reset timestamp (#3425). Deliberately far tighter than the 6h disk-hydration + * horizon in `quota.ts`: shorter than any plausible five-hour burst window, so a persisted + * reading can never strand a recovered account, and long enough that a snapshot taken at + * admission is still fresh when selection reads it. + */ +export const TERMINAL_SHORT_WINDOW_FRESHNESS_MS = 5 * 60_000; +/** How long a transient failure keeps the account out of pool selection. */ +export const CODEX_TRANSIENT_SOFT_AVOID_MS = 30_000; +export const CODEX_TRANSIENT_SOFT_AVOID_ESCALATION_MS = [ + CODEX_TRANSIENT_SOFT_AVOID_MS, + 2 * 60_000, + 10 * 60_000, + 30 * 60_000, +] as const; + +export type CodexUpstreamOutcome = number | "connect_error" | "timeout" | "connect_neutral"; +export type CodexUpstreamOutcomeClass = "success" | "credential" + | "workspace" | "quota" | "transient" | "caller" | "neutral" | "unknown"; +export type CodexCooldownSource = "retry-after" | "reset-derived" | "default"; + +export type CodexUpstreamOutcomeMeta = { + retryAfter?: string | null; + resetAt?: unknown | unknown[]; + now?: number; + /** (provider, host) ledger key for account-neutral reachability failures (#914). */ + hostKey?: string; + /** + * Upstream denial evidence for a 403. A workspace/entitlement denial means the CREDENTIAL + * is fine and the account simply cannot reach this workspace, so it must not be quarantined + * for reauthentication (#1789). Absent evidence keeps the historical credential handling. + */ + denial?: "workspace" | "entitlement"; + /** Stable transport code recorded alongside a neutral host failure. */ + lastFailureCode?: string; + /** Native model selected for this request; used only for confirmed scoped quotas. */ + modelId?: string; + /** When set, clears affinity for this thread immediately on transient failure. */ + threadId?: string | null; + /** + * Suppress Pool rotation and quota/transient affinity mutations for an account-qualified + * request. Credential failures still sweep stale affinities because reauthentication is + * account-wide. + */ + fixedAccount?: boolean; + /** + * Probe lease held by this request, when it was admitted through an active + * quota cooldown. Only the outcome carrying the current lease may clear the + * cooldown (#433). + */ + probeLeaseId?: string; + /** Scope of `probeLeaseId` when it was granted against a model-scoped cooldown. */ + probeQuotaScope?: CodexQuotaScope; + /** + * Already-chosen alternate for same-request 429 retry. When set, promotion + * reuses this account instead of calling {@link pickAlternateCodexAccount} + * again (which would advance a round-robin ring twice). + */ + promoteAccountId?: string; + /** Generation captured when this routed account was selected. */ + writerGeneration?: number; + /** + * Credential generation this request's bearer was read at. Distinct from + * `writerGeneration`, which tracks the config store. + * + * A 401 that arrives after the credential was already replaced is evidence about a + * token nobody is using any more, so it must not quarantine the replacement. Absent + * means the caller cannot supply lineage and the historical unfenced handling stands. + */ + credentialGeneration?: number; +}; + +export function computeCodexUsageScore(quota: { + weeklyPercent?: number; + monthlyPercent?: number; + shortPercent?: number; + shortResetAt?: number; + shortObservedAt?: number; +} | null, plan?: unknown, now: number = Date.now()): number { + if (!quota) return CODEX_UNKNOWN_USAGE_SCORE; + const finite = (value: unknown): value is number => typeof value === "number" && Number.isFinite(value); + const longWindows = isThirtyDayOnlyCodexPlan(plan) + ? [quota.monthlyPercent] + : [quota.weeklyPercent, quota.monthlyPercent]; + const knownLong = longWindows.filter(finite); + // The short burst window only REFINES a known long-window position; it cannot stand in for + // one. A snapshot carrying just `shortPercent: 0` would otherwise score a flat 0 and make an + // account whose weekly/monthly usage is entirely unverified look like the emptiest in the + // pool, so `pickLowestUsageAmong` would send every request to it. Unknown has to stay + // unknown until a governing window is actually observed. + // + // A FULL burst window is the exception (#3029). It is not an optimistic guess about an + // unobserved window — it is a direct observation that the account cannot serve a request + // right now, whatever its monthly position turns out to be. Unknown-means-selectable is + // correct for uncertainty and wrong for a measured refusal: the account stays selected, + // `applyQuotaAutoSwitch` never fires, and the pool wedges on an exhausted credential. + if (knownLong.length === 0) { + return isTerminalShortWindow(quota, now) ? CODEX_EXHAUSTED_USAGE_PERCENT : CODEX_UNKNOWN_USAGE_SCORE; + } + const values = finite(quota.shortPercent) ? [...knownLong, quota.shortPercent] : knownLong; + return Math.max(...values); +} + +/** + * A short-only reading that proves the account is blocked NOW. + * + * Freshness is not optional. `getAccountQuota` performs no expiry check, partial updates + * carry a still-open short tuple forward, and disk hydration accepts a persisted reading for + * hours — so scoring 100 from `shortPercent` alone would keep excluding an account whose + * five-hour window has since reset. Merge no longer carries an elapsed shortResetAt, but an + * explicit incoming elapsed tuple is still stored, and a missing reset cannot be aged there. + * That is #3029 pointed the other way: the issue is that + * an exhausted account stays selected, and "a recovered account stays excluded" trades one + * unusable pool for another. + * + * A reading with no `shortResetAt` cannot be aged, so it stays unknown. The conservative + * direction here is the one that keeps an account selectable: a wrongly-selected account + * fails one request, while a wrongly-excluded one is invisible until someone reads the pool + * by hand. + * + * A missing reset can instead be aged by shortObservedAt (#3425). General updatedAt is not + * sufficient: credit-only updates preserve the old short tuple but advance that timestamp. + * Old disk snapshots without short-window provenance remain unknown. + */ +function isTerminalShortWindow( + quota: { shortPercent?: number; shortResetAt?: number; shortObservedAt?: number }, + now: number, +): boolean { + if (typeof quota.shortPercent !== "number" || !Number.isFinite(quota.shortPercent)) return false; + if (quota.shortPercent < CODEX_EXHAUSTED_USAGE_PERCENT) return false; + const resetAt = quota.shortResetAt; + if (typeof resetAt !== "number" || !Number.isFinite(resetAt) || resetAt <= 0) { + const observedAt = quota.shortObservedAt; + if (typeof observedAt !== "number" || !Number.isFinite(observedAt)) return false; + const age = now - observedAt; + return age >= 0 && age <= TERMINAL_SHORT_WINDOW_FRESHNESS_MS; + } + // Seconds and milliseconds both reach storage, so the split lives in one place next to the + // merge that also ages a stored reset instant (`resetAtToMs`, src/codex/quota.ts). + return resetAtToMs(resetAt) > now; +} + +export function classifyCodexUpstreamOutcome( + outcome: CodexUpstreamOutcome, + denial?: "workspace" | "entitlement", +): CodexUpstreamOutcomeClass { + if (outcome === "connect_neutral") return "neutral"; + if (outcome === "connect_error" || outcome === "timeout") return "transient"; + if (!Number.isFinite(outcome)) return "unknown"; + if (outcome >= 200 && outcome < 300) return "success"; + // Explicit 3xx policy (#914): a redirect response is relayed as-is and is + // never account or host health evidence — it proves the host is reachable + // and says nothing about the credential. Relayed as the neutral class so a + // stray 3xx cannot increment an account's transient streak. + if (outcome >= 300 && outcome < 400) return "neutral"; + // 401 is always a credential problem. A 403 is only a credential problem when nothing + // tells us otherwise: a workspace/entitlement denial (#1789) means the credential is valid + // and the account simply lacks access here, so quarantining it for reauth is wrong advice. + // Absent denial evidence the historical mapping stands, so the change fails safe. + if (outcome === 403 && denial !== undefined) return "workspace"; + if (outcome === 401 || outcome === 403) return "credential"; + // 402 Payment Required is treated as quota exhaustion for pool cooldown/failover + // (same-request alternate retry records this outcome for the depleted account). + if (outcome === 429 || outcome === 402) return "quota"; + if (outcome >= 400 && outcome < 500) return "caller"; + if (outcome >= 500 && outcome < 600) return "transient"; + return "unknown"; +} + +function clampCooldownMs(ms: number): number { + return Math.min(Math.max(ms, 1), CODEX_MAX_QUOTA_COOLDOWN_MS); +} + +export function parseRetryAfterMs(value: string | null | undefined, now = Date.now()): number | undefined { + const text = value?.trim(); + if (!text) return undefined; + if (/^\d+(?:\.\d+)?$/.test(text)) { + const seconds = Number(text); + if (Number.isFinite(seconds) && seconds > 0) return clampCooldownMs(Math.ceil(seconds * 1000)); + } + const timestamp = Date.parse(text); + if (!Number.isFinite(timestamp)) return undefined; + const delay = timestamp - now; + return delay > 0 ? clampCooldownMs(delay) : undefined; +} + +function resetTimestampMs(value: unknown): number | undefined { + const numeric = typeof value === "number" + ? value + : typeof value === "string" && value.trim() !== "" + ? Number(value) + : undefined; + if (typeof numeric !== "number" || !Number.isFinite(numeric) || numeric <= 0) return undefined; + return numeric < 1_000_000_000_000 ? numeric * 1000 : numeric; +} + +export function parseResetCooldownMs(resetAt: unknown | unknown[] | undefined, now = Date.now()): number | undefined { + const values = Array.isArray(resetAt) ? resetAt : [resetAt]; + let best: number | undefined; + for (const value of values) { + const timestamp = resetTimestampMs(value); + if (timestamp === undefined) continue; + const delay = timestamp - now; + if (delay <= 0) continue; + // A far-future reset must not pin the account for the full Retry-After + // ceiling: quota usually frees up well before the advertised window (#433). + const clamped = Math.min(clampCooldownMs(delay), CODEX_MAX_RESET_DERIVED_COOLDOWN_MS); + if (best === undefined || clamped < best) best = clamped; + } + return best; +} + +export function computeQuotaCooldown(meta: CodexUpstreamOutcomeMeta = {}): { + until: number; + source: CodexCooldownSource; +} { + const now = meta.now ?? Date.now(); + const retryAfterMs = parseRetryAfterMs(meta.retryAfter, now); + if (retryAfterMs !== undefined) return { until: now + retryAfterMs, source: "retry-after" }; + const resetCooldownMs = parseResetCooldownMs(meta.resetAt, now); + if (resetCooldownMs !== undefined) return { until: now + resetCooldownMs, source: "reset-derived" }; + return { until: now + CODEX_DEFAULT_QUOTA_COOLDOWN_MS, source: "default" }; +} + +/** + * When the pool should stop preferring an account after it refused on quota. + * + * The earliest window the refusal actually announced, bounded by {@link CODEX_MAX_QUOTA_AVOID_MS}, + * and never shorter than the cooldown the same refusal produced — a Retry-After directive that + * outlasts every announcement still governs. + */ +export function quotaAvoidUntilFor(meta: CodexUpstreamOutcomeMeta, now: number, cooldownUntil: number): number { + const values = Array.isArray(meta.resetAt) ? meta.resetAt : [meta.resetAt]; + let announced: number | undefined; + for (const value of values) { + const timestamp = resetTimestampMs(value); + if (timestamp === undefined) continue; + const delay = timestamp - now; + if (delay <= 0) continue; + const until = now + Math.min(delay, CODEX_MAX_QUOTA_AVOID_MS); + if (announced === undefined || until < announced) announced = until; + } + return Math.max(cooldownUntil, announced ?? 0); +} + +export function computeQuotaCooldownUntil(meta: CodexUpstreamOutcomeMeta = {}): number { + return computeQuotaCooldown(meta).until; +} diff --git a/src/codex/routing/health-store.ts b/src/codex/routing/health-store.ts new file mode 100644 index 0000000000..9c0d922b97 --- /dev/null +++ b/src/codex/routing/health-store.ts @@ -0,0 +1,402 @@ +import { isCodexAccountGenerationLive } from "../account-store"; +import { NATIVE_RESERVE_MODEL } from "../catalog/native-models"; +import { MAIN_CODEX_ACCOUNT_ID } from "../main-account"; +import { POOL_KEY_CODEX } from "../pool-rotation"; +import { isCanonicalOpenAiForwardProvider } from "../../providers/openai-tiers"; +import type { OcxConfig } from "../../types"; +import type { CodexCooldownSource } from "./cooldown-math"; + +export type CodexUpstreamHealth = { + consecutiveFailures: number; + /** Consecutive healthy terminals observed while recovering from escalation level 2+. */ + consecutiveSuccesses?: number; + lastFailureStatus?: number; + lastFailureAt?: number; + /** Hard cooldown (quota 429). Survives a later 2xx; blocks auth + selection. */ + cooldownUntil?: number; + /** + * How long a quota refusal keeps selection away from this account (or this native quota + * group), as opposed to how long it is hard-blocked. + * + * The two are deliberately different lengths. {@link CODEX_MAX_RESET_DERIVED_COOLDOWN_MS} + * caps the hard cooldown at 15 minutes because a reset announcement is advisory and plan + * quota usually frees up before it — an account must stay reachable so the pool can find + * that out (#433). The window the refusal announced is not 15 minutes, though, so once the + * cooldown lapses the account is selectable again while its burst window is still spent, + * and the strategy picks it straight back: this proxy reads a weekly bar a burst limit never + * touches, so a refused account still scores as the coolest in the pool. Every request then + * earns the same 429 until the process restarts, which is the only thing that drops this map. + * + * So the announcement governs avoidance and the cap still governs blocking. Avoidance is soft + * in the {@link softAvoidUntil} sense: it reorders the pool and releases a bound thread, and + * the last-resort paths still reach the account when nothing else can serve, so one pessimistic + * announcement cannot stall routing. + */ + quotaAvoidUntil?: number; + /** When the current cooldown was recorded; origin of the probe interval clock. */ + cooldownSince?: number; + /** + * What produced the cooldown. An explicit Retry-After is a literal retry + * directive and is never probed; a quota resetAt only announces a window + * refresh, so it may be probed early (#433). + */ + cooldownSource?: CodexCooldownSource; + /** + * Bumped on every cooldown write. A probe lease records the generation it was + * issued for so a lease cannot clear a cooldown that a later 429 replaced. + */ + cooldownGeneration?: number; + /** + * Identity of the in-flight probe. A cooled-down account sends no traffic, so + * no organic 2xx can prove recovery; only the outcome carrying this id may + * clear the cooldown. + */ + probeLeaseId?: string; + /** Cooldown generation at the moment the lease was granted. */ + probeLeaseGeneration?: number; + /** Last probe grant or conclusion; paces the probe interval. */ + lastProbeAt?: number; + /** + * Soft avoid after connect_error / timeout / transient 5xx. Cleared on 2xx. + * Blocks pool selection + thread affinity reuse so a sticky session can leave a + * flaky account without throwing CodexAccountCooldownError (hard-only). + */ + softAvoidUntil?: number; + /** + * Credential generation a 401/403 quarantine was derived from (#2892 gap 4). + * + * Provenance lives ON the entry rather than in a side map keyed by account id. A side map spends + * "whatever health is current when the old credential is found dead", which deletes a later + * unrelated entry: a G1 401, then a G2 save, then a genuine G2 503 would lose the 503. Only the + * entry that carries this field can be spent, and any later write simply replaces it. + */ + credentialFailureGeneration?: number; +}; + +const upstreamHealth = new Map(); +/** + * Reset-derived 429s can describe a quota owned by one native model family, + * rather than the whole ChatGPT account. Keep those advisory cooldowns apart + * from account-wide Retry-After/default throttles and transient health. + */ +const quotaScopedHealth = new Map>(); +/** + * Spend a credential-failure health entry whose credential no longer exists (#2892 gap 4). + * + * A 401/403 describes one CREDENTIAL, not an account, and a replacement can land at any point after + * the outcome is recorded — so re-reading the store inside `recordCodexUpstreamOutcome` narrows the + * window without closing it. The reader decides instead, and it may only spend an entry that + * actually carries credential provenance: a later transient or quota write replaces the entry and + * with it the tag, so this can never delete evidence that belongs to a different failure. + */ +export function dropSpentCredentialFailure(accountId: string): void { + const health = upstreamHealth.get(accountId); + const generation = health?.credentialFailureGeneration; + if (health === undefined || generation === undefined) return; + if (isCodexAccountGenerationLive(accountId, generation)) return; + upstreamHealth.delete(accountId); +} +let lastReconciledGeneration = 0; +let liveHealthAccountIds = new Set(); + +/** + * Native Codex quota groups known to be independent upstream. Keep the mapping + * deliberately conservative: unlisted models share the normal native group. + * Add a new explicit group here only when its independent upstream quota is + * confirmed, so shared limits never receive cross-model bypasses. + */ +export type CodexQuotaScope = "shared" | "reserve"; + + +export const NATIVE_MODEL_QUOTA_SCOPES: Readonly> = { + [NATIVE_RESERVE_MODEL]: "reserve", +}; + +export function codexQuotaScopeForModel(modelId: string | undefined): CodexQuotaScope | undefined { + if (!modelId?.trim()) return undefined; + return NATIVE_MODEL_QUOTA_SCOPES[modelId.trim().toLowerCase()] ?? "shared"; +} + +/** Independent quota groups must not mutate the shared active-account cursor. */ +export function isIndependentCodexQuotaScope(quotaScope?: CodexQuotaScope): boolean { + return quotaScope !== undefined && quotaScope !== "shared"; +} + +export function codexPoolKeyForScope(quotaScope?: CodexQuotaScope): string { + return isIndependentCodexQuotaScope(quotaScope) ? `${POOL_KEY_CODEX}:${quotaScope}` : POOL_KEY_CODEX; +} + +export function listLiveCodexAccountIds(config: OcxConfig): ReadonlySet { + const ids = new Set((config.codexAccounts ?? []).map(account => account.id)); + const openai = config.providers.openai; + if (openai && openai.disabled !== true && isCanonicalOpenAiForwardProvider(openai)) { + ids.add(MAIN_CODEX_ACCOUNT_ID); + } + return ids; +} + +export function getCodexUpstreamHealth( + accountId: string, +): CodexUpstreamHealth | null { + dropSpentCredentialFailure(accountId); + return upstreamHealth.get(accountId) ?? null; +} + +export function scopedHealthFor(accountId: string, scope: CodexQuotaScope): CodexUpstreamHealth | undefined { + return quotaScopedHealth.get(accountId)?.get(scope); +} + +export function setScopedHealth(accountId: string, scope: CodexQuotaScope, health: CodexUpstreamHealth): void { + let scopes = quotaScopedHealth.get(accountId); + if (!scopes) { + scopes = new Map(); + quotaScopedHealth.set(accountId, scopes); + } + scopes.set(scope, health); +} + +export function deleteScopedHealth(accountId: string, scope: CodexQuotaScope): void { + const scopes = quotaScopedHealth.get(accountId); + if (!scopes) return; + scopes.delete(scope); + if (scopes.size === 0) quotaScopedHealth.delete(accountId); +} + +/** Live quota-refusal avoidance for an account, including the lane the request belongs to. */ +function codexQuotaAvoidUntil( + accountId: string, + quotaScope: CodexQuotaScope | undefined, + now: number, +): number | null { + const live = (value: number | undefined): number | null => + typeof value === "number" && Number.isFinite(value) && value > now ? value : null; + const account = live(upstreamHealth.get(accountId)?.quotaAvoidUntil); + const scoped = quotaScope === undefined + ? null + : live(scopedHealthFor(accountId, quotaScope)?.quotaAvoidUntil); + if (account === null) return scoped; + return scoped === null ? account : Math.max(account, scoped); +} + +export function isCodexQuotaAvoided( + accountId: string, + quotaScope: CodexQuotaScope | undefined, + now: number, +): boolean { + return codexQuotaAvoidUntil(accountId, quotaScope, now) !== null; +} + +/** + * Hard-cooldown bookkeeping that ordinary success/transient transitions rebuild + * their health object from. Dropping these would let one late unrelated response + * erase a Retry-After source, a cooldown generation, or someone else's live probe. + */ +export function preservedCooldownFields(health: CodexUpstreamHealth | undefined): Partial { + if (!health) return {}; + // `credentialFailureGeneration` is provenance for ONE credential failure, so it must not survive + // into a later transient or quota entry — otherwise that entry inherits the tag and gets spent + // when the old credential dies, deleting evidence that was never about it (#2892 gap 4 review). + const { + consecutiveFailures: _f, consecutiveSuccesses: _s, lastFailureStatus: _st, lastFailureAt: _at, + softAvoidUntil: _sa, credentialFailureGeneration: _cg, ...cooldownFields + } = health; + return cooldownFields; +} + +export function getCodexAccountCooldownUntil(accountId: string, now = Date.now()): number | null { + const cooldownUntil = upstreamHealth.get(accountId)?.cooldownUntil; + return typeof cooldownUntil === "number" && Number.isFinite(cooldownUntil) && cooldownUntil > now ? cooldownUntil : null; +} + +/** Read-only cooldown snapshot for shared OAuth health projection (no write side effects). */ +export function getCodexAccountHealthSnapshot(accountId: string, now = Date.now()): { + cooldownUntil?: number; + cooldownSource?: CodexCooldownSource; +} | null { + const cooldownUntil = getCodexAccountCooldownUntil(accountId, now); + if (cooldownUntil === null) return null; + const source = upstreamHealth.get(accountId)?.cooldownSource; + return { + cooldownUntil, + ...(source ? { cooldownSource: source } : {}), + }; +} + +/** + * Read the cooldown relevant to a routed native model. Account-wide cooldowns + * (Retry-After/default) always win; reset-derived scoped state applies only to + * its confirmed quota group. + */ +export function getCodexQuotaHealthSnapshot( + accountId: string, + quotaScope: CodexQuotaScope | undefined, + now = Date.now(), +): { + cooldownUntil?: number; + cooldownSource?: CodexCooldownSource; + quotaScope?: CodexQuotaScope; +} | null { + const account = getCodexAccountHealthSnapshot(accountId, now); + if (account) return account; + if (!quotaScope) return null; + const scoped = scopedHealthFor(accountId, quotaScope); + const cooldownUntil = scoped?.cooldownUntil; + if (typeof cooldownUntil !== "number" || !Number.isFinite(cooldownUntil) || cooldownUntil <= now) return null; + return { + cooldownUntil, + ...(scoped?.cooldownSource ? { cooldownSource: scoped.cooldownSource } : {}), + quotaScope, + }; +} + +export function isCodexAccountInCooldown(accountId: string, now = Date.now()): boolean { + return getCodexAccountCooldownUntil(accountId, now) !== null; +} + +/** + * Manually lift a hard quota cooldown without touching failure history. + * + * Injected Codex routing makes this proxy the ONLY model path for Codex Desktop, so a + * cooldown that outlives the real upstream limit reads to the user as "the whole app is + * broken" with no escape but editing config.toml. This is that escape hatch. + * + * Deliberately narrow: + * - Failure counters and softAvoid survive. Clearing a cooldown says "the quota window + * moved", not "this account is healthy"; failover must keep its knowledge. + * - Dropping `probeLeaseId` is what stops a stale in-flight probe from later "proving" + * recovery against a NEWER cooldown: {@link ownsProbeLease} needs the id to match. + * `cooldownGeneration` is preserved and bumped as redundancy only — a fresh 429 already + * bumps it in {@link recordCodexUpstreamOutcome}, so the bump here is not load-bearing + * today and is kept so the invariant survives a future change that retains the lease. + * + * Returns false when the account carried neither a live cooldown nor a live avoidance window. + * The window outlives the cooldown by design — the cooldown caps at fifteen minutes and the + * window runs up to six hours — so the moment an operator actually reaches for this escape + * hatch is usually after the cooldown lapsed and only the window is still keeping the account + * out of rotation. Refusing to look at the window then would leave the hatch shut in the one + * case it exists for. + */ +export function clearCodexAccountCooldown(accountId: string, now = Date.now()): boolean { + const clear = (health: CodexUpstreamHealth): CodexUpstreamHealth | null => { + const cooldownUntil = health.cooldownUntil; + const liveCooldown = typeof cooldownUntil === "number" && Number.isFinite(cooldownUntil) && cooldownUntil > now; + const avoidUntil = health.quotaAvoidUntil; + const liveAvoidance = typeof avoidUntil === "number" && Number.isFinite(avoidUntil) && avoidUntil > now; + if (!liveCooldown && !liveAvoidance) return null; + const { + cooldownUntil: _until, + cooldownSince: _since, + cooldownSource: _source, + probeLeaseId: _leaseId, + probeLeaseGeneration: _leaseGeneration, + // Same reasoning as the probe recovery above: "the quota window moved" is a statement + // about the whole refusal, so the avoidance it announced goes with the block it + // produced. Keeping it would leave this escape hatch not escaping, because selection + // would still pass over the account for as long as the announced window runs. + quotaAvoidUntil: _avoid, + ...rest + } = health; + return { + ...rest, + cooldownGeneration: (health.cooldownGeneration ?? 0) + 1, + lastProbeAt: now, + }; + }; + + let cleared = false; + const accountHealth = upstreamHealth.get(accountId); + if (accountHealth) { + const next = clear(accountHealth); + if (next) { + upstreamHealth.set(accountId, next); + cleared = true; + } + } + for (const [scope, health] of quotaScopedHealth.get(accountId) ?? []) { + const next = clear(health); + if (next) { + setScopedHealth(accountId, scope, next); + cleared = true; + } + } + return cleared; +} + +export function getCodexAccountSoftAvoidUntil(accountId: string, now = Date.now()): number | null { + const softAvoidUntil = upstreamHealth.get(accountId)?.softAvoidUntil; + return typeof softAvoidUntil === "number" && Number.isFinite(softAvoidUntil) && softAvoidUntil > now + ? softAvoidUntil + : null; +} + +export function isCodexAccountSoftAvoided(accountId: string, now = Date.now()): boolean { + return getCodexAccountSoftAvoidUntil(accountId, now) !== null; +} + +/** + * Closed package-internal accessors for the account-wide health maps. Selection, + * the probe lease, and the active cursor mutate health only through these; the + * Map bindings themselves never leave this module. + */ +export function getAccountHealth(accountId: string): CodexUpstreamHealth | undefined { + return upstreamHealth.get(accountId); +} + +export function setAccountHealth(accountId: string, health: CodexUpstreamHealth): void { + upstreamHealth.set(accountId, health); +} + +export function deleteAccountHealth(accountId: string): void { + upstreamHealth.delete(accountId); +} + +export function listScopedHealthEntries(accountId: string): Array<[CodexQuotaScope, CodexUpstreamHealth]> { + return [...(quotaScopedHealth.get(accountId) ?? [])]; +} + +export function deleteAllScopedHealth(accountId: string): void { + quotaScopedHealth.delete(accountId); +} + +export function isHealthAccountAdmissible(accountId: string, writerGeneration: number): boolean { + return writerGeneration >= lastReconciledGeneration || liveHealthAccountIds.has(accountId); +} + +export function isHealthGenerationReconciled(generation: number): boolean { + return generation <= lastReconciledGeneration; +} + +export function pruneHealthAccountsForContext(codexAccountIds: ReadonlySet): number { + let removed = 0; + for (const accountId of upstreamHealth.keys()) { + if (codexAccountIds.has(accountId)) continue; + upstreamHealth.delete(accountId); + removed += 1; + } + for (const accountId of quotaScopedHealth.keys()) { + if (codexAccountIds.has(accountId)) continue; + quotaScopedHealth.delete(accountId); + removed += 1; + } + return removed; +} + +export function commitHealthReconcile(generation: number, codexAccountIds: ReadonlySet): void { + liveHealthAccountIds = new Set(codexAccountIds); + lastReconciledGeneration = generation; +} + +export function clearUpstreamHealthState(): void { + upstreamHealth.clear(); + quotaScopedHealth.clear(); +} + +export function resetHealthReconcileState(): void { + lastReconciledGeneration = 0; + liveHealthAccountIds = new Set(); +} + +export function deleteAllHealthForAccount(accountId: string): void { + upstreamHealth.delete(accountId); + quotaScopedHealth.delete(accountId); +} diff --git a/src/codex/routing/probe-lease.ts b/src/codex/routing/probe-lease.ts new file mode 100644 index 0000000000..0ae865ac47 --- /dev/null +++ b/src/codex/routing/probe-lease.ts @@ -0,0 +1,358 @@ +import { randomUUID } from "node:crypto"; +import { isCodexAccountGenerationLive, readCodexAccountRecord, type CodexRefreshProvenance } from "../account-store"; +import { isCodexAccountPaused } from "../account-pause"; +import { isSelectableCodexPoolAccount } from "../account-id"; +import { isAccountNeedsReauth } from "../account-runtime-state"; +import { MAIN_CODEX_ACCOUNT_ID } from "../main-account"; +import type { OcxConfig } from "../../types"; +import { CODEX_QUOTA_PROBE_INTERVAL_MS, type CodexUpstreamOutcomeMeta } from "./cooldown-math"; +import { + deleteScopedHealth, + getAccountHealth, + listScopedHealthEntries, + scopedHealthFor, + setAccountHealth, + setScopedHealth, + type CodexQuotaScope, + type CodexUpstreamHealth, +} from "./health-store"; + +export type CodexQuotaRecoveryProbeClaim = { + accountId: string; + scope?: CodexQuotaScope; + leaseId: string; + cooldownGeneration: number; + credentialGeneration: number; + /** Claim-time `replacedAt`; unchanged after a probe-owned refresh, stamped on external replacement. */ + credentialReplacedAt?: number; +}; + +export type CodexQuotaRecoveryProbeProof = { + credentialGeneration?: number; +}; + +/** + * Grant at most one probe lease per interval for a cooled-down account. + * + * A cooled-down account is short-circuited locally, so it never sends traffic and + * no organic 2xx can prove that upstream quota recovered — the cooldown can only + * end by expiry or a proxy restart (#433). Releasing a single probe breaks that + * deadlock. Explicit Retry-After cooldowns are excluded: those are literal retry + * directives, not window announcements. + * + * Returns the lease id, or null when no probe may go out right now. + */ +export function tryAcquireCodexQuotaProbeLease(accountId: string, now = Date.now()): string | null { + if (!canAcquireCodexQuotaProbeLease(accountId, now)) return null; + const health = getAccountHealth(accountId)!; + const probeLeaseId = randomUUID(); + setAccountHealth(accountId, { + ...health, + probeLeaseId, + probeLeaseGeneration: health.cooldownGeneration ?? 0, + lastProbeAt: now, + }); + return probeLeaseId; +} + +/** Side-effect-free check mirroring {@link tryAcquireCodexQuotaProbeLease} eligibility. */ +export function canAcquireCodexQuotaProbeLease(accountId: string, now = Date.now()): boolean { + return canAcquireQuotaProbeLease(getAccountHealth(accountId), now); +} + +function canAcquireQuotaProbeLease(health: CodexUpstreamHealth | undefined, now: number): boolean { + if (!health) return false; + const cooldownUntil = health.cooldownUntil; + if (typeof cooldownUntil !== "number" || !Number.isFinite(cooldownUntil) || cooldownUntil <= now) return false; + if (health.cooldownSource === "retry-after") return false; + if (health.probeLeaseId !== undefined) return false; + const origin = health.lastProbeAt ?? health.cooldownSince ?? cooldownUntil; + return now - origin >= CODEX_QUOTA_PROBE_INTERVAL_MS; +} + +/** + * Claim due reset-derived cooldown probes without consulting account selection. + * Added Pool credentials only; owned main usage recovery is handled separately. + */ +export function claimDueCodexQuotaRecoveryProbes( + config: OcxConfig, + limit: number, + now = Date.now(), +): CodexQuotaRecoveryProbeClaim[] { + const boundedLimit = Math.max(0, Math.floor(limit)); + if (boundedLimit === 0) return []; + const candidates: Array<{ + accountId: string; + scope?: CodexQuotaScope; + health: CodexUpstreamHealth; + credentialGeneration: number; + credentialReplacedAt?: number; + order: number; + }> = []; + for (const [order, account] of (config.codexAccounts ?? []).entries()) { + if (!isSelectableCodexPoolAccount(account) + || isCodexAccountPaused(config, account.id) + || isAccountNeedsReauth(account.id)) continue; + const record = readCodexAccountRecord(account.id); + if (!record?.credential || record.deletedAt != null) continue; + const due = [ + { scope: undefined, health: getAccountHealth(account.id) }, + ...[...(listScopedHealthEntries(account.id))].map(([scope, health]) => ({ scope, health })), + ].filter((entry): entry is { scope?: CodexQuotaScope; health: CodexUpstreamHealth } => + // Generic WHAM evidence can recover only ordinary quota, never Reserve. + // Do not spend this account's one claim per pass on an independent scope and + // delay the shared scope that the response can actually recover. + (entry.scope === undefined || entry.scope === "shared") + && entry.health?.cooldownSource === "reset-derived" + && canAcquireQuotaProbeLease(entry.health, now)) + .sort((a, b) => + (a.health.lastProbeAt ?? a.health.cooldownSince ?? 0) + - (b.health.lastProbeAt ?? b.health.cooldownSince ?? 0)); + const candidate = due[0]; + if (candidate) candidates.push({ + accountId: account.id, + ...(candidate.scope ? { scope: candidate.scope } : {}), + health: candidate.health, + credentialGeneration: record.generation, + ...(record.replacedAt !== undefined ? { credentialReplacedAt: record.replacedAt } : {}), + order, + }); + } + candidates.sort((a, b) => { + const age = (a.health.lastProbeAt ?? a.health.cooldownSince ?? 0) + - (b.health.lastProbeAt ?? b.health.cooldownSince ?? 0); + return age || a.order - b.order; + }); + return candidates.slice(0, boundedLimit).map(candidate => { + const leaseId = randomUUID(); + const next = { + ...candidate.health, + probeLeaseId: leaseId, + probeLeaseGeneration: candidate.health.cooldownGeneration ?? 0, + lastProbeAt: now, + }; + if (candidate.scope) setScopedHealth(candidate.accountId, candidate.scope, next); + else setAccountHealth(candidate.accountId, next); + return { + accountId: candidate.accountId, + ...(candidate.scope ? { scope: candidate.scope } : {}), + leaseId, + cooldownGeneration: candidate.health.cooldownGeneration ?? 0, + credentialGeneration: candidate.credentialGeneration, + ...(candidate.credentialReplacedAt !== undefined + ? { credentialReplacedAt: candidate.credentialReplacedAt } + : {}), + }; + }); +} + +type CooldownRecoveryLease = Pick; + +export type ManualResetCooldownClaim = + | { kind: "pool"; probe: CodexQuotaRecoveryProbeClaim } + | { kind: "main"; probe: CooldownRecoveryLease }; + +function manualResetAccountEligible(config: OcxConfig, accountId: string): boolean { + return !isCodexAccountPaused(config, accountId) && !isAccountNeedsReauth(accountId) + && (accountId === MAIN_CODEX_ACCOUNT_ID + || (config.codexAccounts ?? []).some(account => account.id === accountId && isSelectableCodexPoolAccount(account))); +} + +/** Explicit reset bypasses probe pacing, never another owner's lease or quota scope. */ +export function claimManualResetCooldowns( + config: OcxConfig, + accountId: string, + now = Date.now(), + expectedPoolGeneration?: number, +): ManualResetCooldownClaim[] { + if (!manualResetAccountEligible(config, accountId)) return []; + const record = accountId === MAIN_CODEX_ACCOUNT_ID ? undefined : readCodexAccountRecord(accountId); + if (accountId !== MAIN_CODEX_ACCOUNT_ID && (!record?.credential || record.deletedAt != null)) return []; + if (record && expectedPoolGeneration !== undefined && record.generation !== expectedPoolGeneration) return []; + const claims: ManualResetCooldownClaim[] = []; + for (const scope of [undefined, "shared"] as const) { + const health = scope ? scopedHealthFor(accountId, scope) : getAccountHealth(accountId); + if (!health || health.cooldownSource !== "reset-derived" || health.probeLeaseId !== undefined + || !Number.isFinite(health.cooldownUntil) || !(health.cooldownUntil! > now)) continue; + const leaseId = randomUUID(); + const cooldownGeneration = health.cooldownGeneration ?? 0; + const next = { ...health, probeLeaseId: leaseId, probeLeaseGeneration: cooldownGeneration, lastProbeAt: now }; + if (scope) setScopedHealth(accountId, scope, next); + else setAccountHealth(accountId, next); + const probe = { accountId, scope, leaseId, cooldownGeneration }; + claims.push(record ? { kind: "pool", probe: { + ...probe, credentialGeneration: record.generation, credentialReplacedAt: record.replacedAt, + } } : { kind: "main", probe }); + } + return claims; +} + +export type ManualResetRefreshLineage = Readonly<{ + fromGeneration: number; + toGeneration: number; + provenance: CodexRefreshProvenance; +}>; + +type ManualResetQuotaProof = CodexQuotaRecoveryProbeProof & { + refreshLineage?: ManualResetRefreshLineage; +}; + +/** Main proof is checked by the already-owned auth operation, never by a Pool record. */ +export function settleManualResetCooldown( + config: OcxConfig, + claim: ManualResetCooldownClaim, + recovered: boolean, + proof: ManualResetQuotaProof = {}, + now = Date.now(), +): boolean { + if (!recovered) return settleCooldownRecoveryLease(claim.probe, false, now); + const eligible = manualResetAccountEligible(config, claim.probe.accountId); + if (claim.kind === "main") return settleCooldownRecoveryLease(claim.probe, eligible, now); + const lineage = proof.refreshLineage; + // Equal wall-clock replacement stamps do not establish ancestry. Manual +1 + // recovery additionally needs the actual forced-refresh result for this edge. + const ownedGeneration = proof.credentialGeneration === claim.probe.credentialGeneration + || (proof.credentialGeneration === claim.probe.credentialGeneration + 1 + && lineage?.fromGeneration === claim.probe.credentialGeneration + && lineage.toGeneration === proof.credentialGeneration + && (lineage.provenance === "self-refresh" || lineage.provenance === "joined-lineage")); + return settleCodexQuotaRecoveryProbe(claim.probe, eligible && ownedGeneration, proof, now); +} + +/** Settle one background recovery claim without mutating account-wide outcome state. */ +export function settleCodexQuotaRecoveryProbe( + claim: CodexQuotaRecoveryProbeClaim, + recovered: boolean, + proof: CodexQuotaRecoveryProbeProof, + now = Date.now(), +): boolean { + const health = claim.scope + ? scopedHealthFor(claim.accountId, claim.scope) + : getAccountHealth(claim.accountId); + if (!health || health.probeLeaseId !== claim.leaseId) return false; + const currentRecord = readCodexAccountRecord(claim.accountId); + const proofGeneration = proof.credentialGeneration; + // A probe-owned token refresh (getValidCodexToken) advances the credential generation by + // exactly one while preserving `replacedAt`; an external credential replacement bumps the + // generation too but stamps a fresh `replacedAt`. Accept the +1 transition only when the + // claim-time lineage is intact AND the generation the fresh quota was proven under is live. + const generationFenced = proofGeneration !== undefined + && (proofGeneration === claim.credentialGeneration + ? isCodexAccountGenerationLive(claim.accountId, proofGeneration) + : proofGeneration === claim.credentialGeneration + 1 + && currentRecord?.replacedAt === claim.credentialReplacedAt + && isCodexAccountGenerationLive(claim.accountId, proofGeneration)); + return settleCooldownRecoveryLease(claim, recovered && generationFenced, now); +} + +function settleCooldownRecoveryLease(claim: CooldownRecoveryLease, recovered: boolean, now: number): boolean { + const health = claim.scope ? scopedHealthFor(claim.accountId, claim.scope) : getAccountHealth(claim.accountId); + if (!health || health.probeLeaseId !== claim.leaseId) return false; + const fenced = (claim.scope === undefined || claim.scope === "shared") + && health.cooldownSource === "reset-derived" + && (health.cooldownGeneration ?? 0) === claim.cooldownGeneration + && (health.probeLeaseGeneration ?? 0) === claim.cooldownGeneration; + if (!recovered || !fenced) { + const released = withProbeLeaseReleased(health, now); + if (claim.scope) setScopedHealth(claim.accountId, claim.scope, released); + else setAccountHealth(claim.accountId, released); + return false; + } + if (claim.scope) { + deleteScopedHealth(claim.accountId, claim.scope); + } else { + const { + cooldownUntil: _until, + cooldownSince: _since, + cooldownSource: _source, + probeLeaseId: _leaseId, + probeLeaseGeneration: _leaseGeneration, + // "The quota window moved" is a statement about the whole refusal, so the avoidance it + // announced goes with the block it produced. Leaving it would make this escape hatch stop + // escaping: the account would still be passed over by every selection it is meant to win. + quotaAvoidUntil: _avoid, + ...rest + } = health; + setAccountHealth(claim.accountId, { + ...rest, + cooldownGeneration: claim.cooldownGeneration + 1, + lastProbeAt: now, + }); + } + return true; +} + +/** Acquire the recovery probe for one confirmed model-specific quota group. */ +export function tryAcquireCodexQuotaScopeProbeLease( + accountId: string, + scope: CodexQuotaScope, + now = Date.now(), +): string | null { + const health = scopedHealthFor(accountId, scope); + if (!canAcquireQuotaProbeLease(health, now)) return null; + const probeLeaseId = randomUUID(); + setScopedHealth(accountId, scope, { + ...health!, + probeLeaseId, + probeLeaseGeneration: health!.cooldownGeneration ?? 0, + lastProbeAt: now, + }); + return probeLeaseId; +} + +/** Side-effect-free check for a confirmed model-specific quota probe. */ +export function canAcquireCodexQuotaScopeProbeLease( + accountId: string, + scope: CodexQuotaScope, + now = Date.now(), +): boolean { + return canAcquireQuotaProbeLease(scopedHealthFor(accountId, scope), now); +} + +/** + * Hand a probe lease back without recording an upstream outcome. Used by paths + * that take a lease and then fail before any request reaches upstream. + */ +export function releaseCodexQuotaProbeLease(accountId: string, leaseId: string, now = Date.now()): void { + const health = getAccountHealth(accountId); + if (!health || health.probeLeaseId !== leaseId) return; + setAccountHealth(accountId, withProbeLeaseReleased(health, now)); +} + +/** Release a model-specific quota probe when the request never reaches upstream. */ +export function releaseCodexQuotaScopeProbeLease( + accountId: string, + scope: CodexQuotaScope, + leaseId: string, + now = Date.now(), +): void { + const health = scopedHealthFor(accountId, scope); + if (!health || health.probeLeaseId !== leaseId) return; + setScopedHealth(accountId, scope, withProbeLeaseReleased(health, now)); +} + +/** + * True when this outcome belongs to the account's in-flight probe. The + * undefined-id guard matters: without it an outcome carrying no lease would match + * an account holding no lease and be mistaken for the probe owner. + */ +export function ownsProbeLease(health: CodexUpstreamHealth | undefined, meta: CodexUpstreamOutcomeMeta): boolean { + return meta.probeLeaseId !== undefined && meta.probeLeaseId === health?.probeLeaseId; +} + +/** + * True when the owning probe may still clear the cooldown. A later 429 bumps the + * generation, so a probe that started under an older cooldown must not erase the + * newer restriction (which may carry an explicit Retry-After). + */ +export function probeMayClearCooldown(health: CodexUpstreamHealth | undefined, meta: CodexUpstreamOutcomeMeta): boolean { + return ownsProbeLease(health, meta) + && (health!.probeLeaseGeneration ?? 0) === (health!.cooldownGeneration ?? 0); +} + +/** Strip the in-flight lease while preserving every hard-cooldown field. */ +export function withProbeLeaseReleased(health: CodexUpstreamHealth, now: number): CodexUpstreamHealth { + const { probeLeaseId: _id, probeLeaseGeneration: _gen, ...rest } = health; + return { ...rest, lastProbeAt: now }; +} diff --git a/src/codex/routing/selection.ts b/src/codex/routing/selection.ts new file mode 100644 index 0000000000..03f562a56c --- /dev/null +++ b/src/codex/routing/selection.ts @@ -0,0 +1,698 @@ +import { isCodexAccountPaused } from "../account-pause"; +import { codexAccountPriorityLookup, pinnedCodexAccountId } from "../account-priority"; +import { isSelectableCodexPoolAccount } from "../account-id"; +import { isAccountNeedsReauth } from "../account-runtime-state"; +import { isCodexAccountUsable, type CodexAccountUsabilityOptions } from "../account-usability"; +import { + normalizeAccountPoolStickyLimit, + normalizeCodexAccountPoolStrategy, + notePoolRotationSuccess, + peekRoundRobinAccount, + pickRoundRobinAccount, + selectPriorityTier, +} from "../pool-rotation"; +import { CODEX_UNKNOWN_USAGE_SCORE, getAccountQuota, resetAtToMs } from "../quota"; +import { codexPlanKey } from "../plan"; +import { MAIN_CODEX_ACCOUNT_ID, getMainAccountPlan, hasMainAccountRefreshGrant } from "../main-account"; +import type { OcxConfig } from "../../types"; +import { CODEX_FAILURE_WINDOW_MS, computeCodexUsageScore } from "./cooldown-math"; +import { + codexPoolKeyForScope, + dropSpentCredentialFailure, + getAccountHealth, + getCodexQuotaHealthSnapshot, + isCodexAccountSoftAvoided, + isCodexQuotaAvoided, + isIndependentCodexQuotaScope, + type CodexQuotaScope, +} from "./health-store"; +import { bindThreadAffinity, type CodexAffinityReason } from "./thread-affinity"; +import { + getEffectiveActiveCodexAccountId, + manualPreferenceBlocks, + promoteActiveCodexAccount, + rememberActiveCodexAccount, + setActiveCodexAccount, +} from "./active-account"; + +/** + * Plan keys the operator excluded from automatic rotation. Absent or empty means no policy, so an + * existing install rotates exactly as before. Compared with `codexPlanKey` because the stored plan + * is an unrestricted provider string whose casing this repository does not control. + */ +function excludedCodexPoolPlanKeys(config: OcxConfig): ReadonlySet | undefined { + const configured = config.codexPool?.excludedPlans; + if (!configured?.length) return undefined; + const keys = configured + .map(plan => codexPlanKey(plan)) + .filter((key): key is string => key !== undefined); + return keys.length > 0 ? new Set(keys) : undefined; +} + +/** + * Whether the operator's plan policy removes this account from automatic selection. + * + * Modelled on pause rather than usability: an excluded account keeps its credential, quota history, + * and affinity, stays visible on the account surface, and is still reachable by explicit account + * selection. Only automatic rotation skips it, which is the distinction #4211 asked for. + * + * It is checked in the same two places pause is checked, and that is not redundancy. The eligible + * list is consulted only when routing picks a NEW account; an already-active or already-affined + * account is served straight from {@link isCodexAccountSelectable}. A lapsed subscription leaves + * behind exactly that account, so a policy that filtered only the eligible list would miss the case + * it exists for. + * + * `__main__` is exempt. {@link getPoolAccountPlanForSelection} withholds the main plan during a + * selection-only drain so routing never reads the fenced native credential for it, so a rule that + * covered main would disagree with itself between drain and ordinary routing. + */ +export function isCodexAccountPlanExcluded( + config: OcxConfig, + accountId: string, + precomputed?: ReadonlySet, +): boolean { + if (accountId === MAIN_CODEX_ACCOUNT_ID) return false; + // Callers that test a whole list pass the set once rather than rebuilding it per row. + const excluded = precomputed ?? excludedCodexPoolPlanKeys(config); + if (!excluded) return false; + const plan = codexPlanKey(getPoolAccountPlan(config, accountId)); + return plan !== undefined && excluded.has(plan); +} + +export function isCodexAccountSelectable( + config: OcxConfig, + accountId: string, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): boolean { + return !isCodexAccountPaused(config, accountId) + && !isCodexAccountPlanExcluded(config, accountId) + && getCodexQuotaHealthSnapshot(accountId, quotaScope, now) === null + && !isCodexQuotaAvoided(accountId, quotaScope, now) + && !isCodexAccountSoftAvoided(accountId, now) + && isCodexAccountUsable(config, accountId, selectionOptions); +} + +/** + * Which guard in {@link isCodexAccountSelectable} refused this account, if any. + * + * Deliberately the same predicates in the same order as that function, because the point is to + * REPORT the guard that actually fired rather than to re-derive a plausible-looking cause. An + * earlier version of the release reason checked only a subset and let a paused, plan-excluded, + * cooled-down or quota-avoided release fall through to a quota fallback, which named something + * routing never used -- a diagnostic that is confidently wrong in exactly the cases an operator + * would consult it for (#4598). + */ +export function codexAccountBlockReason( + config: OcxConfig, + accountId: string, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): CodexAffinityReason | undefined { + if (isCodexAccountPaused(config, accountId)) return "paused"; + if (isCodexAccountPlanExcluded(config, accountId)) return "plan_excluded"; + if (getCodexQuotaHealthSnapshot(accountId, quotaScope, now) !== null) return "cooldown"; + if (isCodexQuotaAvoided(accountId, quotaScope, now)) return "quota_avoided"; + if (isCodexAccountSoftAvoided(accountId, now)) return "transient"; + if (!isCodexAccountUsable(config, accountId, selectionOptions)) return "unusable"; + return undefined; +} + +export function getEligiblePoolAccounts( + config: OcxConfig, + excludeId?: string, + now = Date.now(), + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + skipFailoverReadyCandidates = false, +): readonly string[] { + const excludedPlans = excludedCodexPoolPlanKeys(config); + const ids = (config.codexAccounts ?? []) + .filter(account => isSelectableCodexPoolAccount(account) + && account.id !== excludeId + && !isCodexAccountPaused(config, account.id) + && !isCodexAccountPlanExcluded(config, account.id, excludedPlans) + && !isAccountNeedsReauth(account.id) + && (!skipFailoverReadyCandidates || !shouldFailover(config, account.id, now))) + .filter(account => getCodexQuotaHealthSnapshot(account.id, quotaScope, now) === null) + .filter(account => !isCodexAccountSoftAvoided(account.id, now)) + .filter(account => !isCodexQuotaAvoided(account.id, quotaScope, now)) + .filter(account => isCodexAccountUsable(config, account.id, selectionOptions)) + .map(account => account.id); + // The main Codex account is not stored in config.codexAccounts; include it as a + // first-class rotation candidate when its read-only token is usable (Option A). + if ( + excludeId !== MAIN_CODEX_ACCOUNT_ID + && !isCodexAccountPaused(config, MAIN_CODEX_ACCOUNT_ID) + && (!isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID) || hasMainAccountRefreshGrant()) + && getCodexQuotaHealthSnapshot(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) === null + && !isCodexAccountSoftAvoided(MAIN_CODEX_ACCOUNT_ID, now) + // The main login is not in `config.codexAccounts`, so it never passes through the + // filters above and this is the only place an avoidance window can exclude it. Without + // this the window a refusal announced applies to the pool but not to the account that + // earned it: the cooldown caps at fifteen minutes, the window runs up to six hours, and + // in between the main account returns as a first-class candidate. + && !isCodexQuotaAvoided(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) + && (!skipFailoverReadyCandidates || !shouldFailover(config, MAIN_CODEX_ACCOUNT_ID, now)) + && isCodexAccountUsable(config, MAIN_CODEX_ACCOUNT_ID, selectionOptions) + ) { + ids.unshift(MAIN_CODEX_ACCOUNT_ID); + } + // Single choke point for selection order: every strategy, failover, and preview + // reaches the pool through here, so tiering applies once rather than per picker. + // Eligibility above is unchanged — this only narrows an already-eligible list. + return selectPriorityTier( + ids, + codexAccountPriorityLookup(config), + id => hasCodexQuotaHeadroom(config, id, selectionOptions, now), + pinnedCodexAccountId(config), + ); +} + +function listEligibleCodexAccountIds( + config: OcxConfig, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): readonly string[] { + return getEligiblePoolAccounts(config, undefined, now, quotaScope, selectionOptions); +} + +/** Shared reset timestamps are not evidence for independent model-quota groups. */ +export function accountPoolStrategyForScope(config: OcxConfig, quotaScope?: CodexQuotaScope) { + const strategy = normalizeCodexAccountPoolStrategy(config.accountPoolStrategy); + return strategy === "reset-first" && isIndependentCodexQuotaScope(quotaScope) ? "quota" : strategy; +} + +function stickyLimitForConfig(config: OcxConfig): number { + return normalizeAccountPoolStickyLimit(config.accountPoolStickyLimit); +} + +/** + * Whether an account still has quota to give under the auto-switch threshold. + * + * Fill-first and the priority tier filter share this predicate, and share both of + * its escape hatches. A disabled threshold means only health, pause, and reauth + * may drain an account; unknown usage is a guess, so it must neither force + * fill-first off the active account nor drain a tier that was simply never + * primed. A genuinely exhausted account 429s into cooldown and leaves + * eligibility on its own. + */ +export function hasCodexQuotaHeadroom( + config: OcxConfig, + accountId: string, + selectionOptions?: CodexAccountUsabilityOptions, + now: number = Date.now(), +): boolean { + const threshold = config.autoSwitchThreshold ?? 80; + if (threshold <= 0) return true; + const usage = computeCodexUsageScore( + getAccountQuota(accountId), + getPoolAccountPlanForSelection(config, accountId, selectionOptions), + now, + ); + if (isUnknownUsage(usage)) return true; + return usage < threshold; +} + +/** + * Is a live binding held for its prompt cache? + * + * Unset means yes. Cache affinity shipped as an opt-in flag (#4292) and then #4546 measured + * what the default costs: a pool whose accounts all sit in the 80-99% band hands a bound + * conversation from account to account, and because provider prompt caches are account-isolated + * every hop re-sends the entire prefix. An install that has never heard of this flag is exactly + * the install that gets hurt by it, so the protection cannot be something you have to find. + * + * `false` restores capacity-first routing byte-for-byte. It is a real choice -- a pinned thread + * on a busy account pays latency -- and it stays available; it is just no longer the default. + */ +export function isCacheAffinityEnabled(config: OcxConfig): boolean { + return config.pool?.cacheAffinity !== false; +} + +/** Earliest future shared short/weekly reset; missing evidence and ties use usage order. */ +export function pickResetFirstCodexAccount( + config: OcxConfig, + ids: readonly string[], + now: number, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + const available = ids.filter(id => hasCodexQuotaHeadroom(config, id, selectionOptions, now)); + if (available.length === 0) return pickLowestUsageAmong(config, ids, selectionOptions, now); + let earliest = Number.POSITIVE_INFINITY; + let candidates: string[] = []; + for (const id of available) { + const quota = getAccountQuota(id); + const resets = [quota?.shortResetAt, quota?.weeklyResetAt] + .filter((reset): reset is number => typeof reset === "number" && Number.isFinite(reset)) + .map(resetAtToMs) + .filter(reset => reset > now); + const next = Math.min(...resets); + if (next < earliest) { + earliest = next; + candidates = [id]; + } else if (next === earliest) candidates.push(id); + } + return pickLowestUsageAmong(config, candidates, selectionOptions, now); +} + +/** + * Fill-first: keep selectable active under threshold; otherwise advance to the next + * eligible id in stable sorted order after the current active (wrapping). + */ +function pickFillFirstCodexAccount( + config: OcxConfig, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + const eligible = listEligibleCodexAccountIds(config, now, quotaScope, selectionOptions); + if (eligible.length === 0) return null; + + const active = getEffectiveActiveCodexAccountId(config); + if (active && eligible.includes(active) && hasCodexQuotaHeadroom(config, active, selectionOptions, now)) { + return active; + } + + return pickNextFillFirstCodexAccount(config, active ?? null, eligible, now, selectionOptions); +} + +/** Next eligible account in stable order after `afterId` (wrapping). */ +function pickNextFillFirstCodexAccount( + config: OcxConfig, + afterId: string | null, + eligible: readonly string[] = listEligibleCodexAccountIds(config, Date.now()), + now = Date.now(), + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + if (eligible.length === 0) return null; + const ordered = [...eligible].sort((a, b) => a.localeCompare(b)); + if (!afterId) { + // Prefer an under-threshold account when starting with no active cursor. + for (const id of ordered) { + if (hasCodexQuotaHeadroom(config, id, selectionOptions, now)) return id; + } + return ordered[0] ?? null; + } + + const allConfigured = [ + ...(isCodexAccountUsable(config, MAIN_CODEX_ACCOUNT_ID, selectionOptions) || afterId === MAIN_CODEX_ACCOUNT_ID + ? [MAIN_CODEX_ACCOUNT_ID] + : []), + ...(config.codexAccounts ?? []).filter(account => !account.isMain).map(account => account.id), + ]; + const stableAll = [...new Set(allConfigured)].sort((a, b) => a.localeCompare(b)); + const startIdx = stableAll.indexOf(afterId); + if (startIdx < 0) { + for (const id of ordered) { + if (hasCodexQuotaHeadroom(config, id, selectionOptions, now)) return id; + } + return ordered[0] ?? null; + } + + // Skip successors that are also at/above threshold (known drained usage). + let fallback: string | null = null; + for (let step = 1; step <= stableAll.length; step++) { + const candidate = stableAll[(startIdx + step) % stableAll.length]!; + if (!eligible.includes(candidate)) continue; + if (!fallback) fallback = candidate; + if (hasCodexQuotaHeadroom(config, candidate, selectionOptions, now)) return candidate; + } + return fallback ?? ordered[0] ?? null; +} + +/** + * Unbound new-session pick for round-robin / fill-first. Returns null to fall through + * to the legacy quota path (or when the strategy is quota). + * + * When `commit` is true (resolve path), advances RR state. `commitSharedActive` + * and `commitAffinity` independently control the two cross-request side effects: + * model-scoped entitlement selection can bind a new task without replacing an + * existing task binding or global active choice. Preview remains a dry-run peek. + * + * Automatic strategy picks never sync-write config; only manual selection persists active. + * + * Known limitation (follow-up): when a subagent preview peeks an RR account and the request + * then falls back to a non-Codex provider, the ring is not reserved/committed. Prefer seeding + * the peeked account if that path becomes load-bearing. + */ +export function pickUnboundStrategyAccount( + config: OcxConfig, + threadId: string | null, + now: number, + commit: boolean, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + commitSharedActive = commit, + commitAffinity = commit, +): string | null { + const strategy = accountPoolStrategyForScope(config, quotaScope); + if (strategy === "quota") return null; + const poolKey = codexPoolKeyForScope(quotaScope); + + let picked: string | null = null; + if (strategy === "round-robin") { + const eligible = listEligibleCodexAccountIds(config, now, quotaScope, selectionOptions); + const limit = stickyLimitForConfig(config); + if (!commit) { + return peekRoundRobinAccount(poolKey, eligible, limit); + } + picked = pickRoundRobinAccount(poolKey, eligible, limit); + if (!picked) return null; + if (commitSharedActive) { + if (!isIndependentCodexQuotaScope(quotaScope) + && !manualPreferenceBlocks(codexPoolKeyForScope(quotaScope), picked)) { + rememberActiveCodexAccount(config, picked); + } + } + if (commitAffinity && threadId) bindThreadAffinity(threadId, picked, now, quotaScope); + notePoolRotationSuccess(poolKey, picked, limit); + return picked; + } + + if (strategy === "fill-first" || strategy === "reset-first") { + picked = strategy === "reset-first" + ? pickResetFirstCodexAccount(config, listEligibleCodexAccountIds(config, now, quotaScope, selectionOptions), now, selectionOptions) + : pickFillFirstCodexAccount(config, now, quotaScope, selectionOptions); + if (!picked) return null; + if (commitSharedActive) { + if (!isIndependentCodexQuotaScope(quotaScope) + && !manualPreferenceBlocks(codexPoolKeyForScope(quotaScope), picked)) { + rememberActiveCodexAccount(config, picked); + } + } + if (commitAffinity && threadId) bindThreadAffinity(threadId, picked, now, quotaScope); + return picked; + } + + return null; +} + +export function getPoolAccountPlan(config: OcxConfig, accountId: string): string | undefined { + if (accountId === MAIN_CODEX_ACCOUNT_ID) return getMainAccountPlan(); + return (config.codexAccounts ?? []) + .find(account => isSelectableCodexPoolAccount(account) && account.id === accountId)?.plan; +} + +/** Selection-only main routing must not lazily read the fenced native credential for its plan. */ +export function getPoolAccountPlanForSelection( + config: OcxConfig, + accountId: string, + selectionOptions?: CodexAccountUsabilityOptions, +): string | undefined { + if (accountId === MAIN_CODEX_ACCOUNT_ID && selectionOptions?.nativeMainSelectionOnly === true) { + return undefined; + } + return getPoolAccountPlan(config, accountId); +} + +/** Shared routing state must ignore a request-scoped entitlement roster. */ +export function sharedStateSelectionOptions( + selectionOptions?: CodexAccountUsabilityOptions, +): Pick< + CodexAccountUsabilityOptions, + "nativeMainSelectionOnly" | "isMainAccountTokenLive" +> | undefined { + if (!selectionOptions) return undefined; + return { + ...(selectionOptions.nativeMainSelectionOnly !== undefined + ? { nativeMainSelectionOnly: selectionOptions.nativeMainSelectionOnly } + : {}), + ...(selectionOptions.isMainAccountTokenLive + ? { isMainAccountTokenLive: selectionOptions.isMainAccountTokenLive } + : {}), + }; +} + +export function pickLowerUsageAccount( + config: OcxConfig, + active: string, + activeUsage: number, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + skipFailoverReadyCandidates = false, +): string { + let best = active; + let bestUsage = activeUsage; + for (const id of getEligiblePoolAccounts( + config, + active, + now, + quotaScope, + selectionOptions, + skipFailoverReadyCandidates, + )) { + const usage = computeCodexUsageScore( + getAccountQuota(id), + getPoolAccountPlanForSelection(config, id, selectionOptions), + now, + ); + if (usage < bestUsage) { + best = id; + bestUsage = usage; + } + } + return best; +} + +/** Coolest account in an already-selected candidate list; first index wins ties. */ +export function pickLowestUsageAmong( + config: OcxConfig, + ids: readonly string[], + selectionOptions?: CodexAccountUsabilityOptions, + now: number = Date.now(), +): string | null { + let best: string | null = null; + let bestUsage = Number.POSITIVE_INFINITY; + for (const id of ids) { + const usage = computeCodexUsageScore( + getAccountQuota(id), + getPoolAccountPlanForSelection(config, id, selectionOptions), + now, + ); + if (usage < bestUsage) { + best = id; + bestUsage = usage; + } + } + return best; +} + +export function pickLowestUsageCodexAccount( + config: OcxConfig, + excludeId?: string, + now = Date.now(), + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + return pickLowestUsageAmong( + config, + getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions), + selectionOptions, + now, + ); +} + +/** + * Strategy-aware alternate after a cooled/excluded account (same-request 429 retry + * and active promotion). Quota keeps lowest-usage; fill-first advances stable order; + * round-robin takes the next ring pick (caller should have noted the failure). + */ +export function pickAlternateCodexAccount( + config: OcxConfig, + excludeId: string, + now = Date.now(), + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + const strategy = accountPoolStrategyForScope(config, quotaScope); + // The exclusion is passed into eligibility rather than post-filtered off its + // result: when the excluded account is the only healthy member of the top + // tier, the tier walk must be free to descend instead of selecting that tier + // and then handing back an empty list. + if (strategy === "round-robin") { + const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); + return pickRoundRobinAccount(codexPoolKeyForScope(quotaScope), eligible, stickyLimitForConfig(config)); + } + if (strategy === "fill-first") { + const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); + return pickNextFillFirstCodexAccount(config, excludeId, eligible, now, selectionOptions); + } + if (strategy === "reset-first") { + return pickResetFirstCodexAccount(config, getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions), now, selectionOptions); + } + return pickLowestUsageCodexAccount(config, excludeId, now, quotaScope, selectionOptions); +} + +/** + * The account {@link pickAlternateCodexAccount} WOULD return, without returning it. + * + * Only the round-robin branch has a side effect -- `pickRoundRobinAccount` commits the pick and + * advances the ring -- so every other strategy delegates rather than growing a second copy of + * the selection rule that could drift from it. + * + * This exists because preview and resolve have to agree on the FIRST transient detour, not just + * on later ones. Preview feeds subagent model-availability scoring, so a preview that reported + * the bound account while resolve was about to serve from a cool sibling could retire a model + * over usage the request would never have touched. + */ +export function peekAlternateCodexAccount( + config: OcxConfig, + excludeId: string, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + if (accountPoolStrategyForScope(config, quotaScope) === "round-robin") { + const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); + return peekRoundRobinAccount(codexPoolKeyForScope(quotaScope), eligible, stickyLimitForConfig(config)); + } + return pickAlternateCodexAccount(config, excludeId, now, quotaScope, selectionOptions); +} + +export function isUnknownUsage(usage: number): boolean { + return usage >= CODEX_UNKNOWN_USAGE_SCORE; +} + +/** + * Move an unbound request back up when a higher tier regains headroom — the + * weekly-reset case. Returns null when nothing should change. + * + * Downward moves are deliberately left to {@link applyQuotaAutoSwitch}: this only + * fires when the tier filter has already excluded `active`, and only toward a + * tier that strictly outranks it. Threads bound by affinity never reach here. + */ +export function pickPriorityPreemption( + config: OcxConfig, + active: string, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + const eligible = getEligiblePoolAccounts(config, undefined, now, quotaScope, selectionOptions); + if (eligible.length === 0 || eligible.includes(active)) return null; + const pinned = pinnedCodexAccountId(config); + // A live pin already lowered the tier ceiling; never preempt past an explicit + // operator choice. Same liveness test the tier filter applies, so preview and + // resolve agree even before the pin is garbage-collected. + if ( + pinned !== undefined + && eligible.includes(pinned) + && hasCodexQuotaHeadroom(config, pinned, selectionOptions, now) + ) return null; + const priorityOf = codexAccountPriorityLookup(config); + if (priorityOf(eligible[0]!) <= priorityOf(active)) return null; + // Members without headroom are in the tier only because a sibling has some; + // picking one would hand the request straight back to a drained account. + return pickLowestUsageAmong( + config, + eligible.filter(id => hasCodexQuotaHeadroom(config, id, selectionOptions, now)), + selectionOptions, + now, + ); +} + +export function applyQuotaAutoSwitch( + config: OcxConfig, + active: string, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + commitSharedSelection = true, +): string { + const threshold = config.autoSwitchThreshold ?? 80; + if (threshold <= 0) return active; + const quota = getAccountQuota(active); + const activeUsage = computeCodexUsageScore( + quota, + getPoolAccountPlanForSelection(config, active, selectionOptions), + now, + ); + // Unknown usage is not evidence that a user's explicit selection crossed the + // threshold. Wait for quota priming instead of rotating among guesses. + if (isUnknownUsage(activeUsage)) return active; + if (activeUsage < threshold) return active; + const best = pickLowerUsageAccount(config, active, activeUsage, now, quotaScope, selectionOptions); + if (best !== active) { + if (commitSharedSelection && !isIndependentCodexQuotaScope(quotaScope)) { + setActiveCodexAccount(config, best); + } + return best; + } + + return active; +} + +export function shouldFailover(config: OcxConfig, accountId: string, now: number): boolean { + const threshold = config.upstreamFailoverThreshold ?? 3; + if (threshold <= 0) return false; + dropSpentCredentialFailure(accountId); + const health = getAccountHealth(accountId); + if (health?.lastFailureAt && now - health.lastFailureAt > CODEX_FAILURE_WINDOW_MS) return false; + return !!health && health.consecutiveFailures >= threshold; +} + +export function isHealthySharedCodexSelection( + config: OcxConfig, + accountId: string, + now: number, + quotaScope: CodexQuotaScope | undefined, + selectionOptions: CodexAccountUsabilityOptions | undefined, +): boolean { + return isCodexAccountSelectable(config, accountId, now, quotaScope, selectionOptions) + && hasCodexQuotaHeadroom(config, accountId, selectionOptions, now) + && !shouldFailover(config, accountId, now); +} + +export function strategySelectionOptionsForModelDetour( + config: OcxConfig, + now: number, + quotaScope: CodexQuotaScope | undefined, + selectionOptions: CodexAccountUsabilityOptions | undefined, +): CodexAccountUsabilityOptions | undefined { + if (selectionOptions?.modelEligibleAccountIds === undefined) return selectionOptions; + const sharedSelectionOptions = sharedStateSelectionOptions(selectionOptions) ?? {}; + return { + ...selectionOptions, + modelEligibleAccountIds: new Set( + [...selectionOptions.modelEligibleAccountIds].filter(accountId => + isHealthySharedCodexSelection( + config, + accountId, + now, + quotaScope, + sharedSelectionOptions, + ) + ), + ), + }; +} + +export function applyFailureFailover( + config: OcxConfig, + active: string, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + commitSharedSelection = true, +): string { + if (!shouldFailover(config, active, now)) return active; + const best = pickAlternateCodexAccount(config, active, now, quotaScope, selectionOptions); + if (best) { + // The scope still routes away from the failing account — that is this request's + // own decision — but an independent one must not persist a new shared active + // account. recordCodexUpstreamOutcome only suppresses the promotion it makes at + // the moment of the failure; the streak outlives the soft avoid, so a later + // scoped resolve reaches here with the streak still tripped and would otherwise + // move the shared cursor after all. + if (commitSharedSelection && !isIndependentCodexQuotaScope(quotaScope)) { + promoteActiveCodexAccount(config, best); + } + return best; + } + return active; +} diff --git a/src/codex/routing/thread-affinity.ts b/src/codex/routing/thread-affinity.ts new file mode 100644 index 0000000000..cd493d20f4 --- /dev/null +++ b/src/codex/routing/thread-affinity.ts @@ -0,0 +1,419 @@ +import { isCodexAccountGenerationLive, readCodexAccountRecord } from "../account-store"; +import { MAIN_CODEX_ACCOUNT_ID } from "../main-account"; +import { retainedUtf8Bytes } from "../../lib/admission"; +import type { CodexQuotaScope } from "./health-store"; + +export type ThreadAffinityEntry = { + accountId: string; + generation: number; + createdAt: number; + lastUsedAt: number; + // Last time the bound account's quota threshold was re-evaluated for this + // thread (interval-gated to avoid per-request flapping). See REEVAL_INTERVAL_MS. + lastReevalAt: number; + // When a transient failure streak first forced this thread onto another account + // while the binding was HELD (#4546). Cleared the moment the bound account serves + // again; once it ages past CODEX_TRANSIENT_AFFINITY_HOLD_MS the binding is + // released through the ordinary path instead of detouring forever. + transientHoldSince?: number; + // Which account is serving this thread while its own is held under a transient hold. + // Remembered rather than re-picked per request: under round-robin a fresh pick each turn + // would walk the ring and start cold on every hop, which is the behaviour the hold exists + // to prevent. Cleared with transientHoldSince when the bound account serves again. + transientDetourAccountId?: string; +}; + +export type CodexThreadResolution = + | { status: "selected"; accountId: string; affinity?: CodexAffinityDecision } + | { status: "none"; affinity?: CodexAffinityDecision } + | { status: "expired"; accountId: string; affinity?: CodexAffinityDecision }; + +/** What happened to this thread's binding on this request (#4546). */ +export type CodexAffinityMove = + /** Served by its own bound account, which was healthy. */ + | "reused" + /** Served by its own bound account while something transient was wrong with it. */ + | "held" + /** Served by another account while the binding stayed put. */ + | "detour" + /** The binding was released and a different account took the thread. */ + | "rebound" + /** There was no live binding; this request established one. */ + | "new_bind" + /** The binding was released without a replacement on this request. */ + | "cleared"; + +/** + * Why. A move is the expensive event -- it discards the prompt-cache prefix warmed on the old + * account -- so the operator should not have to infer it from account labels across log lines, + * which is how #4546 had to be diagnosed. + */ +export type CodexAffinityReason = + | "healthy" + | "quota_headroom" + | "quota_refusal" + | "transient" + | "transient_hold_expired" + | "unusable" + | "paused" + | "plan_excluded" + | "cooldown" + | "quota_avoided" + | "generation" + | "expired" + | "model_lane"; + +export interface CodexAffinityDecision { + move: CodexAffinityMove; + reason: CodexAffinityReason; +} + +/** The decision to report once a binding has been released and selection starts over. */ +export function affinityAfterRelease( + threadId: string | null, + releaseReason: CodexAffinityReason | undefined, +): CodexAffinityDecision { + // Reported now, so it must not be reported again by the next request. + clearPendingReleaseReason(threadId); + return releaseReason === undefined + ? { move: "new_bind", reason: "healthy" } + : { move: "rebound", reason: releaseReason }; +} + +/** + * What to report when selection produced no account at all. The binding is gone and nothing took + * it, which is a `cleared`, and the pending reason is deliberately NOT consumed: a no-account + * result reaches no auth context and therefore no usage entry, so the next resolve that does + * produce one is the first place this release can actually be seen. + */ +export function affinityOnNoAccount( + threadId: string | null, + releaseReason: CodexAffinityReason | undefined, +): CodexAffinityDecision | undefined { + if (releaseReason === undefined) return undefined; + // Hand it forward as well as reporting it. A reason derived from the entry this request just + // released lives only in a local, so without this the next resolve finds no entry and no + // pending reason and calls the rebind a fresh healthy bind. + notePendingReleaseReason(threadId, releaseReason); + return { move: "cleared", reason: releaseReason }; +} + +export const CODEX_THREAD_AFFINITY_IDLE_TTL_MS = 24 * 60 * 60_000; +export const CODEX_THREAD_AFFINITY_MAX_ENTRIES = 2048; +const MAX_AFFINITY_COMPONENT_BYTES = 512; +// Min interval between quota threshold re-evaluations for a single bound thread. +// Well under the 5h/weekly quota windows, but enough to stop per-request flapping. +export const CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS = 60_000; + +/** + * How long a live binding outlives a TRANSIENT failure streak on its own account (#4546). + * + * Being unable to send right now is not the same as losing ownership of the conversation. + * A 5xx streak is frequently provider-wide rather than account-specific, and deleting the + * binding for it discards a prompt-cache prefix that the next turn then pays for again -- + * the same cost the quota threshold used to impose, arriving through a different door. + * So the request detours to another account while the binding is held here. + * + * Bounded, because an unbounded hold is its own defect: an account that never recovers + * would keep a thread detouring indefinitely while the conversation's real warm prefix + * accumulates somewhere else. Ten minutes is longer than the whole soft-avoid escalation + * ladder up to its final step, so an ordinary outage resolves inside the hold and a + * genuine one converts to a real rebind instead of a permanent detour. + */ +export const CODEX_TRANSIENT_AFFINITY_HOLD_MS = 10 * 60_000; + +/** + * Requests without a resolved native model retain the historic one-account-per- + * thread behavior. Requests with a known quota scope get an independent + * affinity so a Reserve failover cannot displace the same thread's Terra/Luna + * account (and vice versa). + */ +type BaseThreadAffinityScope = CodexQuotaScope | "legacy"; +type ModelDetourAffinityScope = `model-detour:${BaseThreadAffinityScope}:${string}`; +type ThreadAffinityScope = BaseThreadAffinityScope | ModelDetourAffinityScope; + +function isModelDetourAffinityScope(scope: ThreadAffinityScope): scope is ModelDetourAffinityScope { + return scope.startsWith("model-detour:"); +} +const LEGACY_THREAD_AFFINITY_SCOPE = "legacy" as const; +const threadAccountMap = new Map>(); +let threadAffinityEntryTotal = 0; + +export function clearThreadAccountMap(): void { + threadAccountMap.clear(); + threadAffinityEntryTotal = 0; +} + +export function clearThreadAccountMapForAccount( + accountId: string, + reason: CodexAffinityReason = "unusable", +): void { + for (const [threadId, affinities] of threadAccountMap) { + for (const [scope, entry] of affinities) { + if (entry.accountId === accountId && affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + notePendingReleaseReason(threadId, reason); + } + } + if (affinities.size === 0) threadAccountMap.delete(threadId); + } +} + +/** + * Why a binding was released, held until that thread's next resolve can report it (#4546). + * + * A release and the request that pays for it are two different moments: a 429 clears the pin + * inside the outcome recorder, and the next request arrives with nothing left to explain why it + * is starting cold. Bounded, because it is a diagnostic and must not become a leak. + */ +const pendingReleaseReasons = new Map(); +const MAX_PENDING_RELEASE_REASONS = 4096; + +function notePendingReleaseReason(threadId: string | null, reason: CodexAffinityReason): void { + if (threadId === null) return; + if (!pendingReleaseReasons.has(threadId) && pendingReleaseReasons.size >= MAX_PENDING_RELEASE_REASONS) { + const oldest = pendingReleaseReasons.keys().next(); + if (!oldest.done) pendingReleaseReasons.delete(oldest.value); + } + pendingReleaseReasons.set(threadId, reason); +} + +export function peekPendingReleaseReason(threadId: string | null): CodexAffinityReason | undefined { + if (threadId === null) return undefined; + return pendingReleaseReasons.get(threadId); +} + +/** + * Forget a release only once it has actually been reported. + * + * Consuming it at derivation time lost it whenever selection then failed to produce an account: + * a no-account return carries no payload, so the release went unrecorded and the next successful + * resolve claimed a fresh healthy bind (#4598). A release survives until some resolve reports it. + */ +function clearPendingReleaseReason(threadId: string | null): void { + if (threadId !== null) pendingReleaseReasons.delete(threadId); +} + +function threadAffinityScope(quotaScope?: CodexQuotaScope): BaseThreadAffinityScope { + return quotaScope ?? LEGACY_THREAD_AFFINITY_SCOPE; +} + +function admissibleAffinityComponent(value: string): boolean { + return retainedUtf8Bytes(value) <= MAX_AFFINITY_COMPONENT_BYTES; +} + +function modelDetourAffinityScope( + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): ModelDetourAffinityScope | undefined { + const canonicalModelId = modelId?.trim().toLowerCase(); + if (!canonicalModelId || !admissibleAffinityComponent(canonicalModelId)) return undefined; + return `model-detour:${threadAffinityScope(quotaScope)}:${canonicalModelId}`; +} + +function getThreadAffinityForScope( + threadId: string, + scope: ThreadAffinityScope, +): ThreadAffinityEntry | undefined { + if (!admissibleAffinityComponent(threadId)) return undefined; + return threadAccountMap.get(threadId)?.get(scope); +} + +export function getThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): ThreadAffinityEntry | undefined { + return getThreadAffinityForScope(threadId, threadAffinityScope(quotaScope)); +} + +export function getModelDetourAffinity( + threadId: string, + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): ThreadAffinityEntry | undefined { + const scope = modelDetourAffinityScope(modelId, quotaScope); + return scope ? getThreadAffinityForScope(threadId, scope) : undefined; +} + +function deleteThreadAffinityForScope(threadId: string, scope: ThreadAffinityScope): void { + if (!admissibleAffinityComponent(threadId)) return; + const affinities = threadAccountMap.get(threadId); + if (!affinities) return; + if (affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + } + if (affinities.size === 0) threadAccountMap.delete(threadId); +} + +export function deleteThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): void { + deleteThreadAffinityForScope(threadId, threadAffinityScope(quotaScope)); +} + +export function deleteModelDetourAffinity( + threadId: string, + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): void { + const scope = modelDetourAffinityScope(modelId, quotaScope); + if (scope) deleteThreadAffinityForScope(threadId, scope); +} + +/** Remove only the matching failed account's affinities for one thread. */ +export function deleteThreadAffinitiesForAccount(threadId: string, accountId: string): void { + if (!admissibleAffinityComponent(threadId) || !admissibleAffinityComponent(accountId)) return; + const affinities = threadAccountMap.get(threadId); + if (!affinities) return; + for (const [scope, entry] of affinities) { + if (entry.accountId === accountId && affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + } + } + if (affinities.size === 0) threadAccountMap.delete(threadId); +} + +function threadAffinityEntryCount(): number { + return threadAffinityEntryTotal; +} + +export function isThreadAffinityExpired(entry: ThreadAffinityEntry, now: number): boolean { + return now - entry.lastUsedAt > CODEX_THREAD_AFFINITY_IDLE_TTL_MS; +} + +export function isThreadAffinityGenerationLive(entry: ThreadAffinityEntry): boolean { + if (entry.accountId === MAIN_CODEX_ACCOUNT_ID) return entry.generation === 0; + return isCodexAccountGenerationLive(entry.accountId, entry.generation); +} + +/** Generations this account's affinity entries are bound at. Test observability only. */ +export function debugCodexAffinityGenerations(accountId: string): number[] { + const generations: number[] = []; + for (const affinities of threadAccountMap.values()) { + for (const entry of affinities.values()) { + if (entry.accountId === accountId) generations.push(entry.generation); + } + } + return generations; +} + +/** + * Advance this account's affinity entries from the generation a rejected credential + * was bound under to the generation its own refresh produced. + * + * A 401 refresh-and-replay keeps the request on the same account, but the CAS write + * moves the credential from G to G+1, and {@link isThreadAffinityGenerationLive} + * demands exact equality — so without this the entry the replay just preserved is + * dead on the next request. Not quarantining an account is not the same as keeping + * its affinity. + * + * Lineage is proven by the CALLER, which must pass only a generation its own refresh + * produced. Re-deriving it here from `replacedAt` cannot work: the caller reads that + * field after the refresh and this function would re-read the same record, so the + * comparison is tautological and an external replacement passes it. An external + * replacement must retire the affinity, because that credential may belong to a + * different upstream identity. + */ +export function handOffThreadAffinityGeneration( + accountId: string, + fromGeneration: number, + toGeneration: number, +): boolean { + if (accountId === MAIN_CODEX_ACCOUNT_ID) return false; + if (toGeneration !== fromGeneration + 1) return false; + const record = readCodexAccountRecord(accountId); + if (!record?.credential || record.deletedAt != null) return false; + if (record.generation !== toGeneration) return false; + let handedOff = false; + for (const affinities of threadAccountMap.values()) { + for (const entry of affinities.values()) { + if (entry.accountId !== accountId || entry.generation !== fromGeneration) continue; + entry.generation = toGeneration; + handedOff = true; + } + } + return handedOff; +} + +function pruneExpiredThreadAffinities(now: number): void { + for (const [threadId, affinities] of threadAccountMap) { + for (const [scope, entry] of affinities) { + if (isThreadAffinityExpired(entry, now) && affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + } + } + if (affinities.size === 0) threadAccountMap.delete(threadId); + } +} + +function pruneLruThreadAffinities(): void { + if (threadAffinityEntryCount() <= CODEX_THREAD_AFFINITY_MAX_ENTRIES) return; + while (threadAffinityEntryCount() > CODEX_THREAD_AFFINITY_MAX_ENTRIES) { + let oldestThreadId: string | null = null; + let oldestScope: ThreadAffinityScope | null = null; + let oldestLastUsedAt = Number.POSITIVE_INFINITY; + let oldestIsDetour = false; + for (const [threadId, affinities] of threadAccountMap) { + for (const [scope, entry] of affinities) { + const candidateIsDetour = isModelDetourAffinityScope(scope); + if ( + (candidateIsDetour && !oldestIsDetour) + || (candidateIsDetour === oldestIsDetour && entry.lastUsedAt < oldestLastUsedAt) + ) { + oldestThreadId = threadId; + oldestScope = scope; + oldestLastUsedAt = entry.lastUsedAt; + oldestIsDetour = candidateIsDetour; + } + } + } + if (!oldestThreadId || !oldestScope) return; + deleteThreadAffinityForScope(oldestThreadId, oldestScope); + } +} + +function bindThreadAffinityForScope( + threadId: string, + accountId: string, + now: number, + scope: ThreadAffinityScope, +): void { + if (!admissibleAffinityComponent(threadId) || !admissibleAffinityComponent(accountId)) return; + const record = accountId === MAIN_CODEX_ACCOUNT_ID ? undefined : readCodexAccountRecord(accountId); + if (accountId !== MAIN_CODEX_ACCOUNT_ID && (!record?.credential || record.deletedAt != null)) return; + pruneExpiredThreadAffinities(now); + const affinities = threadAccountMap.get(threadId) ?? new Map(); + const previous = affinities.get(scope); + affinities.set(scope, { + accountId, + generation: accountId === MAIN_CODEX_ACCOUNT_ID ? 0 : record!.generation, + createdAt: previous?.createdAt ?? now, + lastUsedAt: now, + lastReevalAt: now, + }); + if (!previous) threadAffinityEntryTotal += 1; + threadAccountMap.set(threadId, affinities); + pruneLruThreadAffinities(); +} + +export function bindThreadAffinity( + threadId: string, + accountId: string, + now: number, + quotaScope?: CodexQuotaScope, +): void { + bindThreadAffinityForScope(threadId, accountId, now, threadAffinityScope(quotaScope)); +} + +export function bindModelDetourAffinity( + threadId: string, + accountId: string, + now: number, + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): void { + const scope = modelDetourAffinityScope(modelId, quotaScope); + if (scope) bindThreadAffinityForScope(threadId, accountId, now, scope); +} + +/** Read-only view of one thread's scope-keyed affinity entries. */ +export function getThreadAffinityScopes( + threadId: string, +): ReadonlyMap | undefined { + return threadAccountMap.get(threadId); +} diff --git a/src/providers/quota.ts b/src/providers/quota.ts index 386e921391..e02447ffb0 100644 --- a/src/providers/quota.ts +++ b/src/providers/quota.ts @@ -1,3078 +1,102 @@ -import { createHash } from "node:crypto"; -import { - effectiveCodexAuthAccountId, - fetchMainAccountInfoSnapshot, - listCodexAuthAccountsSnapshot, -} from "../codex/auth-api"; -import { withoutRetiredCodexQuota, type StoredAccountQuota } from "../codex/quota"; -import { isMainAccountIdentityGenerationLive } from "../codex/main-account-cache"; -import { MAIN_CODEX_ACCOUNT_ID } from "../codex/main-account"; -import { codexPlanKey } from "../codex/plan"; -import { resolveEnvValue } from "../config"; -import { resolveProviderApiKey } from "./key-store"; -import { getValidAccessToken, getValidAccessTokenForAccount } from "../oauth"; -import { getAccountCredential, getAccountSet, getCredential } from "../oauth/store"; -import { antigravityUserAgent } from "../adapters/client-fingerprint"; -import { isCanonicalOllamaCloudUrl } from "../adapters/ollama-native-url"; -import { DestinationDnsResolutionError } from "../lib/destination-policy"; -import { PinnedHttpError } from "../lib/pinned-http"; -import { ProviderOutboundPolicyError, providerOutboundPost, providerRedirectError, type ProviderOutboundDependencies } from "../lib/provider-outbound"; -import { apiKeyPoolEntryId } from "./api-keys"; -import { fetchMuseKeyQuotaSnapshot } from "./muse-key-quota"; -import { XAI_GROK_CLIENT_VERSION, XAI_GROK_COMPATIBILITY } from "./xai-transport"; -import { getProviderRegistryEntry, providerCodexAccountMode, registryEntryForProviderDestination } from "./registry"; -import type { OcxConfig, OcxProviderConfig } from "../types"; -import { isCanonicalOpenAiForwardProvider, OPENAI_CODEX_PROVIDER_ID } from "./openai-tiers"; -import { - captureConfigGeneration, - sweepExpiredOnWrite, - type GenerationContext, -} from "../lib/state-store-sweeper"; -import { - ACCOUNT_QUOTA_TTL_MS, - asRecord, - CACHE_TTL_MS, - normalizePercent, - normalizeResetAt, - QUOTA_JSON_READ_FAILURE, - readQuotaJson, - REQUEST_TIMEOUT_MS, - toFiniteNumber, -} from "./quota-wire"; -import { - clearCachedProviderQuotas, - providerQuotaRoutingBinding, - replaceCachedProviderQuotas, - type ProviderQuotaRoutingEvidence, -} from "./quota-routing-cache"; -import { - aggregateCodexPoolCapacity, - CODEX_CAPACITY_MAX_QUOTA_AGE_MS, - type CodexCapacityAggregation, - type CodexCapacityQuota, -} from "./codex-capacity"; -import type { - AccountQuotaMode, - QuotaFailureCode, - ProviderQuota, - ProviderQuotaCreditsUsd, - ProviderQuotaWindow, - ProviderRoutingQuota, -} from "./quota-types"; -import { - clearKiroAccountUsageState, - commitKiroAccountUsageState, - fetchKiroUsageSnapshot, - type KiroUsageSnapshot, - kiroUsageContextForAccount, - reconcileKiroAccountUsageState, -} from "./kiro-usage"; -import { - cancelPendingAccountQuotaPersist, - readPersistedAccountQuotas, - schedulePersistAccountQuotas, -} from "./account-quota-disk"; -import { clearProviderApiKeyQuotaCache, mapQuotaRoster, readProviderApiKeyQuotas, type ProviderApiKeyQuota } from "./quota-key-accounts"; - -export type { ProviderQuota, ProviderQuotaCreditsUsd, ProviderQuotaWindow } from "./quota-types"; - -/** Match oauth/index REFRESH_SKEW_MS — use stored access without refresh when still fresh. */ -const ACCOUNT_TOKEN_SKEW_MS = 60_000; -/** Successful provider quota payloads are small; reject oversized or stalled JSON before parsing. */ -export { QUOTA_RESPONSE_MAX_BYTES } from "./quota-wire"; -const KIMI_CODE_BASE_URL = "https://api.kimi.com/coding/v1"; -const KIMI_CODE_USAGE_URL = `${KIMI_CODE_BASE_URL}/usages`; -const COMMAND_CODE_BASE_URL = "https://api.commandcode.ai"; -const COMMAND_CODE_WHOAMI_URL = `${COMMAND_CODE_BASE_URL}/alpha/whoami`; -const COMMAND_CODE_CREDITS_URL = `${COMMAND_CODE_BASE_URL}/alpha/billing/credits`; -const COMMAND_CODE_SUBSCRIPTIONS_URL = `${COMMAND_CODE_BASE_URL}/alpha/billing/subscriptions`; -const COMMAND_CODE_USAGE_URL = `${COMMAND_CODE_BASE_URL}/alpha/usage/summary`; -const A6API_BASE_URL = "https://api.a6api.com"; -const OPENCODE_GO_BASE_URL = "https://opencode.ai/zen/go/v1"; -const OPENCODE_GO_USAGE_URL = `${OPENCODE_GO_BASE_URL}/usage`; -const OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"; -const DEEPSEEK_BASE_URL = "https://api.deepseek.com"; -const CLINE_BASE_URL = "https://api.cline.bot"; -const OLLAMA_CLOUD_BASE_URL = "https://ollama.com"; -const OLLAMA_CLOUD_USAGE_URL = `${OLLAMA_CLOUD_BASE_URL}/api/usage`; -const ZAI_BASE_URL = "https://api.z.ai"; -const ZAI_CN_BASE_URL = "https://open.bigmodel.cn"; -const MINIMAX_REMAINS_URL = "https://www.minimax.io/v1/token_plan/remains"; -const MOONSHOT_BASE_URL = "https://api.moonshot.ai/v1"; -const VENICE_BASE_URL = "https://api.venice.ai/api/v1"; -const SYNTHETIC_BASE_URL = "https://api.synthetic.new/v2"; -const DEEPINFRA_BASE_URL = "https://api.deepinfra.com"; -const NEURALWATT_BASE_URL = "https://api.neuralwatt.com/v1"; -const XAI_BILLING_URL = "https://cli-chat-proxy.grok.com/v1/billing"; -const XAI_CREDITS_URL = `${XAI_BILLING_URL}?format=credits`; -/** Keep a failed probe's previous row at most this long before dropping it. */ -const LAST_GOOD_MAX_AGE_MS = CODEX_CAPACITY_MAX_QUOTA_AGE_MS; -const nativeMainReportGenerations = new WeakMap(); -const accountReportCurrent = new WeakMap boolean>(); -const routingEvidence = new WeakMap(); -let providerQuotaBeforePublishForTests: (() => void | Promise) | null = null; - -/** Test-only seam for identity/config invalidation after probes but before publication. */ -export function setProviderQuotaBeforePublishForTests( - hook: (() => void | Promise) | null, -): void { - providerQuotaBeforePublishForTests = hook; -} -const TERMINAL_QUOTA_FAILURE = Symbol("terminal-quota-failure"); -/** - * The probe succeeded and the upstream authoritatively reported NO model-quota windows. - * - * Distinct from `null`, which means "this probe told us nothing" and deliberately preserves - * the last-good row for up to 30 minutes. Collapsing the two would let a stale report outlive - * the authoritative answer that replaced it: a GLM plan whose payload carries only MCP - * `TIME_LIMIT` rows has no model windows, and the dashboard and quota-aware routing must stop - * showing the previous token windows rather than keep them for another half hour. - * - * Suppression is shared with `TERMINAL_QUOTA_FAILURE`; only the reason differs. - */ -const AUTHORITATIVE_EMPTY_QUOTA = Symbol("authoritative-empty-quota"); -type ProviderQuotaProbeResult = - | ProviderQuotaReport - | null - | typeof TERMINAL_QUOTA_FAILURE - | typeof AUTHORITATIVE_EMPTY_QUOTA; - -export interface ProviderQuotaReport { - provider: string; - label: string; - source: string; - quota: ProviderQuota; - updatedAt: number; - /** Added by the management response projection, never stored on a cached report. */ - routingQuota?: ProviderRoutingQuota; - reverseEngineered?: boolean; - /** - * The row was OBSERVED in-band on a streaming turn rather than probed. - * - * Age means something different for these. A probed provider re-reads on its own TTL, - * so a row older than the last-good bound means the probe is failing and showing it - * would misrepresent a live number. A passive provider publishes no endpoint at all - * (`hasPassiveAccountQuota`), so its last observation is not a stale reading of - * something fresher — it is the only measurement that exists, and dropping it leaves - * the operator with nothing. Consumers that enforce a freshness bound must exempt - * these and state the observation age instead. - */ - observed?: boolean; - aggregation?: CodexCapacityAggregation; -} - -export interface ProviderQuotaResponse { - generatedAt: number; - reports: ProviderQuotaReport[]; -} - -let cache: { key: string; ts: number; response: ProviderQuotaResponse } | null = null; -const inflight = new Map }>(); -/** Bumped on cache clear and on force-refresh start; stale-epoch probes lose commit authority. */ -let invalidationEpoch = 0; - -/** Invalidate the report cache (e.g. after switching a provider's active account). */ -export function clearProviderQuotaCache(): void { - cache = null; - clearCachedProviderQuotas(); - clearProviderApiKeyQuotaCache(); - invalidationEpoch += 1; -} - -function cacheKey(config: OcxConfig): string { - const providers = Object.entries(config.providers) - .map(([name, provider]) => { - const resolvedKey = typeof provider.apiKey === "string" - ? resolveProviderApiKey(provider.apiKey)?.trim() - : undefined; - const activeKeyId = resolvedKey ? apiKeyPoolEntryId(resolvedKey) : "none"; - return `${name}:${provider.adapter}:${provider.authMode ?? "key"}:${providerCodexAccountMode(name, provider) ?? "none"}:${provider.disabled === true ? "off" : "on"}:${provider.baseUrl}:${activeKeyId}`; - }) - .sort() - .join("|"); - return `${config.defaultProvider}|${providers}`; -} - -type CodexAuthAccountsSnapshotPromise = ReturnType; - -function hasCodexPoolProvider(config: OcxConfig): boolean { - return Object.entries(config.providers).some(([name, provider]) => ( - provider.disabled !== true - && isBuiltInChatGptForwardProvider(name, provider) - && providerCodexAccountMode(name, provider) !== "direct" - )); -} - -function quotaSignatureValue(quota: CodexCapacityQuota | null): unknown { - if (!quota) return null; - return { - fiveHourPercent: quota.fiveHourPercent, - fiveHourResetAt: quota.fiveHourResetAt, - weeklyPercent: quota.weeklyPercent, - weeklyResetAt: quota.weeklyResetAt, - monthlyPercent: quota.monthlyPercent, - monthlyResetAt: quota.monthlyResetAt, - updatedAt: quota.updatedAt, - customWindows: [...(quota.customWindows ?? [])] - .map(window => ({ label: window.label, percent: window.percent, resetAt: window.resetAt })) - .sort((a, b) => a.label.localeCompare(b.label)), - }; -} - -function providerQuotaFromCodexQuota( - quota: StoredAccountQuota | Omit | null | undefined, -): CodexCapacityQuota | null { - if (!quota) return null; - // Direct snapshots bypass account DTOs; sanitize here as well as at ingestion. - quota = withoutRetiredCodexQuota(quota); - if (!quota) return null; - const projected: CodexCapacityQuota = { - ...(quota.shortPercent !== undefined ? { fiveHourPercent: quota.shortPercent } : {}), - ...(quota.shortResetAt !== undefined ? { fiveHourResetAt: quota.shortResetAt } : {}), - ...(quota.weeklyPercent !== undefined ? { weeklyPercent: quota.weeklyPercent } : {}), - ...(quota.weeklyResetAt !== undefined ? { weeklyResetAt: quota.weeklyResetAt } : {}), - ...(quota.monthlyPercent !== undefined ? { monthlyPercent: quota.monthlyPercent } : {}), - ...(quota.monthlyResetAt !== undefined ? { monthlyResetAt: quota.monthlyResetAt } : {}), - ...(quota.customWindows !== undefined ? { customWindows: quota.customWindows } : {}), - updatedAt: "updatedAt" in quota ? quota.updatedAt : Date.now(), - }; - return hasQuotaRows(projected) ? projected : null; -} - -/** Hash only presentation-relevant state; account ids and email addresses never enter the key. */ -function cacheKeyWithAggregationState( - config: OcxConfig, - prefetchedSnapshot?: CodexAuthAccountsSnapshotPromise, -): string | Promise { - const base = cacheKey(config); - if (!hasCodexPoolProvider(config)) return base; - return (async () => { - try { - const activeId = effectiveCodexAuthAccountId(config); - const snapshot = await (prefetchedSnapshot ?? listCodexAuthAccountsSnapshot(config, false)); - const rows = snapshot.accounts.map(account => ({ - isMain: account.isMain, - active: account.id === activeId, - plan: codexPlanKey(account.plan) ?? null, - paused: account.paused, - needsReauth: account.needsReauth === true, - quota: quotaSignatureValue(providerQuotaFromCodexQuota(account.quota)), - })); - const canonicalRows = rows.map(row => JSON.stringify(row)).sort(); - const digest = createHash("sha256").update(JSON.stringify(canonicalRows)).digest("hex").slice(0, 24); - return `${base}|codex-pool:${digest}`; - } catch { - return `${base}|codex-pool:unavailable`; - } - })(); -} - -function publicCapacityWindow(window: import("./codex-capacity").CodexCapacityWindowAggregation) { - const { totalWeight: _totalWeight, consumedWeight: _consumedWeight, remainingWeight: _remainingWeight, ...safe } = window; - return safe; -} - -/** Management API metadata intentionally omits configured/weighted unit counts. */ -function publicCapacityAggregation( - aggregation: CodexCapacityAggregation, - presentation: NonNullable, -): CodexCapacityAggregation { - const safeCurrentAccount = presentation === "coverage-only" && aggregation.currentAccount - ? { ...aggregation.currentAccount, quota: null } - : aggregation.currentAccount; - return { - ...aggregation, - presentation, - ...(safeCurrentAccount ? { currentAccount: safeCurrentAccount } : {}), - ...(aggregation.fiveHour ? { fiveHour: publicCapacityWindow(aggregation.fiveHour) } : {}), - ...(aggregation.weekly ? { weekly: publicCapacityWindow(aggregation.weekly) } : {}), - ...(aggregation.monthly ? { monthly: publicCapacityWindow(aggregation.monthly) } : {}), - ...(aggregation.customWindows ? { - customWindows: aggregation.customWindows.map(window => ({ - label: window.label, - ...publicCapacityWindow(window), - })), - } : {}), - }; -} - -function hasQuotaRows(quota: ProviderQuota | null | undefined): quota is ProviderQuota { - if (!quota) return false; - return typeof quota.fiveHourPercent === "number" - || typeof quota.weeklyPercent === "number" - || typeof quota.monthlyPercent === "number" - || quota.creditsUsd?.unlimited === true - || typeof quota.creditsUsd?.percent === "number" - || !!quota.customWindows?.some(window => typeof window.percent === "number"); -} - -function providerLabel(providerId: string): string { - return getProviderRegistryEntry(providerId)?.label ?? providerId; -} - -/** Test-only access to the quota reader's deadline and cancellation contract. */ -export async function readProviderQuotaJsonForTests(response: Response, timeoutMs: number): Promise { - const result = await readQuotaJson(response, timeoutMs); - return result === QUOTA_JSON_READ_FAILURE ? null : result; -} - -function isBuiltInChatGptForwardProvider(name: string, provider: OcxProviderConfig): boolean { - return name === OPENAI_CODEX_PROVIDER_ID && isCanonicalOpenAiForwardProvider(provider); -} - -function isCanonicalA6apiBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === A6API_BASE_URL || normalized === `${A6API_BASE_URL}/v1`; -} - -function isCanonicalOpenCodeGoBaseUrl(baseUrl: string): boolean { - return normalizedBaseUrl(baseUrl) === OPENCODE_GO_BASE_URL; -} - -function isCanonicalOpenRouterBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === OPENROUTER_BASE_URL; -} - -function isCanonicalDeepSeekBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === DEEPSEEK_BASE_URL || normalized === `${DEEPSEEK_BASE_URL}/v1`; -} - -function isCanonicalClineBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === CLINE_BASE_URL || normalized === `${CLINE_BASE_URL}/api/v1`; -} - -function isCanonicalOllamaCloudBaseUrl(baseUrl?: string): boolean { - if (!baseUrl) return false; - try { - return isCanonicalOllamaCloudUrl(baseUrl); - } catch { - return false; - } -} - -function zaiQuotaMonitorHost(baseUrl: string): string | null { - // Admission and destination selection must share one mapping: admitting a new - // international wire must never fall through to the CN host/authentication scheme. - switch (normalizedBaseUrl(baseUrl)) { - case ZAI_BASE_URL: - case `${ZAI_BASE_URL}/api/coding/paas/v4`: - case `${ZAI_BASE_URL}/api/anthropic`: - case `${ZAI_BASE_URL}/api/v1`: - return ZAI_BASE_URL; - case ZAI_CN_BASE_URL: - case `${ZAI_CN_BASE_URL}/api/coding/paas/v4`: - case `${ZAI_CN_BASE_URL}/api/v1`: - return ZAI_CN_BASE_URL; - default: - return null; - } -} - -function isCanonicalZaiBaseUrl(baseUrl: string): boolean { - return zaiQuotaMonitorHost(baseUrl) !== null; -} - -function isCanonicalMinimaxBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === "https://api.minimax.io/v1" || normalized === "https://api.minimaxi.com/v1"; -} - -function isCanonicalMoonshotBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === MOONSHOT_BASE_URL || normalized === "https://api.moonshot.cn/v1"; -} - -function isCanonicalVeniceBaseUrl(baseUrl: string): boolean { - return normalizedBaseUrl(baseUrl) === VENICE_BASE_URL; -} - -function isCanonicalSyntheticBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === SYNTHETIC_BASE_URL || normalized === "https://api.synthetic.new/openai/v1"; -} - -function isCanonicalDeepInfraBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - return normalized === DEEPINFRA_BASE_URL || normalized === `${DEEPINFRA_BASE_URL}/v1/openai`; -} - -function isCanonicalNeuralwattBaseUrl(baseUrl: string): boolean { - return normalizedBaseUrl(baseUrl) === NEURALWATT_BASE_URL; -} - -function a6apiPayload(value: unknown): Record | null { - const body = asRecord(value); - return asRecord(body?.data) ?? body; -} - -function firstFinite(record: Record | null, names: string[]): number | undefined { - if (!record) return undefined; - for (const name of names) { - const value = toFiniteNumber(record[name]); - if (value !== undefined) return value; - } - return undefined; -} - -async function fetchA6apiQuota(provider: string, config: OcxProviderConfig): Promise { - // Never send a configured API key to a lookalike host or through a redirect. - if (!isCanonicalA6apiBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const headers = { Accept: "application/json", Authorization: `Bearer ${apiKey}` } as const; - const [subscriptionResponse, tokenResponse] = await Promise.all([ - fetch(`${A6API_BASE_URL}/dashboard/billing/subscription`, { - headers, redirect: "error", signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }), - fetch(`${A6API_BASE_URL}/api/usage/token/`, { - headers, redirect: "error", signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }), - ]); - if (!subscriptionResponse.ok || !tokenResponse.ok) { - const statuses = [subscriptionResponse.status, tokenResponse.status]; - // 408/429 are transient (timeout/throttle), not invalid-account signals: keep the - // last-good row like 5xx/network failures. 401/403 (bad key) and 404 (contract change) - // stay terminal. - return statuses.some(status => status >= 400 && status < 500 && status !== 429 && status !== 408) - ? TERMINAL_QUOTA_FAILURE - : null; - } - const [subscriptionBody, tokenBody] = await Promise.all([ - readQuotaJson(subscriptionResponse), - readQuotaJson(tokenResponse), - ]); - if (subscriptionBody === QUOTA_JSON_READ_FAILURE || tokenBody === QUOTA_JSON_READ_FAILURE) return null; - const subscription = a6apiPayload(subscriptionBody); - const token = a6apiPayload(tokenBody); - const unlimited = token?.unlimited_quota === true - || token?.unlimited_quota === 1 - || token?.unlimited_quota === "true"; - const normalizedExpiry = normalizeResetAt(token?.expires_at); - const expiry = normalizedExpiry && normalizedExpiry > 0 - ? { expiresAt: normalizedExpiry } - : {}; - if (unlimited) { - // Every row is an API-credit constraint on inference, so the display quota is also - // the routing projection. Passing it explicitly is the opt-in. - const quota: ProviderQuota = { - creditsUsd: { - used: 0, - limit: 0, - remaining: 0, - percent: 0, - unlimited: true, - ...expiry, - }, - customWindows: [{ label: "Unlimited API credits", percent: 0 }], - updatedAt: Date.now(), - }; - return keyReport(provider, "a6api:billing", quota, config, apiKey, quota); - } - const limitUsd = firstFinite(subscription, ["hard_limit_usd"]); - const grantedUnits = firstFinite(token, ["total_granted"]); - const usedUnits = firstFinite(token, ["total_used"]); - const availableUnits = firstFinite(token, ["total_available"]); - const reconciledUnits = usedUnits !== undefined && availableUnits !== undefined - ? usedUnits + availableUnits - : undefined; - const reconciliationTolerance = grantedUnits !== undefined - ? Math.abs(grantedUnits) * 1e-9 - : 0; - if (limitUsd === undefined || grantedUnits === undefined || usedUnits === undefined - || availableUnits === undefined || limitUsd <= 0 || grantedUnits <= 0 - || usedUnits < 0 || availableUnits < 0 - || reconciledUnits === undefined - || Math.abs(reconciledUnits - grantedUnits) > reconciliationTolerance) return TERMINAL_QUOTA_FAILURE; - const usdPerUnit = limitUsd / grantedUnits; - const usedUsd = usedUnits * usdPerUnit; - const remainingUsd = Math.max(0, availableUnits * usdPerUnit); - const percent = normalizePercent((usedUsd / limitUsd) * 100); - if (percent === undefined) return TERMINAL_QUOTA_FAILURE; - const label = `API credits ($${remainingUsd.toFixed(2)} of $${limitUsd.toFixed(2)} remaining)`; - const quota: ProviderQuota = { - creditsUsd: { - used: usedUsd, - limit: limitUsd, - remaining: remainingUsd, - percent, - ...expiry, - }, - customWindows: [{ label, percent }], - updatedAt: Date.now(), - }; - // The credit balance funds inference itself, so display and routing scope agree. - return keyReport(provider, "a6api:billing", quota, config, apiKey, quota); -} - -function parseOpenCodeGoUsageWindow(value: unknown): { percent: number; resetAt?: number } | null { - const row = asRecord(value); - if (!row) return null; - const percent = normalizePercent(row.percent); - if (percent === undefined) return null; - const resetAt = normalizeResetAt(row.resetsAt); - return { percent, ...(resetAt !== undefined ? { resetAt } : {}) }; -} - -async function fetchOpenCodeGoQuota(provider: string, config: OcxProviderConfig): Promise { - // Never send a configured API key when the provider destination is not the built-in Go endpoint. - if (!isCanonicalOpenCodeGoBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(OPENCODE_GO_USAGE_URL, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const usage = asRecord(body?.usage); - if (!usage) return null; - const rolling = parseOpenCodeGoUsageWindow(usage.rolling); - const weekly = parseOpenCodeGoUsageWindow(usage.weekly); - const monthly = parseOpenCodeGoUsageWindow(usage.monthly); - const quota: ProviderQuota = { - ...(rolling ? { - fiveHourPercent: rolling.percent, - ...(rolling.resetAt !== undefined ? { fiveHourResetAt: rolling.resetAt } : {}), - } : {}), - ...(weekly ? { - weeklyPercent: weekly.percent, - ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), - } : {}), - ...(monthly ? { - monthlyPercent: monthly.percent, - ...(monthly.resetAt !== undefined ? { monthlyResetAt: monthly.resetAt } : {}), - } : {}), - updatedAt: Date.now(), - }; - return keyReport(provider, "opencode-go:usage", quota, config, apiKey, quota); -} - -/** - * OpenRouter `GET /api/v1/key` — the key's own credit balance and optional - * per-key spending cap. `limit` is the configured cap (absent = uncapped); - * `usage` is lifetime spend; `limit_remaining` is what is left of the cap. - * When no cap is set there is no hard limit to meter against, so no bar is - * produced — the provider falls back to its documented reference. - */ -async function fetchOpenRouterQuota(provider: string, config: OcxProviderConfig): Promise { - // Never send a configured API key to a lookalike host or through a redirect. - if (!isCanonicalOpenRouterBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(`${OPENROUTER_BASE_URL}/key`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const data = asRecord(body?.data) ?? body; - if (!data) return null; - const limit = toFiniteNumber(data.limit); - const limitRemaining = toFiniteNumber(data.limit_remaining); - const usage = toFiniteNumber(data.usage); - // A successful no-cap response is a DELIBERATE change, not a transient - // failure: the old capped row must be dropped, not preserved as last-good. - if (limit === undefined || limit <= 0) return TERMINAL_QUOTA_FAILURE; - // Prefer the authoritative remaining-cap value when present: `usage` is - // lifetime accumulated spend and overstates a reset or re-capped key. - const used = limitRemaining !== undefined - ? Math.max(0, limit - limitRemaining) - : usage !== undefined && usage >= 0 ? usage : undefined; - if (used === undefined) return null; - const percent = normalizePercent((used / limit) * 100); - if (percent === undefined) return null; - const remaining = Math.max(0, limit - used); - const label = `API credits ($${remaining.toFixed(2)} of $${limit.toFixed(2)} remaining)`; - // The per-key spending cap stops every request this credential can make, so the - // whole report is inference-wide routing evidence. - const quota: ProviderQuota = { - customWindows: [{ label, percent }], - updatedAt: Date.now(), - }; - return keyReport(provider, "openrouter:key-info", quota, config, apiKey, quota); -} - -/** - * DeepSeek `GET /user/balance` — the account's granted + topped-up credit - * balance. The payload places `total_balance` / `granted_balance` inside - * entries of `balance_infos` (one row per currency); the row for the account's - * currency is selected by preference. `granted_balance` is a CURRENT balance - * component, not the original grant ceiling, so no consumed percentage is - * fabricated — the balance is reported as a balance-only window. - */ -async function fetchDeepSeekQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalDeepSeekBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(`${DEEPSEEK_BASE_URL}/user/balance`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - // The payload nests balances under `balance_infos` rows keyed by currency; - // prefer a USD row, then CNY, then the first row that parses. - const infos = Array.isArray(body?.balance_infos) ? body.balance_infos as unknown[] : null; - const rows = infos - ? infos.map((raw): Record | null => asRecord(raw)).filter((r): r is Record => r !== null) - : []; - const pick = (currency: string): Record | null => - rows.find(row => String(row.currency ?? "").toUpperCase() === currency) ?? null; - const preferred = pick("USD") ?? pick("CNY") ?? rows[0] ?? null; - if (!preferred) return null; - const totalBalance = toFiniteNumber(preferred.total_balance); - const grantedBalance = toFiniteNumber(preferred.granted_balance); - const toppedUp = toFiniteNumber(preferred.topped_up_balance); - const balance = totalBalance ?? grantedBalance ?? toppedUp; - if (balance === undefined || balance < 0) return null; - const label = grantedBalance !== undefined && grantedBalance > 0 - ? `API balance ($${balance.toFixed(2)} total, $${grantedBalance.toFixed(2)} granted)` - : `API balance ($${balance.toFixed(2)})`; - return report(provider, "deepseek:balance", { - customWindows: [{ label, percent: 0 }], - updatedAt: Date.now(), - }); -} - -/** - * ClinePass `GET /api/v1/users/me/plan/usage-limits` — the subscription's - * rolling five-hour, weekly, and monthly utilization, matching the existing - * ProviderQuota windows directly. The endpoint 404s (or returns a null plan) - * for accounts without an active ClinePass, which is a no-report, not an error. - */ -async function fetchClineQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalClineBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(`${CLINE_BASE_URL}/api/v1/users/me/plan/usage-limits`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - // 404 = no active plan; a plain "no plan" is a no-report, everything else - // 4xx (except 408/429) is a credential/contract problem. - if (response.status === 404) return null; - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const data = asRecord(body?.data) ?? body; - const limits = Array.isArray(data?.limits) ? data.limits : null; - if (!limits) return null; - const quota: ProviderQuota = { updatedAt: Date.now() }; - let windows = 0; - for (const raw of limits) { - const row = asRecord(raw); - if (!row) continue; - const percent = normalizePercent(row.percentUsed); - if (percent === undefined) continue; - const resetAt = normalizeResetAt(row.resetsAt); - if (row.type === "five_hour") { - quota.fiveHourPercent = percent; - if (resetAt !== undefined) quota.fiveHourResetAt = resetAt; - windows += 1; - } else if (row.type === "weekly") { - quota.weeklyPercent = percent; - if (resetAt !== undefined) quota.weeklyResetAt = resetAt; - windows += 1; - } else if (row.type === "monthly") { - quota.monthlyPercent = percent; - if (resetAt !== undefined) quota.monthlyResetAt = resetAt; - windows += 1; - } - } - return windows > 0 ? keyReport(provider, "cline:plan-usage-limits", quota, config, apiKey, quota) : null; -} - -/** - * Ollama Cloud `GET https://ollama.com/api/usage` — returns account usage. - * Legacy plans report rolling 5-hour `limits.session.usage` and 7-day - * `limits.weekly.usage`. Migrated monthly-credit plans report - * `limits.monthly.usage`. `usage` values are normalized fractions (0..1). - */ -function parseOllamaPercent(usageValue: unknown): number | undefined { - const usage = toFiniteNumber(usageValue); - if (usage === undefined || usage < 0) return undefined; - const percent = Math.round(usage * 10000) / 100; - return normalizePercent(percent); -} - -export function parseOllamaCloudQuota(body: Record | null): ProviderQuota | null { - if (!body) return null; - const limits = asRecord(body.limits); - if (!limits) return null; - - const quota: ProviderQuota = { updatedAt: Date.now() }; - let windows = 0; - - const session = asRecord(limits.session); - if (session) { - const percent = parseOllamaPercent(session.usage); - if (percent !== undefined) { - quota.fiveHourPercent = percent; - windows += 1; - } - } - - const weekly = asRecord(limits.weekly); - if (weekly) { - const percent = parseOllamaPercent(weekly.usage); - if (percent !== undefined) { - quota.weeklyPercent = percent; - windows += 1; - } - } - - const monthly = asRecord(limits.monthly); - if (monthly) { - const percent = parseOllamaPercent(monthly.usage); - if (percent !== undefined) { - quota.monthlyPercent = percent; - windows += 1; - } - } - - return windows > 0 ? quota : null; -} - -async function fetchOllamaCloudQuota(provider: string, config: OcxProviderConfig): Promise { - const effectiveBaseUrl = config.baseUrl ?? getProviderRegistryEntry(provider)?.baseUrl ?? ""; - if (!isCanonicalOllamaCloudBaseUrl(effectiveBaseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(OLLAMA_CLOUD_USAGE_URL, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - if (response.status === 404) return null; - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const quota = parseOllamaCloudQuota(body); - return quota ? keyReport(provider, "ollama-cloud:usage", quota, config, apiKey, quota) : null; -} - -/** - * Z.AI GLM Coding Plan `GET /api/monitor/usage/quota/limit` — the coding-plan - * limits arrive as a `limits` array of `TOKENS_LIMIT` (newer plans call the - * same rows `CREDIT_LIMIT`) and `TIME_LIMIT` rows. `TOKENS_LIMIT`/`CREDIT_LIMIT` - * rows carry the window length as `unit`/`number`: unit 3 is hours (number 5 → - * the rolling five-hour window), unit 6 is weeks (number 1 → the weekly - * window). Every row's `percentage` is the consumed share (falling - * back to `currentValue`/`usage` when absent) and `nextResetTime` (unix ms) - * the window reset. - * - * `TIME_LIMIT` rows are deliberately ignored (issue #1168). They are the shared - * monthly MCP *call* allowance for Web Search / Web Reader / Zread — not a - * model-token budget — and `ProviderQuota.monthlyPercent` is consumed as a - * model-capacity signal: `headroomOf()` in `src/oauth/account-quota-rank.ts` - * takes the MAX across every window, so a user who spent their MCP search - * allowance would be ranked as having no model capacity left, and the dashboard - * would draw a full monthly bar for a plan whose model tokens are untouched. - * A payload carrying only `TIME_LIMIT` rows therefore reports no quota at all, - * which is the honest answer rather than a fabricated one. - */ -export function parseZaiQuotaLimits(data: Record | null): ProviderQuota | null { - const limits = Array.isArray(data?.limits) ? data.limits as unknown[] : null; - if (!limits) return null; - const quota: ProviderQuota = { updatedAt: Date.now() }; - let windows = 0; - for (const raw of limits) { - const row = asRecord(raw); - if (!row) continue; - // Gate on row type before deriving a percentage: an MCP row must not even - // contribute a parsed value to a model-quota report. - if (row.type !== "TOKENS_LIMIT" && row.type !== "CREDIT_LIMIT") continue; - const resetAt = normalizeResetAt(row.nextResetTime); - let percent = normalizePercent(row.percentage); - if (percent === undefined) { - const used = toFiniteNumber(row.currentValue); - const total = toFiniteNumber(row.usage); - if (used !== undefined && total !== undefined && total > 0) { - percent = normalizePercent((used / total) * 100); - } - } - if (percent === undefined) continue; - const unit = toFiniteNumber(row.unit); - const number = toFiniteNumber(row.number); - if (unit === 3 && number === 5) { - quota.fiveHourPercent = percent; - if (resetAt !== undefined) quota.fiveHourResetAt = resetAt; - windows += 1; - } else if (unit === 6 && number === 1) { - quota.weeklyPercent = percent; - if (resetAt !== undefined) quota.weeklyResetAt = resetAt; - windows += 1; - } - } - return windows > 0 ? quota : null; -} - -/** - * Legacy Z.AI payload shape: percent fields with window identifiers directly on - * the data object (optionally nested under `quota`). Kept as a fallback so - * older responses keep rendering when the `limits` array is absent. - */ -function parseZaiQuotaLegacyFields(data: Record | null): ProviderQuota | null { - if (!data) return null; - const quota: ProviderQuota = { updatedAt: Date.now() }; - let windows = 0; - const percentAt = (key: string): number | undefined => { - const value = normalizePercent(data[key]); - if (value !== undefined) return value; - const nested = asRecord(data.quota); - return nested ? normalizePercent(nested[key]) : undefined; - }; - const fiveHour = percentAt("fiveHourPercent") ?? percentAt("fiveHourUsage") ?? percentAt("fiveHourUsed"); - const weekly = percentAt("weeklyPercent") ?? percentAt("weeklyUsage") ?? percentAt("weeklyUsed"); - const monthly = percentAt("monthlyPercent") ?? percentAt("mcpPercent") ?? percentAt("monthlyMCPUsage"); - if (fiveHour !== undefined) { - quota.fiveHourPercent = fiveHour; - windows += 1; - } - if (weekly !== undefined) { - quota.weeklyPercent = weekly; - windows += 1; - } - if (monthly !== undefined) { - quota.monthlyPercent = monthly; - windows += 1; - } - return windows > 0 ? quota : null; -} - -/** - * Fetches the Z.AI GLM Coding Plan quota — on whichever region the provider - * points at (api.z.ai or open.bigmodel.cn). The `limits` array shape is - * preferred; older field-name payloads fall back to the legacy parser. - * - * Authentication differs by host (issue #1168). `api.z.ai` takes the API key as - * a Bearer token per Z.AI's API reference; `open.bigmodel.cn` expects the key - * directly in `Authorization` with no scheme prefix and answers a Bearer header - * with an auth error, which is why BigModel Coding Plan quota never rendered. - * The host is already canonicalized by `isCanonicalZaiBaseUrl` above and - * `redirect: "error"` stays set, so the bare key cannot travel to a lookalike - * host or follow a redirect off-origin. - */ -async function fetchZaiQuota(provider: string, config: OcxProviderConfig): Promise { - const monitorHost = zaiQuotaMonitorHost(config.baseUrl); - if (!monitorHost) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const authorization = monitorHost === ZAI_CN_BASE_URL ? apiKey : `Bearer ${apiKey}`; - const response = await fetch(`${monitorHost}/api/monitor/usage/quota/limit`, { - headers: { Accept: "application/json", Authorization: authorization }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - if (!body || body.success === false) return null; - const data = asRecord(body.data) ?? body; - if (Array.isArray(data?.limits)) { - const quota = parseZaiQuotaLimits(data); - // A well-formed `limits[]` we fully understood is authoritative even when it yields no - // model window — for example a plan reporting only the monthly MCP `TIME_LIMIT` row. - // Returning `null` here would preserve the previous token windows for up to 30 minutes - // and keep quota-aware routing acting on a report the provider has already superseded. - return quota - ? keyReport(provider, "zai:quota-limit", quota, config, apiKey, quota) - : AUTHORITATIVE_EMPTY_QUOTA; - } - const legacy = parseZaiQuotaLegacyFields(data); - if (!legacy) return null; - // The legacy monthly figure also carries MCP usage; it is display evidence, not - // proof that model inference is unavailable. Modern TOKEN_LIMIT rows above are scoped. - const inferenceQuota = { ...legacy }; - delete inferenceQuota.monthlyPercent; - delete inferenceQuota.monthlyResetAt; - return keyReport(provider, "zai:quota-limit", legacy, config, apiKey, inferenceQuota); -} - -/** - * MiniMax Token Plan `GET /v1/token_plan/remains` — the subscription's - * remaining quota as a countdown-time value (ms). The endpoint does not expose - * the plan's total duration, so no percentage is fabricated from a presumed - * window: the remaining time is reported as a duration-only window. When the - * API supplies a total (`total_time` / `plan_duration_ms`), a consumed share - * is derived from it. Region selects the host: `minimax` → www.minimax.io, - * `minimax-cn` → api.minimaxi.com. - */ -async function fetchMinimaxQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalMinimaxBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const cnHost = normalizedBaseUrl(config.baseUrl)?.startsWith("https://api.minimaxi.com"); - const remainsUrl = cnHost ? "https://api.minimaxi.com/v1/token_plan/remains" : MINIMAX_REMAINS_URL; - const response = await fetch(remainsUrl, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - if (!body || body.success === false) return null; - const data = asRecord(body.data) ?? body; - const remainsMs = toFiniteNumber(data.remains_time ?? data.remainsTime); - if (remainsMs === undefined || remainsMs < 0) return null; - const hours = Math.floor(remainsMs / 3_600_000); - const label = `Token Plan remaining (${hours}h)`; - // Only derive a consumed share when the API actually reports the plan total; - // a presumed window (e.g. 30 days) would fabricate utilization. A valid - // response that omits the total after a prior refresh had it is a DELIBERATE - // contract change — the old row must be dropped (terminal), not preserved as - // a transient last-good. - const totalMs = toFiniteNumber(data.total_time ?? data.plan_duration_ms ?? data.total_duration_ms); - if (totalMs === undefined || totalMs <= 0) return TERMINAL_QUOTA_FAILURE; - const consumed = Math.max(0, totalMs - remainsMs); - const percent = normalizePercent((consumed / totalMs) * 100); - if (percent === undefined) return null; - return report(provider, "minimax:token-plan-remains", { - customWindows: [{ label, percent }], - updatedAt: Date.now(), - }); -} - -/** - * Moonshot/Kimi `GET /v1/users/me/balance` — the account's available balance - * (voucher + cash). Renders a single balance window against the sum of - * voucher + cash when positive (there is no per-window rate limit to meter). - */ -async function fetchMoonshotQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalMoonshotBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const host = normalizedBaseUrl(config.baseUrl)?.startsWith("https://api.moonshot.cn") ? "https://api.moonshot.cn/v1" : MOONSHOT_BASE_URL; - const response = await fetch(`${host}/users/me/balance`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const data = asRecord(body?.data) ?? body; - if (!data) return null; - const available = toFiniteNumber(data.available_balance); - const voucher = toFiniteNumber(data.voucher_balance); - const cash = toFiniteNumber(data.cash_balance); - if (available === undefined || available < 0) return null; - // Moonshot exposes no per-window quota ceiling, only a balance — report it - // as a balance-only window (percent 0) rather than a fabricated utilization. - // Currency is host-scoped: China platform (api.moonshot.cn) bills in CNY; - // the international platform (api.moonshot.ai) bills in USD. Do not force - // either side into the other unit — the number is correct, only the unit - // must match the host. - const isChinaHost = host.startsWith("https://api.moonshot.cn"); - const money = (n: number) => isChinaHost ? `¥${n.toFixed(2)}` : `$${n.toFixed(2)}`; - const unit = isChinaHost ? "CNY" : "USD"; - const label = voucher !== undefined && cash !== undefined - ? `Balance (${money(available)} ${unit} available, ${money(voucher)} voucher)` - : `Balance (${money(available)} ${unit} available)`; - return report(provider, "moonshot:balance", { - customWindows: [{ label, percent: 0 }], - updatedAt: Date.now(), - }); -} - -/** - * Venice `GET /api/v1/billing/balance` — DIEM (native credits) or USD balance. - * Shows the remaining balance; epoch allocation progress when present. - */ -async function fetchVeniceQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalVeniceBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(`${VENICE_BASE_URL}/billing/balance`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const data = asRecord(body?.data) ?? body; - if (!data) return null; - const diemBalance = toFiniteNumber(data.balance); - const usdBalance = toFiniteNumber(data.balance_usd); - const epochUsed = toFiniteNumber(data.diem_epoch_used); - const epochAllocated = toFiniteNumber(data.diem_epoch_allocated); - if (diemBalance === undefined && usdBalance === undefined) return null; - const label = diemBalance !== undefined - ? `DIEM balance (${Math.round(diemBalance)})` - : `USD balance ($${usdBalance?.toFixed(2) ?? "?"})`; - if (epochAllocated !== undefined && epochAllocated > 0 && epochUsed !== undefined) { - const percent = normalizePercent((epochUsed / epochAllocated) * 100); - if (percent === undefined) return null; - return report(provider, "venice:billing-balance", { - customWindows: [{ label, percent }], - updatedAt: Date.now(), - }); - } - return report(provider, "venice:billing-balance", { - customWindows: [{ label, percent: 0 }], - updatedAt: Date.now(), - }); -} - -/** - * Synthetic `GET /v2/quotas` — the known quota lanes (rolling 5-hour, - * weekly token, search-hourly) mapped onto the quota windows. - */ -async function fetchSyntheticQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalSyntheticBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(`${SYNTHETIC_BASE_URL}/quotas`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const data = asRecord(body?.data) ?? body; - const quota: ProviderQuota = { updatedAt: Date.now() }; - let windows = 0; - const percentAt = (key: string): number | undefined => { - const value = normalizePercent(data?.[key]); - if (value !== undefined) return value; - const nested = asRecord(data?.quota) ?? asRecord(data?.quotas); - return nested ? normalizePercent(nested[key]) : undefined; - }; - const fiveHour = percentAt("rollingFiveHourLimit"); - const weekly = percentAt("weeklyTokenLimit"); - if (fiveHour !== undefined) { - quota.fiveHourPercent = fiveHour; - windows += 1; - } - if (weekly !== undefined) { - quota.weeklyPercent = weekly; - windows += 1; - } - const search = asRecord(data?.search); - const searchHourly = search ? normalizePercent(search.hourly) : undefined; - if (searchHourly !== undefined) { - quota.customWindows = [...(quota.customWindows ?? []), { label: "Search hourly", percent: searchHourly }]; - windows += 1; - } - const inferenceQuota = { ...quota }; - delete inferenceQuota.customWindows; // search.hourly does not constrain model inference. - return windows > 0 ? keyReport(provider, "synthetic:quotas", quota, config, apiKey, inferenceQuota) : null; -} - -/** - * DeepInfra `GET /payment/checklist?compute_owed=true` — prepaid balance, - * recent spend, spending limit, and suspension state. Renders a balance - * window (prepaid funds are a negative `stripe_balance` → positive available). - */ -async function fetchDeepInfraQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalDeepInfraBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(`${DEEPINFRA_BASE_URL}/payment/checklist?compute_owed=true`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const data = asRecord(body?.data) ?? body; - if (!data) return null; - const stripeBalance = toFiniteNumber(data.stripe_balance); - const spendLimit = toFiniteNumber(data.spending_limit); - const total = toFiniteNumber(data.total_amount_due); - if (stripeBalance === undefined) return null; - // Prepaid funds are negative; a positive value is money owed. - const available = stripeBalance < 0 ? -stripeBalance : 0; - if (spendLimit !== undefined && spendLimit > 0) { - const spent = total !== undefined && total > 0 ? total : Math.max(0, spendLimit - available); - const percent = normalizePercent((spent / spendLimit) * 100); - if (percent === undefined) return null; - return report(provider, "deepinfra:billing-checklist", { - customWindows: [{ label: `Billing cycle spend ($${spent.toFixed(2)} of $${spendLimit.toFixed(2)})`, percent }], - updatedAt: Date.now(), - }); - } - return report(provider, "deepinfra:billing-checklist", { - customWindows: [{ label: `Prepaid balance ($${available.toFixed(2)})`, percent: 0 }], - updatedAt: Date.now(), - }); -} - -/** - * Neuralwatt `GET /v1/quota` — subscription kWh usage (primary window) and - * prepaid USD credit balance (secondary). - */ -async function fetchNeuralwattQuota(provider: string, config: OcxProviderConfig): Promise { - if (!isCanonicalNeuralwattBaseUrl(config.baseUrl)) return null; - const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); - if (!apiKey) return null; - const response = await fetch(`${NEURALWATT_BASE_URL}/quota`, { - headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const body = asRecord(await readQuotaJson(response)); - const data = asRecord(body?.data) ?? body; - const quota: ProviderQuota = { updatedAt: Date.now() }; - let windows = 0; - const subscription = asRecord(data?.subscription); - const kwhUsed = subscription ? toFiniteNumber(subscription.kwh_used) : undefined; - const kwhIncluded = subscription ? toFiniteNumber(subscription.kwh_included) : undefined; - if (kwhUsed !== undefined && kwhIncluded !== undefined && kwhIncluded > 0) { - const percent = normalizePercent((kwhUsed / kwhIncluded) * 100); - if (percent !== undefined) { - quota.fiveHourPercent = percent; - const periodEnd = subscription ? normalizeResetAt(subscription.current_period_end) : undefined; - if (periodEnd !== undefined) quota.fiveHourResetAt = periodEnd; - windows += 1; - } - } - const balance = asRecord(data?.balance); - const totalCredits = balance ? toFiniteNumber(balance.total_credits_usd) : undefined; - const remainingCredits = balance ? toFiniteNumber(balance.credits_remaining_usd) : undefined; - if (totalCredits !== undefined && totalCredits > 0 && remainingCredits !== undefined) { - // Utilization is CONSUMED credits, not the remaining share. - const used = Math.max(0, totalCredits - remainingCredits); - const percent = normalizePercent((used / totalCredits) * 100); - if (percent !== undefined) { - quota.customWindows = [...(quota.customWindows ?? []), { label: "Prepaid credits", percent }]; - windows += 1; - } - } - return windows > 0 ? report(provider, "neuralwatt:quota", quota) : null; -} - -function report( - provider: string, - source: string, - quota: ProviderQuota, - aggregation?: CodexCapacityAggregation, -): ProviderQuotaReport | null { - if (!hasQuotaRows(quota)) return null; - return { - provider, - label: providerLabel(provider), - source, - quota, - updatedAt: quota.updatedAt, - ...(aggregation ? { aggregation } : {}), - }; -} - -/** - * Publish a credential-bound report, and routing evidence only when the producer - * hands over its inference-only projection. - * - * The projection is deliberately not defaulted to the display quota. A producer must - * decide that its rows really do constrain inference on the probed credential; omitting - * the argument leaves the report display-only, so a new producer cannot inherit - * provider-veto authority merely by calling this helper. Ownership alone is not the - * scope decision: providerQuotaRoutingBinding resolving is necessary, never sufficient. - */ -function keyReport( - provider: string, - source: string, - quota: ProviderQuota, - config: OcxProviderConfig, - probedCredential: string, - inferenceQuota?: ProviderQuota, -): ProviderQuotaReport | null { - const result = report(provider, source, quota); - if (!result || !inferenceQuota) return result; - const binding = providerQuotaRoutingBinding(provider, config, probedCredential); - if (binding) routingEvidence.set(result, { quota: inferenceQuota, binding }); - return result; -} - -function tagNativeMainReport( - value: ProviderQuotaReport | null, - generation: number, -): ProviderQuotaReport | null { - if (value) nativeMainReportGenerations.set(value, generation); - return value; -} - -/** - * Test-only seam: publish exactly as a credential-bound producer does, and hand back the - * routing evidence the publication actually attached. - * - * Live producers all pass a projection today, so no probe fixture can prove the OTHER half - * of the contract: that omitting it stays display-only. Routing an omitted argument through - * the real helper keeps that provable, and a re-introduced `= quota` default would be - * observed here (a defaulted parameter also fires for an explicitly undefined argument). - */ -export function publishKeyReportForTests( - provider: string, - source: string, - quota: ProviderQuota, - config: OcxProviderConfig, - probedCredential: string, - inferenceQuota?: ProviderQuota, -): { report: ProviderQuotaReport | null; routing: ProviderQuotaRoutingEvidence | undefined } { - const result = keyReport(provider, source, quota, config, probedCredential, inferenceQuota); - return { report: result, routing: result ? routingEvidence.get(result) : undefined }; -} - -function isProviderQuotaReportCurrent(value: ProviderQuotaReport): boolean { - const generation = nativeMainReportGenerations.get(value); - return (generation === undefined || isMainAccountIdentityGenerationLive(generation)) - && (accountReportCurrent.get(value)?.() ?? true); -} - -async function fetchChatGptForwardQuota( - config: OcxConfig, - provider: string, - providerConfig: OcxProviderConfig, - forceRefresh: boolean, - prefetchedSnapshot?: CodexAuthAccountsSnapshotPromise, -): Promise { - if (providerCodexAccountMode(provider, providerConfig) === "direct") { - const snapshot = await fetchMainAccountInfoSnapshot(forceRefresh); - const quota = providerQuotaFromCodexQuota(snapshot.info.quota); - if (quota) quota.updatedAt = Date.now(); - return quota - ? tagNativeMainReport(report(provider, "chatgpt:wham", quota), snapshot.mainIdentityGeneration) - : null; - } - const snapshot = await (prefetchedSnapshot ?? listCodexAuthAccountsSnapshot(config, forceRefresh)); - const accounts = snapshot.accounts; - const activeId = effectiveCodexAuthAccountId(config); - const capacityAccounts = accounts.map(account => ({ - ...account, - active: account.id === activeId, - quota: providerQuotaFromCodexQuota(account.quota), - })); - const active = capacityAccounts.find(account => account.active) - ?? capacityAccounts.find(account => account.id === MAIN_CODEX_ACCOUNT_ID) - ?? capacityAccounts[0]; - const now = Date.now(); - const capacity = aggregateCodexPoolCapacity(capacityAccounts, now); - if (capacity.aggregation && capacity.quota) { - return tagNativeMainReport( - report( - provider, - "chatgpt:wham", - capacity.quota as ProviderQuota, - publicCapacityAggregation(capacity.aggregation, "aggregate"), - ), - snapshot.mainIdentityGeneration, - ); - } - const activeUsable = !!active && !active.paused && active.needsReauth !== true; - const quota = activeUsable && active?.quota - ? { ...active.quota, updatedAt: active.quota.updatedAt ?? Date.now() } as CodexCapacityQuota - : null; - const quotaFresh = !!quota - && Number.isFinite(quota.updatedAt) - && now - quota.updatedAt < CODEX_CAPACITY_MAX_QUOTA_AGE_MS; - if (quota && quotaFresh) { - const fallback = report( - provider, - "chatgpt:wham", - quota as ProviderQuota, - capacity.aggregation - ? publicCapacityAggregation(capacity.aggregation, "effective-account-fallback") - : undefined, - ); - return tagNativeMainReport(fallback, snapshot.mainIdentityGeneration); - } - if (capacity.aggregation) { - const updatedAt = Date.now(); - return tagNativeMainReport( - { - provider, - label: providerLabel(provider), - source: "chatgpt:wham", - quota: { updatedAt }, - updatedAt, - aggregation: publicCapacityAggregation(capacity.aggregation, "coverage-only"), - }, - snapshot.mainIdentityGeneration, - ); - } - return null; -} - -function centsValue(value: unknown): number | undefined { - const rec = asRecord(value); - return rec ? toFiniteNumber(rec.val) : undefined; -} - -/** Decode JWT payload `sub` for xAI weekly credits when the stored credential lacks accountId. */ -function xaiUserIdFromAccessToken(accessToken: string): string | undefined { - const parts = accessToken.split("."); - if (parts.length < 2 || !parts[1]) return undefined; - try { - const payload = JSON.parse(Buffer.from(parts[1], "base64url").toString("utf8")) as { sub?: unknown }; - return typeof payload.sub === "string" && payload.sub.trim() ? payload.sub.trim() : undefined; - } catch { - return undefined; - } -} - -/** - * Grok Build weekly credits envelope: - * `{ config: { creditUsagePercent?, currentPeriod: { type: "USAGE_PERIOD_TYPE_WEEKLY", end } } }`. - * Omitted percent is treated as 0 (proto3 default). - */ -export function parseXaiCreditsResponse(value: unknown): { percent: number; resetAt?: number } | null { - const body = asRecord(value); - const config = asRecord(body?.config); - if (!config) return null; - const period = asRecord(config.currentPeriod); - if (!period || period.type !== "USAGE_PERIOD_TYPE_WEEKLY") return null; - let percent = 0; - if (config.creditUsagePercent !== undefined) { - const normalized = normalizePercent(config.creditUsagePercent); - if (normalized === undefined) return null; - percent = normalized; - } - const resetAt = normalizeResetAt(period.end); - return { - percent, - ...(resetAt !== undefined ? { resetAt } : {}), - }; -} - -async function fetchXaiWeeklyCredits(accessToken: string, userId: string): Promise { - try { - const response = await fetch(XAI_CREDITS_URL, { - redirect: "error", - headers: { - Accept: "application/json", - Authorization: `Bearer ${accessToken}`, - [XAI_GROK_COMPATIBILITY.headers.tokenAuth]: "xai-grok-cli", - [XAI_GROK_COMPATIBILITY.headers.authenticateResponse]: "authenticate-response", - "x-userid": userId, - [XAI_GROK_COMPATIBILITY.headers.clientVersion]: XAI_GROK_CLIENT_VERSION, - }, - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) return null; - const parsed = parseXaiCreditsResponse(await readQuotaJson(response)); - if (!parsed) return null; - return { - weeklyPercent: parsed.percent, - ...(parsed.resetAt !== undefined ? { weeklyResetAt: parsed.resetAt } : {}), - updatedAt: Date.now(), - }; - } catch { - return null; - } -} - -async function fetchXaiQuota(provider: string, context: { accessToken: string; upstreamAccountId?: string }): Promise { - const { accessToken } = context; - - // Prefer the SuperGrok weekly credits window that actually gates prompting (#1283). - const userId = context.upstreamAccountId?.trim() || xaiUserIdFromAccessToken(accessToken); - if (userId) { - const weekly = await fetchXaiWeeklyCredits(accessToken, userId); - if (weekly) return report(provider, "xai:grok-billing-credits", weekly); - } - - // Legacy monthly dollar pool — retained when weekly is unavailable. - try { - const response = await fetch(XAI_BILLING_URL, { - redirect: "error", - headers: { Accept: "application/json", Authorization: `Bearer ${accessToken}` }, - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) return null; - const body = asRecord(await readQuotaJson(response)); - const config = asRecord(body?.config); - if (!config) return null; - const limitCents = centsValue(config.monthlyLimit); - const usedCents = centsValue(config.used); - if (limitCents === undefined || usedCents === undefined || limitCents <= 0) return null; - const percent = normalizePercent((usedCents / limitCents) * 100); - if (percent === undefined) return null; - return report(provider, "xai:grok-billing", { - monthlyPercent: percent, - monthlyResetAt: normalizeResetAt(config.billingPeriodEnd), - updatedAt: Date.now(), - }); - } catch { - return null; - } -} - -function parseClaudeBucket(value: unknown): { percent?: number; resetAt?: number } | null { - const rec = asRecord(value); - if (!rec) return null; - const percent = normalizePercent(rec.utilization); - const resetAt = normalizeResetAt(rec.resets_at); - if (percent === undefined && resetAt === undefined) return null; - return { percent, resetAt }; -} - -function parseClaudeLimit(value: unknown): { label: string; percent: number; resetAt?: number } | null { - const rec = asRecord(value); - if (!rec) return null; - const percent = normalizePercent(rec.percent); - if (percent === undefined) return null; - const scope = asRecord(rec.scope); - const model = asRecord(scope?.model); - const rawLabel = String(model?.display_name ?? "").trim(); - if (!rawLabel) return null; - const lowerLabel = rawLabel.toLowerCase(); - const label = lowerLabel.includes("fable") ? "Fable" - : lowerLabel.includes("opus") ? "Opus" - : lowerLabel.includes("sonnet") ? "Sonnet" - : rawLabel; - const resetAt = normalizeResetAt(rec.resets_at); - return { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }; -} - -/** Claude's OAuth usage endpoint, probed with ONE account's own bearer token. */ -const anthropicUsageInflight = new Map>(); - -/** - * Anthropic per-credential usage. - * - * This endpoint reports quota only. Its body carries `five_hour`, `seven_day`, the - * model-scoped weekly buckets (`seven_day_fable`/`_opus`/`_sonnet`) and a `limits` array, - * and **no subscription or tier field** — nor does the OAuth token response, which yields only - * `account.uuid` and `account.email_address` (`src/oauth/anthropic.ts`). That is why - * `OAuthAccountSummary.plan` is `null` for Anthropic rather than populated here (#3777); it is - * a missing upstream field, not an unfinished mapping. - * - * A tier must not be inferred from what is here. Percentages are normalized per account, so a - * Max x5 seat at 50% is byte-identical to a Max x20 seat at 50%, and the presence of a - * model-scoped window tracks entitlement rather than seat size. Populate `plan` only when - * upstream returns the tier itself. - */ -async function fetchAnthropicUsageQuota(accessToken: string): Promise { - const joinable = anthropicUsageInflight.get(accessToken); - if (joinable) return joinable; - - const probe = (async (): Promise => { - const response = await fetch("https://api.anthropic.com/api/oauth/usage", { - headers: { - Accept: "application/json, text/plain, */*", - "Content-Type": "application/json", - "User-Agent": "claude-cli/2.1.63 (external, cli)", - "anthropic-beta": "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27,prompt-caching-scope-2026-01-05", - Authorization: `Bearer ${accessToken}`, - }, - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) return null; - const body = asRecord(await readQuotaJson(response)); - if (!body) return null; - const fiveHour = parseClaudeBucket(body.five_hour); - const sevenDay = parseClaudeBucket(body.seven_day); - const fable = parseClaudeBucket(body.seven_day_fable); - const opus = parseClaudeBucket(body.seven_day_opus); - const sonnet = parseClaudeBucket(body.seven_day_sonnet); - const customWindows: ProviderQuotaWindow[] = []; - if (fable?.percent !== undefined) customWindows.push({ label: "Fable", percent: fable.percent, ...(fable.resetAt !== undefined ? { resetAt: fable.resetAt } : {}) }); - if (opus?.percent !== undefined) customWindows.push({ label: "Opus", percent: opus.percent, ...(opus.resetAt !== undefined ? { resetAt: opus.resetAt } : {}) }); - if (sonnet?.percent !== undefined) customWindows.push({ label: "Sonnet", percent: sonnet.percent, ...(sonnet.resetAt !== undefined ? { resetAt: sonnet.resetAt } : {}) }); - const knownLabels = new Set(customWindows.map(window => window.label.toLowerCase())); - const limits = Array.isArray(body.limits) ? body.limits : []; - for (const rawLimit of limits) { - const limitRecord = asRecord(rawLimit); - // `session` and `weekly_all` mirror the canonical five-hour and weekly - // buckets above; only model-scoped weekly limits add a third window. - if (String(limitRecord?.kind ?? "").trim().toLowerCase() !== "weekly_scoped") continue; - const limit = parseClaudeLimit(rawLimit); - if (!limit || knownLabels.has(limit.label.toLowerCase())) continue; - knownLabels.add(limit.label.toLowerCase()); - customWindows.push(limit); - } - const quota: ProviderQuota = { - // Claude's 5-hour window is a first-class rate limit, same as the Codex login 5h/weekly - // rows: report it in the canonical fields so the dashboard renders it with the standard - // "5-hour limit" label and ordering instead of as a generic extra window. - ...(fiveHour?.percent !== undefined ? { fiveHourPercent: fiveHour.percent } : {}), - ...(fiveHour?.resetAt !== undefined ? { fiveHourResetAt: fiveHour.resetAt } : {}), - ...(sevenDay?.percent !== undefined ? { weeklyPercent: sevenDay.percent } : {}), - ...(sevenDay?.resetAt !== undefined ? { weeklyResetAt: sevenDay.resetAt } : {}), - ...(customWindows.length > 0 ? { customWindows } : {}), - updatedAt: Date.now(), - }; - // Empty / schema-changed payloads must not cache as "success with no bars". - return hasQuotaRows(quota) ? quota : null; - })().finally(() => { - if (anthropicUsageInflight.get(accessToken) === probe) anthropicUsageInflight.delete(accessToken); - }); - anthropicUsageInflight.set(accessToken, probe); - return probe; -} - -async function fetchAnthropicQuota(provider: string): Promise { - // Capture the account we intend to probe before awaiting — a mid-flight active - // switch must not seed the wrong account's cache with this response. - const probedAccountId = getAccountSet("anthropic")?.activeAccountId; - const probedAccountKey = probedAccountId ? accountCacheKey("anthropic", probedAccountId) : null; - const writerGeneration = captureConfigGeneration(); - let accessToken: string; - try { - accessToken = await getValidAccessToken("anthropic"); - } catch { - return null; - } - const quota = await fetchAnthropicUsageQuota(accessToken); - if (!quota) return null; - // Share the active-account probe with the per-account cache so Providers-page - // loads do not double-hit Anthropic's rate-limited usage endpoint. - if (probedAccountId && probedAccountKey) { - const stillOwnsToken = getAccountCredential("anthropic", probedAccountId)?.access === accessToken; - if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { - accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota }); - } - } - return report(provider, "anthropic:oauth-usage", quota); -} - -/** - * Provider-level Kiro row: the active account's usage, shown on the Providers page. - * - * The per-account cache is seeded from the same probe so opening that page does not read - * the active account twice, and the account id is captured before the await so a - * concurrent account switch cannot file this answer under the wrong account. - */ -async function fetchKiroQuota(provider: string): Promise { - const probedAccountId = getAccountSet("kiro")?.activeAccountId; - if (!probedAccountId) return null; - const probedAccountKey = accountCacheKey("kiro", probedAccountId); - const writerGeneration = captureConfigGeneration(); - let snapshot: KiroUsageSnapshot | null; - try { - snapshot = await fetchKiroUsageSnapshot(await kiroUsageContextForAccount(probedAccountId)); - } catch { - return null; - } - if (!snapshot) return null; - if (mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { - accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota: snapshot.quota }); - commitKiroAccountUsageState(probedAccountKey, snapshot); - } - return report(provider, "kiro:usage-limits", snapshot.quota); -} - -/** - * Provider-level row probed from the key endpoint, for an account that CAN be probed. - * - * Written through the same account cache the passive path reads, so the measurement - * survives a restart and the per-account rows at oauth-account-routes.ts:313 pick it up - * with no mode change. Deliberately does not flip providerOAuthAccountQuotaMode: that - * mode selects readPassiveProviderAccountQuotas, and the probed per-account path it would - * switch to is gated on supportsPerAccountQuota, which has no meta-muse reader, so the - * GUI account list would go from showing observations to showing nothing. - */ -async function fetchMuseKeyQuota(provider: string): Promise { - const probedAccountId = getAccountSet(provider)?.activeAccountId; - if (!probedAccountId) return null; - const oauthAccessToken = getAccountCredential(provider, probedAccountId)?.muse?.oauthAccessToken; - // An imported or pasted credential has no account token and never will: it is - // capability, not provider id, that decides whether a probe is possible. - if (!oauthAccessToken) return null; - const probedAccountKey = accountCacheKey(provider, probedAccountId); - const writerGeneration = captureConfigGeneration(); - const quota = await fetchMuseKeyQuotaSnapshot(probedAccountId, oauthAccessToken); - if (!quota) return null; - if (mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { - // Hydrate before writing, for the same reason recordPassiveAccountQuota does: - // persistAccountQuotaCache serializes the whole in-memory map. - hydrateAccountQuotaCache(); - accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota }); - persistAccountQuotaCache(); - } - return report(provider, `${provider}:key-endpoint`, quota); -} -/** - * Provider-level row for a passive provider: the ACTIVE account's last observed - * subscription windows, the same shape `fetchAnthropicQuota` and `fetchKiroQuota` - * return. - * - * Cache-only. A dashboard load or `ocx account refresh` must never spend an inference - * turn, so `forceRefresh` does not exist on this path — there is nothing to refresh. - * `report.updatedAt` is the observation time, which is what both GUI surfaces render - * as the relative age of the row. - */ -async function fetchPassiveProviderQuota(provider: string): Promise { - const activeId = getAccountSet(provider)?.activeAccountId; - if (!activeId) return null; - // Idempotent; without it a proxy restart shows nothing until the next streaming turn - // even though the last observation is on disk. - hydrateAccountQuotaCache(); - const entry = accountQuotaCache.get(accountCacheKey(provider, activeId)); - if (!entry?.quota) return null; - const built = report(provider, `${provider}:subscription-observation`, entry.quota); - // Tagged here rather than inside report(), which every probed path shares. - return built ? { ...built, observed: true } : null; -} - -// --------------------------------------------------------------------------- -// Per-account quota (multiauth) -// --------------------------------------------------------------------------- - -/** - * Anthropic and Kiro both report usage per CREDENTIAL, so every logged-in account can be - * probed with its own bearer token — the active-account selection and the local usage log - * are irrelevant here. Mirrors the Codex pool behaviour - * (codex/auth-api.ts:fetchPoolAccountQuota), including a per-account TTL so N accounts cost - * at most N upstream calls per window. `ACCOUNT_QUOTA_TTL_MS` lives in `quota-wire.ts` - * because the Kiro exhaustion reader applies the same staleness bound. - */ -type AccountQuotaCacheEntry = { - ts: number; - quota: ProviderQuota | null; - /** Last probe failed (429 / network / expired login); still may hold last-good quota. */ - unavailable?: true; - quotaFailure?: QuotaFailureCode; - quotaFailureIsCurrent?: () => boolean; - /** Private new-reader identity; never persisted or serialized. */ - identity?: string; - isCurrent?: () => boolean; -}; -/** Expired measurements become unknown; missing reset evidence never implies a fresh allowance. */ -function normalizeAnthropicQuota(quota: ProviderQuota | null | undefined, now: number): ProviderQuota | null { - if (!quota) return null; - const validReset = (resetAt: unknown): resetAt is number => typeof resetAt === "number" - && Number.isFinite(resetAt) && resetAt > 0 && Number.isFinite(new Date(resetAt).getTime()); - let result = quota; - for (const [percent, reset] of [ - ["fiveHourPercent", "fiveHourResetAt"], - ["weeklyPercent", "weeklyResetAt"], - ["monthlyPercent", "monthlyResetAt"], - ] as const) { - const resetAt = quota[reset]; - if (resetAt === undefined) continue; - const valid = validReset(resetAt); - if (valid && resetAt > now) continue; - if (result === quota) result = { ...quota }; - if (valid) delete result[percent]; - delete result[reset]; - } - // Persisted rows validate only the outer quota object, so custom data may be malformed. - if (quota.customWindows !== undefined) { - const windows = Array.isArray(quota.customWindows) ? quota.customWindows : []; - const retained: ProviderQuotaWindow[] = []; - let changed = !Array.isArray(quota.customWindows); - for (const window of windows) { - if (!window || typeof window !== "object" || typeof window.label !== "string" || !window.label.trim() - || typeof window.percent !== "number" || !Number.isFinite(window.percent) - || window.percent < 0 || window.percent > 100) { - changed = true; - continue; - } - if (validReset(window.resetAt) && window.resetAt <= now) { - changed = true; - continue; - } - if (window.resetAt !== undefined && !validReset(window.resetAt)) { - const normalized = { ...window }; - delete normalized.resetAt; - retained.push(normalized); - changed = true; - } else { - retained.push(window); - } - } - if (changed) { - if (result === quota) result = { ...quota }; - if (retained.length) result.customWindows = retained; - else delete result.customWindows; - } - } - return hasQuotaRows(result) ? result : null; -} - -const accountQuotaCache = new Map(); -let explicitAccountEpoch = 0; - -/** - * Seed the cache from the last run, once. - * - * Without this a restart forgets every measurement, so the pool opens its next turn with - * no idea which account has room — the exact blindness pre-dispatch selection exists to - * remove. A hydrated row is still subject to the ordinary TTL, so it orders the first - * request and is replaced by a live probe immediately after. - */ -let diskHydrated = false; -function hydrateAccountQuotaCache(): void { - if (diskHydrated) return; - diskHydrated = true; - for (const [key, quota] of readPersistedAccountQuotas()) { - // Disk stores observation time, not the Anthropic usage probe's clock. - if (!accountQuotaCache.has(key)) { - const anthropic = key.startsWith("anthropic\u0000"); - accountQuotaCache.set(key, { - ts: anthropic ? 0 : quota.updatedAt, - quota: anthropic ? normalizeAnthropicQuota(quota, Date.now()) : quota, - }); - } - } -} - -function persistAccountQuotaCache(): void { - schedulePersistAccountQuotas(function* () { - const now = Date.now(); - for (const [key, entry] of accountQuotaCache) { - const quota = key.startsWith("anthropic\u0000") ? normalizeAnthropicQuota(entry.quota, now) : entry.quota; - if (quota) yield [key, quota] as [string, ProviderQuota]; - } - }); -} -const accountQuotaInflight = new Map>(); -let lastReconciledGeneration = 0; -let liveAccountQuotaKeys = new Set(); -let liveProviderQuotaKeys = new Set(); - -function mayCommitAccountQuotaKey(key: string, writerGeneration: number): boolean { - return writerGeneration >= lastReconciledGeneration || liveAccountQuotaKeys.has(key); -} - -function mayCommitProviderQuotaKey(key: string, writerGeneration: number): boolean { - return writerGeneration >= lastReconciledGeneration || liveProviderQuotaKeys.has(key); -} - -export interface ProviderAccountQuota { - accountId: string; - quota: ProviderQuota | null; - /** Set when the probe could not reach upstream (expired login, 429, network). */ - unavailable?: true; - quotaFailure?: QuotaFailureCode; - quotaFailureIsCurrent?: () => boolean; - isCurrent?: () => boolean; -} - -/** Providers whose per-account quota can be probed. Extend as other OAuth APIs are covered. */ -export function supportsPerAccountQuota(provider: string): boolean { - return provider === "anthropic" || provider === "kiro" || provider === "google-antigravity" - || explicitAccountReader(provider); -} - -function explicitAccountReader(provider: string): boolean { - return provider === "xai" || provider === "cursor" || provider === "kimi" || provider === "command-code"; -} - -export function providerOAuthAccountQuotaMode(provider: string): AccountQuotaMode { - return hasPassiveAccountQuota(provider) ? "passive" : supportsPerAccountQuota(provider) ? "probe" : "unsupported"; -} - -function accountCacheKey(provider: string, accountId: string): string { - return `${provider}\u0000${accountId}`; -} - -/** - * Synchronous last-good per-account quota read for routing. Never probes the network. - * Returns null when nothing is cached (or the cached row has no bars). - */ -export function getCachedProviderAccountQuota(provider: string, accountId: string): ProviderQuota | null { - const entry = accountQuotaCache.get(accountCacheKey(provider, accountId)); - if (entry?.isCurrent && !entry.isCurrent()) return null; - return provider === "anthropic" ? normalizeAnthropicQuota(entry?.quota, Date.now()) : entry?.quota ?? null; -} - -/** Test-only: seed or clear the per-account quota cache without probing upstream. */ -export function setCachedProviderAccountQuotaForTests( - provider: string, - accountId: string, - quota: ProviderQuota | null, -): void { - const key = accountCacheKey(provider, accountId); - if (quota === null) { - accountQuotaCache.delete(key); - return; - } - accountQuotaCache.set(key, { ts: Date.now(), quota }); -} - -/** Unified headers report utilization fractions and epoch-second reset times. */ -function anthropicHeaderResetAt(value: string | null): number | undefined { - const seconds = toFiniteNumber(value); - if (seconds === undefined || seconds <= 0) return undefined; - const timestamp = seconds * 1000; - return Number.isFinite(new Date(timestamp).getTime()) ? timestamp : undefined; -} - -export function parseAnthropicRateLimitHeaders(headers: Headers): ProviderQuota | null { - const fiveHourPercent = normalizeUtilizationFraction(headers.get("anthropic-ratelimit-unified-5h-utilization")); - const weeklyPercent = normalizeUtilizationFraction(headers.get("anthropic-ratelimit-unified-7d-utilization")); - if (fiveHourPercent === undefined && weeklyPercent === undefined) return null; - const fiveHourResetAt = anthropicHeaderResetAt(headers.get("anthropic-ratelimit-unified-5h-reset")); - const weeklyResetAt = anthropicHeaderResetAt(headers.get("anthropic-ratelimit-unified-7d-reset")); - return { - ...(fiveHourPercent !== undefined ? { fiveHourPercent } : {}), - ...(fiveHourPercent !== undefined && fiveHourResetAt !== undefined ? { fiveHourResetAt } : {}), - ...(weeklyPercent !== undefined ? { weeklyPercent } : {}), - ...(weeklyPercent !== undefined && weeklyResetAt !== undefined ? { weeklyResetAt } : {}), - updatedAt: Date.now(), - }; -} - -/** Reject unknown scales; round fraction conversion for persisted/displayed percentages. */ -function normalizeUtilizationFraction(value: string | null): number | undefined { - const numeric = toFiniteNumber(value); - if (numeric === undefined || numeric < 0 || numeric > 1) return undefined; - return Math.round(numeric * 10_000) / 100; -} - -/** - * Merge serving-account observations without advancing the usage probe's clock or - * erasing model-specific windows. The caller owns credential attribution; this guard - * prevents a retired account key from being revived by an older config generation. - */ -export function recordAnthropicAccountQuotaFromHeaders( - accountId: string, - headers: Headers, - writerGeneration: number, -): void { - if (!accountId) return; - const observed = parseAnthropicRateLimitHeaders(headers); - if (!observed) return; - const key = accountCacheKey("anthropic", accountId); - if (!mayCommitAccountQuotaKey(key, writerGeneration)) return; - // Hydrate before writing, for the same reason `recordPassiveAccountQuota` does: this write - // arrives unprompted from the request path, and `persistAccountQuotaCache` serializes the - // whole map. Landing before any reader has hydrated would persist this single row and erase - // every other provider's saved row. - hydrateAccountQuotaCache(); - const previous = accountQuotaCache.get(key); - accountQuotaCache.set(key, { - ...previous, - // Headers do not prove that the last usage probe succeeded. - ts: previous?.ts ?? 0, - quota: normalizeAnthropicQuota({ - ...normalizeAnthropicQuota(previous?.quota, observed.updatedAt), ...observed, - }, observed.updatedAt), - }); - persistAccountQuotaCache(); -} - -/** - * Providers whose per-account quota is OBSERVED in-band, never probed. - * - * Deliberately separate from `supportsPerAccountQuota` rather than folded into it. That - * predicate gates explicit upstream readers. Meta publishes no quota endpoint, so it - * remains a cache-only observation even when every probe reader is account-scoped. - */ -export function hasPassiveAccountQuota(provider: string): boolean { - return provider === "meta-muse"; -} - -/** - * Record a quota observed in-band on a streaming turn. - * - * The CALLER captures `writerGeneration` when it resolves the serving credential, not - * this function at write time. A streaming turn is a long await, and a generation - * captured immediately before the write cannot see a config or account change that - * happened EARLIER in the same turn — which is exactly the case the fence exists for. - */ -export function recordPassiveAccountQuota( - provider: string, - accountId: string, - quota: ProviderQuota, - writerGeneration: number, -): void { - if (!hasPassiveAccountQuota(provider) || !accountId) return; - const key = accountCacheKey(provider, accountId); - if (!mayCommitAccountQuotaKey(key, writerGeneration)) return; - // Hydrate BEFORE writing, not only on the read path. `persistAccountQuotaCache` - // serializes the whole in-memory map, so a passive write that lands before anything - // has read the cache would persist this one row and erase every other provider's - // saved row -- and `diskHydrated` would then stop any later reader from recovering - // them. A probe writer cannot hit this because its own read hydrates first; an - // observation arrives unprompted, so it must hydrate itself. - hydrateAccountQuotaCache(); - accountQuotaCache.set(key, { ts: Date.now(), quota }); - // Persisted so a restart keeps the last observation: with no probe to re-establish it, - // a forgotten row stays forgotten until the user happens to run another streaming turn. - persistAccountQuotaCache(); - // sweepExpiredOnWrite is deliberately NOT called. Existing probe writers call it - // because they run on a poll; this runs on the request path, where a state sweep does - // not belong. Passive rows are still reclaimed by generation reconciliation - // (reconcileProviderAccountQuotaRows) and by the disk reader's age bound. -} - -/** - * Cache-only per-account rows for a passive provider. Never probes, never refreshes. - * - * An account with no observation is OMITTED rather than returned with `quota: null` and - * `unavailable`: that pair means "a probe was attempted and failed", and no probe was - * ever attempted here. A user who has not yet run a streaming turn simply has no - * measurement, which is not an error state. - */ -export function readPassiveProviderAccountQuotas(provider: string): ProviderAccountQuota[] { - if (!hasPassiveAccountQuota(provider)) return []; - // Idempotent, and otherwise only reached from probe paths a passive provider never - // enters — without it a restart shows nothing until the next streaming turn, even - // though the row is sitting on disk. - hydrateAccountQuotaCache(); - const set = getAccountSet(provider); - if (!set) return []; - const rows: ProviderAccountQuota[] = []; - for (const account of set.accounts) { - const entry = accountQuotaCache.get(accountCacheKey(provider, account.id)); - if (entry?.quota) rows.push({ accountId: account.id, quota: entry.quota }); - } - return rows; -} - -export function sweepExpiredProviderAccountQuotaRows(now = Date.now()): number { - let removed = 0; - for (const [key, entry] of accountQuotaCache) { - // Anthropic observations extend retention, never the usage probe's eligibility clock. - const retainedAt = key.startsWith("anthropic\u0000") - ? Math.max(entry.ts, entry.quota?.updatedAt ?? 0) - : entry.ts; - if (retainedAt + ACCOUNT_QUOTA_TTL_MS > now) continue; - accountQuotaCache.delete(key); - removed += 1; - } - return removed; -} - -export function reconcileProviderAccountQuotaRows(context: GenerationContext): number { - if (context.generation <= lastReconciledGeneration) return 0; - let removed = 0; - for (const key of accountQuotaCache.keys()) { - if (context.oauthAccountKeys.has(key)) continue; - accountQuotaCache.delete(key); - removed += 1; - } - // Kiro exhaustion rows are keyed identically, so they retire with their quota row; a - // verdict outliving its account would hand the replacement a cooldown it never earned. - removed += reconcileKiroAccountUsageState(context.oauthAccountKeys); - if (cache) { - const reports = cache.response.reports.filter(report => context.providerNames.has(report.provider)); - removed += cache.response.reports.length - reports.length; - cache = { ...cache, response: { ...cache.response, reports } }; - replaceCachedProviderQuotas(reports, routingEvidence); - } - liveAccountQuotaKeys = new Set(context.oauthAccountKeys); - liveProviderQuotaKeys = new Set(context.providerNames); - lastReconciledGeneration = context.generation; - return removed; -} - -/** Test-only reset so a direct reconcile call in one file cannot leak across files. */ -export function resetProviderQuotaReconcileStateForTests(): void { - lastReconciledGeneration = 0; - liveAccountQuotaKeys = new Set(); - liveProviderQuotaKeys = new Set(); -} - -/** Drop cached per-account rows (all, or just one provider's). */ -export function clearAccountQuotaCache(provider?: string): void { - explicitAccountEpoch += 1; - if (!provider) { - accountQuotaCache.clear(); - accountQuotaInflight.clear(); - clearKiroAccountUsageState(); - // A cleared cache must not be re-seeded from the file it was just cleared of, and any - // pending write of the old rows is abandoned. - diskHydrated = false; - cancelPendingAccountQuotaPersist(); - return; - } - const prefix = `${provider}\u0000`; - for (const key of [...accountQuotaCache.keys()]) { - if (key.startsWith(prefix)) accountQuotaCache.delete(key); - } - clearKiroAccountUsageState(prefix); - // Drop in-flight probes too so a late resolve cannot repopulate after logout/remove. - for (const key of [...accountQuotaInflight.keys()]) { - if (key.startsWith(prefix)) accountQuotaInflight.delete(key); - } - persistAccountQuotaCache(); -} - -/** - * Resolve a bearer for quota probing without silently adopting a newer global - * Claude CLI credential into a background multiauth slot. - * - * - Fresh stored access → use as-is (no refresh). - * - Active account with expired access → normal refresh path. - * - Background `local-cli` with expired access → fail closed (unavailable): - * `getValidAccessTokenForAccount` can persist a mismatched Claude CLI identity. - * - Background ordinary OAuth (`source !== "local-cli"`) → safe to refresh; - * Anthropic's lock only adopts disk credentials for `local-cli` rows. - */ -async function getTokenForAccountQuotaProbe(provider: string, accountId: string): Promise { - const stored = getAccountCredential(provider, accountId); - if (!stored) throw new Error("account credential missing"); - if (stored.expires > Date.now() + ACCOUNT_TOKEN_SKEW_MS) return stored.access; - const activeId = getAccountSet(provider)?.activeAccountId; - if (activeId !== accountId && stored.source === "local-cli") { - throw new Error("background local-cli token expired; skip CLI-adopting refresh for quota probe"); - } - return getValidAccessTokenForAccount(provider, accountId); -} - -function explicitQuotaConfig(provider: string, configured?: OcxProviderConfig): OcxProviderConfig | undefined { - if (configured) return configured; - const entry = getProviderRegistryEntry(provider); - return entry ? { adapter: entry.adapter, baseUrl: entry.baseUrl, authMode: "oauth" } : undefined; -} - -function explicitQuotaIdentity(provider: string, accountId: string, configured?: OcxProviderConfig): string | undefined { - const credential = getAccountCredential(provider, accountId); - const target = explicitQuotaConfig(provider, configured); - if (!credential || !target) return undefined; - return quotaCredentialIdentity(provider, accountId, credential, target); -} - -function quotaCredentialIdentity(provider: string, accountId: string, credential: NonNullable>, target: OcxProviderConfig): string { - return createHash("sha256").update(JSON.stringify([ - provider, accountId, credential.access, credential.refresh, credential.expires, - credential.accountId, credential.projectId, credential.source, - target.adapter, target.baseUrl, target.authMode, target.disabled === true, - ])).digest("hex"); -} - -function explicitQuotaDestination(provider: string, config: OcxProviderConfig): boolean { - if (config.disabled === true || config.authMode !== "oauth") return false; - if (provider === "kimi") return isCanonicalKimiCodeBaseUrl(config.baseUrl); - if (provider === "command-code") return isCanonicalCommandCodeBaseUrl(config.baseUrl); - // These readers use fixed canonical billing origins, never config.baseUrl. - return provider === "xai" || provider === "cursor"; -} - -async function readExplicitAccountQuota(provider: string, accountId: string, configured?: OcxProviderConfig): Promise<{ - result: ProviderQuotaProbeResult; - identity: string | undefined; - isCurrent: () => boolean; -} | null> { - const target = explicitQuotaConfig(provider, configured); - if (!target || !explicitQuotaDestination(provider, target)) return null; - const config = { ...target }; - const epoch = explicitAccountEpoch; - const accessToken = await getTokenForAccountQuotaProbe(provider, accountId); - const credential = getAccountCredential(provider, accountId); - if (!credential || credential.access !== accessToken) return null; - // Pair the post-renewal credential with the destination captured before renewal. - const identity = explicitQuotaIdentity(provider, accountId, config); - const isCurrent = () => epoch === explicitAccountEpoch - && identity === explicitQuotaIdentity(provider, accountId, configured); - if (!isCurrent()) return null; - let result: ProviderQuotaProbeResult; - switch (provider) { - case "xai": result = await fetchXaiQuota(provider, { accessToken, upstreamAccountId: credential.accountId }); break; - case "cursor": result = await fetchCursorQuota(provider, accessToken); break; - case "kimi": result = await fetchKimiQuota(provider, config, accessToken); break; - case "command-code": result = await fetchCommandCodeQuota(provider, config, accessToken); break; - default: return null; - } - return { result, identity, isCurrent }; -} - -async function fetchExplicitAccountQuota(provider: string, accountId: string, force: boolean, configured?: OcxProviderConfig): Promise { - const key = accountCacheKey(provider, accountId); - const identity = explicitQuotaIdentity(provider, accountId, configured); - const previous = accountQuotaCache.get(key); - const cached = identity && previous?.identity === identity && previous.isCurrent?.() ? previous : undefined; - if (!force && cached && Date.now() - cached.ts < ACCOUNT_QUOTA_TTL_MS - && (!cached.quota || Date.now() - cached.quota.updatedAt < LAST_GOOD_MAX_AGE_MS)) return cached; - const flightKey = `${key}\u0000${identity ?? "missing"}`; - const running = accountQuotaInflight.get(flightKey); - if (running) return running; - const epoch = explicitAccountEpoch; - const lastGood = cached?.quota && Date.now() - cached.quota.updatedAt < LAST_GOOD_MAX_AGE_MS ? cached.quota : null; - const flight = (async (): Promise => { - let read: Awaited> = null; - try { read = await readExplicitAccountQuota(provider, accountId, configured); } catch { /* unavailable */ } - const isCurrent = read?.isCurrent ?? (() => epoch === explicitAccountEpoch && !!identity - && identity === explicitQuotaIdentity(provider, accountId, configured)); - const result = read?.result; - const current = epoch === explicitAccountEpoch && isCurrent(); - const quota = current && result && typeof result !== "symbol" ? result.quota : null; - const empty = result === AUTHORITATIVE_EMPTY_QUOTA; - const entry: AccountQuotaCacheEntry = { - ts: Date.now(), - quota: quota ?? (current && result !== TERMINAL_QUOTA_FAILURE && !empty - && lastGood && Date.now() - lastGood.updatedAt < LAST_GOOD_MAX_AGE_MS ? lastGood : null), - ...(!current || (!quota && !empty) ? { unavailable: true as const } : {}), - identity: read?.identity ?? identity, - isCurrent: () => epoch === explicitAccountEpoch && isCurrent(), - }; - if (entry.isCurrent?.()) accountQuotaCache.set(key, entry); - return entry; - })().finally(() => { if (accountQuotaInflight.get(flightKey) === flight) accountQuotaInflight.delete(flightKey); }); - accountQuotaInflight.set(flightKey, flight); - return flight; -} - -async function fetchExplicitCurrentQuota(provider: string, config: OcxProviderConfig, liveConfig: OcxConfig): Promise { - const id = getAccountSet(provider)?.activeAccountId; - if (!id) return null; - const read = await readExplicitAccountQuota(provider, id, config); - if (!read) return null; - const isCurrent = () => liveConfig.providers[provider] === config - && read.isCurrent() && getAccountSet(provider)?.activeAccountId === id; - if (!isCurrent()) return TERMINAL_QUOTA_FAILURE; - if (read.result && typeof read.result !== "symbol") accountReportCurrent.set(read.result, isCurrent); - return read.result; -} - -function antigravityQuotaDiagnosticIdentity(accountId: string, credential = getAccountCredential("google-antigravity", accountId)): string | undefined { - return credential ? quotaCredentialIdentity("google-antigravity", accountId, credential, { - adapter: "google", baseUrl: ANTIGRAVITY_ACCOUNT_QUOTA_BASE, authMode: "oauth", - }) : undefined; -} - -async function fetchAccountQuota( - provider: string, - accountId: string, - forceRefresh: boolean, - providerConfig?: OcxProviderConfig, -): Promise { - if (!supportsPerAccountQuota(provider)) return { ts: Date.now(), quota: null, unavailable: true }; - if (explicitAccountReader(provider)) return fetchExplicitAccountQuota(provider, accountId, forceRefresh, providerConfig); - if (provider === "anthropic") hydrateAccountQuotaCache(); - const key = accountCacheKey(provider, accountId); - const writerGeneration = captureConfigGeneration(); - const cached = accountQuotaCache.get(key); - if (!forceRefresh && cached && Date.now() - cached.ts < ACCOUNT_QUOTA_TTL_MS) { - if (provider === "google-antigravity" && cached.quotaFailure && cached.quotaFailureIsCurrent?.() !== true) return { ...cached, quotaFailure: undefined }; - return provider === "anthropic" ? { ...cached, quota: normalizeAnthropicQuota(cached.quota, Date.now()) } : cached; - } - const joinable = accountQuotaInflight.get(key); - if (joinable) return joinable; - - const epoch = explicitAccountEpoch; - const probe = (async (): Promise => { - let diagnosticIdentity: string | undefined; - let quotaFailure: QuotaFailureCode | undefined; - const quotaFailureIsCurrent = () => { - try { return epoch === explicitAccountEpoch && diagnosticIdentity !== undefined && diagnosticIdentity === antigravityQuotaDiagnosticIdentity(accountId); } - catch { return false; } - }; - const diagnosticFields = () => quotaFailure && quotaFailureIsCurrent() ? { quotaFailure, quotaFailureIsCurrent } : {}; - try { - if (provider === "google-antigravity") diagnosticIdentity = antigravityQuotaDiagnosticIdentity(accountId); - let quota: ProviderQuota | null; - let kiroSnapshot: KiroUsageSnapshot | null = null; - if (provider === "kiro") { - // Kiro resolves the bearer and its routing metadata from ONE account-scoped - // snapshot. It deliberately does not use getTokenForAccountQuotaProbe: that - // helper refuses to refresh a background `local-cli` slot because Anthropic's - // lock can adopt a mismatched Claude CLI identity, but Kiro marks every - // CLI-imported credential `local-cli`, so the same rule would blank the quota of - // every inactive pool account the moment its token expired. - kiroSnapshot = await fetchKiroUsageSnapshot(await kiroUsageContextForAccount(accountId)); - quota = kiroSnapshot?.quota ?? null; - } else { - const token = await getTokenForAccountQuotaProbe(provider, accountId); - if (provider === "google-antigravity") { - // Per-account Gem/Cla windows (#1082). The project id is part of the stored - // credential; without it the probe cannot be made, and that is "unavailable", - // never 0%. - const credential = getAccountCredential(provider, accountId); - diagnosticIdentity = credential?.access === token ? antigravityQuotaDiagnosticIdentity(accountId, credential) : undefined; - if (!diagnosticIdentity || !credential?.projectId) throw new Error("antigravity account unavailable"); - const result = await probeAntigravityUsageQuota(token, credential.projectId); - quota = result.kind === "available" ? result.quota : null; - if (result.kind === "unavailable") quotaFailure = result.failure; - } else if (provider === "anthropic") { - quota = await fetchAnthropicUsageQuota(token); - } else { - return { ts: Date.now(), quota: null, unavailable: true }; - } - } - if (!quota) { - // Preserve last-good bars and mark unavailable; advance TTL so failures - // negative-cache instead of re-probing on every GUI poll. - const entry: AccountQuotaCacheEntry = { - ts: Date.now(), - // Settle once for all joiners against observations committed during the probe. - quota: provider === "anthropic" - ? normalizeAnthropicQuota(accountQuotaCache.get(key)?.quota, Date.now()) : cached?.quota ?? null, - unavailable: true, - ...diagnosticFields(), - }; - if (mayCommitAccountQuotaKey(key, writerGeneration)) { - accountQuotaCache.set(key, entry); - if (provider === "kiro") commitKiroAccountUsageState(key, null); - sweepExpiredOnWrite(entry.ts); - } - return entry; - } - const entry: AccountQuotaCacheEntry = { - ts: Date.now(), quota: provider === "anthropic" ? normalizeAnthropicQuota(quota, Date.now()) : quota, - }; - if (mayCommitAccountQuotaKey(key, writerGeneration)) { - accountQuotaCache.set(key, entry); - // Exhaustion state rides the SAME commit guard as the quota row: a probe from a - // superseded config generation must not publish either half. - if (provider === "kiro") commitKiroAccountUsageState(key, kiroSnapshot); - sweepExpiredOnWrite(entry.ts); - } - return entry; - } catch { - if (provider === "google-antigravity") quotaFailure = "account_unavailable"; - const entry: AccountQuotaCacheEntry = { - ts: Date.now(), - quota: provider === "anthropic" - ? normalizeAnthropicQuota(accountQuotaCache.get(key)?.quota, Date.now()) : cached?.quota ?? null, - unavailable: true, - ...diagnosticFields(), - }; - if (mayCommitAccountQuotaKey(key, writerGeneration)) { - accountQuotaCache.set(key, entry); - sweepExpiredOnWrite(entry.ts); - } - return entry; - } - })().finally(() => { - if (accountQuotaInflight.get(key) === probe) accountQuotaInflight.delete(key); - }); - accountQuotaInflight.set(key, probe); - return probe; -} - -/** - * Per-account quota rows for a provider's logged-in accounts. Probes run in parallel; a - * single failing account never blocks the others. - */ -export async function fetchProviderAccountQuotas( - provider: string, - forceRefresh = false, - providerConfig?: OcxProviderConfig, -): Promise { - if (!supportsPerAccountQuota(provider)) return []; - const set = getAccountSet(provider); - if (!set) return []; - return mapQuotaRoster(set.accounts, async account => { - const entry = await fetchAccountQuota(provider, account.id, forceRefresh, providerConfig); - const result: ProviderAccountQuota = { - accountId: account.id, - quota: provider === "anthropic" ? normalizeAnthropicQuota(entry.quota, Date.now()) : entry.quota, - ...(entry.unavailable ? { unavailable: true as const } : {}), - ...(entry.unavailable && entry.quotaFailure && entry.quotaFailureIsCurrent?.() === true ? { quotaFailure: entry.quotaFailure } : {}), - }; - if (entry.quotaFailureIsCurrent) Object.defineProperty(result, "quotaFailureIsCurrent", { value: entry.quotaFailureIsCurrent }); - if (!explicitAccountReader(provider)) return result; - const identity = entry.identity; - Object.defineProperty(result, "isCurrent", { value: () => { - if (entry.isCurrent) return entry.isCurrent(); - const credential = getAccountCredential(provider, account.id); - return !!credential && (!identity || explicitQuotaIdentity(provider, account.id, providerConfig) === identity); - } }); - return result; - }); -} - -function normalizedBaseUrl(value: string): string | null { - try { - const url = new URL(value); - if (url.username || url.password || url.search || url.hash) return null; - return `${url.origin.toLowerCase()}${url.pathname.replace(/\/+$/, "")}`; - } catch { - return null; - } -} - -function quotaResetAt(row: Record): number | undefined { - return normalizeResetAt(row.resetTime ?? row.resetAt ?? row.reset_time ?? row.reset_at); -} - -function isCanonicalKimiCodeBaseUrl(baseUrl: string): boolean { - return normalizedBaseUrl(baseUrl) === KIMI_CODE_BASE_URL; -} - -function isCanonicalCommandCodeBaseUrl(baseUrl: string): boolean { - const normalized = normalizedBaseUrl(baseUrl); - // OAuth preset points at the API root; the Provider-API preset at /provider/v1. - return normalized === COMMAND_CODE_BASE_URL || normalized === `${COMMAND_CODE_BASE_URL}/provider/v1`; -} - -/** Prefer the nested `data` shell when the outer object is only an envelope. */ -function unwrapKimiQuotaPayload(value: unknown): Record | null { - const body = asRecord(value); - if (!body) return null; - const nested = asRecord(body.data); - if (!nested) return body; - // A null/non-usable outer field is a placeholder, not data — an envelope like - // { usage: null, data: { usage: {...} } } must still unwrap to the nested payload. - const usable = (field: unknown): boolean => field !== undefined && field !== null; - const outerHasUsage = usable(body.usage) || usable(body.limits) || usable(body.totalQuota); - const nestedHasUsage = usable(nested.usage) || usable(nested.limits) || usable(nested.totalQuota); - return !outerHasUsage && nestedHasUsage ? nested : body; -} - -function kimiLimitLabel(item: Record, detail: Record): string { - return [item.name, item.title, item.scope, detail.name, detail.title] - .filter((value): value is string => typeof value === "string") - .join(" ") - .toLowerCase(); -} - -function parseKimiQuotaRow(value: unknown, resetFallback?: Record): { percent: number; resetAt?: number } | null { - const row = asRecord(value); - if (!row) return null; - const resetAt = quotaResetAt(row) ?? (resetFallback ? quotaResetAt(resetFallback) : undefined); - const limit = toFiniteNumber(row.limit); - if (limit !== undefined && limit > 0) { - let used = toFiniteNumber(row.used); - if (used === undefined) { - const remaining = toFiniteNumber(row.remaining); - if (remaining !== undefined) used = limit - remaining; - } - if (used !== undefined) { - const percent = normalizePercent((used / limit) * 100); - if (percent !== undefined) return { percent, ...(resetAt !== undefined ? { resetAt } : {}) }; - } - } - // Some payloads expose utilisation directly when limit/used arithmetic is absent. - const direct = normalizePercent(row.utilization ?? row.percent ?? row.usedPercent ?? row.used_percent); - return direct === undefined ? null : { percent: direct, ...(resetAt !== undefined ? { resetAt } : {}) }; -} - -function isKimiFiveHourLimit(item: Record, detail: Record, window: Record): boolean { - const duration = toFiniteNumber(window.duration ?? item.duration ?? detail.duration); - const unit = String(window.timeUnit ?? item.timeUnit ?? detail.timeUnit ?? "").toUpperCase(); - if ((unit.includes("MINUTE") && duration === 300) || (unit.includes("HOUR") && duration === 5)) return true; - return /(^|\b)5\s*(?:h|hour)/.test(kimiLimitLabel(item, detail)); -} - -function isKimiWeeklyLimit(item: Record, detail: Record, window: Record): boolean { - const duration = toFiniteNumber(window.duration ?? item.duration ?? detail.duration); - const unit = String(window.timeUnit ?? item.timeUnit ?? detail.timeUnit ?? "").toUpperCase(); - if ((unit.includes("DAY") && duration === 7) || (unit.includes("HOUR") && duration === 168)) return true; - return /weekly|7\s*(?:d|day)/.test(kimiLimitLabel(item, detail)); -} - -function parseKimiQuotaPayload(value: unknown): ProviderQuota | null { - const body = unwrapKimiQuotaPayload(value); - if (!body) return null; - let weekly = parseKimiQuotaRow(body.usage); - const total = parseKimiQuotaRow(body.totalQuota); - let fiveHour: { percent: number; resetAt?: number } | null = null; - if (Array.isArray(body.limits)) { - for (const rawItem of body.limits) { - const item = asRecord(rawItem); - if (!item) continue; - const detail = asRecord(item.detail) ?? item; - const window = asRecord(item.window) ?? {}; - if (!fiveHour && isKimiFiveHourLimit(item, detail, window)) { - fiveHour = parseKimiQuotaRow(detail, window); - } - if (!weekly && isKimiWeeklyLimit(item, detail, window)) { - weekly = parseKimiQuotaRow(detail, window); - } - if (fiveHour && weekly) break; - } - } - const quota: ProviderQuota = { - ...(fiveHour ? { - fiveHourPercent: fiveHour.percent, - ...(fiveHour.resetAt !== undefined ? { fiveHourResetAt: fiveHour.resetAt } : {}), - } : {}), - ...(weekly ? { - weeklyPercent: weekly.percent, - ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), - } : {}), - ...(total ? { customWindows: [{ label: "Total subscription credits", percent: total.percent, ...(total.resetAt !== undefined ? { resetAt: total.resetAt } : {}) }] } : {}), - updatedAt: Date.now(), - }; - return hasQuotaRows(quota) ? quota : null; -} - -async function resolveKimiQuotaBearer(config: OcxProviderConfig, accountId?: string): Promise { - if (config.authMode === "oauth") { - try { - return accountId ? await getTokenForAccountQuotaProbe("kimi", accountId) : null; - } catch { - return null; - } - } - // ACTIVE key only: silently walking apiKeyPool when the primary env reference is - // unresolved would render a quota bar for a DIFFERENT account than the one routing - // requests — a wrong meter is worse than no meter. - const primary = resolveProviderApiKey(config.apiKey)?.trim(); - return primary || null; -} - -async function fetchKimiQuota(provider: string, config: OcxProviderConfig, accessToken: string): Promise { - // Never release credentials to a user-edited or lookalike provider host. - if (!isCanonicalKimiCodeBaseUrl(config.baseUrl)) return null; - if (!accessToken) return null; - const response = await fetch(KIMI_CODE_USAGE_URL, { - headers: { Accept: "application/json", Authorization: `Bearer ${accessToken}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) return null; - const quota = parseKimiQuotaPayload(await readQuotaJson(response)); - return quota ? keyReport(provider, "kimi:usages", quota, config, accessToken, quota) : null; -} - -/** - * Command Code rolling window: `{ cap, used, resetAt }` off /alpha/billing/credits, - * normalized to a percent with an optional reset timestamp. - */ -function parseCommandCodeWindow(value: unknown): { percent: number; resetAt?: number } | null { - const row = asRecord(value); - if (!row) return null; - const cap = toFiniteNumber(row.cap); - const used = toFiniteNumber(row.used); - if (cap === undefined || used === undefined || cap <= 0 || used < 0) return null; - const percent = normalizePercent((used / cap) * 100); - if (percent === undefined) return null; - const resetAt = quotaResetAt(row); - return { percent, ...(resetAt !== undefined ? { resetAt } : {}) }; -} - -/** Soft-fail GET returning a parsed record, or null when unavailable. */ -async function fetchCommandCodeJson(url: string, bearer: string): Promise | null> { - try { - const response = await fetch(url, { - headers: { Accept: "application/json", Authorization: `Bearer ${bearer}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) return null; - return asRecord(await readQuotaJson(response)); - } catch { - return null; - } -} - -/** - * Soft-fail period spend (used) against the remaining credit pools → creditsUsd. - * Period scoping: `since=` keeps spend aligned with the - * pools' billing cycle, and `currentPeriodEnd` becomes expiresAt. - */ -async function fetchCommandCodeSpend( - bearer: string, - credits: Record | null, - orgQuery: string, -): Promise { - if (!credits) return undefined; - const subscriptionBody = await fetchCommandCodeJson(`${COMMAND_CODE_SUBSCRIPTIONS_URL}${orgQuery}`, bearer); - const subscription = asRecord(subscriptionBody?.data) ?? subscriptionBody; - const periodStart = typeof subscription?.currentPeriodStart === "string" ? subscription.currentPeriodStart.trim() : ""; - // Unscoped /usage/summary is lifetime spend; mixing it with current-cycle - // remaining pools produces a wrong percent. Omit creditsUsd until a period exists. - if (!periodStart) return undefined; - const sinceQuery = `${orgQuery ? "&" : "?"}since=${encodeURIComponent(periodStart)}`; - const expiresAt = normalizeResetAt(subscription?.currentPeriodEnd); - const summaryBody = await fetchCommandCodeJson(`${COMMAND_CODE_USAGE_URL}${orgQuery}${sinceQuery}`, bearer); - const summary = asRecord(summaryBody?.data) ?? summaryBody; - const used = toFiniteNumber(summary?.totalCost) ?? toFiniteNumber(summary?.totalMonthlyCredits); - if (used === undefined || used < 0) return undefined; - const pools = [credits.monthlyCredits, credits.purchasedCredits, credits.freeCredits] - .map(value => toFiniteNumber(value)) - .filter((value): value is number => value !== undefined); - // Field presence is what separates a real balance from absent data: an exhausted - // all-zero account still reports remaining=0, while no remaining-credit field at - // all means there is nothing to meter. - if (pools.length === 0) return undefined; - const remaining = pools.reduce((sum, value) => sum + Math.max(0, value ?? 0), 0); - const limit = used + remaining; - const percent = normalizePercent(limit > 0 ? (used / limit) * 100 : 0); - // Purchased credits roll over past the subscription period end, so an expiry is - // only truthful when the aggregate contains no non-expiring purchased pool. - const purchased = toFiniteNumber(credits.purchasedCredits) ?? 0; - return percent === undefined - ? undefined - : { - used, - limit, - remaining, - percent, - ...(expiresAt !== undefined && purchased <= 0 ? { expiresAt } : {}), - }; -} - -/** OAuth access token or ACTIVE Provider-API key for the Command Code quota probe. */ -async function resolveCommandCodeQuotaBearer(config: OcxProviderConfig, accountId?: string): Promise { - if (config.authMode === "oauth") { - try { - return accountId ? await getTokenForAccountQuotaProbe("command-code", accountId) : null; - } catch { - return null; - } - } - // ACTIVE key only: a quota bar for a different account than the one routing - // requests is a wrong meter, not a helpful one. - return resolveProviderApiKey(config.apiKey)?.trim() || null; -} - -/** - * Command Code `GET /alpha/billing/credits` — the same Bearer surface the CLI's - * usage view uses (windowLimits.fiveHour / windowLimits.weekly), plus soft - * whoami (team orgId scoping) and subscription-scoped spend for creditsUsd. - */ -async function fetchCommandCodeQuota(provider: string, config: OcxProviderConfig, bearer: string): Promise { - // Never release credentials to a user-edited or lookalike provider host. - if (!isCanonicalCommandCodeBaseUrl(config.baseUrl)) return null; - if (!bearer) return null; - const whoamiBody = await fetchCommandCodeJson(COMMAND_CODE_WHOAMI_URL, bearer); - const whoami = asRecord(whoamiBody?.data) ?? whoamiBody; - const org = asRecord(whoami?.org); - const orgId = typeof org?.id === "string" && org.id.trim() ? org.id.trim() : null; - const orgQuery = orgId ? `?orgId=${encodeURIComponent(orgId)}` : ""; - const response = await fetch(`${COMMAND_CODE_CREDITS_URL}${orgQuery}`, { - headers: { Accept: "application/json", Authorization: `Bearer ${bearer}` }, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) { - return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 - ? TERMINAL_QUOTA_FAILURE - : null; - } - const raw = asRecord(await readQuotaJson(response)); - const body = asRecord(raw?.data) ?? raw; - const credits = asRecord(body?.credits); - const limits = asRecord(body?.windowLimits); - if (!credits && !limits) return null; - const fiveHour = parseCommandCodeWindow(limits?.fiveHour); - const weekly = parseCommandCodeWindow(limits?.weekly); - const creditsUsd = await fetchCommandCodeSpend(bearer, credits, orgQuery); - const quota: ProviderQuota = { - ...(fiveHour ? { - fiveHourPercent: fiveHour.percent, - ...(fiveHour.resetAt !== undefined ? { fiveHourResetAt: fiveHour.resetAt } : {}), - } : {}), - ...(weekly ? { - weeklyPercent: weekly.percent, - ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), - } : {}), - ...(creditsUsd ? { creditsUsd } : {}), - updatedAt: Date.now(), - }; - // Rolling windows and the credit balance both gate inference on this bearer. - return keyReport(provider, "command-code:credits", quota, config, bearer, quota); -} - -/** Cursor included usage via api2.cursor.sh (Bearer from OAuth) — unofficial, may change. */ -async function fetchCursorQuota(provider: string, accessToken: string): Promise { - - const authHeaders = { - Accept: "application/json", - Authorization: `Bearer ${accessToken}`, - "User-Agent": "opencodex-quota", - } as const; - - // Prefer dashboard period usage (Pro/Team/Ultra spend allowance in USD cents). - // Field names follow Cursor's Connect RPC shape (limit/remaining/includedSpend), not usedCents. - try { - const periodRes = await fetch("https://api2.cursor.sh/aiserver.v1.DashboardService/GetCurrentPeriodUsage", { - method: "POST", - redirect: "error", - headers: { - ...authHeaders, - "Content-Type": "application/json", - "Connect-Protocol-Version": "1", - }, - body: "{}", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (periodRes.ok) { - const body = asRecord(await readQuotaJson(periodRes)); - const planUsage = asRecord(body?.planUsage); - if (planUsage) { - const resetAt = normalizeResetAt(body?.billingCycleEnd ?? planUsage.billingCycleEnd ?? body?.periodEnd); - - // Primary meter: overall included allowance (Cursor Settings → Usage total %). - // autoPercentUsed / apiPercentUsed are secondary pools and must not replace the total. - const limit = toFiniteNumber(planUsage.limit ?? planUsage.limitCents ?? planUsage.totalLimitCents); - const remaining = toFiniteNumber(planUsage.remaining ?? planUsage.remainingCents); - const includedSpend = toFiniteNumber(planUsage.includedSpend ?? planUsage.usedCents ?? planUsage.used); - const totalSpend = toFiniteNumber(planUsage.totalSpend); - let used: number | undefined; - if (includedSpend !== undefined) used = includedSpend; - else if (limit !== undefined && remaining !== undefined) used = Math.max(0, limit - remaining); - else if (totalSpend !== undefined) used = totalSpend; - const totalPercent = normalizePercent(planUsage.totalPercentUsed ?? planUsage.percentUsed) - ?? (limit !== undefined && limit > 0 && used !== undefined - ? normalizePercent((used / limit) * 100) - : undefined); - - const autoPercent = normalizePercent(planUsage.autoPercentUsed); - const apiPercent = normalizePercent(planUsage.apiPercentUsed); - const customWindows: ProviderQuotaWindow[] = []; - if (autoPercent !== undefined) { - customWindows.push({ - label: "First-party models", - percent: autoPercent, - ...(resetAt !== undefined ? { resetAt } : {}), - }); - } - if (apiPercent !== undefined) { - customWindows.push({ - label: "API usage", - percent: apiPercent, - ...(resetAt !== undefined ? { resetAt } : {}), - }); - } - - if (totalPercent !== undefined || customWindows.length > 0) { - const built = report(provider, "cursor:period-usage", { - ...(totalPercent !== undefined ? { - monthlyPercent: totalPercent, - ...(resetAt !== undefined ? { monthlyResetAt: resetAt } : {}), - } : {}), - ...(customWindows.length > 0 ? { customWindows } : {}), - updatedAt: Date.now(), - }); - if (built) return { ...built, reverseEngineered: true }; - } - } - } - } catch { - /* fall through */ - } - - // /api/usage/summary — same host, sometimes richer than /auth/usage for Team plans. - try { - const summaryRes = await fetch("https://api2.cursor.sh/api/usage/summary", { - headers: authHeaders, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (summaryRes.ok) { - const body = asRecord(await readQuotaJson(summaryRes)); - const individual = asRecord(body?.individualUsage); - const plan = asRecord(individual?.plan); - if (plan) { - const used = toFiniteNumber(plan.used); - const limit = toFiniteNumber(plan.limit); - const percent = normalizePercent(plan.totalPercentUsed) - ?? (used !== undefined && limit !== undefined && limit > 0 - ? normalizePercent((used / limit) * 100) - : undefined); - if (percent !== undefined) { - const built = report(provider, "cursor:usage-summary", { - monthlyPercent: percent, - monthlyResetAt: normalizeResetAt(body?.billingCycleEnd), - updatedAt: Date.now(), - }); - if (built) return { ...built, reverseEngineered: true }; - } - } - } - } catch { - /* fall through to /auth/usage */ - } - - const response = await fetch("https://api2.cursor.sh/auth/usage", { - headers: authHeaders, - redirect: "error", - signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }); - if (!response.ok) return null; - const body = asRecord(await readQuotaJson(response)); - if (!body) return null; - - // Prefer the gpt-4 bucket (historical "fast requests"); else first model with used+limit. - let used: number | undefined; - let limit: number | undefined; - const gpt4 = asRecord(body["gpt-4"]); - if (gpt4) { - used = toFiniteNumber(gpt4.numRequests ?? gpt4.used); - limit = toFiniteNumber(gpt4.maxRequestUsage ?? gpt4.limit ?? gpt4.maxRequests); - } - if (used === undefined || limit === undefined || limit <= 0) { - for (const [key, value] of Object.entries(body)) { - if (key === "startOfMonth" || key === "billingCycleStart") continue; - const bucket = asRecord(value); - if (!bucket) continue; - const bucketUsed = toFiniteNumber(bucket.numRequests ?? bucket.used); - const bucketLimit = toFiniteNumber(bucket.maxRequestUsage ?? bucket.limit ?? bucket.maxRequests); - if (bucketUsed !== undefined && bucketLimit !== undefined && bucketLimit > 0) { - used = bucketUsed; - limit = bucketLimit; - break; - } - } - } - if (used === undefined || limit === undefined || limit <= 0) return null; - const percent = normalizePercent((used / limit) * 100); - if (percent === undefined) return null; - const startOfMonth = normalizeResetAt(body.startOfMonth ?? body.billingCycleStart); - // Next reset = same day next month, computed in UTC to avoid timezone-shifted rollover. - const monthlyResetAt = startOfMonth !== undefined - ? (() => { - const start = new Date(startOfMonth); - return Date.UTC(start.getUTCFullYear(), start.getUTCMonth() + 1, start.getUTCDate()); - })() - : undefined; - const built = report(provider, "cursor:auth-usage", { - monthlyPercent: percent, - ...(monthlyResetAt !== undefined ? { monthlyResetAt } : {}), - updatedAt: Date.now(), - }); - return built ? { ...built, reverseEngineered: true } : null; -} - -function quotaInfoEntries(modelInfo: Record): Record[] { - const entries: Record[] = []; - const add = (value: unknown, tier?: string) => { - const rec = asRecord(value); - if (!rec) return; - entries.push(tier ? { ...rec, tier } : rec); - }; - const addArray = (value: unknown) => { - if (!Array.isArray(value)) return; - for (const entry of value) add(entry); - }; - - if (Array.isArray(modelInfo.quotaInfo)) addArray(modelInfo.quotaInfo); - else add(modelInfo.quotaInfo); - addArray(modelInfo.quotaInfos); - - const byTier = asRecord(modelInfo.quotaInfoByTier); - if (byTier) { - for (const [tier, value] of Object.entries(byTier)) { - if (Array.isArray(value)) { - for (const entry of value) add(entry, tier); - } else { - add(value, tier); - } - } - } - return entries; -} - -function classifyAntigravityFamily(modelId: string, modelInfo: Record, quotaInfo: Record): "Gem" | "Cla" | null { - const displayName = typeof modelInfo.displayName === "string" ? modelInfo.displayName : ""; - const tier = typeof quotaInfo.tier === "string" ? quotaInfo.tier : ""; - const haystack = `${modelId} ${displayName} ${tier}`.toLowerCase(); - if (haystack.includes("gemini")) return "Gem"; - if (haystack.includes("claude") || haystack.includes("opus") || haystack.includes("sonnet") || haystack.includes("gpt-oss") || haystack.includes("gpt_oss")) return "Cla"; - return null; -} - -function antigravityUsedPercent(quotaInfo: Record): number | undefined { - const target = asRecord(quotaInfo.remaining) ?? quotaInfo; - const remaining = normalizePercent(toFiniteNumber(target.remainingFraction) !== undefined - ? toFiniteNumber(target.remainingFraction)! * 100 - : toFiniteNumber(target.remainingPercentage) !== undefined - ? toFiniteNumber(target.remainingPercentage)! * 100 - : undefined); - if (remaining === undefined) return undefined; - return normalizePercent(100 - remaining); -} - -/** Gem/Cla windows from a `fetchAvailableModels` body; shared by the provider and account probes. */ -function antigravityWindowsFromModels(body: Record | null): ProviderQuotaWindow[] { - const models = asRecord(body?.models); - if (!models) return []; - - const windows = new Map(); - for (const [modelId, rawModelInfo] of Object.entries(models)) { - const modelInfo = asRecord(rawModelInfo); - if (!modelInfo) continue; - for (const quotaInfo of quotaInfoEntries(modelInfo)) { - const label = classifyAntigravityFamily(modelId, modelInfo, quotaInfo); - if (!label || windows.has(label)) continue; - const percent = antigravityUsedPercent(quotaInfo); - if (percent === undefined) continue; - windows.set(label, { - label, - percent, - ...(normalizeResetAt(quotaInfo.resetTime) !== undefined ? { resetAt: normalizeResetAt(quotaInfo.resetTime) } : {}), - }); - } - } - - const customWindows = ["Gem", "Cla"].flatMap(label => { - const window = windows.get(label); - return window ? [window] : []; - }); - return customWindows; -} - -/** - * Parse Google Antigravity quota from `v1internal:retrieveUserQuotaSummary`. - * Groups contain Gemini models and Claude/3P models, each with 5h and weekly limit buckets. - */ -function parseAntigravityQuotaSummary(body: Record | null): ProviderQuota | null { - const groups = Array.isArray(body?.groups) ? (body.groups as unknown[]) : []; - if (groups.length === 0) return null; - - const customWindowsMap = new Map(); - - for (const rawGroup of groups) { - const group = asRecord(rawGroup); - if (!group) continue; - const groupName = `${typeof group.displayName === "string" ? group.displayName : ""} ${typeof group.description === "string" ? group.description : ""}`.toLowerCase(); - const isGemini = groupName.includes("gemini"); - const isClaude = groupName.includes("claude") || groupName.includes("3p") || groupName.includes("gpt"); - - const buckets = Array.isArray(group.buckets) ? (group.buckets as unknown[]) : []; - for (const rawBucket of buckets) { - const bucket = asRecord(rawBucket); - if (!bucket) continue; - const windowStr = `${typeof bucket.window === "string" ? bucket.window : ""} ${typeof bucket.bucketId === "string" ? bucket.bucketId : ""} ${typeof bucket.displayName === "string" ? bucket.displayName : ""}`.toLowerCase(); - const percent = antigravityUsedPercent(bucket); - if (percent === undefined) continue; - const resetAt = normalizeResetAt(bucket.resetTime); - - const isWeekly = windowStr.includes("week"); - const is5h = windowStr.includes("5h") || windowStr.includes("five"); - - if (isGemini) { - const label = is5h ? "Gem" : isWeekly ? "Gem (Weekly)" : ""; - if (label && !customWindowsMap.has(label)) { - customWindowsMap.set(label, { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }); - } - } else if (isClaude) { - const label = is5h ? "Cla" : isWeekly ? "Cla (Weekly)" : ""; - if (label && !customWindowsMap.has(label)) { - customWindowsMap.set(label, { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }); - } - } else { - const baseLabel = typeof group.displayName === "string" ? group.displayName : "Other"; - const label = isWeekly ? `${baseLabel} (Weekly)` : baseLabel; - if (!customWindowsMap.has(label)) { - customWindowsMap.set(label, { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }); - } - } - } - } - - const PREFERRED_ORDER = ["Gem", "Gem (Weekly)", "Cla", "Cla (Weekly)"]; - const customWindows = Array.from(customWindowsMap.values()).sort((a, b) => { - const ia = PREFERRED_ORDER.indexOf(a.label); - const ib = PREFERRED_ORDER.indexOf(b.label); - if (ia !== -1 && ib !== -1) return ia - ib; - if (ia !== -1) return -1; - if (ib !== -1) return 1; - return a.label.localeCompare(b.label); - }); - - if (customWindows.length === 0) { - return null; - } - - return { - customWindows, - updatedAt: Date.now(), - }; -} - -const ANTIGRAVITY_ACCOUNT_QUOTA_BASE = "https://daily-cloudcode-pa.googleapis.com"; -const ANTIGRAVITY_QUOTA_SUMMARY_URL = `${ANTIGRAVITY_ACCOUNT_QUOTA_BASE}/v1internal:retrieveUserQuotaSummary`; -const ANTIGRAVITY_QUOTA_MODELS_URL = `${ANTIGRAVITY_ACCOUNT_QUOTA_BASE}/v1internal:fetchAvailableModels`; - -/** Only these fixed accounting destinations may use transparent Fake-IP DNS. */ -export function isCanonicalAntigravityQuotaUrl(name: string, url: string): boolean { - return name === "google-antigravity" - && (url === ANTIGRAVITY_QUOTA_SUMMARY_URL || url === ANTIGRAVITY_QUOTA_MODELS_URL); -} - -let antigravityOutboundDependencies: ProviderOutboundDependencies = { - isCanonicalUrl: isCanonicalAntigravityQuotaUrl, -}; - -/** Test seam: inject resolver/pinned transport for provider and per-account probes. */ -export function setAntigravityAccountQuotaTransportForTests(dependencies: ProviderOutboundDependencies | null): void { - antigravityOutboundDependencies = { ...dependencies, isCanonicalUrl: isCanonicalAntigravityQuotaUrl }; -} - -/** - * Per-account Antigravity quota (#1082). Always probes Google's own Cloud Code Assist host - * through the pinned provider-outbound transport: a configured `baseUrl` is a routing choice - * for requests, not a second source of Google's accounting for a stored credential, and fixing - * the destination keeps the `provider\0accountId` cache identity exact across config changes. - * A redirect or non-2xx yields null (unavailable), never a partial row. - */ -type AntigravityQuotaProbeResult = - | { kind: "available"; quota: ProviderQuota; source: "google-antigravity:retrieveUserQuotaSummary" | "google-antigravity:fetchAvailableModels" } - | { kind: "unavailable"; failure: QuotaFailureCode; legacy: { kind: "null" } | { kind: "throw"; error: unknown } }; - -function quotaTransportFailure(error: unknown): QuotaFailureCode { - if (error instanceof ProviderOutboundPolicyError) return "destination_blocked"; - if (error instanceof DestinationDnsResolutionError) return "dns_failed"; - if (error instanceof PinnedHttpError) return error.code === "output_byte_limit" ? "response_unusable" : "timeout"; - if (error instanceof DOMException && error.name === "TimeoutError") return "timeout"; - return "transport_error"; -} - -function quotaHttpFailure(status: number): QuotaFailureCode { - if (status >= 300 && status < 400) return "redirect_blocked"; - if (status === 401 || status === 403) return "access_denied"; - if (status === 429) return "rate_limited"; - return "upstream_error"; -} - -function unavailableAntigravityQuota(failure: QuotaFailureCode): AntigravityQuotaProbeResult { - return { kind: "unavailable", failure, legacy: { kind: "null" } }; -} - -/** - * Prefer a summary network-policy diagnosis over a vaguer fallback. A blocked - * destination is an actionable local-network fact, while "upstream_error" tells - * the operator to go look at Google. A successful models probe still clears - * the first failure completely. - */ -function antigravityUnavailableFailure( - summaryFailure: QuotaFailureCode | undefined, - fallbackFailure: QuotaFailureCode, -): QuotaFailureCode { - if ( - (summaryFailure === "destination_blocked" || summaryFailure === "dns_failed") - && fallbackFailure !== "destination_blocked" - && fallbackFailure !== "dns_failed" - ) { - return summaryFailure; - } - return fallbackFailure; -} - -async function probeAntigravityUsageQuota(accessToken: string, projectId: string): Promise { - const fetchQuota = (url: string) => providerOutboundPost("google-antigravity", { baseUrl: ANTIGRAVITY_ACCOUNT_QUOTA_BASE }, url, { - headers: { - Accept: "application/json", "Content-Type": "application/json", - "User-Agent": antigravityUserAgent(), Authorization: `Bearer ${accessToken}`, - }, - body: JSON.stringify({ project: projectId }), signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), - }, antigravityOutboundDependencies); - let summaryFailure: QuotaFailureCode | undefined; - try { - const response = await fetchQuota(ANTIGRAVITY_QUOTA_SUMMARY_URL); - if (await providerRedirectError(response, ANTIGRAVITY_QUOTA_SUMMARY_URL)) return unavailableAntigravityQuota("redirect_blocked"); - if (response.status === 401 || response.status === 403) return unavailableAntigravityQuota("access_denied"); - if (response.ok) { - const quota = parseAntigravityQuotaSummary(asRecord(await readQuotaJson(response))); - if (quota) return { kind: "available", quota, source: "google-antigravity:retrieveUserQuotaSummary" }; - } - } catch (error) { - // Existing behavior: summary transport/parse failure may recover through the models probe. - summaryFailure = quotaTransportFailure(error); - } - try { - const response = await fetchQuota(ANTIGRAVITY_QUOTA_MODELS_URL); - if (await providerRedirectError(response, ANTIGRAVITY_QUOTA_MODELS_URL)) { - return unavailableAntigravityQuota(antigravityUnavailableFailure(summaryFailure, "redirect_blocked")); - } - if (!response.ok) { - return unavailableAntigravityQuota(antigravityUnavailableFailure(summaryFailure, quotaHttpFailure(response.status))); - } - const customWindows = antigravityWindowsFromModels(asRecord(await readQuotaJson(response))); - if (!customWindows.length) { - return unavailableAntigravityQuota(antigravityUnavailableFailure(summaryFailure, "response_unusable")); - } - return { kind: "available", quota: { customWindows, updatedAt: Date.now() }, source: "google-antigravity:fetchAvailableModels" }; - } catch (error) { - // The public compatibility wrapper still rejects this exact fallback error; it never enters a DTO. - return { - kind: "unavailable", - failure: antigravityUnavailableFailure(summaryFailure, quotaTransportFailure(error)), - legacy: { kind: "throw", error }, - }; - } -} - -export async function fetchAntigravityUsageQuota(accessToken: string, projectId: string): Promise { - const result = await probeAntigravityUsageQuota(accessToken, projectId); - if (result.kind === "available") return result.quota; - if (result.legacy.kind === "throw") throw result.legacy.error; - return null; -} - -async function fetchAntigravityQuota(provider: string): Promise { - const credential = getCredential("google-antigravity"); - if (!credential?.projectId) return null; - let accessToken: string; - try { accessToken = await getValidAccessToken("google-antigravity"); } catch { return null; } - const result = await probeAntigravityUsageQuota(accessToken, credential.projectId); - if (result.kind === "available") return report(provider, result.source, result.quota); - if (result.legacy.kind === "throw") throw result.legacy.error; - return null; -} -type KeyQuotaReader = (name: string, provider: OcxProviderConfig) => Promise; - -/** Same selector drives cheap capabilities and uncached reads; never resolves credentials. */ -function keyQuotaReaderForProvider(name: string, provider: OcxProviderConfig): KeyQuotaReader | null { - if (provider.disabled === true || (provider.authMode ?? "key") !== "key") return null; - if (isCanonicalKimiCodeBaseUrl(provider.baseUrl)) { - return async (id, config) => { - const bearer = await resolveKimiQuotaBearer(config); - return bearer ? fetchKimiQuota(id, config, bearer) : null; - }; - } - if (name === "commandcode" && isCanonicalCommandCodeBaseUrl(provider.baseUrl)) { - return async (id, config) => { - const bearer = await resolveCommandCodeQuotaBearer(config); - return bearer ? fetchCommandCodeQuota(id, config, bearer) : null; - }; - } - if (registryEntryForProviderDestination(provider)?.id === "opencode-go") return fetchOpenCodeGoQuota; - if (isCanonicalA6apiBaseUrl(provider.baseUrl)) return fetchA6apiQuota; - if (name === "openrouter" && isCanonicalOpenRouterBaseUrl(provider.baseUrl)) return fetchOpenRouterQuota; - if (name === "deepseek" && isCanonicalDeepSeekBaseUrl(provider.baseUrl)) return fetchDeepSeekQuota; - if (name === "cline-pass" && isCanonicalClineBaseUrl(provider.baseUrl)) return fetchClineQuota; - if (isCanonicalOllamaCloudBaseUrl(provider.baseUrl ?? getProviderRegistryEntry(name)?.baseUrl)) return fetchOllamaCloudQuota; - // #4201: the Responses preset is the same domestic GLM Coding Plan subscription on the OpenAI - // Responses wire, so it reads the same monitor endpoint. Eligibility stays a name list AND the - // canonical-URL guard: the guard is what keeps BigModel's bare-key Authorization from reaching a - // lookalike host, so a same-named custom destination still dispatches nothing. - if (["zai", "glm", "glm-cn", "zhipu-bigmodel-coding", "zhipu-bigmodel-responses"].includes(name) && isCanonicalZaiBaseUrl(provider.baseUrl)) return fetchZaiQuota; - if (["minimax", "minimax-cn"].includes(name) && isCanonicalMinimaxBaseUrl(provider.baseUrl)) return fetchMinimaxQuota; - if (name === "moonshot" && isCanonicalMoonshotBaseUrl(provider.baseUrl)) return fetchMoonshotQuota; - if (name === "venice" && isCanonicalVeniceBaseUrl(provider.baseUrl)) return fetchVeniceQuota; - if (name === "synthetic" && isCanonicalSyntheticBaseUrl(provider.baseUrl)) return fetchSyntheticQuota; - if (name === "deepinfra" && isCanonicalDeepInfraBaseUrl(provider.baseUrl)) return fetchDeepInfraQuota; - if (name === "neuralwatt" && isCanonicalNeuralwattBaseUrl(provider.baseUrl)) return fetchNeuralwattQuota; - return null; -} +import { listCodexAuthAccountsSnapshot } from "../codex/auth-api"; +import { resolveEnvValue } from "../config"; +import { getAccountCredential, getAccountSet } from "../oauth/store"; +import { apiKeyPoolEntryId } from "./api-keys"; +import { captureConfigGeneration, sweepExpiredOnWrite } from "../lib/state-store-sweeper"; +import { ACCOUNT_QUOTA_TTL_MS, CACHE_TTL_MS } from "./quota-wire"; +import { replaceCachedProviderQuotas } from "./quota-routing-cache"; +import { + commitKiroAccountUsageState, + fetchKiroUsageSnapshot, + type KiroUsageSnapshot, + kiroUsageContextForAccount, +} from "./kiro-usage"; +import { mapQuotaRoster, readProviderApiKeyQuotas, type ProviderApiKeyQuota } from "./quota-key-accounts"; +import type { OcxConfig, OcxProviderConfig } from "../types"; +import type { ProviderQuota, QuotaFailureCode } from "./quota-types"; +import { + accountReportCurrent, + AUTHORITATIVE_EMPTY_QUOTA, + bumpProviderQuotaInvalidationEpoch, + cacheKeyWithAggregationState, + getProviderQuotaReportCache, + hasCodexPoolProvider, + inflight, + invalidationEpoch, + isBuiltInChatGptForwardProvider, + isProviderQuotaReportCurrent, + LAST_GOOD_MAX_AGE_MS, + providerQuotaBeforePublishForTests, + routingEvidence, + setProviderQuotaReportCache, + TERMINAL_QUOTA_FAILURE, + type CodexAuthAccountsSnapshotPromise, + type ProviderQuotaProbeResult, + type ProviderQuotaReport, + type ProviderQuotaResponse, +} from "./quota/report-cache"; +import { + accountCacheKey, + accountQuotaCache, + accountQuotaInflight, + explicitAccountEpoch, + explicitAccountReader, + explicitQuotaConfig, + explicitQuotaDestination, + explicitQuotaIdentity, + getTokenForAccountQuotaProbe, + hasPassiveAccountQuota, + hydrateAccountQuotaCache, + mayCommitAccountQuotaKey, + mayCommitProviderQuotaKey, + normalizeAnthropicQuota, + supportsPerAccountQuota, + type AccountQuotaCacheEntry, + type ProviderAccountQuota, +} from "./quota/account-cache"; +import { + fetchAnthropicQuota, + fetchAnthropicUsageQuota, + fetchChatGptForwardQuota, + fetchCursorQuota, + fetchKiroQuota, + fetchMuseKeyQuota, + fetchPassiveProviderQuota, + fetchXaiQuota, +} from "./quota/vendor-probes-oauth"; +import { fetchCommandCodeQuota, fetchKimiQuota, keyQuotaReaderForProvider } from "./quota/vendor-probes-key"; +import { antigravityQuotaDiagnosticIdentity, fetchAntigravityQuota, probeAntigravityUsageQuota } from "./quota/antigravity"; -export function providerApiKeyQuotaMode(name: string, provider: OcxProviderConfig): AccountQuotaMode { - return keyQuotaReaderForProvider(name, provider) ? "probe" : "unsupported"; -} +export type { ProviderQuota, ProviderQuotaCreditsUsd, ProviderQuotaWindow } from "./quota-types"; +export { QUOTA_RESPONSE_MAX_BYTES } from "./quota-wire"; +export { + clearProviderQuotaCache, + publishKeyReportForTests, + readProviderQuotaJsonForTests, + setProviderQuotaBeforePublishForTests, + type ProviderQuotaReport, + type ProviderQuotaResponse, +} from "./quota/report-cache"; +export { + clearAccountQuotaCache, + getCachedProviderAccountQuota, + hasPassiveAccountQuota, + parseAnthropicRateLimitHeaders, + providerOAuthAccountQuotaMode, + readPassiveProviderAccountQuotas, + recordAnthropicAccountQuotaFromHeaders, + recordPassiveAccountQuota, + reconcileProviderAccountQuotaRows, + resetProviderQuotaReconcileStateForTests, + setCachedProviderAccountQuotaForTests, + supportsPerAccountQuota, + sweepExpiredProviderAccountQuotaRows, + type ProviderAccountQuota, +} from "./quota/account-cache"; +export { fetchAntigravityUsageQuota, isCanonicalAntigravityQuotaUrl, setAntigravityAccountQuotaTransportForTests } from "./quota/antigravity"; +export { parseOllamaCloudQuota, parseZaiQuotaLimits, providerApiKeyQuotaMode } from "./quota/vendor-probes-key"; +export { parseXaiCreditsResponse } from "./quota/vendor-probes-oauth"; export async function fetchProviderApiKeyQuotas(config: OcxConfig, name: string, forceRefresh = false): Promise { const provider = config.providers[name]; @@ -3220,19 +244,21 @@ export async function fetchProviderQuotaReports(config: OcxConfig, forceRefresh // by construction and never becomes fresher on its own. Without the exemption a single // configured passive provider makes this predicate permanently false, so every dashboard // poll re-probes every OTHER provider upstream instead of serving the 5-minute cache. - const cacheFresh = cache && cache.key === key && now - cache.ts < CACHE_TTL_MS - && cache.response.reports.every(item => + const currentCache = getProviderQuotaReportCache(); + const cacheFresh = currentCache && currentCache.key === key && now - currentCache.ts < CACHE_TTL_MS + && currentCache.response.reports.every(item => (item.observed === true || now - item.updatedAt < LAST_GOOD_MAX_AGE_MS) && isProviderQuotaReportCurrent(item)); - if (!forceRefresh && cacheFresh) return cache!.response; + if (!forceRefresh && cacheFresh) return currentCache!.response; const joinable = inflight.get(key); if (!forceRefresh && joinable && joinable.epoch === invalidationEpoch) return joinable.promise; // A forced probe takes commit authority: older in-flight probes must not overwrite its result. - if (forceRefresh) invalidationEpoch += 1; + if (forceRefresh) bumpProviderQuotaInvalidationEpoch(); const epoch = invalidationEpoch; const promise = (async (): Promise => { - const previous = cache && cache.key === key ? cache.response.reports : []; + const previousCache = getProviderQuotaReportCache(); + const previous = previousCache && previousCache.key === key ? previousCache.response.reports : []; const probeResults = await Promise.all( Object.entries(config.providers).map(([name, provider]) => ( maybeFetchProviderQuota(name, provider, config, forceRefresh, prefetchedCodexSnapshot) @@ -3296,7 +322,7 @@ export async function fetchProviderQuotaReports(config: OcxConfig, forceRefresh && generationMismatchedProviders.size === 0 ) { const reports = response.reports.filter(item => mayCommitProviderQuotaKey(item.provider, writerGeneration)); - cache = { key, ts: Date.now(), response: { ...response, reports } }; + setProviderQuotaReportCache({ key, ts: Date.now(), response: { ...response, reports } }); replaceCachedProviderQuotas(reports, routingEvidence); notifyProviderQuotaSnapshot(reports, config); } @@ -3311,3 +337,222 @@ export async function fetchProviderQuotaReports(config: OcxConfig, forceRefresh if (inflight.get(key) === entry) inflight.delete(key); } } + +async function readExplicitAccountQuota(provider: string, accountId: string, configured?: OcxProviderConfig): Promise<{ + result: ProviderQuotaProbeResult; + identity: string | undefined; + isCurrent: () => boolean; +} | null> { + const target = explicitQuotaConfig(provider, configured); + if (!target || !explicitQuotaDestination(provider, target)) return null; + const config = { ...target }; + const epoch = explicitAccountEpoch; + const accessToken = await getTokenForAccountQuotaProbe(provider, accountId); + const credential = getAccountCredential(provider, accountId); + if (!credential || credential.access !== accessToken) return null; + // Pair the post-renewal credential with the destination captured before renewal. + const identity = explicitQuotaIdentity(provider, accountId, config); + const isCurrent = () => epoch === explicitAccountEpoch + && identity === explicitQuotaIdentity(provider, accountId, configured); + if (!isCurrent()) return null; + let result: ProviderQuotaProbeResult; + switch (provider) { + case "xai": result = await fetchXaiQuota(provider, { accessToken, upstreamAccountId: credential.accountId }); break; + case "cursor": result = await fetchCursorQuota(provider, accessToken); break; + case "kimi": result = await fetchKimiQuota(provider, config, accessToken); break; + case "command-code": result = await fetchCommandCodeQuota(provider, config, accessToken); break; + default: return null; + } + return { result, identity, isCurrent }; +} + +async function fetchExplicitAccountQuota(provider: string, accountId: string, force: boolean, configured?: OcxProviderConfig): Promise { + const key = accountCacheKey(provider, accountId); + const identity = explicitQuotaIdentity(provider, accountId, configured); + const previous = accountQuotaCache.get(key); + const cached = identity && previous?.identity === identity && previous.isCurrent?.() ? previous : undefined; + if (!force && cached && Date.now() - cached.ts < ACCOUNT_QUOTA_TTL_MS + && (!cached.quota || Date.now() - cached.quota.updatedAt < LAST_GOOD_MAX_AGE_MS)) return cached; + const flightKey = `${key}\u0000${identity ?? "missing"}`; + const running = accountQuotaInflight.get(flightKey); + if (running) return running; + const epoch = explicitAccountEpoch; + const lastGood = cached?.quota && Date.now() - cached.quota.updatedAt < LAST_GOOD_MAX_AGE_MS ? cached.quota : null; + const flight = (async (): Promise => { + let read: Awaited> = null; + try { read = await readExplicitAccountQuota(provider, accountId, configured); } catch { /* unavailable */ } + const isCurrent = read?.isCurrent ?? (() => epoch === explicitAccountEpoch && !!identity + && identity === explicitQuotaIdentity(provider, accountId, configured)); + const result = read?.result; + const current = epoch === explicitAccountEpoch && isCurrent(); + const quota = current && result && typeof result !== "symbol" ? result.quota : null; + const empty = result === AUTHORITATIVE_EMPTY_QUOTA; + const entry: AccountQuotaCacheEntry = { + ts: Date.now(), + quota: quota ?? (current && result !== TERMINAL_QUOTA_FAILURE && !empty + && lastGood && Date.now() - lastGood.updatedAt < LAST_GOOD_MAX_AGE_MS ? lastGood : null), + ...(!current || (!quota && !empty) ? { unavailable: true as const } : {}), + identity: read?.identity ?? identity, + isCurrent: () => epoch === explicitAccountEpoch && isCurrent(), + }; + if (entry.isCurrent?.()) accountQuotaCache.set(key, entry); + return entry; + })().finally(() => { if (accountQuotaInflight.get(flightKey) === flight) accountQuotaInflight.delete(flightKey); }); + accountQuotaInflight.set(flightKey, flight); + return flight; +} + +async function fetchExplicitCurrentQuota(provider: string, config: OcxProviderConfig, liveConfig: OcxConfig): Promise { + const id = getAccountSet(provider)?.activeAccountId; + if (!id) return null; + const read = await readExplicitAccountQuota(provider, id, config); + if (!read) return null; + const isCurrent = () => liveConfig.providers[provider] === config + && read.isCurrent() && getAccountSet(provider)?.activeAccountId === id; + if (!isCurrent()) return TERMINAL_QUOTA_FAILURE; + if (read.result && typeof read.result !== "symbol") accountReportCurrent.set(read.result, isCurrent); + return read.result; +} + + +async function fetchAccountQuota( + provider: string, + accountId: string, + forceRefresh: boolean, + providerConfig?: OcxProviderConfig, +): Promise { + if (!supportsPerAccountQuota(provider)) return { ts: Date.now(), quota: null, unavailable: true }; + if (explicitAccountReader(provider)) return fetchExplicitAccountQuota(provider, accountId, forceRefresh, providerConfig); + if (provider === "anthropic") hydrateAccountQuotaCache(); + const key = accountCacheKey(provider, accountId); + const writerGeneration = captureConfigGeneration(); + const cached = accountQuotaCache.get(key); + if (!forceRefresh && cached && Date.now() - cached.ts < ACCOUNT_QUOTA_TTL_MS) { + if (provider === "google-antigravity" && cached.quotaFailure && cached.quotaFailureIsCurrent?.() !== true) return { ...cached, quotaFailure: undefined }; + return provider === "anthropic" ? { ...cached, quota: normalizeAnthropicQuota(cached.quota, Date.now()) } : cached; + } + const joinable = accountQuotaInflight.get(key); + if (joinable) return joinable; + + const epoch = explicitAccountEpoch; + const probe = (async (): Promise => { + let diagnosticIdentity: string | undefined; + let quotaFailure: QuotaFailureCode | undefined; + const quotaFailureIsCurrent = () => { + try { return epoch === explicitAccountEpoch && diagnosticIdentity !== undefined && diagnosticIdentity === antigravityQuotaDiagnosticIdentity(accountId); } + catch { return false; } + }; + const diagnosticFields = () => quotaFailure && quotaFailureIsCurrent() ? { quotaFailure, quotaFailureIsCurrent } : {}; + try { + if (provider === "google-antigravity") diagnosticIdentity = antigravityQuotaDiagnosticIdentity(accountId); + let quota: ProviderQuota | null; + let kiroSnapshot: KiroUsageSnapshot | null = null; + if (provider === "kiro") { + // Kiro resolves the bearer and its routing metadata from ONE account-scoped + // snapshot. It deliberately does not use getTokenForAccountQuotaProbe: that + // helper refuses to refresh a background `local-cli` slot because Anthropic's + // lock can adopt a mismatched Claude CLI identity, but Kiro marks every + // CLI-imported credential `local-cli`, so the same rule would blank the quota of + // every inactive pool account the moment its token expired. + kiroSnapshot = await fetchKiroUsageSnapshot(await kiroUsageContextForAccount(accountId)); + quota = kiroSnapshot?.quota ?? null; + } else { + const token = await getTokenForAccountQuotaProbe(provider, accountId); + if (provider === "google-antigravity") { + // Per-account Gem/Cla windows (#1082). The project id is part of the stored + // credential; without it the probe cannot be made, and that is "unavailable", + // never 0%. + const credential = getAccountCredential(provider, accountId); + diagnosticIdentity = credential?.access === token ? antigravityQuotaDiagnosticIdentity(accountId, credential) : undefined; + if (!diagnosticIdentity || !credential?.projectId) throw new Error("antigravity account unavailable"); + const result = await probeAntigravityUsageQuota(token, credential.projectId); + quota = result.kind === "available" ? result.quota : null; + if (result.kind === "unavailable") quotaFailure = result.failure; + } else if (provider === "anthropic") { + quota = await fetchAnthropicUsageQuota(token); + } else { + return { ts: Date.now(), quota: null, unavailable: true }; + } + } + if (!quota) { + // Preserve last-good bars and mark unavailable; advance TTL so failures + // negative-cache instead of re-probing on every GUI poll. + const entry: AccountQuotaCacheEntry = { + ts: Date.now(), + // Settle once for all joiners against observations committed during the probe. + quota: provider === "anthropic" + ? normalizeAnthropicQuota(accountQuotaCache.get(key)?.quota, Date.now()) : cached?.quota ?? null, + unavailable: true, + ...diagnosticFields(), + }; + if (mayCommitAccountQuotaKey(key, writerGeneration)) { + accountQuotaCache.set(key, entry); + if (provider === "kiro") commitKiroAccountUsageState(key, null); + sweepExpiredOnWrite(entry.ts); + } + return entry; + } + const entry: AccountQuotaCacheEntry = { + ts: Date.now(), quota: provider === "anthropic" ? normalizeAnthropicQuota(quota, Date.now()) : quota, + }; + if (mayCommitAccountQuotaKey(key, writerGeneration)) { + accountQuotaCache.set(key, entry); + // Exhaustion state rides the SAME commit guard as the quota row: a probe from a + // superseded config generation must not publish either half. + if (provider === "kiro") commitKiroAccountUsageState(key, kiroSnapshot); + sweepExpiredOnWrite(entry.ts); + } + return entry; + } catch { + if (provider === "google-antigravity") quotaFailure = "account_unavailable"; + const entry: AccountQuotaCacheEntry = { + ts: Date.now(), + quota: provider === "anthropic" + ? normalizeAnthropicQuota(accountQuotaCache.get(key)?.quota, Date.now()) : cached?.quota ?? null, + unavailable: true, + ...diagnosticFields(), + }; + if (mayCommitAccountQuotaKey(key, writerGeneration)) { + accountQuotaCache.set(key, entry); + sweepExpiredOnWrite(entry.ts); + } + return entry; + } + })().finally(() => { + if (accountQuotaInflight.get(key) === probe) accountQuotaInflight.delete(key); + }); + accountQuotaInflight.set(key, probe); + return probe; +} + +/** + * Per-account quota rows for a provider's logged-in accounts. Probes run in parallel; a + * single failing account never blocks the others. + */ +export async function fetchProviderAccountQuotas( + provider: string, + forceRefresh = false, + providerConfig?: OcxProviderConfig, +): Promise { + if (!supportsPerAccountQuota(provider)) return []; + const set = getAccountSet(provider); + if (!set) return []; + return mapQuotaRoster(set.accounts, async account => { + const entry = await fetchAccountQuota(provider, account.id, forceRefresh, providerConfig); + const result: ProviderAccountQuota = { + accountId: account.id, + quota: provider === "anthropic" ? normalizeAnthropicQuota(entry.quota, Date.now()) : entry.quota, + ...(entry.unavailable ? { unavailable: true as const } : {}), + ...(entry.unavailable && entry.quotaFailure && entry.quotaFailureIsCurrent?.() === true ? { quotaFailure: entry.quotaFailure } : {}), + }; + if (entry.quotaFailureIsCurrent) Object.defineProperty(result, "quotaFailureIsCurrent", { value: entry.quotaFailureIsCurrent }); + if (!explicitAccountReader(provider)) return result; + const identity = entry.identity; + Object.defineProperty(result, "isCurrent", { value: () => { + if (entry.isCurrent) return entry.isCurrent(); + const credential = getAccountCredential(provider, account.id); + return !!credential && (!identity || explicitQuotaIdentity(provider, account.id, providerConfig) === identity); + } }); + return result; + }); +} diff --git a/src/providers/quota/account-cache.ts b/src/providers/quota/account-cache.ts new file mode 100644 index 0000000000..93c437a61d --- /dev/null +++ b/src/providers/quota/account-cache.ts @@ -0,0 +1,441 @@ +import { createHash } from "node:crypto"; +import { getValidAccessTokenForAccount } from "../../oauth"; +import { getAccountCredential, getAccountSet } from "../../oauth/store"; +import type { GenerationContext } from "../../lib/state-store-sweeper"; +import { ACCOUNT_QUOTA_TTL_MS, toFiniteNumber } from "../quota-wire"; +import { clearKiroAccountUsageState, reconcileKiroAccountUsageState } from "../kiro-usage"; +import { cancelPendingAccountQuotaPersist, readPersistedAccountQuotas, schedulePersistAccountQuotas } from "../account-quota-disk"; +import { replaceCachedProviderQuotas } from "../quota-routing-cache"; +import { getProviderRegistryEntry } from "../registry"; +import { getProviderQuotaReportCache, hasQuotaRows, routingEvidence, setProviderQuotaReportCache } from "./report-cache"; +import { isCanonicalCommandCodeBaseUrl, isCanonicalKimiCodeBaseUrl } from "./vendor-probes-key"; +import type { AccountQuotaMode, ProviderQuota, ProviderQuotaWindow, QuotaFailureCode } from "../quota-types"; +import type { OcxConfig, OcxProviderConfig } from "../../types"; + +/** Match oauth/index REFRESH_SKEW_MS — use stored access without refresh when still fresh. */ +const ACCOUNT_TOKEN_SKEW_MS = 60_000; + +/** + * Anthropic and Kiro both report usage per CREDENTIAL, so every logged-in account can be + * probed with its own bearer token — the active-account selection and the local usage log + * are irrelevant here. Mirrors the Codex pool behaviour + * (codex/auth-api.ts:fetchPoolAccountQuota), including a per-account TTL so N accounts cost + * at most N upstream calls per window. `ACCOUNT_QUOTA_TTL_MS` lives in `quota-wire.ts` + * because the Kiro exhaustion reader applies the same staleness bound. + */ +export type AccountQuotaCacheEntry = { + ts: number; + quota: ProviderQuota | null; + /** Last probe failed (429 / network / expired login); still may hold last-good quota. */ + unavailable?: true; + quotaFailure?: QuotaFailureCode; + quotaFailureIsCurrent?: () => boolean; + /** Private new-reader identity; never persisted or serialized. */ + identity?: string; + isCurrent?: () => boolean; +}; +/** Expired measurements become unknown; missing reset evidence never implies a fresh allowance. */ +export function normalizeAnthropicQuota(quota: ProviderQuota | null | undefined, now: number): ProviderQuota | null { + if (!quota) return null; + const validReset = (resetAt: unknown): resetAt is number => typeof resetAt === "number" + && Number.isFinite(resetAt) && resetAt > 0 && Number.isFinite(new Date(resetAt).getTime()); + let result = quota; + for (const [percent, reset] of [ + ["fiveHourPercent", "fiveHourResetAt"], + ["weeklyPercent", "weeklyResetAt"], + ["monthlyPercent", "monthlyResetAt"], + ] as const) { + const resetAt = quota[reset]; + if (resetAt === undefined) continue; + const valid = validReset(resetAt); + if (valid && resetAt > now) continue; + if (result === quota) result = { ...quota }; + if (valid) delete result[percent]; + delete result[reset]; + } + // Persisted rows validate only the outer quota object, so custom data may be malformed. + if (quota.customWindows !== undefined) { + const windows = Array.isArray(quota.customWindows) ? quota.customWindows : []; + const retained: ProviderQuotaWindow[] = []; + let changed = !Array.isArray(quota.customWindows); + for (const window of windows) { + if (!window || typeof window !== "object" || typeof window.label !== "string" || !window.label.trim() + || typeof window.percent !== "number" || !Number.isFinite(window.percent) + || window.percent < 0 || window.percent > 100) { + changed = true; + continue; + } + if (validReset(window.resetAt) && window.resetAt <= now) { + changed = true; + continue; + } + if (window.resetAt !== undefined && !validReset(window.resetAt)) { + const normalized = { ...window }; + delete normalized.resetAt; + retained.push(normalized); + changed = true; + } else { + retained.push(window); + } + } + if (changed) { + if (result === quota) result = { ...quota }; + if (retained.length) result.customWindows = retained; + else delete result.customWindows; + } + } + return hasQuotaRows(result) ? result : null; +} + +export const accountQuotaCache = new Map(); +export let explicitAccountEpoch = 0; + +/** + * Seed the cache from the last run, once. + * + * Without this a restart forgets every measurement, so the pool opens its next turn with + * no idea which account has room — the exact blindness pre-dispatch selection exists to + * remove. A hydrated row is still subject to the ordinary TTL, so it orders the first + * request and is replaced by a live probe immediately after. + */ +let diskHydrated = false; +export function hydrateAccountQuotaCache(): void { + if (diskHydrated) return; + diskHydrated = true; + for (const [key, quota] of readPersistedAccountQuotas()) { + // Disk stores observation time, not the Anthropic usage probe's clock. + if (!accountQuotaCache.has(key)) { + const anthropic = key.startsWith("anthropic\u0000"); + accountQuotaCache.set(key, { + ts: anthropic ? 0 : quota.updatedAt, + quota: anthropic ? normalizeAnthropicQuota(quota, Date.now()) : quota, + }); + } + } +} + +export function persistAccountQuotaCache(): void { + schedulePersistAccountQuotas(function* () { + const now = Date.now(); + for (const [key, entry] of accountQuotaCache) { + const quota = key.startsWith("anthropic\u0000") ? normalizeAnthropicQuota(entry.quota, now) : entry.quota; + if (quota) yield [key, quota] as [string, ProviderQuota]; + } + }); +} +export const accountQuotaInflight = new Map>(); +let lastReconciledGeneration = 0; +let liveAccountQuotaKeys = new Set(); +let liveProviderQuotaKeys = new Set(); + +export function mayCommitAccountQuotaKey(key: string, writerGeneration: number): boolean { + return writerGeneration >= lastReconciledGeneration || liveAccountQuotaKeys.has(key); +} + +export function mayCommitProviderQuotaKey(key: string, writerGeneration: number): boolean { + return writerGeneration >= lastReconciledGeneration || liveProviderQuotaKeys.has(key); +} + +export interface ProviderAccountQuota { + accountId: string; + quota: ProviderQuota | null; + /** Set when the probe could not reach upstream (expired login, 429, network). */ + unavailable?: true; + quotaFailure?: QuotaFailureCode; + quotaFailureIsCurrent?: () => boolean; + isCurrent?: () => boolean; +} + +/** Providers whose per-account quota can be probed. Extend as other OAuth APIs are covered. */ +export function supportsPerAccountQuota(provider: string): boolean { + return provider === "anthropic" || provider === "kiro" || provider === "google-antigravity" + || explicitAccountReader(provider); +} + +export function explicitAccountReader(provider: string): boolean { + return provider === "xai" || provider === "cursor" || provider === "kimi" || provider === "command-code"; +} + +export function providerOAuthAccountQuotaMode(provider: string): AccountQuotaMode { + return hasPassiveAccountQuota(provider) ? "passive" : supportsPerAccountQuota(provider) ? "probe" : "unsupported"; +} + +export function accountCacheKey(provider: string, accountId: string): string { + return `${provider}\u0000${accountId}`; +} + +/** + * Synchronous last-good per-account quota read for routing. Never probes the network. + * Returns null when nothing is cached (or the cached row has no bars). + */ +export function getCachedProviderAccountQuota(provider: string, accountId: string): ProviderQuota | null { + const entry = accountQuotaCache.get(accountCacheKey(provider, accountId)); + if (entry?.isCurrent && !entry.isCurrent()) return null; + return provider === "anthropic" ? normalizeAnthropicQuota(entry?.quota, Date.now()) : entry?.quota ?? null; +} + +/** Test-only: seed or clear the per-account quota cache without probing upstream. */ +export function setCachedProviderAccountQuotaForTests( + provider: string, + accountId: string, + quota: ProviderQuota | null, +): void { + const key = accountCacheKey(provider, accountId); + if (quota === null) { + accountQuotaCache.delete(key); + return; + } + accountQuotaCache.set(key, { ts: Date.now(), quota }); +} + +/** Unified headers report utilization fractions and epoch-second reset times. */ +function anthropicHeaderResetAt(value: string | null): number | undefined { + const seconds = toFiniteNumber(value); + if (seconds === undefined || seconds <= 0) return undefined; + const timestamp = seconds * 1000; + return Number.isFinite(new Date(timestamp).getTime()) ? timestamp : undefined; +} + +export function parseAnthropicRateLimitHeaders(headers: Headers): ProviderQuota | null { + const fiveHourPercent = normalizeUtilizationFraction(headers.get("anthropic-ratelimit-unified-5h-utilization")); + const weeklyPercent = normalizeUtilizationFraction(headers.get("anthropic-ratelimit-unified-7d-utilization")); + if (fiveHourPercent === undefined && weeklyPercent === undefined) return null; + const fiveHourResetAt = anthropicHeaderResetAt(headers.get("anthropic-ratelimit-unified-5h-reset")); + const weeklyResetAt = anthropicHeaderResetAt(headers.get("anthropic-ratelimit-unified-7d-reset")); + return { + ...(fiveHourPercent !== undefined ? { fiveHourPercent } : {}), + ...(fiveHourPercent !== undefined && fiveHourResetAt !== undefined ? { fiveHourResetAt } : {}), + ...(weeklyPercent !== undefined ? { weeklyPercent } : {}), + ...(weeklyPercent !== undefined && weeklyResetAt !== undefined ? { weeklyResetAt } : {}), + updatedAt: Date.now(), + }; +} + +/** Reject unknown scales; round fraction conversion for persisted/displayed percentages. */ +function normalizeUtilizationFraction(value: string | null): number | undefined { + const numeric = toFiniteNumber(value); + if (numeric === undefined || numeric < 0 || numeric > 1) return undefined; + return Math.round(numeric * 10_000) / 100; +} + +/** + * Merge serving-account observations without advancing the usage probe's clock or + * erasing model-specific windows. The caller owns credential attribution; this guard + * prevents a retired account key from being revived by an older config generation. + */ +export function recordAnthropicAccountQuotaFromHeaders( + accountId: string, + headers: Headers, + writerGeneration: number, +): void { + if (!accountId) return; + const observed = parseAnthropicRateLimitHeaders(headers); + if (!observed) return; + const key = accountCacheKey("anthropic", accountId); + if (!mayCommitAccountQuotaKey(key, writerGeneration)) return; + // Hydrate before writing, for the same reason `recordPassiveAccountQuota` does: this write + // arrives unprompted from the request path, and `persistAccountQuotaCache` serializes the + // whole map. Landing before any reader has hydrated would persist this single row and erase + // every other provider's saved row. + hydrateAccountQuotaCache(); + const previous = accountQuotaCache.get(key); + accountQuotaCache.set(key, { + ...previous, + // Headers do not prove that the last usage probe succeeded. + ts: previous?.ts ?? 0, + quota: normalizeAnthropicQuota({ + ...normalizeAnthropicQuota(previous?.quota, observed.updatedAt), ...observed, + }, observed.updatedAt), + }); + persistAccountQuotaCache(); +} + +/** + * Providers whose per-account quota is OBSERVED in-band, never probed. + * + * Deliberately separate from `supportsPerAccountQuota` rather than folded into it. That + * predicate gates explicit upstream readers. Meta publishes no quota endpoint, so it + * remains a cache-only observation even when every probe reader is account-scoped. + */ +export function hasPassiveAccountQuota(provider: string): boolean { + return provider === "meta-muse"; +} + +/** + * Record a quota observed in-band on a streaming turn. + * + * The CALLER captures `writerGeneration` when it resolves the serving credential, not + * this function at write time. A streaming turn is a long await, and a generation + * captured immediately before the write cannot see a config or account change that + * happened EARLIER in the same turn — which is exactly the case the fence exists for. + */ +export function recordPassiveAccountQuota( + provider: string, + accountId: string, + quota: ProviderQuota, + writerGeneration: number, +): void { + if (!hasPassiveAccountQuota(provider) || !accountId) return; + const key = accountCacheKey(provider, accountId); + if (!mayCommitAccountQuotaKey(key, writerGeneration)) return; + // Hydrate BEFORE writing, not only on the read path. `persistAccountQuotaCache` + // serializes the whole in-memory map, so a passive write that lands before anything + // has read the cache would persist this one row and erase every other provider's + // saved row -- and `diskHydrated` would then stop any later reader from recovering + // them. A probe writer cannot hit this because its own read hydrates first; an + // observation arrives unprompted, so it must hydrate itself. + hydrateAccountQuotaCache(); + accountQuotaCache.set(key, { ts: Date.now(), quota }); + // Persisted so a restart keeps the last observation: with no probe to re-establish it, + // a forgotten row stays forgotten until the user happens to run another streaming turn. + persistAccountQuotaCache(); + // sweepExpiredOnWrite is deliberately NOT called. Existing probe writers call it + // because they run on a poll; this runs on the request path, where a state sweep does + // not belong. Passive rows are still reclaimed by generation reconciliation + // (reconcileProviderAccountQuotaRows) and by the disk reader's age bound. +} + +/** + * Cache-only per-account rows for a passive provider. Never probes, never refreshes. + * + * An account with no observation is OMITTED rather than returned with `quota: null` and + * `unavailable`: that pair means "a probe was attempted and failed", and no probe was + * ever attempted here. A user who has not yet run a streaming turn simply has no + * measurement, which is not an error state. + */ +export function readPassiveProviderAccountQuotas(provider: string): ProviderAccountQuota[] { + if (!hasPassiveAccountQuota(provider)) return []; + // Idempotent, and otherwise only reached from probe paths a passive provider never + // enters — without it a restart shows nothing until the next streaming turn, even + // though the row is sitting on disk. + hydrateAccountQuotaCache(); + const set = getAccountSet(provider); + if (!set) return []; + const rows: ProviderAccountQuota[] = []; + for (const account of set.accounts) { + const entry = accountQuotaCache.get(accountCacheKey(provider, account.id)); + if (entry?.quota) rows.push({ accountId: account.id, quota: entry.quota }); + } + return rows; +} + +export function sweepExpiredProviderAccountQuotaRows(now = Date.now()): number { + let removed = 0; + for (const [key, entry] of accountQuotaCache) { + // Anthropic observations extend retention, never the usage probe's eligibility clock. + const retainedAt = key.startsWith("anthropic\u0000") + ? Math.max(entry.ts, entry.quota?.updatedAt ?? 0) + : entry.ts; + if (retainedAt + ACCOUNT_QUOTA_TTL_MS > now) continue; + accountQuotaCache.delete(key); + removed += 1; + } + return removed; +} + +export function reconcileProviderAccountQuotaRows(context: GenerationContext): number { + if (context.generation <= lastReconciledGeneration) return 0; + let removed = 0; + for (const key of accountQuotaCache.keys()) { + if (context.oauthAccountKeys.has(key)) continue; + accountQuotaCache.delete(key); + removed += 1; + } + // Kiro exhaustion rows are keyed identically, so they retire with their quota row; a + // verdict outliving its account would hand the replacement a cooldown it never earned. + removed += reconcileKiroAccountUsageState(context.oauthAccountKeys); + const cachedReports = getProviderQuotaReportCache(); + if (cachedReports) { + const reports = cachedReports.response.reports.filter(report => context.providerNames.has(report.provider)); + removed += cachedReports.response.reports.length - reports.length; + setProviderQuotaReportCache({ ...cachedReports, response: { ...cachedReports.response, reports } }); + replaceCachedProviderQuotas(reports, routingEvidence); + } + liveAccountQuotaKeys = new Set(context.oauthAccountKeys); + liveProviderQuotaKeys = new Set(context.providerNames); + lastReconciledGeneration = context.generation; + return removed; +} + +/** Test-only reset so a direct reconcile call in one file cannot leak across files. */ +export function resetProviderQuotaReconcileStateForTests(): void { + lastReconciledGeneration = 0; + liveAccountQuotaKeys = new Set(); + liveProviderQuotaKeys = new Set(); +} + +/** Drop cached per-account rows (all, or just one provider's). */ +export function clearAccountQuotaCache(provider?: string): void { + explicitAccountEpoch += 1; + if (!provider) { + accountQuotaCache.clear(); + accountQuotaInflight.clear(); + clearKiroAccountUsageState(); + // A cleared cache must not be re-seeded from the file it was just cleared of, and any + // pending write of the old rows is abandoned. + diskHydrated = false; + cancelPendingAccountQuotaPersist(); + return; + } + const prefix = `${provider}\u0000`; + for (const key of [...accountQuotaCache.keys()]) { + if (key.startsWith(prefix)) accountQuotaCache.delete(key); + } + clearKiroAccountUsageState(prefix); + // Drop in-flight probes too so a late resolve cannot repopulate after logout/remove. + for (const key of [...accountQuotaInflight.keys()]) { + if (key.startsWith(prefix)) accountQuotaInflight.delete(key); + } + persistAccountQuotaCache(); +} + +/** + * Resolve a bearer for quota probing without silently adopting a newer global + * Claude CLI credential into a background multiauth slot. + * + * - Fresh stored access → use as-is (no refresh). + * - Active account with expired access → normal refresh path. + * - Background `local-cli` with expired access → fail closed (unavailable): + * `getValidAccessTokenForAccount` can persist a mismatched Claude CLI identity. + * - Background ordinary OAuth (`source !== "local-cli"`) → safe to refresh; + * Anthropic's lock only adopts disk credentials for `local-cli` rows. + */ +export async function getTokenForAccountQuotaProbe(provider: string, accountId: string): Promise { + const stored = getAccountCredential(provider, accountId); + if (!stored) throw new Error("account credential missing"); + if (stored.expires > Date.now() + ACCOUNT_TOKEN_SKEW_MS) return stored.access; + const activeId = getAccountSet(provider)?.activeAccountId; + if (activeId !== accountId && stored.source === "local-cli") { + throw new Error("background local-cli token expired; skip CLI-adopting refresh for quota probe"); + } + return getValidAccessTokenForAccount(provider, accountId); +} + +export function explicitQuotaConfig(provider: string, configured?: OcxProviderConfig): OcxProviderConfig | undefined { + if (configured) return configured; + const entry = getProviderRegistryEntry(provider); + return entry ? { adapter: entry.adapter, baseUrl: entry.baseUrl, authMode: "oauth" } : undefined; +} + +export function explicitQuotaIdentity(provider: string, accountId: string, configured?: OcxProviderConfig): string | undefined { + const credential = getAccountCredential(provider, accountId); + const target = explicitQuotaConfig(provider, configured); + if (!credential || !target) return undefined; + return quotaCredentialIdentity(provider, accountId, credential, target); +} + +export function quotaCredentialIdentity(provider: string, accountId: string, credential: NonNullable>, target: OcxProviderConfig): string { + return createHash("sha256").update(JSON.stringify([ + provider, accountId, credential.access, credential.refresh, credential.expires, + credential.accountId, credential.projectId, credential.source, + target.adapter, target.baseUrl, target.authMode, target.disabled === true, + ])).digest("hex"); +} + +export function explicitQuotaDestination(provider: string, config: OcxProviderConfig): boolean { + if (config.disabled === true || config.authMode !== "oauth") return false; + if (provider === "kimi") return isCanonicalKimiCodeBaseUrl(config.baseUrl); + if (provider === "command-code") return isCanonicalCommandCodeBaseUrl(config.baseUrl); + // These readers use fixed canonical billing origins, never config.baseUrl. + return provider === "xai" || provider === "cursor"; +} diff --git a/src/providers/quota/antigravity.ts b/src/providers/quota/antigravity.ts new file mode 100644 index 0000000000..bafdc1eee0 --- /dev/null +++ b/src/providers/quota/antigravity.ts @@ -0,0 +1,295 @@ +import { antigravityUserAgent } from "../../adapters/client-fingerprint"; +import { DestinationDnsResolutionError } from "../../lib/destination-policy"; +import { PinnedHttpError } from "../../lib/pinned-http"; +import { ProviderOutboundPolicyError, providerOutboundPost, providerRedirectError, type ProviderOutboundDependencies } from "../../lib/provider-outbound"; +import { getValidAccessToken } from "../../oauth"; +import { getAccountCredential, getCredential } from "../../oauth/store"; +import { asRecord, normalizePercent, normalizeResetAt, readQuotaJson, REQUEST_TIMEOUT_MS, toFiniteNumber } from "../quota-wire"; +import { report, type ProviderQuotaReport } from "./report-cache"; +import { quotaCredentialIdentity } from "./account-cache"; +import type { ProviderQuota, ProviderQuotaWindow, QuotaFailureCode } from "../quota-types"; + +export function antigravityQuotaDiagnosticIdentity(accountId: string, credential = getAccountCredential("google-antigravity", accountId)): string | undefined { + return credential ? quotaCredentialIdentity("google-antigravity", accountId, credential, { + adapter: "google", baseUrl: ANTIGRAVITY_ACCOUNT_QUOTA_BASE, authMode: "oauth", + }) : undefined; +} + + +function quotaInfoEntries(modelInfo: Record): Record[] { + const entries: Record[] = []; + const add = (value: unknown, tier?: string) => { + const rec = asRecord(value); + if (!rec) return; + entries.push(tier ? { ...rec, tier } : rec); + }; + const addArray = (value: unknown) => { + if (!Array.isArray(value)) return; + for (const entry of value) add(entry); + }; + + if (Array.isArray(modelInfo.quotaInfo)) addArray(modelInfo.quotaInfo); + else add(modelInfo.quotaInfo); + addArray(modelInfo.quotaInfos); + + const byTier = asRecord(modelInfo.quotaInfoByTier); + if (byTier) { + for (const [tier, value] of Object.entries(byTier)) { + if (Array.isArray(value)) { + for (const entry of value) add(entry, tier); + } else { + add(value, tier); + } + } + } + return entries; +} + +function classifyAntigravityFamily(modelId: string, modelInfo: Record, quotaInfo: Record): "Gem" | "Cla" | null { + const displayName = typeof modelInfo.displayName === "string" ? modelInfo.displayName : ""; + const tier = typeof quotaInfo.tier === "string" ? quotaInfo.tier : ""; + const haystack = `${modelId} ${displayName} ${tier}`.toLowerCase(); + if (haystack.includes("gemini")) return "Gem"; + if (haystack.includes("claude") || haystack.includes("opus") || haystack.includes("sonnet") || haystack.includes("gpt-oss") || haystack.includes("gpt_oss")) return "Cla"; + return null; +} + +function antigravityUsedPercent(quotaInfo: Record): number | undefined { + const target = asRecord(quotaInfo.remaining) ?? quotaInfo; + const remaining = normalizePercent(toFiniteNumber(target.remainingFraction) !== undefined + ? toFiniteNumber(target.remainingFraction)! * 100 + : toFiniteNumber(target.remainingPercentage) !== undefined + ? toFiniteNumber(target.remainingPercentage)! * 100 + : undefined); + if (remaining === undefined) return undefined; + return normalizePercent(100 - remaining); +} + +/** Gem/Cla windows from a `fetchAvailableModels` body; shared by the provider and account probes. */ +function antigravityWindowsFromModels(body: Record | null): ProviderQuotaWindow[] { + const models = asRecord(body?.models); + if (!models) return []; + + const windows = new Map(); + for (const [modelId, rawModelInfo] of Object.entries(models)) { + const modelInfo = asRecord(rawModelInfo); + if (!modelInfo) continue; + for (const quotaInfo of quotaInfoEntries(modelInfo)) { + const label = classifyAntigravityFamily(modelId, modelInfo, quotaInfo); + if (!label || windows.has(label)) continue; + const percent = antigravityUsedPercent(quotaInfo); + if (percent === undefined) continue; + windows.set(label, { + label, + percent, + ...(normalizeResetAt(quotaInfo.resetTime) !== undefined ? { resetAt: normalizeResetAt(quotaInfo.resetTime) } : {}), + }); + } + } + + const customWindows = ["Gem", "Cla"].flatMap(label => { + const window = windows.get(label); + return window ? [window] : []; + }); + return customWindows; +} + +/** + * Parse Google Antigravity quota from `v1internal:retrieveUserQuotaSummary`. + * Groups contain Gemini models and Claude/3P models, each with 5h and weekly limit buckets. + */ +function parseAntigravityQuotaSummary(body: Record | null): ProviderQuota | null { + const groups = Array.isArray(body?.groups) ? (body.groups as unknown[]) : []; + if (groups.length === 0) return null; + + const customWindowsMap = new Map(); + + for (const rawGroup of groups) { + const group = asRecord(rawGroup); + if (!group) continue; + const groupName = `${typeof group.displayName === "string" ? group.displayName : ""} ${typeof group.description === "string" ? group.description : ""}`.toLowerCase(); + const isGemini = groupName.includes("gemini"); + const isClaude = groupName.includes("claude") || groupName.includes("3p") || groupName.includes("gpt"); + + const buckets = Array.isArray(group.buckets) ? (group.buckets as unknown[]) : []; + for (const rawBucket of buckets) { + const bucket = asRecord(rawBucket); + if (!bucket) continue; + const windowStr = `${typeof bucket.window === "string" ? bucket.window : ""} ${typeof bucket.bucketId === "string" ? bucket.bucketId : ""} ${typeof bucket.displayName === "string" ? bucket.displayName : ""}`.toLowerCase(); + const percent = antigravityUsedPercent(bucket); + if (percent === undefined) continue; + const resetAt = normalizeResetAt(bucket.resetTime); + + const isWeekly = windowStr.includes("week"); + const is5h = windowStr.includes("5h") || windowStr.includes("five"); + + if (isGemini) { + const label = is5h ? "Gem" : isWeekly ? "Gem (Weekly)" : ""; + if (label && !customWindowsMap.has(label)) { + customWindowsMap.set(label, { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }); + } + } else if (isClaude) { + const label = is5h ? "Cla" : isWeekly ? "Cla (Weekly)" : ""; + if (label && !customWindowsMap.has(label)) { + customWindowsMap.set(label, { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }); + } + } else { + const baseLabel = typeof group.displayName === "string" ? group.displayName : "Other"; + const label = isWeekly ? `${baseLabel} (Weekly)` : baseLabel; + if (!customWindowsMap.has(label)) { + customWindowsMap.set(label, { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }); + } + } + } + } + + const PREFERRED_ORDER = ["Gem", "Gem (Weekly)", "Cla", "Cla (Weekly)"]; + const customWindows = Array.from(customWindowsMap.values()).sort((a, b) => { + const ia = PREFERRED_ORDER.indexOf(a.label); + const ib = PREFERRED_ORDER.indexOf(b.label); + if (ia !== -1 && ib !== -1) return ia - ib; + if (ia !== -1) return -1; + if (ib !== -1) return 1; + return a.label.localeCompare(b.label); + }); + + if (customWindows.length === 0) { + return null; + } + + return { + customWindows, + updatedAt: Date.now(), + }; +} + +const ANTIGRAVITY_ACCOUNT_QUOTA_BASE = "https://daily-cloudcode-pa.googleapis.com"; +const ANTIGRAVITY_QUOTA_SUMMARY_URL = `${ANTIGRAVITY_ACCOUNT_QUOTA_BASE}/v1internal:retrieveUserQuotaSummary`; +const ANTIGRAVITY_QUOTA_MODELS_URL = `${ANTIGRAVITY_ACCOUNT_QUOTA_BASE}/v1internal:fetchAvailableModels`; + +/** Only these fixed accounting destinations may use transparent Fake-IP DNS. */ +export function isCanonicalAntigravityQuotaUrl(name: string, url: string): boolean { + return name === "google-antigravity" + && (url === ANTIGRAVITY_QUOTA_SUMMARY_URL || url === ANTIGRAVITY_QUOTA_MODELS_URL); +} + +let antigravityOutboundDependencies: ProviderOutboundDependencies = { + isCanonicalUrl: isCanonicalAntigravityQuotaUrl, +}; + +/** Test seam: inject resolver/pinned transport for provider and per-account probes. */ +export function setAntigravityAccountQuotaTransportForTests(dependencies: ProviderOutboundDependencies | null): void { + antigravityOutboundDependencies = { ...dependencies, isCanonicalUrl: isCanonicalAntigravityQuotaUrl }; +} + +/** + * Per-account Antigravity quota (#1082). Always probes Google's own Cloud Code Assist host + * through the pinned provider-outbound transport: a configured `baseUrl` is a routing choice + * for requests, not a second source of Google's accounting for a stored credential, and fixing + * the destination keeps the `provider\0accountId` cache identity exact across config changes. + * A redirect or non-2xx yields null (unavailable), never a partial row. + */ +type AntigravityQuotaProbeResult = + | { kind: "available"; quota: ProviderQuota; source: "google-antigravity:retrieveUserQuotaSummary" | "google-antigravity:fetchAvailableModels" } + | { kind: "unavailable"; failure: QuotaFailureCode; legacy: { kind: "null" } | { kind: "throw"; error: unknown } }; + +function quotaTransportFailure(error: unknown): QuotaFailureCode { + if (error instanceof ProviderOutboundPolicyError) return "destination_blocked"; + if (error instanceof DestinationDnsResolutionError) return "dns_failed"; + if (error instanceof PinnedHttpError) return error.code === "output_byte_limit" ? "response_unusable" : "timeout"; + if (error instanceof DOMException && error.name === "TimeoutError") return "timeout"; + return "transport_error"; +} + +function quotaHttpFailure(status: number): QuotaFailureCode { + if (status >= 300 && status < 400) return "redirect_blocked"; + if (status === 401 || status === 403) return "access_denied"; + if (status === 429) return "rate_limited"; + return "upstream_error"; +} + +function unavailableAntigravityQuota(failure: QuotaFailureCode): AntigravityQuotaProbeResult { + return { kind: "unavailable", failure, legacy: { kind: "null" } }; +} + +/** + * Prefer a summary network-policy diagnosis over a vaguer fallback. A blocked + * destination is an actionable local-network fact, while "upstream_error" tells + * the operator to go look at Google. A successful models probe still clears + * the first failure completely. + */ +function antigravityUnavailableFailure( + summaryFailure: QuotaFailureCode | undefined, + fallbackFailure: QuotaFailureCode, +): QuotaFailureCode { + if ( + (summaryFailure === "destination_blocked" || summaryFailure === "dns_failed") + && fallbackFailure !== "destination_blocked" + && fallbackFailure !== "dns_failed" + ) { + return summaryFailure; + } + return fallbackFailure; +} + +export async function probeAntigravityUsageQuota(accessToken: string, projectId: string): Promise { + const fetchQuota = (url: string) => providerOutboundPost("google-antigravity", { baseUrl: ANTIGRAVITY_ACCOUNT_QUOTA_BASE }, url, { + headers: { + Accept: "application/json", "Content-Type": "application/json", + "User-Agent": antigravityUserAgent(), Authorization: `Bearer ${accessToken}`, + }, + body: JSON.stringify({ project: projectId }), signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }, antigravityOutboundDependencies); + let summaryFailure: QuotaFailureCode | undefined; + try { + const response = await fetchQuota(ANTIGRAVITY_QUOTA_SUMMARY_URL); + if (await providerRedirectError(response, ANTIGRAVITY_QUOTA_SUMMARY_URL)) return unavailableAntigravityQuota("redirect_blocked"); + if (response.status === 401 || response.status === 403) return unavailableAntigravityQuota("access_denied"); + if (response.ok) { + const quota = parseAntigravityQuotaSummary(asRecord(await readQuotaJson(response))); + if (quota) return { kind: "available", quota, source: "google-antigravity:retrieveUserQuotaSummary" }; + } + } catch (error) { + // Existing behavior: summary transport/parse failure may recover through the models probe. + summaryFailure = quotaTransportFailure(error); + } + try { + const response = await fetchQuota(ANTIGRAVITY_QUOTA_MODELS_URL); + if (await providerRedirectError(response, ANTIGRAVITY_QUOTA_MODELS_URL)) { + return unavailableAntigravityQuota(antigravityUnavailableFailure(summaryFailure, "redirect_blocked")); + } + if (!response.ok) { + return unavailableAntigravityQuota(antigravityUnavailableFailure(summaryFailure, quotaHttpFailure(response.status))); + } + const customWindows = antigravityWindowsFromModels(asRecord(await readQuotaJson(response))); + if (!customWindows.length) { + return unavailableAntigravityQuota(antigravityUnavailableFailure(summaryFailure, "response_unusable")); + } + return { kind: "available", quota: { customWindows, updatedAt: Date.now() }, source: "google-antigravity:fetchAvailableModels" }; + } catch (error) { + // The public compatibility wrapper still rejects this exact fallback error; it never enters a DTO. + return { + kind: "unavailable", + failure: antigravityUnavailableFailure(summaryFailure, quotaTransportFailure(error)), + legacy: { kind: "throw", error }, + }; + } +} + +export async function fetchAntigravityUsageQuota(accessToken: string, projectId: string): Promise { + const result = await probeAntigravityUsageQuota(accessToken, projectId); + if (result.kind === "available") return result.quota; + if (result.legacy.kind === "throw") throw result.legacy.error; + return null; +} + +export async function fetchAntigravityQuota(provider: string): Promise { + const credential = getCredential("google-antigravity"); + if (!credential?.projectId) return null; + let accessToken: string; + try { accessToken = await getValidAccessToken("google-antigravity"); } catch { return null; } + const result = await probeAntigravityUsageQuota(accessToken, credential.projectId); + if (result.kind === "available") return report(provider, result.source, result.quota); + if (result.legacy.kind === "throw") throw result.legacy.error; + return null; +} diff --git a/src/providers/quota/report-cache.ts b/src/providers/quota/report-cache.ts new file mode 100644 index 0000000000..44010ebbf7 --- /dev/null +++ b/src/providers/quota/report-cache.ts @@ -0,0 +1,320 @@ +import { createHash } from "node:crypto"; +import { effectiveCodexAuthAccountId, listCodexAuthAccountsSnapshot } from "../../codex/auth-api"; +import { withoutRetiredCodexQuota, type StoredAccountQuota } from "../../codex/quota"; +import { isMainAccountIdentityGenerationLive } from "../../codex/main-account-cache"; +import { codexPlanKey } from "../../codex/plan"; +import { resolveProviderApiKey } from "../key-store"; +import { apiKeyPoolEntryId } from "../api-keys"; +import { getProviderRegistryEntry, providerCodexAccountMode } from "../registry"; +import { isCanonicalOpenAiForwardProvider, OPENAI_CODEX_PROVIDER_ID } from "../openai-tiers"; +import { CODEX_CAPACITY_MAX_QUOTA_AGE_MS, type CodexCapacityAggregation, type CodexCapacityQuota } from "../codex-capacity"; +import { clearCachedProviderQuotas, providerQuotaRoutingBinding, type ProviderQuotaRoutingEvidence } from "../quota-routing-cache"; +import { clearProviderApiKeyQuotaCache } from "../quota-key-accounts"; +import { QUOTA_JSON_READ_FAILURE, readQuotaJson } from "../quota-wire"; +import type { OcxConfig, OcxProviderConfig } from "../../types"; +import type { ProviderQuota, ProviderRoutingQuota } from "../quota-types"; + +/** Keep a failed probe's previous row at most this long before dropping it. */ +export const LAST_GOOD_MAX_AGE_MS = CODEX_CAPACITY_MAX_QUOTA_AGE_MS; +const nativeMainReportGenerations = new WeakMap(); +export const accountReportCurrent = new WeakMap boolean>(); +export const routingEvidence = new WeakMap(); +export let providerQuotaBeforePublishForTests: (() => void | Promise) | null = null; + +/** Test-only seam for identity/config invalidation after probes but before publication. */ +export function setProviderQuotaBeforePublishForTests( + hook: (() => void | Promise) | null, +): void { + providerQuotaBeforePublishForTests = hook; +} +export const TERMINAL_QUOTA_FAILURE = Symbol("terminal-quota-failure"); +/** + * The probe succeeded and the upstream authoritatively reported NO model-quota windows. + * + * Distinct from `null`, which means "this probe told us nothing" and deliberately preserves + * the last-good row for up to 30 minutes. Collapsing the two would let a stale report outlive + * the authoritative answer that replaced it: a GLM plan whose payload carries only MCP + * `TIME_LIMIT` rows has no model windows, and the dashboard and quota-aware routing must stop + * showing the previous token windows rather than keep them for another half hour. + * + * Suppression is shared with `TERMINAL_QUOTA_FAILURE`; only the reason differs. + */ +export const AUTHORITATIVE_EMPTY_QUOTA = Symbol("authoritative-empty-quota"); +export type ProviderQuotaProbeResult = + | ProviderQuotaReport + | null + | typeof TERMINAL_QUOTA_FAILURE + | typeof AUTHORITATIVE_EMPTY_QUOTA; + +export interface ProviderQuotaReport { + provider: string; + label: string; + source: string; + quota: ProviderQuota; + updatedAt: number; + /** Added by the management response projection, never stored on a cached report. */ + routingQuota?: ProviderRoutingQuota; + reverseEngineered?: boolean; + /** + * The row was OBSERVED in-band on a streaming turn rather than probed. + * + * Age means something different for these. A probed provider re-reads on its own TTL, + * so a row older than the last-good bound means the probe is failing and showing it + * would misrepresent a live number. A passive provider publishes no endpoint at all + * (`hasPassiveAccountQuota`), so its last observation is not a stale reading of + * something fresher — it is the only measurement that exists, and dropping it leaves + * the operator with nothing. Consumers that enforce a freshness bound must exempt + * these and state the observation age instead. + */ + observed?: boolean; + aggregation?: CodexCapacityAggregation; +} + +export interface ProviderQuotaResponse { + generatedAt: number; + reports: ProviderQuotaReport[]; +} + +let cache: { key: string; ts: number; response: ProviderQuotaResponse } | null = null; +export const inflight = new Map }>(); +/** Bumped on cache clear and on force-refresh start; stale-epoch probes lose commit authority. */ +export let invalidationEpoch = 0; + +/** Owner-module accessors: cache reassignment stays inside this file. */ +export function getProviderQuotaReportCache(): { key: string; ts: number; response: ProviderQuotaResponse } | null { + return cache; +} + +export function setProviderQuotaReportCache(next: { key: string; ts: number; response: ProviderQuotaResponse } | null): void { + cache = next; +} + +export function bumpProviderQuotaInvalidationEpoch(): void { + invalidationEpoch += 1; +} + +/** Invalidate the report cache (e.g. after switching a provider's active account). */ +export function clearProviderQuotaCache(): void { + cache = null; + clearCachedProviderQuotas(); + clearProviderApiKeyQuotaCache(); + invalidationEpoch += 1; +} + +function cacheKey(config: OcxConfig): string { + const providers = Object.entries(config.providers) + .map(([name, provider]) => { + const resolvedKey = typeof provider.apiKey === "string" + ? resolveProviderApiKey(provider.apiKey)?.trim() + : undefined; + const activeKeyId = resolvedKey ? apiKeyPoolEntryId(resolvedKey) : "none"; + return `${name}:${provider.adapter}:${provider.authMode ?? "key"}:${providerCodexAccountMode(name, provider) ?? "none"}:${provider.disabled === true ? "off" : "on"}:${provider.baseUrl}:${activeKeyId}`; + }) + .sort() + .join("|"); + return `${config.defaultProvider}|${providers}`; +} + +export type CodexAuthAccountsSnapshotPromise = ReturnType; + +export function hasCodexPoolProvider(config: OcxConfig): boolean { + return Object.entries(config.providers).some(([name, provider]) => ( + provider.disabled !== true + && isBuiltInChatGptForwardProvider(name, provider) + && providerCodexAccountMode(name, provider) !== "direct" + )); +} + +function quotaSignatureValue(quota: CodexCapacityQuota | null): unknown { + if (!quota) return null; + return { + fiveHourPercent: quota.fiveHourPercent, + fiveHourResetAt: quota.fiveHourResetAt, + weeklyPercent: quota.weeklyPercent, + weeklyResetAt: quota.weeklyResetAt, + monthlyPercent: quota.monthlyPercent, + monthlyResetAt: quota.monthlyResetAt, + updatedAt: quota.updatedAt, + customWindows: [...(quota.customWindows ?? [])] + .map(window => ({ label: window.label, percent: window.percent, resetAt: window.resetAt })) + .sort((a, b) => a.label.localeCompare(b.label)), + }; +} + +export function providerQuotaFromCodexQuota( + quota: StoredAccountQuota | Omit | null | undefined, +): CodexCapacityQuota | null { + if (!quota) return null; + // Direct snapshots bypass account DTOs; sanitize here as well as at ingestion. + quota = withoutRetiredCodexQuota(quota); + if (!quota) return null; + const projected: CodexCapacityQuota = { + ...(quota.shortPercent !== undefined ? { fiveHourPercent: quota.shortPercent } : {}), + ...(quota.shortResetAt !== undefined ? { fiveHourResetAt: quota.shortResetAt } : {}), + ...(quota.weeklyPercent !== undefined ? { weeklyPercent: quota.weeklyPercent } : {}), + ...(quota.weeklyResetAt !== undefined ? { weeklyResetAt: quota.weeklyResetAt } : {}), + ...(quota.monthlyPercent !== undefined ? { monthlyPercent: quota.monthlyPercent } : {}), + ...(quota.monthlyResetAt !== undefined ? { monthlyResetAt: quota.monthlyResetAt } : {}), + ...(quota.customWindows !== undefined ? { customWindows: quota.customWindows } : {}), + updatedAt: "updatedAt" in quota ? quota.updatedAt : Date.now(), + }; + return hasQuotaRows(projected) ? projected : null; +} + +/** Hash only presentation-relevant state; account ids and email addresses never enter the key. */ +export function cacheKeyWithAggregationState( + config: OcxConfig, + prefetchedSnapshot?: CodexAuthAccountsSnapshotPromise, +): string | Promise { + const base = cacheKey(config); + if (!hasCodexPoolProvider(config)) return base; + return (async () => { + try { + const activeId = effectiveCodexAuthAccountId(config); + const snapshot = await (prefetchedSnapshot ?? listCodexAuthAccountsSnapshot(config, false)); + const rows = snapshot.accounts.map(account => ({ + isMain: account.isMain, + active: account.id === activeId, + plan: codexPlanKey(account.plan) ?? null, + paused: account.paused, + needsReauth: account.needsReauth === true, + quota: quotaSignatureValue(providerQuotaFromCodexQuota(account.quota)), + })); + const canonicalRows = rows.map(row => JSON.stringify(row)).sort(); + const digest = createHash("sha256").update(JSON.stringify(canonicalRows)).digest("hex").slice(0, 24); + return `${base}|codex-pool:${digest}`; + } catch { + return `${base}|codex-pool:unavailable`; + } + })(); +} + +function publicCapacityWindow(window: import("../codex-capacity").CodexCapacityWindowAggregation) { + const { totalWeight: _totalWeight, consumedWeight: _consumedWeight, remainingWeight: _remainingWeight, ...safe } = window; + return safe; +} + +/** Management API metadata intentionally omits configured/weighted unit counts. */ +export function publicCapacityAggregation( + aggregation: CodexCapacityAggregation, + presentation: NonNullable, +): CodexCapacityAggregation { + const safeCurrentAccount = presentation === "coverage-only" && aggregation.currentAccount + ? { ...aggregation.currentAccount, quota: null } + : aggregation.currentAccount; + return { + ...aggregation, + presentation, + ...(safeCurrentAccount ? { currentAccount: safeCurrentAccount } : {}), + ...(aggregation.fiveHour ? { fiveHour: publicCapacityWindow(aggregation.fiveHour) } : {}), + ...(aggregation.weekly ? { weekly: publicCapacityWindow(aggregation.weekly) } : {}), + ...(aggregation.monthly ? { monthly: publicCapacityWindow(aggregation.monthly) } : {}), + ...(aggregation.customWindows ? { + customWindows: aggregation.customWindows.map(window => ({ + label: window.label, + ...publicCapacityWindow(window), + })), + } : {}), + }; +} + +export function hasQuotaRows(quota: ProviderQuota | null | undefined): quota is ProviderQuota { + if (!quota) return false; + return typeof quota.fiveHourPercent === "number" + || typeof quota.weeklyPercent === "number" + || typeof quota.monthlyPercent === "number" + || quota.creditsUsd?.unlimited === true + || typeof quota.creditsUsd?.percent === "number" + || !!quota.customWindows?.some(window => typeof window.percent === "number"); +} + +export function providerLabel(providerId: string): string { + return getProviderRegistryEntry(providerId)?.label ?? providerId; +} + +/** Test-only access to the quota reader's deadline and cancellation contract. */ +export async function readProviderQuotaJsonForTests(response: Response, timeoutMs: number): Promise { + const result = await readQuotaJson(response, timeoutMs); + return result === QUOTA_JSON_READ_FAILURE ? null : result; +} + +export function isBuiltInChatGptForwardProvider(name: string, provider: OcxProviderConfig): boolean { + return name === OPENAI_CODEX_PROVIDER_ID && isCanonicalOpenAiForwardProvider(provider); +} + +export function report( + provider: string, + source: string, + quota: ProviderQuota, + aggregation?: CodexCapacityAggregation, +): ProviderQuotaReport | null { + if (!hasQuotaRows(quota)) return null; + return { + provider, + label: providerLabel(provider), + source, + quota, + updatedAt: quota.updatedAt, + ...(aggregation ? { aggregation } : {}), + }; +} + +/** + * Publish a credential-bound report, and routing evidence only when the producer + * hands over its inference-only projection. + * + * The projection is deliberately not defaulted to the display quota. A producer must + * decide that its rows really do constrain inference on the probed credential; omitting + * the argument leaves the report display-only, so a new producer cannot inherit + * provider-veto authority merely by calling this helper. Ownership alone is not the + * scope decision: providerQuotaRoutingBinding resolving is necessary, never sufficient. + */ +export function keyReport( + provider: string, + source: string, + quota: ProviderQuota, + config: OcxProviderConfig, + probedCredential: string, + inferenceQuota?: ProviderQuota, +): ProviderQuotaReport | null { + const result = report(provider, source, quota); + if (!result || !inferenceQuota) return result; + const binding = providerQuotaRoutingBinding(provider, config, probedCredential); + if (binding) routingEvidence.set(result, { quota: inferenceQuota, binding }); + return result; +} + +export function tagNativeMainReport( + value: ProviderQuotaReport | null, + generation: number, +): ProviderQuotaReport | null { + if (value) nativeMainReportGenerations.set(value, generation); + return value; +} + +/** + * Test-only seam: publish exactly as a credential-bound producer does, and hand back the + * routing evidence the publication actually attached. + * + * Live producers all pass a projection today, so no probe fixture can prove the OTHER half + * of the contract: that omitting it stays display-only. Routing an omitted argument through + * the real helper keeps that provable, and a re-introduced `= quota` default would be + * observed here (a defaulted parameter also fires for an explicitly undefined argument). + */ +export function publishKeyReportForTests( + provider: string, + source: string, + quota: ProviderQuota, + config: OcxProviderConfig, + probedCredential: string, + inferenceQuota?: ProviderQuota, +): { report: ProviderQuotaReport | null; routing: ProviderQuotaRoutingEvidence | undefined } { + const result = keyReport(provider, source, quota, config, probedCredential, inferenceQuota); + return { report: result, routing: result ? routingEvidence.get(result) : undefined }; +} + +export function isProviderQuotaReportCurrent(value: ProviderQuotaReport): boolean { + const generation = nativeMainReportGenerations.get(value); + return (generation === undefined || isMainAccountIdentityGenerationLive(generation)) + && (accountReportCurrent.get(value)?.() ?? true); +} diff --git a/src/providers/quota/vendor-probes-key.ts b/src/providers/quota/vendor-probes-key.ts new file mode 100644 index 0000000000..17594f548e --- /dev/null +++ b/src/providers/quota/vendor-probes-key.ts @@ -0,0 +1,1243 @@ +import { resolveProviderApiKey } from "../key-store"; +import { getProviderRegistryEntry, registryEntryForProviderDestination } from "../registry"; +import { isCanonicalOllamaCloudUrl } from "../../adapters/ollama-native-url"; +import { QUOTA_JSON_READ_FAILURE, asRecord, normalizePercent, normalizeResetAt, readQuotaJson, REQUEST_TIMEOUT_MS, toFiniteNumber } from "../quota-wire"; +import { + AUTHORITATIVE_EMPTY_QUOTA, + hasQuotaRows, + keyReport, + report, + TERMINAL_QUOTA_FAILURE, + type ProviderQuotaProbeResult, + type ProviderQuotaReport, +} from "./report-cache"; +import { getTokenForAccountQuotaProbe } from "./account-cache"; +import type { AccountQuotaMode, ProviderQuota, ProviderQuotaCreditsUsd } from "../quota-types"; +import type { OcxProviderConfig } from "../../types"; + +const KIMI_CODE_BASE_URL = "https://api.kimi.com/coding/v1"; +const KIMI_CODE_USAGE_URL = `${KIMI_CODE_BASE_URL}/usages`; +const COMMAND_CODE_BASE_URL = "https://api.commandcode.ai"; +const COMMAND_CODE_WHOAMI_URL = `${COMMAND_CODE_BASE_URL}/alpha/whoami`; +const COMMAND_CODE_CREDITS_URL = `${COMMAND_CODE_BASE_URL}/alpha/billing/credits`; +const COMMAND_CODE_SUBSCRIPTIONS_URL = `${COMMAND_CODE_BASE_URL}/alpha/billing/subscriptions`; +const COMMAND_CODE_USAGE_URL = `${COMMAND_CODE_BASE_URL}/alpha/usage/summary`; +const A6API_BASE_URL = "https://api.a6api.com"; +const OPENCODE_GO_BASE_URL = "https://opencode.ai/zen/go/v1"; +const OPENCODE_GO_USAGE_URL = `${OPENCODE_GO_BASE_URL}/usage`; +const OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"; +const DEEPSEEK_BASE_URL = "https://api.deepseek.com"; +const CLINE_BASE_URL = "https://api.cline.bot"; +const OLLAMA_CLOUD_BASE_URL = "https://ollama.com"; +const OLLAMA_CLOUD_USAGE_URL = `${OLLAMA_CLOUD_BASE_URL}/api/usage`; +const ZAI_BASE_URL = "https://api.z.ai"; +const ZAI_CN_BASE_URL = "https://open.bigmodel.cn"; +const MINIMAX_REMAINS_URL = "https://www.minimax.io/v1/token_plan/remains"; +const MOONSHOT_BASE_URL = "https://api.moonshot.ai/v1"; +const VENICE_BASE_URL = "https://api.venice.ai/api/v1"; +const SYNTHETIC_BASE_URL = "https://api.synthetic.new/v2"; +const DEEPINFRA_BASE_URL = "https://api.deepinfra.com"; +const NEURALWATT_BASE_URL = "https://api.neuralwatt.com/v1"; + + +function isCanonicalA6apiBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === A6API_BASE_URL || normalized === `${A6API_BASE_URL}/v1`; +} + +function isCanonicalOpenCodeGoBaseUrl(baseUrl: string): boolean { + return normalizedBaseUrl(baseUrl) === OPENCODE_GO_BASE_URL; +} + +function isCanonicalOpenRouterBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === OPENROUTER_BASE_URL; +} + +function isCanonicalDeepSeekBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === DEEPSEEK_BASE_URL || normalized === `${DEEPSEEK_BASE_URL}/v1`; +} + +function isCanonicalClineBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === CLINE_BASE_URL || normalized === `${CLINE_BASE_URL}/api/v1`; +} + +function isCanonicalOllamaCloudBaseUrl(baseUrl?: string): boolean { + if (!baseUrl) return false; + try { + return isCanonicalOllamaCloudUrl(baseUrl); + } catch { + return false; + } +} + +function zaiQuotaMonitorHost(baseUrl: string): string | null { + // Admission and destination selection must share one mapping: admitting a new + // international wire must never fall through to the CN host/authentication scheme. + switch (normalizedBaseUrl(baseUrl)) { + case ZAI_BASE_URL: + case `${ZAI_BASE_URL}/api/coding/paas/v4`: + case `${ZAI_BASE_URL}/api/anthropic`: + case `${ZAI_BASE_URL}/api/v1`: + return ZAI_BASE_URL; + case ZAI_CN_BASE_URL: + case `${ZAI_CN_BASE_URL}/api/coding/paas/v4`: + case `${ZAI_CN_BASE_URL}/api/v1`: + return ZAI_CN_BASE_URL; + default: + return null; + } +} + +function isCanonicalZaiBaseUrl(baseUrl: string): boolean { + return zaiQuotaMonitorHost(baseUrl) !== null; +} + +function isCanonicalMinimaxBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === "https://api.minimax.io/v1" || normalized === "https://api.minimaxi.com/v1"; +} + +function isCanonicalMoonshotBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === MOONSHOT_BASE_URL || normalized === "https://api.moonshot.cn/v1"; +} + +function isCanonicalVeniceBaseUrl(baseUrl: string): boolean { + return normalizedBaseUrl(baseUrl) === VENICE_BASE_URL; +} + +function isCanonicalSyntheticBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === SYNTHETIC_BASE_URL || normalized === "https://api.synthetic.new/openai/v1"; +} + +function isCanonicalDeepInfraBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + return normalized === DEEPINFRA_BASE_URL || normalized === `${DEEPINFRA_BASE_URL}/v1/openai`; +} + +function isCanonicalNeuralwattBaseUrl(baseUrl: string): boolean { + return normalizedBaseUrl(baseUrl) === NEURALWATT_BASE_URL; +} + +function a6apiPayload(value: unknown): Record | null { + const body = asRecord(value); + return asRecord(body?.data) ?? body; +} + +function firstFinite(record: Record | null, names: string[]): number | undefined { + if (!record) return undefined; + for (const name of names) { + const value = toFiniteNumber(record[name]); + if (value !== undefined) return value; + } + return undefined; +} + +async function fetchA6apiQuota(provider: string, config: OcxProviderConfig): Promise { + // Never send a configured API key to a lookalike host or through a redirect. + if (!isCanonicalA6apiBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const headers = { Accept: "application/json", Authorization: `Bearer ${apiKey}` } as const; + const [subscriptionResponse, tokenResponse] = await Promise.all([ + fetch(`${A6API_BASE_URL}/dashboard/billing/subscription`, { + headers, redirect: "error", signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }), + fetch(`${A6API_BASE_URL}/api/usage/token/`, { + headers, redirect: "error", signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }), + ]); + if (!subscriptionResponse.ok || !tokenResponse.ok) { + const statuses = [subscriptionResponse.status, tokenResponse.status]; + // 408/429 are transient (timeout/throttle), not invalid-account signals: keep the + // last-good row like 5xx/network failures. 401/403 (bad key) and 404 (contract change) + // stay terminal. + return statuses.some(status => status >= 400 && status < 500 && status !== 429 && status !== 408) + ? TERMINAL_QUOTA_FAILURE + : null; + } + const [subscriptionBody, tokenBody] = await Promise.all([ + readQuotaJson(subscriptionResponse), + readQuotaJson(tokenResponse), + ]); + if (subscriptionBody === QUOTA_JSON_READ_FAILURE || tokenBody === QUOTA_JSON_READ_FAILURE) return null; + const subscription = a6apiPayload(subscriptionBody); + const token = a6apiPayload(tokenBody); + const unlimited = token?.unlimited_quota === true + || token?.unlimited_quota === 1 + || token?.unlimited_quota === "true"; + const normalizedExpiry = normalizeResetAt(token?.expires_at); + const expiry = normalizedExpiry && normalizedExpiry > 0 + ? { expiresAt: normalizedExpiry } + : {}; + if (unlimited) { + // Every row is an API-credit constraint on inference, so the display quota is also + // the routing projection. Passing it explicitly is the opt-in. + const quota: ProviderQuota = { + creditsUsd: { + used: 0, + limit: 0, + remaining: 0, + percent: 0, + unlimited: true, + ...expiry, + }, + customWindows: [{ label: "Unlimited API credits", percent: 0 }], + updatedAt: Date.now(), + }; + return keyReport(provider, "a6api:billing", quota, config, apiKey, quota); + } + const limitUsd = firstFinite(subscription, ["hard_limit_usd"]); + const grantedUnits = firstFinite(token, ["total_granted"]); + const usedUnits = firstFinite(token, ["total_used"]); + const availableUnits = firstFinite(token, ["total_available"]); + const reconciledUnits = usedUnits !== undefined && availableUnits !== undefined + ? usedUnits + availableUnits + : undefined; + const reconciliationTolerance = grantedUnits !== undefined + ? Math.abs(grantedUnits) * 1e-9 + : 0; + if (limitUsd === undefined || grantedUnits === undefined || usedUnits === undefined + || availableUnits === undefined || limitUsd <= 0 || grantedUnits <= 0 + || usedUnits < 0 || availableUnits < 0 + || reconciledUnits === undefined + || Math.abs(reconciledUnits - grantedUnits) > reconciliationTolerance) return TERMINAL_QUOTA_FAILURE; + const usdPerUnit = limitUsd / grantedUnits; + const usedUsd = usedUnits * usdPerUnit; + const remainingUsd = Math.max(0, availableUnits * usdPerUnit); + const percent = normalizePercent((usedUsd / limitUsd) * 100); + if (percent === undefined) return TERMINAL_QUOTA_FAILURE; + const label = `API credits ($${remainingUsd.toFixed(2)} of $${limitUsd.toFixed(2)} remaining)`; + const quota: ProviderQuota = { + creditsUsd: { + used: usedUsd, + limit: limitUsd, + remaining: remainingUsd, + percent, + ...expiry, + }, + customWindows: [{ label, percent }], + updatedAt: Date.now(), + }; + // The credit balance funds inference itself, so display and routing scope agree. + return keyReport(provider, "a6api:billing", quota, config, apiKey, quota); +} + +function parseOpenCodeGoUsageWindow(value: unknown): { percent: number; resetAt?: number } | null { + const row = asRecord(value); + if (!row) return null; + const percent = normalizePercent(row.percent); + if (percent === undefined) return null; + const resetAt = normalizeResetAt(row.resetsAt); + return { percent, ...(resetAt !== undefined ? { resetAt } : {}) }; +} + +async function fetchOpenCodeGoQuota(provider: string, config: OcxProviderConfig): Promise { + // Never send a configured API key when the provider destination is not the built-in Go endpoint. + if (!isCanonicalOpenCodeGoBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(OPENCODE_GO_USAGE_URL, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const usage = asRecord(body?.usage); + if (!usage) return null; + const rolling = parseOpenCodeGoUsageWindow(usage.rolling); + const weekly = parseOpenCodeGoUsageWindow(usage.weekly); + const monthly = parseOpenCodeGoUsageWindow(usage.monthly); + const quota: ProviderQuota = { + ...(rolling ? { + fiveHourPercent: rolling.percent, + ...(rolling.resetAt !== undefined ? { fiveHourResetAt: rolling.resetAt } : {}), + } : {}), + ...(weekly ? { + weeklyPercent: weekly.percent, + ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), + } : {}), + ...(monthly ? { + monthlyPercent: monthly.percent, + ...(monthly.resetAt !== undefined ? { monthlyResetAt: monthly.resetAt } : {}), + } : {}), + updatedAt: Date.now(), + }; + return keyReport(provider, "opencode-go:usage", quota, config, apiKey, quota); +} + +/** + * OpenRouter `GET /api/v1/key` — the key's own credit balance and optional + * per-key spending cap. `limit` is the configured cap (absent = uncapped); + * `usage` is lifetime spend; `limit_remaining` is what is left of the cap. + * When no cap is set there is no hard limit to meter against, so no bar is + * produced — the provider falls back to its documented reference. + */ +async function fetchOpenRouterQuota(provider: string, config: OcxProviderConfig): Promise { + // Never send a configured API key to a lookalike host or through a redirect. + if (!isCanonicalOpenRouterBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(`${OPENROUTER_BASE_URL}/key`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const data = asRecord(body?.data) ?? body; + if (!data) return null; + const limit = toFiniteNumber(data.limit); + const limitRemaining = toFiniteNumber(data.limit_remaining); + const usage = toFiniteNumber(data.usage); + // A successful no-cap response is a DELIBERATE change, not a transient + // failure: the old capped row must be dropped, not preserved as last-good. + if (limit === undefined || limit <= 0) return TERMINAL_QUOTA_FAILURE; + // Prefer the authoritative remaining-cap value when present: `usage` is + // lifetime accumulated spend and overstates a reset or re-capped key. + const used = limitRemaining !== undefined + ? Math.max(0, limit - limitRemaining) + : usage !== undefined && usage >= 0 ? usage : undefined; + if (used === undefined) return null; + const percent = normalizePercent((used / limit) * 100); + if (percent === undefined) return null; + const remaining = Math.max(0, limit - used); + const label = `API credits ($${remaining.toFixed(2)} of $${limit.toFixed(2)} remaining)`; + // The per-key spending cap stops every request this credential can make, so the + // whole report is inference-wide routing evidence. + const quota: ProviderQuota = { + customWindows: [{ label, percent }], + updatedAt: Date.now(), + }; + return keyReport(provider, "openrouter:key-info", quota, config, apiKey, quota); +} + +/** + * DeepSeek `GET /user/balance` — the account's granted + topped-up credit + * balance. The payload places `total_balance` / `granted_balance` inside + * entries of `balance_infos` (one row per currency); the row for the account's + * currency is selected by preference. `granted_balance` is a CURRENT balance + * component, not the original grant ceiling, so no consumed percentage is + * fabricated — the balance is reported as a balance-only window. + */ +async function fetchDeepSeekQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalDeepSeekBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(`${DEEPSEEK_BASE_URL}/user/balance`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + // The payload nests balances under `balance_infos` rows keyed by currency; + // prefer a USD row, then CNY, then the first row that parses. + const infos = Array.isArray(body?.balance_infos) ? body.balance_infos as unknown[] : null; + const rows = infos + ? infos.map((raw): Record | null => asRecord(raw)).filter((r): r is Record => r !== null) + : []; + const pick = (currency: string): Record | null => + rows.find(row => String(row.currency ?? "").toUpperCase() === currency) ?? null; + const preferred = pick("USD") ?? pick("CNY") ?? rows[0] ?? null; + if (!preferred) return null; + const totalBalance = toFiniteNumber(preferred.total_balance); + const grantedBalance = toFiniteNumber(preferred.granted_balance); + const toppedUp = toFiniteNumber(preferred.topped_up_balance); + const balance = totalBalance ?? grantedBalance ?? toppedUp; + if (balance === undefined || balance < 0) return null; + const label = grantedBalance !== undefined && grantedBalance > 0 + ? `API balance ($${balance.toFixed(2)} total, $${grantedBalance.toFixed(2)} granted)` + : `API balance ($${balance.toFixed(2)})`; + return report(provider, "deepseek:balance", { + customWindows: [{ label, percent: 0 }], + updatedAt: Date.now(), + }); +} + +/** + * ClinePass `GET /api/v1/users/me/plan/usage-limits` — the subscription's + * rolling five-hour, weekly, and monthly utilization, matching the existing + * ProviderQuota windows directly. The endpoint 404s (or returns a null plan) + * for accounts without an active ClinePass, which is a no-report, not an error. + */ +async function fetchClineQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalClineBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(`${CLINE_BASE_URL}/api/v1/users/me/plan/usage-limits`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + // 404 = no active plan; a plain "no plan" is a no-report, everything else + // 4xx (except 408/429) is a credential/contract problem. + if (response.status === 404) return null; + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const data = asRecord(body?.data) ?? body; + const limits = Array.isArray(data?.limits) ? data.limits : null; + if (!limits) return null; + const quota: ProviderQuota = { updatedAt: Date.now() }; + let windows = 0; + for (const raw of limits) { + const row = asRecord(raw); + if (!row) continue; + const percent = normalizePercent(row.percentUsed); + if (percent === undefined) continue; + const resetAt = normalizeResetAt(row.resetsAt); + if (row.type === "five_hour") { + quota.fiveHourPercent = percent; + if (resetAt !== undefined) quota.fiveHourResetAt = resetAt; + windows += 1; + } else if (row.type === "weekly") { + quota.weeklyPercent = percent; + if (resetAt !== undefined) quota.weeklyResetAt = resetAt; + windows += 1; + } else if (row.type === "monthly") { + quota.monthlyPercent = percent; + if (resetAt !== undefined) quota.monthlyResetAt = resetAt; + windows += 1; + } + } + return windows > 0 ? keyReport(provider, "cline:plan-usage-limits", quota, config, apiKey, quota) : null; +} + +/** + * Ollama Cloud `GET https://ollama.com/api/usage` — returns account usage. + * Legacy plans report rolling 5-hour `limits.session.usage` and 7-day + * `limits.weekly.usage`. Migrated monthly-credit plans report + * `limits.monthly.usage`. `usage` values are normalized fractions (0..1). + */ +function parseOllamaPercent(usageValue: unknown): number | undefined { + const usage = toFiniteNumber(usageValue); + if (usage === undefined || usage < 0) return undefined; + const percent = Math.round(usage * 10000) / 100; + return normalizePercent(percent); +} + +export function parseOllamaCloudQuota(body: Record | null): ProviderQuota | null { + if (!body) return null; + const limits = asRecord(body.limits); + if (!limits) return null; + + const quota: ProviderQuota = { updatedAt: Date.now() }; + let windows = 0; + + const session = asRecord(limits.session); + if (session) { + const percent = parseOllamaPercent(session.usage); + if (percent !== undefined) { + quota.fiveHourPercent = percent; + windows += 1; + } + } + + const weekly = asRecord(limits.weekly); + if (weekly) { + const percent = parseOllamaPercent(weekly.usage); + if (percent !== undefined) { + quota.weeklyPercent = percent; + windows += 1; + } + } + + const monthly = asRecord(limits.monthly); + if (monthly) { + const percent = parseOllamaPercent(monthly.usage); + if (percent !== undefined) { + quota.monthlyPercent = percent; + windows += 1; + } + } + + return windows > 0 ? quota : null; +} + +async function fetchOllamaCloudQuota(provider: string, config: OcxProviderConfig): Promise { + const effectiveBaseUrl = config.baseUrl ?? getProviderRegistryEntry(provider)?.baseUrl ?? ""; + if (!isCanonicalOllamaCloudBaseUrl(effectiveBaseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(OLLAMA_CLOUD_USAGE_URL, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + if (response.status === 404) return null; + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const quota = parseOllamaCloudQuota(body); + return quota ? keyReport(provider, "ollama-cloud:usage", quota, config, apiKey, quota) : null; +} + +/** + * Z.AI GLM Coding Plan `GET /api/monitor/usage/quota/limit` — the coding-plan + * limits arrive as a `limits` array of `TOKENS_LIMIT` (newer plans call the + * same rows `CREDIT_LIMIT`) and `TIME_LIMIT` rows. `TOKENS_LIMIT`/`CREDIT_LIMIT` + * rows carry the window length as `unit`/`number`: unit 3 is hours (number 5 → + * the rolling five-hour window), unit 6 is weeks (number 1 → the weekly + * window). Every row's `percentage` is the consumed share (falling + * back to `currentValue`/`usage` when absent) and `nextResetTime` (unix ms) + * the window reset. + * + * `TIME_LIMIT` rows are deliberately ignored (issue #1168). They are the shared + * monthly MCP *call* allowance for Web Search / Web Reader / Zread — not a + * model-token budget — and `ProviderQuota.monthlyPercent` is consumed as a + * model-capacity signal: `headroomOf()` in `src/oauth/account-quota-rank.ts` + * takes the MAX across every window, so a user who spent their MCP search + * allowance would be ranked as having no model capacity left, and the dashboard + * would draw a full monthly bar for a plan whose model tokens are untouched. + * A payload carrying only `TIME_LIMIT` rows therefore reports no quota at all, + * which is the honest answer rather than a fabricated one. + */ +export function parseZaiQuotaLimits(data: Record | null): ProviderQuota | null { + const limits = Array.isArray(data?.limits) ? data.limits as unknown[] : null; + if (!limits) return null; + const quota: ProviderQuota = { updatedAt: Date.now() }; + let windows = 0; + for (const raw of limits) { + const row = asRecord(raw); + if (!row) continue; + // Gate on row type before deriving a percentage: an MCP row must not even + // contribute a parsed value to a model-quota report. + if (row.type !== "TOKENS_LIMIT" && row.type !== "CREDIT_LIMIT") continue; + const resetAt = normalizeResetAt(row.nextResetTime); + let percent = normalizePercent(row.percentage); + if (percent === undefined) { + const used = toFiniteNumber(row.currentValue); + const total = toFiniteNumber(row.usage); + if (used !== undefined && total !== undefined && total > 0) { + percent = normalizePercent((used / total) * 100); + } + } + if (percent === undefined) continue; + const unit = toFiniteNumber(row.unit); + const number = toFiniteNumber(row.number); + if (unit === 3 && number === 5) { + quota.fiveHourPercent = percent; + if (resetAt !== undefined) quota.fiveHourResetAt = resetAt; + windows += 1; + } else if (unit === 6 && number === 1) { + quota.weeklyPercent = percent; + if (resetAt !== undefined) quota.weeklyResetAt = resetAt; + windows += 1; + } + } + return windows > 0 ? quota : null; +} + +/** + * Legacy Z.AI payload shape: percent fields with window identifiers directly on + * the data object (optionally nested under `quota`). Kept as a fallback so + * older responses keep rendering when the `limits` array is absent. + */ +function parseZaiQuotaLegacyFields(data: Record | null): ProviderQuota | null { + if (!data) return null; + const quota: ProviderQuota = { updatedAt: Date.now() }; + let windows = 0; + const percentAt = (key: string): number | undefined => { + const value = normalizePercent(data[key]); + if (value !== undefined) return value; + const nested = asRecord(data.quota); + return nested ? normalizePercent(nested[key]) : undefined; + }; + const fiveHour = percentAt("fiveHourPercent") ?? percentAt("fiveHourUsage") ?? percentAt("fiveHourUsed"); + const weekly = percentAt("weeklyPercent") ?? percentAt("weeklyUsage") ?? percentAt("weeklyUsed"); + const monthly = percentAt("monthlyPercent") ?? percentAt("mcpPercent") ?? percentAt("monthlyMCPUsage"); + if (fiveHour !== undefined) { + quota.fiveHourPercent = fiveHour; + windows += 1; + } + if (weekly !== undefined) { + quota.weeklyPercent = weekly; + windows += 1; + } + if (monthly !== undefined) { + quota.monthlyPercent = monthly; + windows += 1; + } + return windows > 0 ? quota : null; +} + +/** + * Fetches the Z.AI GLM Coding Plan quota — on whichever region the provider + * points at (api.z.ai or open.bigmodel.cn). The `limits` array shape is + * preferred; older field-name payloads fall back to the legacy parser. + * + * Authentication differs by host (issue #1168). `api.z.ai` takes the API key as + * a Bearer token per Z.AI's API reference; `open.bigmodel.cn` expects the key + * directly in `Authorization` with no scheme prefix and answers a Bearer header + * with an auth error, which is why BigModel Coding Plan quota never rendered. + * The host is already canonicalized by `isCanonicalZaiBaseUrl` above and + * `redirect: "error"` stays set, so the bare key cannot travel to a lookalike + * host or follow a redirect off-origin. + */ +async function fetchZaiQuota(provider: string, config: OcxProviderConfig): Promise { + const monitorHost = zaiQuotaMonitorHost(config.baseUrl); + if (!monitorHost) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const authorization = monitorHost === ZAI_CN_BASE_URL ? apiKey : `Bearer ${apiKey}`; + const response = await fetch(`${monitorHost}/api/monitor/usage/quota/limit`, { + headers: { Accept: "application/json", Authorization: authorization }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + if (!body || body.success === false) return null; + const data = asRecord(body.data) ?? body; + if (Array.isArray(data?.limits)) { + const quota = parseZaiQuotaLimits(data); + // A well-formed `limits[]` we fully understood is authoritative even when it yields no + // model window — for example a plan reporting only the monthly MCP `TIME_LIMIT` row. + // Returning `null` here would preserve the previous token windows for up to 30 minutes + // and keep quota-aware routing acting on a report the provider has already superseded. + return quota + ? keyReport(provider, "zai:quota-limit", quota, config, apiKey, quota) + : AUTHORITATIVE_EMPTY_QUOTA; + } + const legacy = parseZaiQuotaLegacyFields(data); + if (!legacy) return null; + // The legacy monthly figure also carries MCP usage; it is display evidence, not + // proof that model inference is unavailable. Modern TOKEN_LIMIT rows above are scoped. + const inferenceQuota = { ...legacy }; + delete inferenceQuota.monthlyPercent; + delete inferenceQuota.monthlyResetAt; + return keyReport(provider, "zai:quota-limit", legacy, config, apiKey, inferenceQuota); +} + +/** + * MiniMax Token Plan `GET /v1/token_plan/remains` — the subscription's + * remaining quota as a countdown-time value (ms). The endpoint does not expose + * the plan's total duration, so no percentage is fabricated from a presumed + * window: the remaining time is reported as a duration-only window. When the + * API supplies a total (`total_time` / `plan_duration_ms`), a consumed share + * is derived from it. Region selects the host: `minimax` → www.minimax.io, + * `minimax-cn` → api.minimaxi.com. + */ +async function fetchMinimaxQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalMinimaxBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const cnHost = normalizedBaseUrl(config.baseUrl)?.startsWith("https://api.minimaxi.com"); + const remainsUrl = cnHost ? "https://api.minimaxi.com/v1/token_plan/remains" : MINIMAX_REMAINS_URL; + const response = await fetch(remainsUrl, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + if (!body || body.success === false) return null; + const data = asRecord(body.data) ?? body; + const remainsMs = toFiniteNumber(data.remains_time ?? data.remainsTime); + if (remainsMs === undefined || remainsMs < 0) return null; + const hours = Math.floor(remainsMs / 3_600_000); + const label = `Token Plan remaining (${hours}h)`; + // Only derive a consumed share when the API actually reports the plan total; + // a presumed window (e.g. 30 days) would fabricate utilization. A valid + // response that omits the total after a prior refresh had it is a DELIBERATE + // contract change — the old row must be dropped (terminal), not preserved as + // a transient last-good. + const totalMs = toFiniteNumber(data.total_time ?? data.plan_duration_ms ?? data.total_duration_ms); + if (totalMs === undefined || totalMs <= 0) return TERMINAL_QUOTA_FAILURE; + const consumed = Math.max(0, totalMs - remainsMs); + const percent = normalizePercent((consumed / totalMs) * 100); + if (percent === undefined) return null; + return report(provider, "minimax:token-plan-remains", { + customWindows: [{ label, percent }], + updatedAt: Date.now(), + }); +} + +/** + * Moonshot/Kimi `GET /v1/users/me/balance` — the account's available balance + * (voucher + cash). Renders a single balance window against the sum of + * voucher + cash when positive (there is no per-window rate limit to meter). + */ +async function fetchMoonshotQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalMoonshotBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const host = normalizedBaseUrl(config.baseUrl)?.startsWith("https://api.moonshot.cn") ? "https://api.moonshot.cn/v1" : MOONSHOT_BASE_URL; + const response = await fetch(`${host}/users/me/balance`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const data = asRecord(body?.data) ?? body; + if (!data) return null; + const available = toFiniteNumber(data.available_balance); + const voucher = toFiniteNumber(data.voucher_balance); + const cash = toFiniteNumber(data.cash_balance); + if (available === undefined || available < 0) return null; + // Moonshot exposes no per-window quota ceiling, only a balance — report it + // as a balance-only window (percent 0) rather than a fabricated utilization. + // Currency is host-scoped: China platform (api.moonshot.cn) bills in CNY; + // the international platform (api.moonshot.ai) bills in USD. Do not force + // either side into the other unit — the number is correct, only the unit + // must match the host. + const isChinaHost = host.startsWith("https://api.moonshot.cn"); + const money = (n: number) => isChinaHost ? `¥${n.toFixed(2)}` : `$${n.toFixed(2)}`; + const unit = isChinaHost ? "CNY" : "USD"; + const label = voucher !== undefined && cash !== undefined + ? `Balance (${money(available)} ${unit} available, ${money(voucher)} voucher)` + : `Balance (${money(available)} ${unit} available)`; + return report(provider, "moonshot:balance", { + customWindows: [{ label, percent: 0 }], + updatedAt: Date.now(), + }); +} + +/** + * Venice `GET /api/v1/billing/balance` — DIEM (native credits) or USD balance. + * Shows the remaining balance; epoch allocation progress when present. + */ +async function fetchVeniceQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalVeniceBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(`${VENICE_BASE_URL}/billing/balance`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const data = asRecord(body?.data) ?? body; + if (!data) return null; + const diemBalance = toFiniteNumber(data.balance); + const usdBalance = toFiniteNumber(data.balance_usd); + const epochUsed = toFiniteNumber(data.diem_epoch_used); + const epochAllocated = toFiniteNumber(data.diem_epoch_allocated); + if (diemBalance === undefined && usdBalance === undefined) return null; + const label = diemBalance !== undefined + ? `DIEM balance (${Math.round(diemBalance)})` + : `USD balance ($${usdBalance?.toFixed(2) ?? "?"})`; + if (epochAllocated !== undefined && epochAllocated > 0 && epochUsed !== undefined) { + const percent = normalizePercent((epochUsed / epochAllocated) * 100); + if (percent === undefined) return null; + return report(provider, "venice:billing-balance", { + customWindows: [{ label, percent }], + updatedAt: Date.now(), + }); + } + return report(provider, "venice:billing-balance", { + customWindows: [{ label, percent: 0 }], + updatedAt: Date.now(), + }); +} + +/** + * Synthetic `GET /v2/quotas` — the known quota lanes (rolling 5-hour, + * weekly token, search-hourly) mapped onto the quota windows. + */ +async function fetchSyntheticQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalSyntheticBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(`${SYNTHETIC_BASE_URL}/quotas`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const data = asRecord(body?.data) ?? body; + const quota: ProviderQuota = { updatedAt: Date.now() }; + let windows = 0; + const percentAt = (key: string): number | undefined => { + const value = normalizePercent(data?.[key]); + if (value !== undefined) return value; + const nested = asRecord(data?.quota) ?? asRecord(data?.quotas); + return nested ? normalizePercent(nested[key]) : undefined; + }; + const fiveHour = percentAt("rollingFiveHourLimit"); + const weekly = percentAt("weeklyTokenLimit"); + if (fiveHour !== undefined) { + quota.fiveHourPercent = fiveHour; + windows += 1; + } + if (weekly !== undefined) { + quota.weeklyPercent = weekly; + windows += 1; + } + const search = asRecord(data?.search); + const searchHourly = search ? normalizePercent(search.hourly) : undefined; + if (searchHourly !== undefined) { + quota.customWindows = [...(quota.customWindows ?? []), { label: "Search hourly", percent: searchHourly }]; + windows += 1; + } + const inferenceQuota = { ...quota }; + delete inferenceQuota.customWindows; // search.hourly does not constrain model inference. + return windows > 0 ? keyReport(provider, "synthetic:quotas", quota, config, apiKey, inferenceQuota) : null; +} + +/** + * DeepInfra `GET /payment/checklist?compute_owed=true` — prepaid balance, + * recent spend, spending limit, and suspension state. Renders a balance + * window (prepaid funds are a negative `stripe_balance` → positive available). + */ +async function fetchDeepInfraQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalDeepInfraBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(`${DEEPINFRA_BASE_URL}/payment/checklist?compute_owed=true`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const data = asRecord(body?.data) ?? body; + if (!data) return null; + const stripeBalance = toFiniteNumber(data.stripe_balance); + const spendLimit = toFiniteNumber(data.spending_limit); + const total = toFiniteNumber(data.total_amount_due); + if (stripeBalance === undefined) return null; + // Prepaid funds are negative; a positive value is money owed. + const available = stripeBalance < 0 ? -stripeBalance : 0; + if (spendLimit !== undefined && spendLimit > 0) { + const spent = total !== undefined && total > 0 ? total : Math.max(0, spendLimit - available); + const percent = normalizePercent((spent / spendLimit) * 100); + if (percent === undefined) return null; + return report(provider, "deepinfra:billing-checklist", { + customWindows: [{ label: `Billing cycle spend ($${spent.toFixed(2)} of $${spendLimit.toFixed(2)})`, percent }], + updatedAt: Date.now(), + }); + } + return report(provider, "deepinfra:billing-checklist", { + customWindows: [{ label: `Prepaid balance ($${available.toFixed(2)})`, percent: 0 }], + updatedAt: Date.now(), + }); +} + +/** + * Neuralwatt `GET /v1/quota` — subscription kWh usage (primary window) and + * prepaid USD credit balance (secondary). + */ +async function fetchNeuralwattQuota(provider: string, config: OcxProviderConfig): Promise { + if (!isCanonicalNeuralwattBaseUrl(config.baseUrl)) return null; + const apiKey = resolveProviderApiKey(config.apiKey)?.trim(); + if (!apiKey) return null; + const response = await fetch(`${NEURALWATT_BASE_URL}/quota`, { + headers: { Accept: "application/json", Authorization: `Bearer ${apiKey}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const body = asRecord(await readQuotaJson(response)); + const data = asRecord(body?.data) ?? body; + const quota: ProviderQuota = { updatedAt: Date.now() }; + let windows = 0; + const subscription = asRecord(data?.subscription); + const kwhUsed = subscription ? toFiniteNumber(subscription.kwh_used) : undefined; + const kwhIncluded = subscription ? toFiniteNumber(subscription.kwh_included) : undefined; + if (kwhUsed !== undefined && kwhIncluded !== undefined && kwhIncluded > 0) { + const percent = normalizePercent((kwhUsed / kwhIncluded) * 100); + if (percent !== undefined) { + quota.fiveHourPercent = percent; + const periodEnd = subscription ? normalizeResetAt(subscription.current_period_end) : undefined; + if (periodEnd !== undefined) quota.fiveHourResetAt = periodEnd; + windows += 1; + } + } + const balance = asRecord(data?.balance); + const totalCredits = balance ? toFiniteNumber(balance.total_credits_usd) : undefined; + const remainingCredits = balance ? toFiniteNumber(balance.credits_remaining_usd) : undefined; + if (totalCredits !== undefined && totalCredits > 0 && remainingCredits !== undefined) { + // Utilization is CONSUMED credits, not the remaining share. + const used = Math.max(0, totalCredits - remainingCredits); + const percent = normalizePercent((used / totalCredits) * 100); + if (percent !== undefined) { + quota.customWindows = [...(quota.customWindows ?? []), { label: "Prepaid credits", percent }]; + windows += 1; + } + } + return windows > 0 ? report(provider, "neuralwatt:quota", quota) : null; +} + + +function normalizedBaseUrl(value: string): string | null { + try { + const url = new URL(value); + if (url.username || url.password || url.search || url.hash) return null; + return `${url.origin.toLowerCase()}${url.pathname.replace(/\/+$/, "")}`; + } catch { + return null; + } +} + +function quotaResetAt(row: Record): number | undefined { + return normalizeResetAt(row.resetTime ?? row.resetAt ?? row.reset_time ?? row.reset_at); +} + +export function isCanonicalKimiCodeBaseUrl(baseUrl: string): boolean { + return normalizedBaseUrl(baseUrl) === KIMI_CODE_BASE_URL; +} + +export function isCanonicalCommandCodeBaseUrl(baseUrl: string): boolean { + const normalized = normalizedBaseUrl(baseUrl); + // OAuth preset points at the API root; the Provider-API preset at /provider/v1. + return normalized === COMMAND_CODE_BASE_URL || normalized === `${COMMAND_CODE_BASE_URL}/provider/v1`; +} + +/** Prefer the nested `data` shell when the outer object is only an envelope. */ +function unwrapKimiQuotaPayload(value: unknown): Record | null { + const body = asRecord(value); + if (!body) return null; + const nested = asRecord(body.data); + if (!nested) return body; + // A null/non-usable outer field is a placeholder, not data — an envelope like + // { usage: null, data: { usage: {...} } } must still unwrap to the nested payload. + const usable = (field: unknown): boolean => field !== undefined && field !== null; + const outerHasUsage = usable(body.usage) || usable(body.limits) || usable(body.totalQuota); + const nestedHasUsage = usable(nested.usage) || usable(nested.limits) || usable(nested.totalQuota); + return !outerHasUsage && nestedHasUsage ? nested : body; +} + +function kimiLimitLabel(item: Record, detail: Record): string { + return [item.name, item.title, item.scope, detail.name, detail.title] + .filter((value): value is string => typeof value === "string") + .join(" ") + .toLowerCase(); +} + +function parseKimiQuotaRow(value: unknown, resetFallback?: Record): { percent: number; resetAt?: number } | null { + const row = asRecord(value); + if (!row) return null; + const resetAt = quotaResetAt(row) ?? (resetFallback ? quotaResetAt(resetFallback) : undefined); + const limit = toFiniteNumber(row.limit); + if (limit !== undefined && limit > 0) { + let used = toFiniteNumber(row.used); + if (used === undefined) { + const remaining = toFiniteNumber(row.remaining); + if (remaining !== undefined) used = limit - remaining; + } + if (used !== undefined) { + const percent = normalizePercent((used / limit) * 100); + if (percent !== undefined) return { percent, ...(resetAt !== undefined ? { resetAt } : {}) }; + } + } + // Some payloads expose utilisation directly when limit/used arithmetic is absent. + const direct = normalizePercent(row.utilization ?? row.percent ?? row.usedPercent ?? row.used_percent); + return direct === undefined ? null : { percent: direct, ...(resetAt !== undefined ? { resetAt } : {}) }; +} + +function isKimiFiveHourLimit(item: Record, detail: Record, window: Record): boolean { + const duration = toFiniteNumber(window.duration ?? item.duration ?? detail.duration); + const unit = String(window.timeUnit ?? item.timeUnit ?? detail.timeUnit ?? "").toUpperCase(); + if ((unit.includes("MINUTE") && duration === 300) || (unit.includes("HOUR") && duration === 5)) return true; + return /(^|\b)5\s*(?:h|hour)/.test(kimiLimitLabel(item, detail)); +} + +function isKimiWeeklyLimit(item: Record, detail: Record, window: Record): boolean { + const duration = toFiniteNumber(window.duration ?? item.duration ?? detail.duration); + const unit = String(window.timeUnit ?? item.timeUnit ?? detail.timeUnit ?? "").toUpperCase(); + if ((unit.includes("DAY") && duration === 7) || (unit.includes("HOUR") && duration === 168)) return true; + return /weekly|7\s*(?:d|day)/.test(kimiLimitLabel(item, detail)); +} + +function parseKimiQuotaPayload(value: unknown): ProviderQuota | null { + const body = unwrapKimiQuotaPayload(value); + if (!body) return null; + let weekly = parseKimiQuotaRow(body.usage); + const total = parseKimiQuotaRow(body.totalQuota); + let fiveHour: { percent: number; resetAt?: number } | null = null; + if (Array.isArray(body.limits)) { + for (const rawItem of body.limits) { + const item = asRecord(rawItem); + if (!item) continue; + const detail = asRecord(item.detail) ?? item; + const window = asRecord(item.window) ?? {}; + if (!fiveHour && isKimiFiveHourLimit(item, detail, window)) { + fiveHour = parseKimiQuotaRow(detail, window); + } + if (!weekly && isKimiWeeklyLimit(item, detail, window)) { + weekly = parseKimiQuotaRow(detail, window); + } + if (fiveHour && weekly) break; + } + } + const quota: ProviderQuota = { + ...(fiveHour ? { + fiveHourPercent: fiveHour.percent, + ...(fiveHour.resetAt !== undefined ? { fiveHourResetAt: fiveHour.resetAt } : {}), + } : {}), + ...(weekly ? { + weeklyPercent: weekly.percent, + ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), + } : {}), + ...(total ? { customWindows: [{ label: "Total subscription credits", percent: total.percent, ...(total.resetAt !== undefined ? { resetAt: total.resetAt } : {}) }] } : {}), + updatedAt: Date.now(), + }; + return hasQuotaRows(quota) ? quota : null; +} + +async function resolveKimiQuotaBearer(config: OcxProviderConfig, accountId?: string): Promise { + if (config.authMode === "oauth") { + try { + return accountId ? await getTokenForAccountQuotaProbe("kimi", accountId) : null; + } catch { + return null; + } + } + // ACTIVE key only: silently walking apiKeyPool when the primary env reference is + // unresolved would render a quota bar for a DIFFERENT account than the one routing + // requests — a wrong meter is worse than no meter. + const primary = resolveProviderApiKey(config.apiKey)?.trim(); + return primary || null; +} + +export async function fetchKimiQuota(provider: string, config: OcxProviderConfig, accessToken: string): Promise { + // Never release credentials to a user-edited or lookalike provider host. + if (!isCanonicalKimiCodeBaseUrl(config.baseUrl)) return null; + if (!accessToken) return null; + const response = await fetch(KIMI_CODE_USAGE_URL, { + headers: { Accept: "application/json", Authorization: `Bearer ${accessToken}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) return null; + const quota = parseKimiQuotaPayload(await readQuotaJson(response)); + return quota ? keyReport(provider, "kimi:usages", quota, config, accessToken, quota) : null; +} + +/** + * Command Code rolling window: `{ cap, used, resetAt }` off /alpha/billing/credits, + * normalized to a percent with an optional reset timestamp. + */ +function parseCommandCodeWindow(value: unknown): { percent: number; resetAt?: number } | null { + const row = asRecord(value); + if (!row) return null; + const cap = toFiniteNumber(row.cap); + const used = toFiniteNumber(row.used); + if (cap === undefined || used === undefined || cap <= 0 || used < 0) return null; + const percent = normalizePercent((used / cap) * 100); + if (percent === undefined) return null; + const resetAt = quotaResetAt(row); + return { percent, ...(resetAt !== undefined ? { resetAt } : {}) }; +} + +/** Soft-fail GET returning a parsed record, or null when unavailable. */ +async function fetchCommandCodeJson(url: string, bearer: string): Promise | null> { + try { + const response = await fetch(url, { + headers: { Accept: "application/json", Authorization: `Bearer ${bearer}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) return null; + return asRecord(await readQuotaJson(response)); + } catch { + return null; + } +} + +/** + * Soft-fail period spend (used) against the remaining credit pools → creditsUsd. + * Period scoping: `since=` keeps spend aligned with the + * pools' billing cycle, and `currentPeriodEnd` becomes expiresAt. + */ +async function fetchCommandCodeSpend( + bearer: string, + credits: Record | null, + orgQuery: string, +): Promise { + if (!credits) return undefined; + const subscriptionBody = await fetchCommandCodeJson(`${COMMAND_CODE_SUBSCRIPTIONS_URL}${orgQuery}`, bearer); + const subscription = asRecord(subscriptionBody?.data) ?? subscriptionBody; + const periodStart = typeof subscription?.currentPeriodStart === "string" ? subscription.currentPeriodStart.trim() : ""; + // Unscoped /usage/summary is lifetime spend; mixing it with current-cycle + // remaining pools produces a wrong percent. Omit creditsUsd until a period exists. + if (!periodStart) return undefined; + const sinceQuery = `${orgQuery ? "&" : "?"}since=${encodeURIComponent(periodStart)}`; + const expiresAt = normalizeResetAt(subscription?.currentPeriodEnd); + const summaryBody = await fetchCommandCodeJson(`${COMMAND_CODE_USAGE_URL}${orgQuery}${sinceQuery}`, bearer); + const summary = asRecord(summaryBody?.data) ?? summaryBody; + const used = toFiniteNumber(summary?.totalCost) ?? toFiniteNumber(summary?.totalMonthlyCredits); + if (used === undefined || used < 0) return undefined; + const pools = [credits.monthlyCredits, credits.purchasedCredits, credits.freeCredits] + .map(value => toFiniteNumber(value)) + .filter((value): value is number => value !== undefined); + // Field presence is what separates a real balance from absent data: an exhausted + // all-zero account still reports remaining=0, while no remaining-credit field at + // all means there is nothing to meter. + if (pools.length === 0) return undefined; + const remaining = pools.reduce((sum, value) => sum + Math.max(0, value ?? 0), 0); + const limit = used + remaining; + const percent = normalizePercent(limit > 0 ? (used / limit) * 100 : 0); + // Purchased credits roll over past the subscription period end, so an expiry is + // only truthful when the aggregate contains no non-expiring purchased pool. + const purchased = toFiniteNumber(credits.purchasedCredits) ?? 0; + return percent === undefined + ? undefined + : { + used, + limit, + remaining, + percent, + ...(expiresAt !== undefined && purchased <= 0 ? { expiresAt } : {}), + }; +} + +/** OAuth access token or ACTIVE Provider-API key for the Command Code quota probe. */ +async function resolveCommandCodeQuotaBearer(config: OcxProviderConfig, accountId?: string): Promise { + if (config.authMode === "oauth") { + try { + return accountId ? await getTokenForAccountQuotaProbe("command-code", accountId) : null; + } catch { + return null; + } + } + // ACTIVE key only: a quota bar for a different account than the one routing + // requests is a wrong meter, not a helpful one. + return resolveProviderApiKey(config.apiKey)?.trim() || null; +} + +/** + * Command Code `GET /alpha/billing/credits` — the same Bearer surface the CLI's + * usage view uses (windowLimits.fiveHour / windowLimits.weekly), plus soft + * whoami (team orgId scoping) and subscription-scoped spend for creditsUsd. + */ +export async function fetchCommandCodeQuota(provider: string, config: OcxProviderConfig, bearer: string): Promise { + // Never release credentials to a user-edited or lookalike provider host. + if (!isCanonicalCommandCodeBaseUrl(config.baseUrl)) return null; + if (!bearer) return null; + const whoamiBody = await fetchCommandCodeJson(COMMAND_CODE_WHOAMI_URL, bearer); + const whoami = asRecord(whoamiBody?.data) ?? whoamiBody; + const org = asRecord(whoami?.org); + const orgId = typeof org?.id === "string" && org.id.trim() ? org.id.trim() : null; + const orgQuery = orgId ? `?orgId=${encodeURIComponent(orgId)}` : ""; + const response = await fetch(`${COMMAND_CODE_CREDITS_URL}${orgQuery}`, { + headers: { Accept: "application/json", Authorization: `Bearer ${bearer}` }, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + return response.status >= 400 && response.status < 500 && response.status !== 408 && response.status !== 429 + ? TERMINAL_QUOTA_FAILURE + : null; + } + const raw = asRecord(await readQuotaJson(response)); + const body = asRecord(raw?.data) ?? raw; + const credits = asRecord(body?.credits); + const limits = asRecord(body?.windowLimits); + if (!credits && !limits) return null; + const fiveHour = parseCommandCodeWindow(limits?.fiveHour); + const weekly = parseCommandCodeWindow(limits?.weekly); + const creditsUsd = await fetchCommandCodeSpend(bearer, credits, orgQuery); + const quota: ProviderQuota = { + ...(fiveHour ? { + fiveHourPercent: fiveHour.percent, + ...(fiveHour.resetAt !== undefined ? { fiveHourResetAt: fiveHour.resetAt } : {}), + } : {}), + ...(weekly ? { + weeklyPercent: weekly.percent, + ...(weekly.resetAt !== undefined ? { weeklyResetAt: weekly.resetAt } : {}), + } : {}), + ...(creditsUsd ? { creditsUsd } : {}), + updatedAt: Date.now(), + }; + // Rolling windows and the credit balance both gate inference on this bearer. + return keyReport(provider, "command-code:credits", quota, config, bearer, quota); +} + + +type KeyQuotaReader = (name: string, provider: OcxProviderConfig) => Promise; + +/** Same selector drives cheap capabilities and uncached reads; never resolves credentials. */ +export function keyQuotaReaderForProvider(name: string, provider: OcxProviderConfig): KeyQuotaReader | null { + if (provider.disabled === true || (provider.authMode ?? "key") !== "key") return null; + if (isCanonicalKimiCodeBaseUrl(provider.baseUrl)) { + return async (id, config) => { + const bearer = await resolveKimiQuotaBearer(config); + return bearer ? fetchKimiQuota(id, config, bearer) : null; + }; + } + if (name === "commandcode" && isCanonicalCommandCodeBaseUrl(provider.baseUrl)) { + return async (id, config) => { + const bearer = await resolveCommandCodeQuotaBearer(config); + return bearer ? fetchCommandCodeQuota(id, config, bearer) : null; + }; + } + if (registryEntryForProviderDestination(provider)?.id === "opencode-go") return fetchOpenCodeGoQuota; + if (isCanonicalA6apiBaseUrl(provider.baseUrl)) return fetchA6apiQuota; + if (name === "openrouter" && isCanonicalOpenRouterBaseUrl(provider.baseUrl)) return fetchOpenRouterQuota; + if (name === "deepseek" && isCanonicalDeepSeekBaseUrl(provider.baseUrl)) return fetchDeepSeekQuota; + if (name === "cline-pass" && isCanonicalClineBaseUrl(provider.baseUrl)) return fetchClineQuota; + if (isCanonicalOllamaCloudBaseUrl(provider.baseUrl ?? getProviderRegistryEntry(name)?.baseUrl)) return fetchOllamaCloudQuota; + // #4201: the Responses preset is the same domestic GLM Coding Plan subscription on the OpenAI + // Responses wire, so it reads the same monitor endpoint. Eligibility stays a name list AND the + // canonical-URL guard: the guard is what keeps BigModel's bare-key Authorization from reaching a + // lookalike host, so a same-named custom destination still dispatches nothing. + if (["zai", "glm", "glm-cn", "zhipu-bigmodel-coding", "zhipu-bigmodel-responses"].includes(name) && isCanonicalZaiBaseUrl(provider.baseUrl)) return fetchZaiQuota; + if (["minimax", "minimax-cn"].includes(name) && isCanonicalMinimaxBaseUrl(provider.baseUrl)) return fetchMinimaxQuota; + if (name === "moonshot" && isCanonicalMoonshotBaseUrl(provider.baseUrl)) return fetchMoonshotQuota; + if (name === "venice" && isCanonicalVeniceBaseUrl(provider.baseUrl)) return fetchVeniceQuota; + if (name === "synthetic" && isCanonicalSyntheticBaseUrl(provider.baseUrl)) return fetchSyntheticQuota; + if (name === "deepinfra" && isCanonicalDeepInfraBaseUrl(provider.baseUrl)) return fetchDeepInfraQuota; + if (name === "neuralwatt" && isCanonicalNeuralwattBaseUrl(provider.baseUrl)) return fetchNeuralwattQuota; + return null; +} + +export function providerApiKeyQuotaMode(name: string, provider: OcxProviderConfig): AccountQuotaMode { + return keyQuotaReaderForProvider(name, provider) ? "probe" : "unsupported"; +} diff --git a/src/providers/quota/vendor-probes-oauth.ts b/src/providers/quota/vendor-probes-oauth.ts new file mode 100644 index 0000000000..7b9c9e7df5 --- /dev/null +++ b/src/providers/quota/vendor-probes-oauth.ts @@ -0,0 +1,590 @@ +import { effectiveCodexAuthAccountId, fetchMainAccountInfoSnapshot, listCodexAuthAccountsSnapshot } from "../../codex/auth-api"; +import { MAIN_CODEX_ACCOUNT_ID } from "../../codex/main-account"; +import { getValidAccessToken } from "../../oauth"; +import { getAccountCredential, getAccountSet } from "../../oauth/store"; +import { fetchMuseKeyQuotaSnapshot } from "../muse-key-quota"; +import { XAI_GROK_CLIENT_VERSION, XAI_GROK_COMPATIBILITY } from "../xai-transport"; +import { + commitKiroAccountUsageState, + fetchKiroUsageSnapshot, + type KiroUsageSnapshot, + kiroUsageContextForAccount, +} from "../kiro-usage"; +import { captureConfigGeneration } from "../../lib/state-store-sweeper"; +import { aggregateCodexPoolCapacity, CODEX_CAPACITY_MAX_QUOTA_AGE_MS, type CodexCapacityQuota } from "../codex-capacity"; +import { asRecord, normalizePercent, normalizeResetAt, readQuotaJson, REQUEST_TIMEOUT_MS, toFiniteNumber } from "../quota-wire"; +import { providerCodexAccountMode } from "../registry"; +import { + hasQuotaRows, + providerLabel, + providerQuotaFromCodexQuota, + publicCapacityAggregation, + report, + tagNativeMainReport, + type CodexAuthAccountsSnapshotPromise, + type ProviderQuotaReport, +} from "./report-cache"; +import { + accountCacheKey, + accountQuotaCache, + hydrateAccountQuotaCache, + mayCommitAccountQuotaKey, + persistAccountQuotaCache, +} from "./account-cache"; +import type { OcxConfig, OcxProviderConfig } from "../../types"; +import type { ProviderQuota, ProviderQuotaWindow } from "../quota-types"; + +const XAI_BILLING_URL = "https://cli-chat-proxy.grok.com/v1/billing"; +const XAI_CREDITS_URL = `${XAI_BILLING_URL}?format=credits`; + +export async function fetchChatGptForwardQuota( + config: OcxConfig, + provider: string, + providerConfig: OcxProviderConfig, + forceRefresh: boolean, + prefetchedSnapshot?: CodexAuthAccountsSnapshotPromise, +): Promise { + if (providerCodexAccountMode(provider, providerConfig) === "direct") { + const snapshot = await fetchMainAccountInfoSnapshot(forceRefresh); + const quota = providerQuotaFromCodexQuota(snapshot.info.quota); + if (quota) quota.updatedAt = Date.now(); + return quota + ? tagNativeMainReport(report(provider, "chatgpt:wham", quota), snapshot.mainIdentityGeneration) + : null; + } + const snapshot = await (prefetchedSnapshot ?? listCodexAuthAccountsSnapshot(config, forceRefresh)); + const accounts = snapshot.accounts; + const activeId = effectiveCodexAuthAccountId(config); + const capacityAccounts = accounts.map(account => ({ + ...account, + active: account.id === activeId, + quota: providerQuotaFromCodexQuota(account.quota), + })); + const active = capacityAccounts.find(account => account.active) + ?? capacityAccounts.find(account => account.id === MAIN_CODEX_ACCOUNT_ID) + ?? capacityAccounts[0]; + const now = Date.now(); + const capacity = aggregateCodexPoolCapacity(capacityAccounts, now); + if (capacity.aggregation && capacity.quota) { + return tagNativeMainReport( + report( + provider, + "chatgpt:wham", + capacity.quota as ProviderQuota, + publicCapacityAggregation(capacity.aggregation, "aggregate"), + ), + snapshot.mainIdentityGeneration, + ); + } + const activeUsable = !!active && !active.paused && active.needsReauth !== true; + const quota = activeUsable && active?.quota + ? { ...active.quota, updatedAt: active.quota.updatedAt ?? Date.now() } as CodexCapacityQuota + : null; + const quotaFresh = !!quota + && Number.isFinite(quota.updatedAt) + && now - quota.updatedAt < CODEX_CAPACITY_MAX_QUOTA_AGE_MS; + if (quota && quotaFresh) { + const fallback = report( + provider, + "chatgpt:wham", + quota as ProviderQuota, + capacity.aggregation + ? publicCapacityAggregation(capacity.aggregation, "effective-account-fallback") + : undefined, + ); + return tagNativeMainReport(fallback, snapshot.mainIdentityGeneration); + } + if (capacity.aggregation) { + const updatedAt = Date.now(); + return tagNativeMainReport( + { + provider, + label: providerLabel(provider), + source: "chatgpt:wham", + quota: { updatedAt }, + updatedAt, + aggregation: publicCapacityAggregation(capacity.aggregation, "coverage-only"), + }, + snapshot.mainIdentityGeneration, + ); + } + return null; +} + +function centsValue(value: unknown): number | undefined { + const rec = asRecord(value); + return rec ? toFiniteNumber(rec.val) : undefined; +} + +/** Decode JWT payload `sub` for xAI weekly credits when the stored credential lacks accountId. */ +function xaiUserIdFromAccessToken(accessToken: string): string | undefined { + const parts = accessToken.split("."); + if (parts.length < 2 || !parts[1]) return undefined; + try { + const payload = JSON.parse(Buffer.from(parts[1], "base64url").toString("utf8")) as { sub?: unknown }; + return typeof payload.sub === "string" && payload.sub.trim() ? payload.sub.trim() : undefined; + } catch { + return undefined; + } +} + +/** + * Grok Build weekly credits envelope: + * `{ config: { creditUsagePercent?, currentPeriod: { type: "USAGE_PERIOD_TYPE_WEEKLY", end } } }`. + * Omitted percent is treated as 0 (proto3 default). + */ +export function parseXaiCreditsResponse(value: unknown): { percent: number; resetAt?: number } | null { + const body = asRecord(value); + const config = asRecord(body?.config); + if (!config) return null; + const period = asRecord(config.currentPeriod); + if (!period || period.type !== "USAGE_PERIOD_TYPE_WEEKLY") return null; + let percent = 0; + if (config.creditUsagePercent !== undefined) { + const normalized = normalizePercent(config.creditUsagePercent); + if (normalized === undefined) return null; + percent = normalized; + } + const resetAt = normalizeResetAt(period.end); + return { + percent, + ...(resetAt !== undefined ? { resetAt } : {}), + }; +} + +async function fetchXaiWeeklyCredits(accessToken: string, userId: string): Promise { + try { + const response = await fetch(XAI_CREDITS_URL, { + redirect: "error", + headers: { + Accept: "application/json", + Authorization: `Bearer ${accessToken}`, + [XAI_GROK_COMPATIBILITY.headers.tokenAuth]: "xai-grok-cli", + [XAI_GROK_COMPATIBILITY.headers.authenticateResponse]: "authenticate-response", + "x-userid": userId, + [XAI_GROK_COMPATIBILITY.headers.clientVersion]: XAI_GROK_CLIENT_VERSION, + }, + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) return null; + const parsed = parseXaiCreditsResponse(await readQuotaJson(response)); + if (!parsed) return null; + return { + weeklyPercent: parsed.percent, + ...(parsed.resetAt !== undefined ? { weeklyResetAt: parsed.resetAt } : {}), + updatedAt: Date.now(), + }; + } catch { + return null; + } +} + +export async function fetchXaiQuota(provider: string, context: { accessToken: string; upstreamAccountId?: string }): Promise { + const { accessToken } = context; + + // Prefer the SuperGrok weekly credits window that actually gates prompting (#1283). + const userId = context.upstreamAccountId?.trim() || xaiUserIdFromAccessToken(accessToken); + if (userId) { + const weekly = await fetchXaiWeeklyCredits(accessToken, userId); + if (weekly) return report(provider, "xai:grok-billing-credits", weekly); + } + + // Legacy monthly dollar pool — retained when weekly is unavailable. + try { + const response = await fetch(XAI_BILLING_URL, { + redirect: "error", + headers: { Accept: "application/json", Authorization: `Bearer ${accessToken}` }, + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) return null; + const body = asRecord(await readQuotaJson(response)); + const config = asRecord(body?.config); + if (!config) return null; + const limitCents = centsValue(config.monthlyLimit); + const usedCents = centsValue(config.used); + if (limitCents === undefined || usedCents === undefined || limitCents <= 0) return null; + const percent = normalizePercent((usedCents / limitCents) * 100); + if (percent === undefined) return null; + return report(provider, "xai:grok-billing", { + monthlyPercent: percent, + monthlyResetAt: normalizeResetAt(config.billingPeriodEnd), + updatedAt: Date.now(), + }); + } catch { + return null; + } +} + +function parseClaudeBucket(value: unknown): { percent?: number; resetAt?: number } | null { + const rec = asRecord(value); + if (!rec) return null; + const percent = normalizePercent(rec.utilization); + const resetAt = normalizeResetAt(rec.resets_at); + if (percent === undefined && resetAt === undefined) return null; + return { percent, resetAt }; +} + +function parseClaudeLimit(value: unknown): { label: string; percent: number; resetAt?: number } | null { + const rec = asRecord(value); + if (!rec) return null; + const percent = normalizePercent(rec.percent); + if (percent === undefined) return null; + const scope = asRecord(rec.scope); + const model = asRecord(scope?.model); + const rawLabel = String(model?.display_name ?? "").trim(); + if (!rawLabel) return null; + const lowerLabel = rawLabel.toLowerCase(); + const label = lowerLabel.includes("fable") ? "Fable" + : lowerLabel.includes("opus") ? "Opus" + : lowerLabel.includes("sonnet") ? "Sonnet" + : rawLabel; + const resetAt = normalizeResetAt(rec.resets_at); + return { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) }; +} + +/** Claude's OAuth usage endpoint, probed with ONE account's own bearer token. */ +const anthropicUsageInflight = new Map>(); + +/** + * Anthropic per-credential usage. + * + * This endpoint reports quota only. Its body carries `five_hour`, `seven_day`, the + * model-scoped weekly buckets (`seven_day_fable`/`_opus`/`_sonnet`) and a `limits` array, + * and **no subscription or tier field** — nor does the OAuth token response, which yields only + * `account.uuid` and `account.email_address` (`src/oauth/anthropic.ts`). That is why + * `OAuthAccountSummary.plan` is `null` for Anthropic rather than populated here (#3777); it is + * a missing upstream field, not an unfinished mapping. + * + * A tier must not be inferred from what is here. Percentages are normalized per account, so a + * Max x5 seat at 50% is byte-identical to a Max x20 seat at 50%, and the presence of a + * model-scoped window tracks entitlement rather than seat size. Populate `plan` only when + * upstream returns the tier itself. + */ +export async function fetchAnthropicUsageQuota(accessToken: string): Promise { + const joinable = anthropicUsageInflight.get(accessToken); + if (joinable) return joinable; + + const probe = (async (): Promise => { + const response = await fetch("https://api.anthropic.com/api/oauth/usage", { + headers: { + Accept: "application/json, text/plain, */*", + "Content-Type": "application/json", + "User-Agent": "claude-cli/2.1.63 (external, cli)", + "anthropic-beta": "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27,prompt-caching-scope-2026-01-05", + Authorization: `Bearer ${accessToken}`, + }, + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) return null; + const body = asRecord(await readQuotaJson(response)); + if (!body) return null; + const fiveHour = parseClaudeBucket(body.five_hour); + const sevenDay = parseClaudeBucket(body.seven_day); + const fable = parseClaudeBucket(body.seven_day_fable); + const opus = parseClaudeBucket(body.seven_day_opus); + const sonnet = parseClaudeBucket(body.seven_day_sonnet); + const customWindows: ProviderQuotaWindow[] = []; + if (fable?.percent !== undefined) customWindows.push({ label: "Fable", percent: fable.percent, ...(fable.resetAt !== undefined ? { resetAt: fable.resetAt } : {}) }); + if (opus?.percent !== undefined) customWindows.push({ label: "Opus", percent: opus.percent, ...(opus.resetAt !== undefined ? { resetAt: opus.resetAt } : {}) }); + if (sonnet?.percent !== undefined) customWindows.push({ label: "Sonnet", percent: sonnet.percent, ...(sonnet.resetAt !== undefined ? { resetAt: sonnet.resetAt } : {}) }); + const knownLabels = new Set(customWindows.map(window => window.label.toLowerCase())); + const limits = Array.isArray(body.limits) ? body.limits : []; + for (const rawLimit of limits) { + const limitRecord = asRecord(rawLimit); + // `session` and `weekly_all` mirror the canonical five-hour and weekly + // buckets above; only model-scoped weekly limits add a third window. + if (String(limitRecord?.kind ?? "").trim().toLowerCase() !== "weekly_scoped") continue; + const limit = parseClaudeLimit(rawLimit); + if (!limit || knownLabels.has(limit.label.toLowerCase())) continue; + knownLabels.add(limit.label.toLowerCase()); + customWindows.push(limit); + } + const quota: ProviderQuota = { + // Claude's 5-hour window is a first-class rate limit, same as the Codex login 5h/weekly + // rows: report it in the canonical fields so the dashboard renders it with the standard + // "5-hour limit" label and ordering instead of as a generic extra window. + ...(fiveHour?.percent !== undefined ? { fiveHourPercent: fiveHour.percent } : {}), + ...(fiveHour?.resetAt !== undefined ? { fiveHourResetAt: fiveHour.resetAt } : {}), + ...(sevenDay?.percent !== undefined ? { weeklyPercent: sevenDay.percent } : {}), + ...(sevenDay?.resetAt !== undefined ? { weeklyResetAt: sevenDay.resetAt } : {}), + ...(customWindows.length > 0 ? { customWindows } : {}), + updatedAt: Date.now(), + }; + // Empty / schema-changed payloads must not cache as "success with no bars". + return hasQuotaRows(quota) ? quota : null; + })().finally(() => { + if (anthropicUsageInflight.get(accessToken) === probe) anthropicUsageInflight.delete(accessToken); + }); + anthropicUsageInflight.set(accessToken, probe); + return probe; +} + +export async function fetchAnthropicQuota(provider: string): Promise { + // Capture the account we intend to probe before awaiting — a mid-flight active + // switch must not seed the wrong account's cache with this response. + const probedAccountId = getAccountSet("anthropic")?.activeAccountId; + const probedAccountKey = probedAccountId ? accountCacheKey("anthropic", probedAccountId) : null; + const writerGeneration = captureConfigGeneration(); + let accessToken: string; + try { + accessToken = await getValidAccessToken("anthropic"); + } catch { + return null; + } + const quota = await fetchAnthropicUsageQuota(accessToken); + if (!quota) return null; + // Share the active-account probe with the per-account cache so Providers-page + // loads do not double-hit Anthropic's rate-limited usage endpoint. + if (probedAccountId && probedAccountKey) { + const stillOwnsToken = getAccountCredential("anthropic", probedAccountId)?.access === accessToken; + if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { + accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota }); + } + } + return report(provider, "anthropic:oauth-usage", quota); +} + +/** + * Provider-level Kiro row: the active account's usage, shown on the Providers page. + * + * The per-account cache is seeded from the same probe so opening that page does not read + * the active account twice, and the account id is captured before the await so a + * concurrent account switch cannot file this answer under the wrong account. + */ +export async function fetchKiroQuota(provider: string): Promise { + const probedAccountId = getAccountSet("kiro")?.activeAccountId; + if (!probedAccountId) return null; + const probedAccountKey = accountCacheKey("kiro", probedAccountId); + const writerGeneration = captureConfigGeneration(); + let snapshot: KiroUsageSnapshot | null; + try { + snapshot = await fetchKiroUsageSnapshot(await kiroUsageContextForAccount(probedAccountId)); + } catch { + return null; + } + if (!snapshot) return null; + if (mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { + accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota: snapshot.quota }); + commitKiroAccountUsageState(probedAccountKey, snapshot); + } + return report(provider, "kiro:usage-limits", snapshot.quota); +} + +/** + * Provider-level row probed from the key endpoint, for an account that CAN be probed. + * + * Written through the same account cache the passive path reads, so the measurement + * survives a restart and the per-account rows at oauth-account-routes.ts:313 pick it up + * with no mode change. Deliberately does not flip providerOAuthAccountQuotaMode: that + * mode selects readPassiveProviderAccountQuotas, and the probed per-account path it would + * switch to is gated on supportsPerAccountQuota, which has no meta-muse reader, so the + * GUI account list would go from showing observations to showing nothing. + */ +export async function fetchMuseKeyQuota(provider: string): Promise { + const probedAccountId = getAccountSet(provider)?.activeAccountId; + if (!probedAccountId) return null; + const oauthAccessToken = getAccountCredential(provider, probedAccountId)?.muse?.oauthAccessToken; + // An imported or pasted credential has no account token and never will: it is + // capability, not provider id, that decides whether a probe is possible. + if (!oauthAccessToken) return null; + const probedAccountKey = accountCacheKey(provider, probedAccountId); + const writerGeneration = captureConfigGeneration(); + const quota = await fetchMuseKeyQuotaSnapshot(probedAccountId, oauthAccessToken); + if (!quota) return null; + if (mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) { + // Hydrate before writing, for the same reason recordPassiveAccountQuota does: + // persistAccountQuotaCache serializes the whole in-memory map. + hydrateAccountQuotaCache(); + accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota }); + persistAccountQuotaCache(); + } + return report(provider, `${provider}:key-endpoint`, quota); +} +/** + * Provider-level row for a passive provider: the ACTIVE account's last observed + * subscription windows, the same shape `fetchAnthropicQuota` and `fetchKiroQuota` + * return. + * + * Cache-only. A dashboard load or `ocx account refresh` must never spend an inference + * turn, so `forceRefresh` does not exist on this path — there is nothing to refresh. + * `report.updatedAt` is the observation time, which is what both GUI surfaces render + * as the relative age of the row. + */ +export async function fetchPassiveProviderQuota(provider: string): Promise { + const activeId = getAccountSet(provider)?.activeAccountId; + if (!activeId) return null; + // Idempotent; without it a proxy restart shows nothing until the next streaming turn + // even though the last observation is on disk. + hydrateAccountQuotaCache(); + const entry = accountQuotaCache.get(accountCacheKey(provider, activeId)); + if (!entry?.quota) return null; + const built = report(provider, `${provider}:subscription-observation`, entry.quota); + // Tagged here rather than inside report(), which every probed path shares. + return built ? { ...built, observed: true } : null; +} + +// --------------------------------------------------------------------------- +// Per-account quota (multiauth) +// --------------------------------------------------------------------------- + + +/** Cursor included usage via api2.cursor.sh (Bearer from OAuth) — unofficial, may change. */ +export async function fetchCursorQuota(provider: string, accessToken: string): Promise { + + const authHeaders = { + Accept: "application/json", + Authorization: `Bearer ${accessToken}`, + "User-Agent": "opencodex-quota", + } as const; + + // Prefer dashboard period usage (Pro/Team/Ultra spend allowance in USD cents). + // Field names follow Cursor's Connect RPC shape (limit/remaining/includedSpend), not usedCents. + try { + const periodRes = await fetch("https://api2.cursor.sh/aiserver.v1.DashboardService/GetCurrentPeriodUsage", { + method: "POST", + redirect: "error", + headers: { + ...authHeaders, + "Content-Type": "application/json", + "Connect-Protocol-Version": "1", + }, + body: "{}", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (periodRes.ok) { + const body = asRecord(await readQuotaJson(periodRes)); + const planUsage = asRecord(body?.planUsage); + if (planUsage) { + const resetAt = normalizeResetAt(body?.billingCycleEnd ?? planUsage.billingCycleEnd ?? body?.periodEnd); + + // Primary meter: overall included allowance (Cursor Settings → Usage total %). + // autoPercentUsed / apiPercentUsed are secondary pools and must not replace the total. + const limit = toFiniteNumber(planUsage.limit ?? planUsage.limitCents ?? planUsage.totalLimitCents); + const remaining = toFiniteNumber(planUsage.remaining ?? planUsage.remainingCents); + const includedSpend = toFiniteNumber(planUsage.includedSpend ?? planUsage.usedCents ?? planUsage.used); + const totalSpend = toFiniteNumber(planUsage.totalSpend); + let used: number | undefined; + if (includedSpend !== undefined) used = includedSpend; + else if (limit !== undefined && remaining !== undefined) used = Math.max(0, limit - remaining); + else if (totalSpend !== undefined) used = totalSpend; + const totalPercent = normalizePercent(planUsage.totalPercentUsed ?? planUsage.percentUsed) + ?? (limit !== undefined && limit > 0 && used !== undefined + ? normalizePercent((used / limit) * 100) + : undefined); + + const autoPercent = normalizePercent(planUsage.autoPercentUsed); + const apiPercent = normalizePercent(planUsage.apiPercentUsed); + const customWindows: ProviderQuotaWindow[] = []; + if (autoPercent !== undefined) { + customWindows.push({ + label: "First-party models", + percent: autoPercent, + ...(resetAt !== undefined ? { resetAt } : {}), + }); + } + if (apiPercent !== undefined) { + customWindows.push({ + label: "API usage", + percent: apiPercent, + ...(resetAt !== undefined ? { resetAt } : {}), + }); + } + + if (totalPercent !== undefined || customWindows.length > 0) { + const built = report(provider, "cursor:period-usage", { + ...(totalPercent !== undefined ? { + monthlyPercent: totalPercent, + ...(resetAt !== undefined ? { monthlyResetAt: resetAt } : {}), + } : {}), + ...(customWindows.length > 0 ? { customWindows } : {}), + updatedAt: Date.now(), + }); + if (built) return { ...built, reverseEngineered: true }; + } + } + } + } catch { + /* fall through */ + } + + // /api/usage/summary — same host, sometimes richer than /auth/usage for Team plans. + try { + const summaryRes = await fetch("https://api2.cursor.sh/api/usage/summary", { + headers: authHeaders, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (summaryRes.ok) { + const body = asRecord(await readQuotaJson(summaryRes)); + const individual = asRecord(body?.individualUsage); + const plan = asRecord(individual?.plan); + if (plan) { + const used = toFiniteNumber(plan.used); + const limit = toFiniteNumber(plan.limit); + const percent = normalizePercent(plan.totalPercentUsed) + ?? (used !== undefined && limit !== undefined && limit > 0 + ? normalizePercent((used / limit) * 100) + : undefined); + if (percent !== undefined) { + const built = report(provider, "cursor:usage-summary", { + monthlyPercent: percent, + monthlyResetAt: normalizeResetAt(body?.billingCycleEnd), + updatedAt: Date.now(), + }); + if (built) return { ...built, reverseEngineered: true }; + } + } + } + } catch { + /* fall through to /auth/usage */ + } + + const response = await fetch("https://api2.cursor.sh/auth/usage", { + headers: authHeaders, + redirect: "error", + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + if (!response.ok) return null; + const body = asRecord(await readQuotaJson(response)); + if (!body) return null; + + // Prefer the gpt-4 bucket (historical "fast requests"); else first model with used+limit. + let used: number | undefined; + let limit: number | undefined; + const gpt4 = asRecord(body["gpt-4"]); + if (gpt4) { + used = toFiniteNumber(gpt4.numRequests ?? gpt4.used); + limit = toFiniteNumber(gpt4.maxRequestUsage ?? gpt4.limit ?? gpt4.maxRequests); + } + if (used === undefined || limit === undefined || limit <= 0) { + for (const [key, value] of Object.entries(body)) { + if (key === "startOfMonth" || key === "billingCycleStart") continue; + const bucket = asRecord(value); + if (!bucket) continue; + const bucketUsed = toFiniteNumber(bucket.numRequests ?? bucket.used); + const bucketLimit = toFiniteNumber(bucket.maxRequestUsage ?? bucket.limit ?? bucket.maxRequests); + if (bucketUsed !== undefined && bucketLimit !== undefined && bucketLimit > 0) { + used = bucketUsed; + limit = bucketLimit; + break; + } + } + } + if (used === undefined || limit === undefined || limit <= 0) return null; + const percent = normalizePercent((used / limit) * 100); + if (percent === undefined) return null; + const startOfMonth = normalizeResetAt(body.startOfMonth ?? body.billingCycleStart); + // Next reset = same day next month, computed in UTC to avoid timezone-shifted rollover. + const monthlyResetAt = startOfMonth !== undefined + ? (() => { + const start = new Date(startOfMonth); + return Date.UTC(start.getUTCFullYear(), start.getUTCMonth() + 1, start.getUTCDate()); + })() + : undefined; + const built = report(provider, "cursor:auth-usage", { + monthlyPercent: percent, + ...(monthlyResetAt !== undefined ? { monthlyResetAt } : {}), + updatedAt: Date.now(), + }); + return built ? { ...built, reverseEngineered: true } : null; +} diff --git a/structure/gui-and-management-api.md b/structure/gui-and-management-api.md index 0496735405..fbbe949174 100644 --- a/structure/gui-and-management-api.md +++ b/structure/gui-and-management-api.md @@ -499,7 +499,7 @@ untouched. ## Z.ai quota destination ownership -`src/providers/quota.ts` uses one exact normalized-base mapping for both Z.ai quota +`src/providers/quota/vendor-probes-key.ts` uses one exact normalized-base mapping for both Z.ai quota eligibility and monitor selection. International root, coding Chat, Anthropic and Responses bases use `api.z.ai` with Bearer authentication. Existing BigModel CN root, coding Chat and Responses bases use `open.bigmodel.cn` with the raw key. Unsupported diff --git a/structure/providers/openai-tiers.md b/structure/providers/openai-tiers.md index a1824746dc..7a36c2414b 100644 --- a/structure/providers/openai-tiers.md +++ b/structure/providers/openai-tiers.md @@ -448,7 +448,7 @@ Listener startup diagnostics follow [the runtime lifecycle contract](../runtime. ## Automatic pool plan exclusions -`src/codex/routing.ts` applies optional `codexPool.excludedPlans` to both candidate selection and existing active/affined accounts. An all-excluded pool returns no automatic candidate, including preview and configured-account fallback. Native main remains exempt and unknown plans remain eligible. Explicit account-qualified routes retain pause, credential and entitlement checks while bypassing only this automatic policy. +`src/codex/routing/selection.ts` applies optional `codexPool.excludedPlans` to both candidate selection and existing active/affined accounts. An all-excluded pool returns no automatic candidate, including preview and configured-account fallback. Native main remains exempt and unknown plans remain eligible. Explicit account-qualified routes retain pause, credential and entitlement checks while bypassing only this automatic policy. `src/codex/auth-api.ts` projects `selectionExcludedReason: "plan_excluded"` and `selectionExcludedPlan` from the routing config, even when a newer display-only WHAM plan could not be persisted. The dashboard and account CLI show the policy reason separately from credential health; renewal clears the derived fields. The automatic next-session action and badge are omitted for excluded rows. ## Paginated history writer boundary @@ -515,7 +515,7 @@ The history read API reports a median effective token estimate and interval samp ## Reset-first account ordering -`src/codex/routing.ts` supports Codex-only `accountPoolStrategy: "reset-first"`. For new shared-quota assignments it chooses the earliest future short/weekly reset after existing eligibility, priority and usage-threshold filtering; ties and absent/elapsed deadlines use the existing usage order. Seconds and milliseconds are normalized with `resetAtToMs`. Threshold zero disables usage filtering while retaining reset ordering. Monthly deadlines do not order this strategy. +`src/codex/routing/selection.ts` supports Codex-only `accountPoolStrategy: "reset-first"`. For new shared-quota assignments it chooses the earliest future short/weekly reset after existing eligibility, priority and usage-threshold filtering; ties and absent/elapsed deadlines use the existing usage order. Seconds and milliseconds are normalized with `resetAtToMs`. Threshold zero disables usage filtering while retaining reset ordering. Monthly deadlines do not order this strategy. Live bindings obey the cache-affinity release policy: `pool.cacheAffinity` is on by default, so threshold crossing alone retains a healthy account. A bound thread that does leave may move only onto an account with genuine quota headroom and strictly lower usage. Manual preference, scoped health and shared-cursor guards remain authoritative. Set the flag false to restore threshold rebinding of bound tasks. Independent `spark`/`reserve` quota scopes resolve reset-first to existing quota selection because shared reset timestamps do not describe those windows. The configured value stays unchanged. diff --git a/structure/runtime.md b/structure/runtime.md index e4a50ba3f8..49a2fc3f2f 100644 --- a/structure/runtime.md +++ b/structure/runtime.md @@ -339,7 +339,7 @@ Automatic Codex pool selection and account status share the [plan exclusion cont `src/web-search/loop.ts` makes at most one extra answer attempt after a clean forced-answer terminal with no visible output or tool call. The recovery has no tools and reuses gathered search results. Malformed calls fail before refusal/truncation passthrough, and well-formed recognized refusal/truncation terminals pass through unchanged, including empty or partial answers. The extra generation may incur provider usage. ## Scoped provider quota for Combo selection -`src/providers/quota.ts` publishes routing evidence only when a producer explicitly supplies its +`src/providers/quota/report-cache.ts` publishes routing evidence only when a producer explicitly supplies its inference-wide projection. A matching credential alone does not grant veto authority. Display-only account, model-group, search and legacy MCP windows remain visible but cannot exclude a provider. The private WeakMap binds provider name, adapter, destination and captured credential; neither diff --git a/tests/config/config-save-boundary.test.ts b/tests/config/config-save-boundary.test.ts index 36dac68303..5b51e8182c 100644 --- a/tests/config/config-save-boundary.test.ts +++ b/tests/config/config-save-boundary.test.ts @@ -20,6 +20,7 @@ const GUARDED_FILES = [ "providers/api-keys.ts", // request-path + management key pool "providers/key-failover.ts", // 429 rotation, reached mid-turn with no user action "codex/routing.ts", // account auto-switch during a turn + "codex/routing/active-account.ts", // setActiveCodexAccount moved here in the routing split "codex/auth-api.ts", // runtime account/quota persistence "cli/claude-desktop.ts", // CLI against a running service "server/management-api.ts", diff --git a/tests/providers/provider-quota.test.ts b/tests/providers/provider-quota.test.ts index b388889a94..0c9c8b210c 100644 --- a/tests/providers/provider-quota.test.ts +++ b/tests/providers/provider-quota.test.ts @@ -119,8 +119,8 @@ afterEach(() => { describe("fetchProviderQuotaReports", () => { test("provider quota probes have no direct Response.json calls", () => { - const source = readFileSync(repoPath("src/providers/quota.ts"), "utf8"); - expect(source).not.toMatch(/\.\s*json\s*\(/); + // Probes live in leaves now; the facade alone no longer holds one. + for (const p of ["quota.ts", "quota/vendor-probes-key.ts", "quota/vendor-probes-oauth.ts", "quota/antigravity.ts"]) expect(readFileSync(repoPath(`src/providers/${p}`), "utf8")).not.toMatch(/\.\s*json\s*\(/); }); test("quota JSON reading cancels a body that stalls before its first byte", async () => { diff --git a/tests/usage/quota-reset-detector.test.ts b/tests/usage/quota-reset-detector.test.ts index 3ce61d7b7b..f369db1a26 100644 --- a/tests/usage/quota-reset-detector.test.ts +++ b/tests/usage/quota-reset-detector.test.ts @@ -117,7 +117,7 @@ describe("quota reset detection", () => { }); test("sentinel reset clocks are ignored rather than read as 1970", () => { - // src/providers/quota.ts:279 and src/codex/quota.ts:192 disagree on whether 0 survives, + // src/providers/quota/account-cache.ts and src/codex/quota.ts disagree on whether 0 survives, // so the detector re-checks: a 0 deadline must not read as a long-passed one. expect(detect({ percent: 90, resetAt: 0 }, { percent: 88, resetAt: 0 })).toBeNull(); });