diff --git a/docs-site/src/content/docs/guides/providers.md b/docs-site/src/content/docs/guides/providers.md index 2e70c4dae3..4615c0726f 100644 --- a/docs-site/src/content/docs/guides/providers.md +++ b/docs-site/src/content/docs/guides/providers.md @@ -127,7 +127,10 @@ Terminal refresh failures mark the account as needing reauthentication instead o **Cooldowns (Codex pool).** Upstream `429` / quota responses set a hard cooldown from `Retry-After`, quota `reset` headers (capped), or a short default backoff. Accounts on an explicit `Retry-After` cooldown are not probed early; reset-derived cooldowns may receive a paced probe lease -so recovery can be detected without flooding the provider. +so recovery can be detected without flooding the provider. Reset-derived native-model cooldowns +also preserve known independent quota groups: `gpt-5.3-codex-spark` does not prevent the same account +from trying the shared GPT-5.6 Terra/Luna quota, while models in that shared group still protect one +another. Explicit `Retry-After` and default cooldowns always remain account-wide. **Session affinity.** Codex thread→account affinity is process-local (in-memory only; not persisted across proxy restarts). On credential failures (`401` / `403`) the account is quarantined for diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index d15c8c4b11..aacea92ee7 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -9,13 +9,16 @@ import { isCodexAccountUsable } from "./account-usability"; import { reconcileMainCodexAccountRuntimeState } from "./account-lifecycle"; import { MAIN_CODEX_ACCOUNT_ID, getMainAccountToken } from "./main-account"; import { - getCodexAccountHealthSnapshot, + codexQuotaScopeForModel, + getCodexQuotaHealthSnapshot, releaseCodexQuotaProbeLease, + releaseCodexQuotaScopeProbeLease, tryAcquireCodexQuotaProbeLease, + tryAcquireCodexQuotaScopeProbeLease, pickAlternateCodexAccount, resolveCodexAccountForThreadDetailed, } from "./routing"; -import type { CodexCooldownSource } from "./routing"; +import type { CodexCooldownSource, CodexQuotaScope } from "./routing"; import { maskAccountId } from "../lib/privacy"; import { formatErrorResponse } from "../bridge"; import { getAccountQuota } from "./quota"; @@ -34,8 +37,12 @@ export type CodexAuthContext = * Set when this request was admitted through an active quota cooldown as * the account's single probe. Must be echoed into the upstream outcome so * only this request can clear the cooldown (#433). - */ + */ probeLeaseId?: string; + /** Native model quota group selected for this request, when known. */ + quotaScope?: CodexQuotaScope; + /** Scope that owns `probeLeaseId`, when it is a scoped recovery probe. */ + probeQuotaScope?: CodexQuotaScope; } | { // Main Codex account participating in rotation: token injected from ~/.codex/auth.json @@ -46,6 +53,8 @@ export type CodexAuthContext = chatgptAccountId: string; /** See `pool.probeLeaseId`. */ probeLeaseId?: string; + quotaScope?: CodexQuotaScope; + probeQuotaScope?: CodexQuotaScope; }; /** Probe lease carried by this context, when it holds one. */ @@ -53,13 +62,20 @@ export function codexProbeLeaseId(ctx: CodexAuthContext | undefined): string | u return ctx?.kind === "pool" || ctx?.kind === "main-pool" ? ctx.probeLeaseId : undefined; } +/** Scope of a lease carried by this context, when it probes a model-specific quota. */ +export function codexProbeQuotaScope(ctx: CodexAuthContext | undefined): CodexQuotaScope | undefined { + return ctx?.kind === "pool" || ctx?.kind === "main-pool" ? ctx.probeQuotaScope : undefined; +} + /** * Hand back a probe lease for a request that will not reach upstream. Safe to * call with a context that holds no lease. */ export function releaseCodexAuthContextProbeLease(ctx: CodexAuthContext | undefined): void { const leaseId = codexProbeLeaseId(ctx); - if (ctx && leaseId) releaseCodexQuotaProbeLease(ctx.accountId!, leaseId); + if (!ctx || ctx.kind === "main" || !leaseId) return; + if (ctx.probeQuotaScope) releaseCodexQuotaScopeProbeLease(ctx.accountId!, ctx.probeQuotaScope, leaseId); + else releaseCodexQuotaProbeLease(ctx.accountId!, leaseId); } export type OcxRuntimeProviderConfig = OcxProviderConfig & { @@ -99,13 +115,20 @@ export class CodexAccountCooldownError extends Error { accountId: string; cooldownUntil: number; cooldownSource?: CodexCooldownSource; + quotaScope?: CodexQuotaScope; - constructor(accountId: string, cooldownUntil: number, cooldownSource?: CodexCooldownSource) { + constructor( + accountId: string, + cooldownUntil: number, + cooldownSource?: CodexCooldownSource, + quotaScope?: CodexQuotaScope, + ) { super("Selected Codex account is cooling down"); this.name = "CodexAccountCooldownError"; this.accountId = accountId; this.cooldownUntil = cooldownUntil; this.cooldownSource = cooldownSource; + this.quotaScope = quotaScope; } } @@ -127,7 +150,12 @@ export function cooldownAccountLabel(accountId: string): string { */ export function cooldownErrorMessage(err: CodexAccountCooldownError): string { const until = new Date(err.cooldownUntil).toISOString(); - return `Selected Codex account (${cooldownAccountLabel(err.accountId)}) is cooling down until ${until}` + const scope = err.quotaScope === "spark" + ? "Spark quota" + : err.quotaScope === "shared" + ? "shared native quota" + : null; + return `Selected Codex account (${cooldownAccountLabel(err.accountId)})${scope ? ` ${scope} is` : " is"} cooling down until ${until}` + ` (source: ${err.cooldownSource ?? "default"}).` + ` Run 'ocx account list openai' to find the id, then` + ` 'ocx account clear-cooldown openai ' to lift it, or switch accounts with 'ocx account use openai '.`; @@ -157,6 +185,8 @@ export function shouldMarkAccountNeedsReauthForCodexAuthFailure(cause: unknown): export interface ResolveCodexAuthContextOptions { excludeAccountId?: string; + /** Final native model selected for this request, used to select its quota group. */ + modelId?: string; } export async function resolveCodexAuthContext( @@ -171,16 +201,17 @@ export async function resolveCodexAuthContext( } reconcileMainCodexAccountRuntimeState(); const threadId = headers.get("x-codex-parent-thread-id"); + const quotaScope = codexQuotaScopeForModel(options.modelId); const resolution = options.excludeAccountId ? (() => { - const accountId = pickAlternateCodexAccount(config, options.excludeAccountId!); + const accountId = pickAlternateCodexAccount(config, options.excludeAccountId!, Date.now(), quotaScope); return accountId ? { status: "selected" as const, accountId } : { status: "none" as const }; })() - : resolveCodexAccountForThreadDetailed(threadId, config); + : resolveCodexAccountForThreadDetailed(threadId, config, Date.now(), quotaScope); if (resolution.status === "expired") throw new CodexThreadAffinityExpiredError(resolution.accountId); - const accountId = resolution.status === "selected" ? resolution.accountId : null; + let accountId = resolution.status === "selected" ? resolution.accountId : null; if (!accountId) throw new CodexPoolAuthenticationError(); // Lazy prime: if the selected account has no quota yet, the pool is likely // unprimed (dashboard never opened, or startup prime was blocked). Kick a @@ -194,15 +225,21 @@ export async function resolveCodexAuthContext( } // Snapshot (not just the deadline) so a refused request can report WHY it is cooled: // a literal Retry-After reads very differently to a user than a reset-derived guess. - const cooldown = getCodexAccountHealthSnapshot(accountId); + const cooldown = getCodexQuotaHealthSnapshot(accountId, quotaScope); const cooldownUntil = cooldown?.cooldownUntil; // A cooled-down account never sends traffic, so upstream recovery can never be // observed and the cooldown outlives the real limit. Admit one probe per // interval; its outcome decides whether the cooldown ends (#433). let probeLeaseId: string | undefined; + let probeQuotaScope: CodexQuotaScope | undefined; if (cooldownUntil) { - probeLeaseId = tryAcquireCodexQuotaProbeLease(accountId) ?? undefined; - if (!probeLeaseId) throw new CodexAccountCooldownError(accountId, cooldownUntil, cooldown?.cooldownSource); + probeQuotaScope = cooldown?.quotaScope; + probeLeaseId = probeQuotaScope + ? tryAcquireCodexQuotaScopeProbeLease(accountId, probeQuotaScope) ?? undefined + : tryAcquireCodexQuotaProbeLease(accountId) ?? undefined; + if (!probeLeaseId) { + throw new CodexAccountCooldownError(accountId, cooldownUntil, cooldown?.cooldownSource, cooldown?.quotaScope); + } } if (accountId === MAIN_CODEX_ACCOUNT_ID) { @@ -210,7 +247,8 @@ export async function resolveCodexAuthContext( const token = getMainAccountToken(); if (!token) { // Nothing will reach upstream, so give the probe back instead of burning it. - if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); + if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); + else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); throw new CodexPoolAuthenticationError(); } return { @@ -218,7 +256,9 @@ export async function resolveCodexAuthContext( accountId, accessToken: token.accessToken, chatgptAccountId: token.chatgptAccountId, + ...(quotaScope ? { quotaScope } : {}), ...(probeLeaseId ? { probeLeaseId } : {}), + ...(probeQuotaScope ? { probeQuotaScope } : {}), }; } @@ -230,10 +270,13 @@ export async function resolveCodexAuthContext( generation: token.generation, accessToken: token.accessToken, chatgptAccountId: token.chatgptAccountId, + ...(quotaScope ? { quotaScope } : {}), ...(probeLeaseId ? { probeLeaseId } : {}), + ...(probeQuotaScope ? { probeQuotaScope } : {}), }; } catch (cause) { - if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); + if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); + else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); if (shouldMarkAccountNeedsReauthForCodexAuthFailure(cause)) { markAccountNeedsReauth(accountId); } @@ -245,9 +288,9 @@ export function assertCodexAuthContextNotCooled(ctx: CodexAuthContext | undefine if (ctx?.kind !== "pool" && ctx?.kind !== "main-pool") return; // A context holding the probe lease was deliberately admitted through the cooldown. if (ctx.probeLeaseId) return; - const cooldown = getCodexAccountHealthSnapshot(ctx.accountId); + const cooldown = getCodexQuotaHealthSnapshot(ctx.accountId, ctx.quotaScope); if (cooldown?.cooldownUntil) { - throw new CodexAccountCooldownError(ctx.accountId, cooldown.cooldownUntil, cooldown.cooldownSource); + throw new CodexAccountCooldownError(ctx.accountId, cooldown.cooldownUntil, cooldown.cooldownSource, cooldown.quotaScope); } } diff --git a/src/codex/routing.ts b/src/codex/routing.ts index f127b12def..a54eb90104 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -33,7 +33,6 @@ export type CodexThreadResolution = | { status: "none" } | { status: "expired"; accountId: string }; -const threadAccountMap = new Map(); /** * Process-local cursor for automatic RR/fill-first (and quota-429 when not * sync-writing) picks. Keeps unrelated `saveConfig` from persisting transient @@ -108,22 +107,76 @@ export const CODEX_THREAD_AFFINITY_MAX_ENTRIES = 2048; export const CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS = 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>(); export type CodexUpstreamOutcome = number | "connect_error" | "timeout"; export type CodexUpstreamOutcomeClass = "success" | "credential" | "quota" | "transient" | "caller" | "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" | "spark"; + +/** + * 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 Spark failover cannot displace the same thread's Terra/Luna + * account (and vice versa). + */ +type ThreadAffinityScope = CodexQuotaScope | "legacy"; +const LEGACY_THREAD_AFFINITY_SCOPE = "legacy" as const; +const threadAccountMap = new Map>(); + +const NATIVE_MODEL_QUOTA_SCOPES: Readonly> = { + "gpt-5.3-codex-spark": "spark", +}; + +// A thread can have one legacy binding plus one binding for each known scope. +// This upper-bound guard avoids an exact map scan until it can be over capacity. +const MAX_THREAD_AFFINITY_SCOPES = new Set([ + LEGACY_THREAD_AFFINITY_SCOPE, + "shared", + ...Object.values(NATIVE_MODEL_QUOTA_SCOPES), +]).size; + +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; + /** 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; /** * 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} @@ -142,18 +195,23 @@ export function clearThreadAccountMap(): void { } export function clearThreadAccountMapForAccount(accountId: string): void { - for (const [threadId, entry] of threadAccountMap) { - if (entry.accountId === accountId) threadAccountMap.delete(threadId); + for (const [threadId, affinities] of threadAccountMap) { + for (const [scope, entry] of affinities) { + if (entry.accountId === accountId) affinities.delete(scope); + } + if (affinities.size === 0) threadAccountMap.delete(threadId); } } export function clearCodexUpstreamHealth(): void { upstreamHealth.clear(); + quotaScopedHealth.clear(); runtimeActiveCodexAccountId = undefined; } export function clearCodexUpstreamHealthForAccount(accountId: string): void { upstreamHealth.delete(accountId); + quotaScopedHealth.delete(accountId); } export function getCodexUpstreamHealth( @@ -162,6 +220,26 @@ export function getCodexUpstreamHealth( 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; @@ -276,7 +354,10 @@ export function tryAcquireCodexQuotaProbeLease(accountId: string, now = Date.now /** Side-effect-free check mirroring {@link tryAcquireCodexQuotaProbeLease} eligibility. */ export function canAcquireCodexQuotaProbeLease(accountId: string, now = Date.now()): boolean { - const health = upstreamHealth.get(accountId); + 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; @@ -286,6 +367,24 @@ export function canAcquireCodexQuotaProbeLease(accountId: string, now = Date.now return now - origin >= CODEX_QUOTA_PROBE_INTERVAL_MS; } +/** 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; +} + /** * 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. @@ -296,6 +395,18 @@ export function releaseCodexQuotaProbeLease(accountId: string, leaseId: string, 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 @@ -341,6 +452,11 @@ export function resetCodexRoutingForManualSelection(accountId: string): void { // 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); + } + } const current = upstreamHealth.get(accountId); if (!current) return; const preserved = preservedCooldownFields(current); @@ -367,6 +483,33 @@ export function getCodexAccountHealthSnapshot(accountId: string, now = Date.now( }; } +/** + * 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; } @@ -390,24 +533,41 @@ export function isCodexAccountInCooldown(accountId: string, now = Date.now()): b * Returns false when the account carried no live cooldown (already expired or never set). */ export function clearCodexAccountCooldown(accountId: string, now = Date.now()): boolean { - const health = upstreamHealth.get(accountId); - if (!health) return false; - const cooldownUntil = health.cooldownUntil; - if (typeof cooldownUntil !== "number" || !Number.isFinite(cooldownUntil) || cooldownUntil <= now) return false; - const { - cooldownUntil: _until, - cooldownSince: _since, - cooldownSource: _source, - probeLeaseId: _leaseId, - probeLeaseGeneration: _leaseGeneration, - ...rest - } = health; - upstreamHealth.set(accountId, { - ...rest, - cooldownGeneration: (health.cooldownGeneration ?? 0) + 1, - lastProbeAt: now, - }); - return true; + const clear = (health: CodexUpstreamHealth): CodexUpstreamHealth | null => { + const cooldownUntil = health.cooldownUntil; + if (typeof cooldownUntil !== "number" || !Number.isFinite(cooldownUntil) || cooldownUntil <= now) return null; + const { + cooldownUntil: _until, + cooldownSince: _since, + cooldownSource: _source, + probeLeaseId: _leaseId, + probeLeaseGeneration: _leaseGeneration, + ...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 { @@ -421,12 +581,48 @@ export function isCodexAccountSoftAvoided(accountId: string, now = Date.now()): return getCodexAccountSoftAvoidUntil(accountId, now) !== null; } -function isCodexAccountSelectable(config: OcxConfig, accountId: string, now: number): boolean { - return !isCodexAccountInCooldown(accountId, now) +function isCodexAccountSelectable( + config: OcxConfig, + accountId: string, + now: number, + quotaScope?: CodexQuotaScope, +): boolean { + return getCodexQuotaHealthSnapshot(accountId, quotaScope, now) === null && !isCodexAccountSoftAvoided(accountId, now) && isCodexAccountUsable(config, accountId); } +function threadAffinityScope(quotaScope?: CodexQuotaScope): ThreadAffinityScope { + return quotaScope ?? LEGACY_THREAD_AFFINITY_SCOPE; +} + +function getThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): ThreadAffinityEntry | undefined { + return threadAccountMap.get(threadId)?.get(threadAffinityScope(quotaScope)); +} + +function deleteThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): void { + const affinities = threadAccountMap.get(threadId); + if (!affinities) return; + affinities.delete(threadAffinityScope(quotaScope)); + if (affinities.size === 0) threadAccountMap.delete(threadId); +} + +/** Remove only the matching failed account's affinities for one thread. */ +function deleteThreadAffinitiesForAccount(threadId: string, accountId: string): void { + const affinities = threadAccountMap.get(threadId); + if (!affinities) return; + for (const [scope, entry] of affinities) { + if (entry.accountId === accountId) affinities.delete(scope); + } + if (affinities.size === 0) threadAccountMap.delete(threadId); +} + +function threadAffinityEntryCount(): number { + let count = 0; + for (const affinities of threadAccountMap.values()) count += affinities.size; + return count; +} + function isThreadAffinityExpired(entry: ThreadAffinityEntry, now: number): boolean { return now - entry.lastUsedAt > CODEX_THREAD_AFFINITY_IDLE_TTL_MS; } @@ -437,45 +633,66 @@ function isThreadAffinityGenerationLive(entry: ThreadAffinityEntry): boolean { } function pruneExpiredThreadAffinities(now: number): void { - for (const [threadId, entry] of threadAccountMap) { - if (isThreadAffinityExpired(entry, now)) threadAccountMap.delete(threadId); + for (const [threadId, affinities] of threadAccountMap) { + for (const [scope, entry] of affinities) { + if (isThreadAffinityExpired(entry, now)) affinities.delete(scope); + } + if (affinities.size === 0) threadAccountMap.delete(threadId); } } function pruneLruThreadAffinities(): void { - while (threadAccountMap.size > CODEX_THREAD_AFFINITY_MAX_ENTRIES) { + if (threadAccountMap.size * MAX_THREAD_AFFINITY_SCOPES <= 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; - for (const [threadId, entry] of threadAccountMap) { - if (entry.lastUsedAt < oldestLastUsedAt) { - oldestThreadId = threadId; - oldestLastUsedAt = entry.lastUsedAt; + for (const [threadId, affinities] of threadAccountMap) { + for (const [scope, entry] of affinities) { + if (entry.lastUsedAt < oldestLastUsedAt) { + oldestThreadId = threadId; + oldestScope = scope; + oldestLastUsedAt = entry.lastUsedAt; + } } } - if (!oldestThreadId) return; - threadAccountMap.delete(oldestThreadId); + if (!oldestThreadId || !oldestScope) return; + deleteThreadAffinity(oldestThreadId, oldestScope === LEGACY_THREAD_AFFINITY_SCOPE ? undefined : oldestScope); } } -function bindThreadAffinity(threadId: string, accountId: string, now: number): void { +function bindThreadAffinity( + threadId: string, + accountId: string, + now: number, + quotaScope?: CodexQuotaScope, +): void { 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 previous = threadAccountMap.get(threadId); - threadAccountMap.set(threadId, { + const scope = threadAffinityScope(quotaScope); + 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, }); + threadAccountMap.set(threadId, affinities); pruneLruThreadAffinities(); } -function getEligiblePoolAccounts(config: OcxConfig, excludeId?: string, now = Date.now()): string[] { +function getEligiblePoolAccounts( + config: OcxConfig, + excludeId?: string, + now = Date.now(), + quotaScope?: CodexQuotaScope, +): string[] { const ids = (config.codexAccounts ?? []) .filter(account => !account.isMain && account.id !== excludeId && !isAccountNeedsReauth(account.id)) - .filter(account => !isCodexAccountInCooldown(account.id, now)) + .filter(account => getCodexQuotaHealthSnapshot(account.id, quotaScope, now) === null) .filter(account => !isCodexAccountSoftAvoided(account.id, now)) .filter(account => isCodexAccountUsable(config, account.id)) .map(account => account.id); @@ -484,7 +701,7 @@ function getEligiblePoolAccounts(config: OcxConfig, excludeId?: string, now = Da if ( excludeId !== MAIN_CODEX_ACCOUNT_ID && !isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID) - && !isCodexAccountInCooldown(MAIN_CODEX_ACCOUNT_ID, now) + && getCodexQuotaHealthSnapshot(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) === null && !isCodexAccountSoftAvoided(MAIN_CODEX_ACCOUNT_ID, now) && isCodexAccountUsable(config, MAIN_CODEX_ACCOUNT_ID) ) { @@ -493,8 +710,12 @@ function getEligiblePoolAccounts(config: OcxConfig, excludeId?: string, now = Da return ids; } -function listEligibleCodexAccountIds(config: OcxConfig, now: number): string[] { - return getEligiblePoolAccounts(config, undefined, now); +function listEligibleCodexAccountIds( + config: OcxConfig, + now: number, + quotaScope?: CodexQuotaScope, +): string[] { + return getEligiblePoolAccounts(config, undefined, now, quotaScope); } function stickyLimitForConfig(config: OcxConfig): number { @@ -514,8 +735,12 @@ function isActiveUnderFillFirstThreshold(config: OcxConfig, accountId: string): * 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): string | null { - const eligible = listEligibleCodexAccountIds(config, now); +function pickFillFirstCodexAccount( + config: OcxConfig, + now: number, + quotaScope?: CodexQuotaScope, +): string | null { + const eligible = listEligibleCodexAccountIds(config, now, quotaScope); if (eligible.length === 0) return null; const active = getEffectiveActiveCodexAccountId(config); @@ -589,31 +814,33 @@ function pickUnboundStrategyAccount( threadId: string | null, now: number, commit: boolean, + quotaScope?: CodexQuotaScope, ): string | null { const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); if (strategy === "quota") return null; + const poolKey = codexPoolKeyForScope(quotaScope); let picked: string | null = null; if (strategy === "round-robin") { - const eligible = listEligibleCodexAccountIds(config, now); + const eligible = listEligibleCodexAccountIds(config, now, quotaScope); const limit = stickyLimitForConfig(config); if (!commit) { - return peekRoundRobinAccount(POOL_KEY_CODEX, eligible, limit); + return peekRoundRobinAccount(poolKey, eligible, limit); } - picked = pickRoundRobinAccount(POOL_KEY_CODEX, eligible, limit); + picked = pickRoundRobinAccount(poolKey, eligible, limit); if (!picked) return null; - rememberActiveCodexAccount(config, picked); - if (threadId) bindThreadAffinity(threadId, picked, now); - notePoolRotationSuccess(POOL_KEY_CODEX, picked, limit); + if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, picked); + if (threadId) bindThreadAffinity(threadId, picked, now, quotaScope); + notePoolRotationSuccess(poolKey, picked, limit); return picked; } if (strategy === "fill-first") { - picked = pickFillFirstCodexAccount(config, now); + picked = pickFillFirstCodexAccount(config, now, quotaScope); if (!picked) return null; if (commit) { - rememberActiveCodexAccount(config, picked); - if (threadId) bindThreadAffinity(threadId, picked, now); + if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, picked); + if (threadId) bindThreadAffinity(threadId, picked, now, quotaScope); } return picked; } @@ -626,10 +853,16 @@ export function getPoolAccountPlan(config: OcxConfig, accountId: string): string return (config.codexAccounts ?? []).find(account => !account.isMain && account.id === accountId)?.plan; } -function pickLowerUsageAccount(config: OcxConfig, active: string, activeUsage: number, now: number): string { +function pickLowerUsageAccount( + config: OcxConfig, + active: string, + activeUsage: number, + now: number, + quotaScope?: CodexQuotaScope, +): string { let best = active; let bestUsage = activeUsage; - for (const id of getEligiblePoolAccounts(config, active, now)) { + for (const id of getEligiblePoolAccounts(config, active, now, quotaScope)) { const usage = computeCodexUsageScore(getAccountQuota(id), getPoolAccountPlan(config, id)); if (usage < bestUsage) { best = id; @@ -639,10 +872,15 @@ function pickLowerUsageAccount(config: OcxConfig, active: string, activeUsage: n return best; } -export function pickLowestUsageCodexAccount(config: OcxConfig, excludeId?: string, now = Date.now()): string | null { +export function pickLowestUsageCodexAccount( + config: OcxConfig, + excludeId?: string, + now = Date.now(), + quotaScope?: CodexQuotaScope, +): string | null { let best: string | null = null; let bestUsage = Number.POSITIVE_INFINITY; - for (const id of getEligiblePoolAccounts(config, excludeId, now)) { + for (const id of getEligiblePoolAccounts(config, excludeId, now, quotaScope)) { const usage = computeCodexUsageScore(getAccountQuota(id), getPoolAccountPlan(config, id)); if (usage < bestUsage) { best = id; @@ -661,17 +899,18 @@ export function pickAlternateCodexAccount( config: OcxConfig, excludeId: string, now = Date.now(), + quotaScope?: CodexQuotaScope, ): string | null { const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); if (strategy === "round-robin") { - const eligible = listEligibleCodexAccountIds(config, now).filter(id => id !== excludeId); - return pickRoundRobinAccount(POOL_KEY_CODEX, eligible, stickyLimitForConfig(config)); + const eligible = listEligibleCodexAccountIds(config, now, quotaScope).filter(id => id !== excludeId); + return pickRoundRobinAccount(codexPoolKeyForScope(quotaScope), eligible, stickyLimitForConfig(config)); } if (strategy === "fill-first") { - const eligible = listEligibleCodexAccountIds(config, now).filter(id => id !== excludeId); + const eligible = listEligibleCodexAccountIds(config, now, quotaScope).filter(id => id !== excludeId); return pickNextFillFirstCodexAccount(config, excludeId, eligible, now); } - return pickLowestUsageCodexAccount(config, excludeId, now); + return pickLowestUsageCodexAccount(config, excludeId, now, quotaScope); } /** Effective active: automatic runtime cursor, else operator/persisted selection. */ @@ -708,7 +947,12 @@ function isUnknownUsage(usage: number): boolean { return usage >= CODEX_UNKNOWN_USAGE_SCORE; } -function applyQuotaAutoSwitch(config: OcxConfig, active: string, now: number): string { +function applyQuotaAutoSwitch( + config: OcxConfig, + active: string, + now: number, + quotaScope?: CodexQuotaScope, +): string { const threshold = config.autoSwitchThreshold ?? 80; if (threshold <= 0) return active; const quota = getAccountQuota(active); @@ -717,9 +961,9 @@ function applyQuotaAutoSwitch(config: OcxConfig, active: string, now: number): s // 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); + const best = pickLowerUsageAccount(config, active, activeUsage, now, quotaScope); if (best !== active) { - setActiveCodexAccount(config, best); + if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, best); return best; } @@ -748,8 +992,9 @@ export function resolveCodexAccountForThread( threadId: string | null, config: OcxConfig, now = Date.now(), + quotaScope?: CodexQuotaScope, ): string | null { - const resolution = resolveCodexAccountForThreadDetailed(threadId, config, now); + const resolution = resolveCodexAccountForThreadDetailed(threadId, config, now, quotaScope); return resolution.status === "selected" ? resolution.accountId : null; } @@ -765,13 +1010,14 @@ export function previewCodexAccountForRequest( threadId: string | null, config: OcxConfig, now = Date.now(), + quotaScope?: CodexQuotaScope, ): string | null { - if (threadId && threadAccountMap.has(threadId)) { - const entry = threadAccountMap.get(threadId)!; + const entry = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; + if (threadId && entry) { if ( !isThreadAffinityExpired(entry, now) && isThreadAffinityGenerationLive(entry) - && isCodexAccountSelectable(config, entry.accountId, now) + && isCodexAccountSelectable(config, entry.accountId, now, quotaScope) && !shouldFailover(config, entry.accountId, now) ) { // Quota strategy only: non-quota strategies keep affinity for ongoing threads @@ -785,7 +1031,7 @@ export function previewCodexAccountForRequest( getPoolAccountPlan(config, entry.accountId), ); if (!isUnknownUsage(usage) && usage >= threshold) { - const best = pickLowerUsageAccount(config, entry.accountId, usage, now); + const best = pickLowerUsageAccount(config, entry.accountId, usage, now, quotaScope); if (best !== entry.accountId) return best; } } @@ -795,15 +1041,15 @@ export function previewCodexAccountForRequest( // Stale/unusable affinity is ignored for preview (no map mutation). } - const strategyPick = pickUnboundStrategyAccount(config, threadId, now, false); + const strategyPick = pickUnboundStrategyAccount(config, threadId, now, false, quotaScope); if (strategyPick) return strategyPick; let active = getEffectiveActiveCodexAccountId(config) ?? null; if (!active) { - return pickLowestUsageCodexAccount(config, undefined, now); + return pickLowestUsageCodexAccount(config, undefined, now, quotaScope); } - if (!isCodexAccountSelectable(config, active, now)) { - const fallback = pickLowestUsageCodexAccount(config, active, now); + if (!isCodexAccountSelectable(config, active, now, quotaScope)) { + const fallback = pickLowestUsageCodexAccount(config, active, now, quotaScope); if (fallback) active = fallback; else if (hasConfiguredPoolAccount(config, active)) return active; else return null; @@ -813,17 +1059,17 @@ export function previewCodexAccountForRequest( if (threshold > 0) { const usage = computeCodexUsageScore(getAccountQuota(active), getPoolAccountPlan(config, active)); if (!isUnknownUsage(usage) && usage >= threshold) { - active = pickLowerUsageAccount(config, active, usage, now); + active = pickLowerUsageAccount(config, active, usage, now, quotaScope); } } if (shouldFailover(config, active, now)) { - const best = pickLowestUsageCodexAccount(config, active, now); + const best = pickLowestUsageCodexAccount(config, active, now, quotaScope); if (best) active = best; } if (!isCodexAccountUsable(config, active)) { return hasConfiguredPoolAccount(config, active) ? active : null; } - if (isCodexAccountInCooldown(active, now)) { + if (getCodexQuotaHealthSnapshot(active, quotaScope, now)) { return hasConfiguredPoolAccount(config, active) ? active : null; } return active; @@ -833,16 +1079,17 @@ export function resolveCodexAccountForThreadDetailed( threadId: string | null, config: OcxConfig, now = Date.now(), + quotaScope?: CodexQuotaScope, ): CodexThreadResolution { - if (threadId && threadAccountMap.has(threadId)) { - const entry = threadAccountMap.get(threadId)!; + const entry = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; + if (threadId && entry) { if (isThreadAffinityExpired(entry, now)) { - threadAccountMap.delete(threadId); + deleteThreadAffinity(threadId, quotaScope); return { status: "expired", accountId: entry.accountId }; } if ( isThreadAffinityGenerationLive(entry) - && isCodexAccountSelectable(config, entry.accountId, now) + && isCodexAccountSelectable(config, entry.accountId, now, quotaScope) // Affined threads must leave a failing account once the streak trips failover // (soft-avoid covers the first-hit case; this catches post-avoid residual streaks). && !shouldFailover(config, entry.accountId, now) @@ -869,10 +1116,10 @@ export function resolveCodexAccountForThreadDetailed( if (overThreshold || now - entry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) { entry.lastReevalAt = now; if (overThreshold) { - const best = pickLowerUsageAccount(config, entry.accountId, usage, now); + const best = pickLowerUsageAccount(config, entry.accountId, usage, now, quotaScope); if (best !== entry.accountId) { - setActiveCodexAccount(config, best); - bindThreadAffinity(threadId, best, now); // rebinds + resets clocks + if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, best); + bindThreadAffinity(threadId, best, now, quotaScope); // rebinds + resets clocks return { status: "selected", accountId: best }; } } @@ -880,23 +1127,23 @@ export function resolveCodexAccountForThreadDetailed( } return { status: "selected", accountId: entry.accountId }; } - threadAccountMap.delete(threadId); + deleteThreadAffinity(threadId, quotaScope); } - const strategyPick = pickUnboundStrategyAccount(config, threadId, now, true); + const strategyPick = pickUnboundStrategyAccount(config, threadId, now, true, quotaScope); if (strategyPick) return { status: "selected", accountId: strategyPick }; let active = getEffectiveActiveCodexAccountId(config); if (!active) { - const selected = pickLowestUsageCodexAccount(config, undefined, now); + const selected = pickLowestUsageCodexAccount(config, undefined, now, quotaScope); if (!selected) return { status: "none" }; - setActiveCodexAccount(config, selected); + if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, selected); active = selected; } - if (!isCodexAccountSelectable(config, active, now)) { - const fallback = pickLowestUsageCodexAccount(config, active, now); + if (!isCodexAccountSelectable(config, active, now, quotaScope)) { + const fallback = pickLowestUsageCodexAccount(config, active, now, quotaScope); if (fallback) { - setActiveCodexAccount(config, fallback); + if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, fallback); active = fallback; } else if (hasConfiguredPoolAccount(config, active)) { return { status: "selected", accountId: active }; @@ -904,15 +1151,15 @@ export function resolveCodexAccountForThreadDetailed( return { status: "none" }; } } - active = applyQuotaAutoSwitch(config, active, now); + active = applyQuotaAutoSwitch(config, active, now, quotaScope); active = applyFailureFailover(config, active, now); if (!isCodexAccountUsable(config, active)) { return hasConfiguredPoolAccount(config, active) ? { status: "selected", accountId: active } : { status: "none" }; } - if (isCodexAccountInCooldown(active, now)) { + if (getCodexQuotaHealthSnapshot(active, quotaScope, now)) { return hasConfiguredPoolAccount(config, active) ? { status: "selected", accountId: active } : { status: "none" }; } - if (threadId) bindThreadAffinity(threadId, active, now); + if (threadId) bindThreadAffinity(threadId, active, now, quotaScope); return { status: "selected", accountId: active }; } @@ -925,7 +1172,18 @@ export function recordCodexUpstreamOutcome( if (!accountId) return; const now = meta.now ?? Date.now(); const outcomeClass = classifyCodexUpstreamOutcome(outcome); + const quotaScope = codexQuotaScopeForModel(meta.modelId); if (outcomeClass === "success") { + const scopedProbe = meta.probeQuotaScope + ? scopedHealthFor(accountId, meta.probeQuotaScope) + : undefined; + if (scopedProbe && meta.probeQuotaScope) { + if (scopedProbe.cooldownUntil && probeMayClearCooldown(scopedProbe, meta)) { + deleteScopedHealth(accountId, meta.probeQuotaScope); + } else if (ownsProbeLease(scopedProbe, meta)) { + setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); + } + } const current = upstreamHealth.get(accountId); const cooldownUntil = getCodexAccountCooldownUntil(accountId, now); // A leased probe that is still on its own cooldown generation proves the @@ -961,6 +1219,12 @@ export function recordCodexUpstreamOutcome( // 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 scopedProbe = meta.probeQuotaScope + ? scopedHealthFor(accountId, meta.probeQuotaScope) + : undefined; + if (scopedProbe && meta.probeQuotaScope && ownsProbeLease(scopedProbe, meta)) { + setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); + } if (ownsProbeLease(current, meta)) { upstreamHealth.set(accountId, withProbeLeaseReleased(current!, now)); } @@ -976,14 +1240,65 @@ export function recordCodexUpstreamOutcome( lastFailureStatus, lastFailureAt: now, }); + quotaScopedHealth.delete(accountId); markAccountNeedsReauth(accountId); clearThreadAccountMapForAccount(accountId); return; } if (outcomeClass === "quota") { - const prior = upstreamHealth.get(accountId); const { until, source } = computeQuotaCooldown(meta); + // A reset timestamp is an advisory quota-window announcement. When the + // selected native model belongs to a confirmed independent group, preserve + // it there so a different group (Spark versus the shared native quota) can + // still reach upstream. Explicit Retry-After/default 429s remain account-wide. + if (source === "reset-derived" && quotaScope) { + const prior = scopedHealthFor(accountId, quotaScope); + const cooldownGeneration = (prior?.cooldownGeneration ?? 0) + 1; + const ownsLease = meta.probeQuotaScope === quotaScope && ownsProbeLease(prior, meta); + setScopedHealth(accountId, quotaScope, { + consecutiveFailures: 0, + lastFailureStatus, + lastFailureAt: now, + cooldownUntil: until, + cooldownSince: now, + cooldownSource: source, + cooldownGeneration, + ...(ownsLease + ? { lastProbeAt: now } + : { + ...(prior?.probeLeaseId !== undefined ? { probeLeaseId: prior.probeLeaseId } : {}), + ...(prior?.probeLeaseGeneration !== undefined ? { probeLeaseGeneration: prior.probeLeaseGeneration } : {}), + ...(prior?.lastProbeAt !== undefined ? { lastProbeAt: prior.lastProbeAt } : {}), + }), + }); + // The shared native scope is the existing account-wide native behavior: + // threads must leave it and new requests should prefer an eligible account. + // Spark remains isolated so a same-account Terra/Luna combo fallback can run. + if (quotaScope === "shared") { + clearThreadAccountMapForAccount(accountId); + notePoolRotationFailure(POOL_KEY_CODEX, accountId); + if (getEffectiveActiveCodexAccountId(config) === accountId) { + // Same-request 429 retry already picked via excludeAccountId — reuse it so + // round-robin does not advance the ring a second time. + const reused = meta.promoteAccountId && meta.promoteAccountId !== accountId + ? meta.promoteAccountId + : null; + const fallback = reused ?? pickAlternateCodexAccount(config, accountId, now, quotaScope); + if (fallback) promoteActiveCodexAccount(config, fallback); + } + } + return; + } + + // A scoped probe that received an account-wide throttle is no longer live. + const scopedProbe = meta.probeQuotaScope + ? scopedHealthFor(accountId, meta.probeQuotaScope) + : undefined; + if (scopedProbe && meta.probeQuotaScope && ownsProbeLease(scopedProbe, meta)) { + setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); + } + const prior = upstreamHealth.get(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; @@ -1022,6 +1337,12 @@ export function recordCodexUpstreamOutcome( // transient (connect_error / timeout / 5xx) const current = upstreamHealth.get(accountId); + const scopedProbe = meta.probeQuotaScope + ? scopedHealthFor(accountId, meta.probeQuotaScope) + : undefined; + if (scopedProbe && meta.probeQuotaScope && ownsProbeLease(scopedProbe, meta)) { + setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); + } // A transient failure concludes an owning probe; an unrelated 5xx must not // consume someone else's live lease or drop hard-cooldown bookkeeping (#433). const transientBase = ownsProbeLease(current, meta) ? withProbeLeaseReleased(current!, now) : current; @@ -1055,8 +1376,7 @@ export function recordCodexUpstreamOutcome( // must not delete a newer healthy binding to account B (race: T→A, A fails, // T→B, late A failure must not delete B's mapping). if (failoverReady && meta.threadId) { - const bound = threadAccountMap.get(meta.threadId); - if (bound?.accountId === accountId) threadAccountMap.delete(meta.threadId); + deleteThreadAffinitiesForAccount(meta.threadId, accountId); } // Once the account is past the failover streak, clear every thread still pinned // to it — matching 429 affinity behavior so "continue" cannot stay on a bad peer. diff --git a/src/server/responses/compact.ts b/src/server/responses/compact.ts index a567c5d833..3d536afd17 100644 --- a/src/server/responses/compact.ts +++ b/src/server/responses/compact.ts @@ -51,6 +51,7 @@ import { isCodexAuthContextUsable, resolveCodexAuthContext, codexProbeLeaseId, + codexProbeQuotaScope, type CodexAuthContext, } from "../../codex/auth-context"; import { @@ -217,7 +218,7 @@ export async function handleResponsesCompact( const headers = new Headers({ "content-type": "application/json" }); try { if (route.codexAccountMode) { - authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode); + authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode, { modelId: selectedModelId }); const selected = headersForCodexAuthContext(req.headers, authCtx); compactProvider = applyCodexAuthContextToProvider(route.provider, authCtx, route.codexAccountMode); for (const name of FORWARD_HEADERS) { @@ -255,12 +256,17 @@ export async function handleResponsesCompact( const compactUrl = `${base}/responses/compact`; const compactThreadId = req.headers.get("x-codex-parent-thread-id"); const connectMs = config.connectTimeoutMs ?? 200_000; - const recordCompactPoolOutcome = (outcome: CodexUpstreamOutcome, meta: { retryAfter?: string | null } = {}) => { + const recordCompactPoolOutcome = ( + outcome: CodexUpstreamOutcome, + meta: { retryAfter?: string | null; resetAt?: unknown | unknown[] } = {}, + ) => { if (!usesCodexForwardPoolAuth(authCtx, route.provider)) return; recordCodexUpstreamOutcome(config, authCtx.accountId, outcome, { ...meta, threadId: compactThreadId, + modelId: selectedModelId, probeLeaseId: codexProbeLeaseId(authCtx), + probeQuotaScope: codexProbeQuotaScope(authCtx), }); }; let upstream: Response; @@ -283,28 +289,30 @@ export async function handleResponsesCompact( { abortSignal: req.signal, label: safeHostLabel(compactUrl) }, ); } catch (err) { - if (req.signal.aborted) return formatErrorResponse(499, "client_cancelled", "Client cancelled compact request"); + if (req.signal.aborted) { + recordCompactPoolOutcome(499); + return formatErrorResponse(499, "client_cancelled", "Client cancelled compact request"); + } const outcome = err instanceof Error && err.name === "TimeoutError" ? "timeout" : "connect_error"; recordCompactPoolOutcome(outcome); return formatErrorResponse(502, "upstream_error", "Failed to connect to compact upstream"); } const retryAfter = upstream.headers.get("retry-after"); + const resetAt = [ + upstream.headers.get("x-codex-primary-reset-at"), + upstream.headers.get("x-codex-secondary-reset-at"), + upstream.headers.get("x-codex-tertiary-reset-at"), + ].filter(Boolean); const buffered = await bufferCompactResponse(upstream, req.signal); // Record pool health only after the body is fully delivered (or definitively failed). // A premature 200 would clear soft-avoid while the client still sees a buffer 502. if (buffered.status === 499) { + recordCompactPoolOutcome(499); return buffered; } - if (upstream.ok && buffered.status >= 500) { - // The upstream account returned 200 — it is healthy. The buffering failure - // (oversized body exceeding COMPACT_RESPONSE_MAX_BYTES, or a rare mid-read - // reset on a small JSON payload) is a local proxy issue, not account flakiness. - // Record the upstream status so a deterministic payload-size limit does not - // soft-avoid a healthy account and rotate a thread for 30s. - recordCompactPoolOutcome(upstream.status, { retryAfter }); - } else { - recordCompactPoolOutcome(upstream.status, { retryAfter }); - } + // Always record the real upstream status: a local buffering failure after a + // 200 upstream response must not soft-avoid a healthy account or rotate a thread. + recordCompactPoolOutcome(upstream.status, { retryAfter, resetAt }); return buffered; } diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 819ce34275..74ab53725c 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -67,6 +67,7 @@ import { isCodexAuthContextUsable, resolveCodexAuthContext, codexProbeLeaseId, + codexProbeQuotaScope, releaseCodexAuthContextProbeLease, stripCodexRuntimeProviderFields, type CodexAuthContext, @@ -320,7 +321,7 @@ async function retryCodexPoolOnAlternateAccount( req.headers, config, "pool", - { excludeAccountId: firstAuthCtx.accountId }, + { excludeAccountId: firstAuthCtx.accountId, modelId: route.modelId }, ); } catch (error) { if ( @@ -342,7 +343,9 @@ async function retryCodexPoolOnAlternateAccount( recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, outcomeStatus, { ...quotaMeta, threadId: req.headers.get("x-codex-parent-thread-id"), + modelId: route.modelId, probeLeaseId: codexProbeLeaseId(firstAuthCtx), + probeQuotaScope: codexProbeQuotaScope(firstAuthCtx), // Retry already advanced the RR ring via excludeAccountId — reuse for promotion. ...(retryAuthCtx.accountId ? { promoteAccountId: retryAuthCtx.accountId } : {}), }); @@ -403,6 +406,7 @@ export function codexForwardTerminalOutcomeRecorder( config: OcxConfig, authCtx: CodexAuthContext, provider: OcxProviderConfig, + modelId?: string, logCtx?: RequestLogContext, threadId?: string | null, ): ((status: ResponsesTerminalStatus, httpStatusOverride?: number) => void) | undefined { @@ -414,7 +418,9 @@ export function codexForwardTerminalOutcomeRecorder( // prior soft-avoid so a healthy account isn't stuck avoided. recordCodexUpstreamOutcome(config, authCtx.accountId, 200, { threadId, + modelId, probeLeaseId: codexProbeLeaseId(authCtx), + probeQuotaScope: codexProbeQuotaScope(authCtx), }); return; } @@ -432,7 +438,9 @@ export function codexForwardTerminalOutcomeRecorder( : (httpStatusOverride ?? logCtx?.terminalHttpStatus ?? 502); recordCodexUpstreamOutcome(config, authCtx.accountId, outcome, { threadId, + modelId, probeLeaseId: codexProbeLeaseId(authCtx), + probeQuotaScope: codexProbeQuotaScope(authCtx), }); }; } @@ -668,7 +676,7 @@ async function resolveResponsesCodexAuth( if (route.codexAccountMode === "direct") validateForwardAdmissionCredential(req.headers, config); let authCtx: CodexAuthContext; if (route.codexAccountMode) { - authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode); + authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode, { modelId: route.modelId }); options.onCodexAuthContextResolved?.(authCtx); } else { authCtx = { kind: "main", accountId: null }; @@ -1428,7 +1436,9 @@ export async function handleResponses( if (usesCodexForwardPoolAuth(authCtx, route.provider)) { recordCodexUpstreamOutcome(config, authCtx.accountId, outcome, { threadId: req.headers.get("x-codex-parent-thread-id"), + modelId: route.modelId, probeLeaseId: codexProbeLeaseId(authCtx), + probeQuotaScope: codexProbeQuotaScope(authCtx), }); } const msg = outcome === "timeout" @@ -1514,6 +1524,7 @@ export async function handleResponses( config, authCtx, route.provider, + route.modelId, logCtx, req.headers.get("x-codex-parent-thread-id"), ); @@ -1553,7 +1564,9 @@ export async function handleResponses( recordCodexUpstreamOutcome(config, authCtx.accountId, upstreamResponse.status, { ...quotaMeta, threadId: req.headers.get("x-codex-parent-thread-id"), + modelId: route.modelId, probeLeaseId: codexProbeLeaseId(authCtx), + probeQuotaScope: codexProbeQuotaScope(authCtx), }); } } diff --git a/structure/08_openai-provider-tiers.md b/structure/08_openai-provider-tiers.md index a0da134877..a0e403eb59 100644 --- a/structure/08_openai-provider-tiers.md +++ b/structure/08_openai-provider-tiers.md @@ -18,6 +18,11 @@ engine. Direct short-circuits that engine before pool state is read or mutated a current caller/main-login bearer. Neither mode may fall through to `openai-apikey`, and the API provider may not fall through to Codex-login credentials. +An explicit `Retry-After` or an unclassified quota 429 is account-wide. A reset-derived native-model +429 is advisory and remains within its confirmed quota group: `gpt-5.3-codex-spark` is separate from +the shared native group (including GPT-5.6 Terra/Luna). This allows a same-account combo to test an +independent quota without allowing fallbacks that share the exhausted quota. + ```text gpt-5.6-sol # openai; Pool or Direct follows the provider option openai-apikey/gpt-5.6-sol # OpenAI API key diff --git a/tests/codex-auth-context.test.ts b/tests/codex-auth-context.test.ts index cd7e7d9bc5..98c6dcfcc0 100644 --- a/tests/codex-auth-context.test.ts +++ b/tests/codex-auth-context.test.ts @@ -257,6 +257,160 @@ describe("Codex auth context", () => { } }); + test("reset-derived native cooldowns stay within their confirmed quota group", async () => { + const originalNow = Date.now; + const now = 1_800_000_000_000; + const cfg = config(); + const headers = new Headers({ authorization: "Bearer main_token" }); + saveCodexAccountCredential("pool-a", { + accessToken: "pool_token", + refreshToken: "pool_refresh", + expiresAt: now + 24 * 60 * 60_000, + chatgptAccountId: "pool_acc", + }); + try { + Date.now = () => now; + const resetAt = Math.floor((now + 4 * 24 * 60 * 60_000) / 1000); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + now, + resetAt, + modelId: "gpt-5.3-codex-spark", + }); + + // Spark owns a separate quota, so Terra can use the same account. + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.6-terra" })) + .resolves.toMatchObject({ kind: "pool", accountId: "pool-a" }); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.3-codex-spark" })) + .rejects.toBeInstanceOf(CodexAccountCooldownError); + + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + now, + resetAt, + modelId: "gpt-5.6-terra", + }); + + // Terra and Luna stay in the shared native quota group, while Spark keeps + // its independent cooldown instead of being overwritten by Terra's 429. + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.6-luna" })) + .rejects.toBeInstanceOf(CodexAccountCooldownError); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.3-codex-spark" })) + .rejects.toBeInstanceOf(CodexAccountCooldownError); + + // An explicit retry directive is still account-wide, regardless of the + // originating model's otherwise independent quota group. + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + now, + retryAfter: "60", + modelId: "gpt-5.3-codex-spark", + }); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.6-terra" })) + .rejects.toBeInstanceOf(CodexAccountCooldownError); + } finally { + Date.now = originalNow; + } + }); + + test("a scoped cooldown uses another account without moving the independent native scope", async () => { + const originalNow = Date.now; + const now = 1_800_000_000_000; + const cfg = config(); + cfg.codexAccounts?.push({ + id: "pool-b", + email: "pool-b@example.test", + isMain: false, + chatgptAccountId: "pool_b_acc", + }); + const headers = new Headers({ + authorization: "Bearer main_token", + "x-codex-parent-thread-id": "independent-scope-thread", + }); + for (const accountId of ["pool-a", "pool-b"]) { + saveCodexAccountCredential(accountId, { + accessToken: `${accountId}-token`, + refreshToken: `${accountId}-refresh`, + expiresAt: now + 24 * 60 * 60_000, + chatgptAccountId: `${accountId}-acc`, + }); + } + try { + Date.now = () => now; + // Establish the shared-scope binding first. The Spark fallback below must + // create a second binding rather than replacing this one. + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.6-terra" })) + .resolves.toMatchObject({ kind: "pool", accountId: "pool-a" }); + + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + now, + resetAt: Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000), + modelId: "gpt-5.3-codex-spark", + }); + + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.3-codex-spark" })) + .resolves.toMatchObject({ kind: "pool", accountId: "pool-b" }); + expect(cfg.activeCodexAccountId).toBe("pool-a"); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.6-terra" })) + .resolves.toMatchObject({ kind: "pool", accountId: "pool-a" }); + // This second Spark request proves routing retained the peer choice for + // the Spark affinity instead of relying on an auth-layer substitution. + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.3-codex-spark" })) + .resolves.toMatchObject({ kind: "pool", accountId: "pool-b" }); + } finally { + Date.now = originalNow; + } + }); + + test("a successful Spark recovery probe leaves the shared native cooldown intact", async () => { + const originalNow = Date.now; + const now = 1_800_000_000_000; + const cfg = config(); + const headers = new Headers({ authorization: "Bearer main_token" }); + saveCodexAccountCredential("pool-a", { + accessToken: "pool_token", + refreshToken: "pool_refresh", + expiresAt: now + 24 * 60 * 60_000, + chatgptAccountId: "pool_acc", + }); + try { + Date.now = () => now; + const resetAt = Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + now, + resetAt, + modelId: "gpt-5.3-codex-spark", + }); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + now, + resetAt, + modelId: "gpt-5.6-terra", + }); + + const probeAt = now + CODEX_QUOTA_PROBE_INTERVAL_MS; + Date.now = () => probeAt; + const sparkProbe = await resolveCodexAuthContext(headers, cfg, "pool", { + modelId: "gpt-5.3-codex-spark", + }); + expect(sparkProbe).toMatchObject({ + kind: "pool", + probeQuotaScope: "spark", + }); + + recordCodexUpstreamOutcome(cfg, "pool-a", 200, { + now: probeAt + 1, + modelId: "gpt-5.3-codex-spark", + probeLeaseId: (sparkProbe as { probeLeaseId?: string }).probeLeaseId, + probeQuotaScope: (sparkProbe as { probeQuotaScope?: "spark" }).probeQuotaScope, + }); + + Date.now = () => probeAt + 1; + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.3-codex-spark" })) + .resolves.toMatchObject({ kind: "pool", accountId: "pool-a" }); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { modelId: "gpt-5.6-luna" })) + .resolves.toMatchObject({ kind: "pool", probeQuotaScope: "shared" }); + } finally { + Date.now = originalNow; + } + }); + test("expired thread affinity fails closed instead of falling back to main auth", async () => { const now = 1_800_000_000_000; saveCodexAccountCredential("pool-a", { @@ -424,6 +578,20 @@ describe("cooldown error surface", () => { expect(message).toContain("account-…3c21"); }); + test("message identifies a model-scoped cooldown without implying an account-wide block", () => { + const err = new CodexAccountCooldownError( + "acct_9f3c21", + Date.parse("2026-07-26T10:00:00.000Z"), + "reset-derived", + "spark", + ); + + const message = cooldownErrorMessage(err); + + expect(message).toContain("Spark quota is cooling down"); + expect(message).not.toContain("Selected Codex account (account-…3c21) is cooling down"); + }); + test("the main login renders as the alias users actually type", () => { const err = new CodexAccountCooldownError(MAIN_CODEX_ACCOUNT_ID, Date.now() + 60_000); diff --git a/tests/codex-pool-rotation.test.ts b/tests/codex-pool-rotation.test.ts index 16a7f32eac..f16f84dfbe 100644 --- a/tests/codex-pool-rotation.test.ts +++ b/tests/codex-pool-rotation.test.ts @@ -156,6 +156,30 @@ describe("accountPoolStrategy new-session routing", () => { expect(new Set(picks).size).toBe(3); }); + test("an independent native scope does not advance the shared round-robin cursor", () => { + const config = makeThreeAccountConfig({ + accountPoolStrategy: "round-robin", + accountPoolStickyLimit: 1, + }); + const now = 1_800_000_000_000; + updateAccountQuota("a", 10); + updateAccountQuota("b", 10); + updateAccountQuota("c", 10); + + expect(resolveCodexAccountForThread(null, config, now, "shared")).toBe("a"); + recordCodexUpstreamOutcome(config, "a", 429, { + now: now + 1, + resetAt: Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000), + modelId: "gpt-5.3-codex-spark", + }); + + // Spark skips A in its own ring. The next shared request still takes B, + // as if the Spark selection had never advanced the shared ring. + expect(resolveCodexAccountForThread(null, config, now + 2, "spark")).toBe("b"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("a"); + expect(resolveCodexAccountForThread(null, config, now + 3, "shared")).toBe("b"); + }); + test("affinity still wins over round-robin", () => { const config = makeThreeAccountConfig({ accountPoolStrategy: "round-robin" }); updateAccountQuota("a", 10); @@ -395,6 +419,38 @@ describe("accountPoolStrategy new-session routing", () => { expect(pickAlternateCodexAccount(config, "a")).toBe("b"); }); + test("scoped reset 429s retain strategy while excluding only the affected native quota", () => { + const config = makeThreeAccountConfig({ + accountPoolStrategy: "fill-first", + activeCodexAccountId: "a", + autoSwitchThreshold: 80, + }); + updateAccountQuota("a", 10); + updateAccountQuota("b", 50); + updateAccountQuota("c", 5); + const now = Date.now(); + const resetAt = Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000); + + recordCodexUpstreamOutcome(config, "b", 429, { + now, + resetAt, + modelId: "gpt-5.3-codex-spark", + }); + + // Fill-first would normally advance a → b, but b is unavailable only to Spark. + expect(pickAlternateCodexAccount(config, "a", now + 1, "spark")).toBe("c"); + expect(pickAlternateCodexAccount(config, "a", now + 1, "shared")).toBe("b"); + + recordCodexUpstreamOutcome(config, "a", 429, { + now, + resetAt, + modelId: "gpt-5.6-terra", + promoteAccountId: "b", + }); + expect(getEffectiveActiveCodexAccountId(config)).toBe("b"); + expect(config.activeCodexAccountId).toBe("a"); + }); + test("RR 429 promotes via ring, not lowest usage", () => { const config = makeThreeAccountConfig({ accountPoolStrategy: "round-robin", diff --git a/tests/codex-routing.test.ts b/tests/codex-routing.test.ts index ce1e14cbb4..8e11154eed 100644 --- a/tests/codex-routing.test.ts +++ b/tests/codex-routing.test.ts @@ -16,6 +16,8 @@ import { clearThreadAccountMapForAccount, computeCodexUsageScore, getCodexAccountCooldownUntil, + getEffectiveActiveCodexAccountId, + getCodexQuotaHealthSnapshot, getCodexAccountSoftAvoidUntil, getCodexUpstreamHealth, isCodexAccountInCooldown, @@ -314,6 +316,47 @@ describe("codex routing", () => { expect(resolveCodexAccountForThread("quota-next", config)).toBe("b"); }); + test("shared native reset cooldown clears affinity and rotates the active account", () => { + const config = makeConfig(); + const now = 1_800_000_000_000; + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + expect(resolveCodexAccountForThread("shared-quota-existing", config, now)).toBe("a"); + + recordCodexUpstreamOutcome(config, "a", 429, { + now, + resetAt: Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000), + modelId: "gpt-5.6-terra", + }); + + expect(config.activeCodexAccountId).toBe("b"); + expect(resolveCodexAccountForThread("shared-quota-existing", config, now + 1)).toBe("b"); + }); + + test("independent native quota scopes keep separate thread affinities", () => { + const config = makeConfig(); + const now = 1_800_000_000_000; + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + + // A known shared-model request binds A for this thread. + expect(resolveCodexAccountForThread("scoped-thread", config, now, "shared")).toBe("a"); + + recordCodexUpstreamOutcome(config, "a", 429, { + now: now + 1, + resetAt: Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000), + modelId: "gpt-5.3-codex-spark", + }); + + // Spark sees its scoped cooldown and binds B without moving the global + // active account or the same thread's shared-scope affinity. + expect(resolveCodexAccountForThread("scoped-thread", config, now + 2, "spark")).toBe("b"); + expect(config.activeCodexAccountId).toBe("a"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("a"); + expect(resolveCodexAccountForThread("scoped-thread", config, now + 3, "shared")).toBe("a"); + expect(resolveCodexAccountForThread("scoped-thread", config, now + 4, "spark")).toBe("b"); + }); + test("2xx responses clear transient failures without clearing an unexpired cooldown", () => { const config = makeConfig(); const now = 1_800_000_000_000; @@ -469,6 +512,28 @@ describe("codex routing", () => { expect(health?.lastFailureStatus).toBe(429); }); + test("clearCodexAccountCooldown lifts every live native-model cooldown", () => { + const config = makeConfig(); + const now = 1_800_000_000_000; + const resetAt = Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000); + recordCodexUpstreamOutcome(config, "a", 429, { + now, + resetAt, + modelId: "gpt-5.3-codex-spark", + }); + recordCodexUpstreamOutcome(config, "a", 429, { + now, + resetAt, + modelId: "gpt-5.6-terra", + }); + + expect(getCodexQuotaHealthSnapshot("a", "spark", now + 1)).not.toBeNull(); + expect(getCodexQuotaHealthSnapshot("a", "shared", now + 1)).not.toBeNull(); + expect(clearCodexAccountCooldown("a", now + 1)).toBe(true); + expect(getCodexQuotaHealthSnapshot("a", "spark", now + 1)).toBeNull(); + expect(getCodexQuotaHealthSnapshot("a", "shared", now + 1)).toBeNull(); + }); + test("clearing is a no-op without a live cooldown", () => { const config = makeConfig(); const now = 1_800_000_000_000; @@ -771,6 +836,27 @@ describe("codex routing", () => { expect(resolveCodexAccountForThread("lru-0", config, now + CODEX_THREAD_AFFINITY_MAX_ENTRIES + 2)).toBe("b"); }); + test("thread affinity LRU cap includes legacy and native quota scopes", () => { + const config = makeConfig(); + const now = 1_800_000_000_000; + const threads = Math.floor(CODEX_THREAD_AFFINITY_MAX_ENTRIES / 3) + 1; + + for (let i = 0; i < threads; i++) { + const threadId = `scoped-lru-${i}`; + expect(resolveCodexAccountForThread(threadId, config, now + i * 3)).toBe("a"); + expect(resolveCodexAccountForThread(threadId, config, now + i * 3 + 1, "shared")).toBe("a"); + expect(resolveCodexAccountForThread(threadId, config, now + i * 3 + 2, "spark")).toBe("a"); + } + + // The oldest legacy entry was evicted, while the same thread's later + // shared and Spark entries remain independently affined to A. + config.activeCodexAccountId = "b"; + const after = now + threads * 3; + expect(resolveCodexAccountForThread("scoped-lru-0", config, after, "shared")).toBe("a"); + expect(resolveCodexAccountForThread("scoped-lru-0", config, after + 1, "spark")).toBe("a"); + expect(resolveCodexAccountForThread("scoped-lru-0", config, after + 2)).toBe("b"); + }); + test("generation mismatch invalidates a mapped thread before reuse", () => { const config = makeConfig(); updateAccountQuota("a", 10); diff --git a/tests/responses-compaction-routing.test.ts b/tests/responses-compaction-routing.test.ts index 4a6da3c0ed..cb50a15fd6 100644 --- a/tests/responses-compaction-routing.test.ts +++ b/tests/responses-compaction-routing.test.ts @@ -5,7 +5,20 @@ * fatals on a compaction turn that came back as an ordinary message. */ import { afterEach, describe, expect, test } from "bun:test"; -import { handleResponses } from "../src/server/responses"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { handleResponses, handleResponsesCompact } from "../src/server/responses"; +import { saveCodexAccountCredential } from "../src/codex/account-store"; +import { + CODEX_QUOTA_PROBE_INTERVAL_MS, + clearCodexUpstreamHealth, + recordCodexUpstreamOutcome, +} from "../src/codex/routing"; +import { + releaseCodexAuthContextProbeLease, + resolveCodexAuthContext, +} from "../src/codex/auth-context"; import { supportsNativeResponsesCompactEndpoint } from "../src/providers/openai-tiers"; import type { OcxConfig, OcxProviderConfig } from "../src/types"; @@ -30,11 +43,33 @@ function keyProviderConfig(overrides: Partial = {}): OcxConfi } as unknown as OcxConfig; } -function compactionRequest(body: Record): Request { +function nativePoolConfig(): OcxConfig { + return { + defaultProvider: "openai", + activeCodexAccountId: "pool-a", + providers: { + openai: { + adapter: "openai-responses", + baseUrl: "https://chatgpt.com/backend-api/codex", + authMode: "forward", + codexAccountMode: "pool", + }, + }, + codexAccounts: [{ + id: "pool-a", + email: "pool@example.test", + isMain: false, + chatgptAccountId: "pool_acc", + }], + } as OcxConfig; +} + +function compactionRequest(body: Record, signal?: AbortSignal): Request { return new Request("http://localhost/v1/responses", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(body), + signal, }); } @@ -112,6 +147,201 @@ describe("supportsNativeResponsesCompactEndpoint (#422)", () => { }); }); +describe("native Codex pool compaction", () => { + test("keeps a Spark reset cooldown separate from a Terra compact request (#590)", async () => { + const testDir = mkdtempSync(join(tmpdir(), "ocx-compact-scope-")); + const previousOpencodexHome = process.env.OPENCODEX_HOME; + const previousCodexHome = process.env.CODEX_HOME; + const config = nativePoolConfig(); + const resetAt = Math.floor((Date.now() + 4 * 24 * 60 * 60_000) / 1_000); + let sparkPhase = true; + try { + process.env.OPENCODEX_HOME = testDir; + process.env.CODEX_HOME = testDir; + clearCodexUpstreamHealth(); + saveCodexAccountCredential("pool-a", { + accessToken: "pool-access-token", + refreshToken: "pool-refresh-token", + expiresAt: Date.now() + 300_000, + chatgptAccountId: "pool_acc", + }); + globalThis.fetch = (async () => { + if (sparkPhase) { + return Response.json({ error: { message: "Spark quota exhausted" } }, { + status: 429, + headers: { "x-codex-primary-reset-at": String(resetAt) }, + }); + } + return jsonResponse(completedPayload("Terra compact response")); + }) as typeof fetch; + const spark = await handleResponsesCompact( + compactionRequest(baseCompactionBody({ model: "gpt-5.3-codex-spark" })), + config, + { model: "", provider: "" }, + ); + expect(spark.status).toBe(429); + + sparkPhase = false; + const cooledSpark = await handleResponsesCompact( + compactionRequest(baseCompactionBody({ model: "gpt-5.3-codex-spark" })), + config, + { model: "", provider: "" }, + ); + expect(cooledSpark.status).toBe(429); + + const terra = await handleResponsesCompact( + compactionRequest(baseCompactionBody({ model: "gpt-5.6-terra" })), + config, + { model: "", provider: "" }, + ); + expect(terra.status).toBe(200); + } finally { + globalThis.fetch = originalFetch; + clearCodexUpstreamHealth(); + rmSync(testDir, { recursive: true, force: true }); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + if (previousCodexHome === undefined) delete process.env.CODEX_HOME; + else process.env.CODEX_HOME = previousCodexHome; + } + }); + + test("a cancelled Spark recovery probe releases its compact lease (#590)", async () => { + const testDir = mkdtempSync(join(tmpdir(), "ocx-compact-probe-")); + const previousOpencodexHome = process.env.OPENCODEX_HOME; + const previousCodexHome = process.env.CODEX_HOME; + const originalNow = Date.now; + const now = 1_800_000_000_000; + const probeAt = now + CODEX_QUOTA_PROBE_INTERVAL_MS; + const config = nativePoolConfig(); + const abort = new AbortController(); + let markReadStarted!: () => void; + let releaseBody!: () => void; + const readStarted = new Promise(resolve => { markReadStarted = resolve; }); + const bodyReleased = new Promise(resolve => { releaseBody = resolve; }); + try { + process.env.OPENCODEX_HOME = testDir; + process.env.CODEX_HOME = testDir; + Date.now = () => now; + clearCodexUpstreamHealth(); + saveCodexAccountCredential("pool-a", { + accessToken: "pool-access-token", + refreshToken: "pool-refresh-token", + expiresAt: now + 30 * 60_000, + chatgptAccountId: "pool_acc", + }); + recordCodexUpstreamOutcome(config, "pool-a", 429, { + now, + resetAt: Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000), + modelId: "gpt-5.3-codex-spark", + }); + Date.now = () => probeAt; + globalThis.fetch = (async () => new Response(new ReadableStream({ + async pull(controller) { + markReadStarted(); + await bodyReleased; + controller.enqueue(new TextEncoder().encode("{\"partial\":")); + controller.close(); + }, + }), { status: 200, headers: { "content-type": "application/json" } })) as typeof fetch; + const pending = handleResponsesCompact( + compactionRequest(baseCompactionBody({ model: "gpt-5.3-codex-spark" }), abort.signal), + config, + { model: "", provider: "" }, + ); + await readStarted; + abort.abort(); + releaseBody(); + const cancelled = await pending; + expect(cancelled.status).toBe(499); + + Date.now = () => probeAt + CODEX_QUOTA_PROBE_INTERVAL_MS; + const nextProbe = await resolveCodexAuthContext( + new Headers({ authorization: "Bearer main-token" }), + config, + "pool", + { modelId: "gpt-5.3-codex-spark" }, + ); + expect(nextProbe).toMatchObject({ probeQuotaScope: "spark" }); + releaseCodexAuthContextProbeLease(nextProbe); + } finally { + Date.now = originalNow; + globalThis.fetch = originalFetch; + clearCodexUpstreamHealth(); + rmSync(testDir, { recursive: true, force: true }); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + if (previousCodexHome === undefined) delete process.env.CODEX_HOME; + else process.env.CODEX_HOME = previousCodexHome; + } + }); + + test("a Spark recovery probe releases its compact lease when connect is cancelled (#590)", async () => { + const testDir = mkdtempSync(join(tmpdir(), "ocx-compact-connect-probe-")); + const previousOpencodexHome = process.env.OPENCODEX_HOME; + const previousCodexHome = process.env.CODEX_HOME; + const originalNow = Date.now; + const now = 1_800_000_000_000; + const probeAt = now + CODEX_QUOTA_PROBE_INTERVAL_MS; + const config = nativePoolConfig(); + const abort = new AbortController(); + let markFetchStarted!: () => void; + const fetchStarted = new Promise(resolve => { markFetchStarted = resolve; }); + try { + process.env.OPENCODEX_HOME = testDir; + process.env.CODEX_HOME = testDir; + Date.now = () => now; + clearCodexUpstreamHealth(); + saveCodexAccountCredential("pool-a", { + accessToken: "pool-access-token", + refreshToken: "pool-refresh-token", + expiresAt: now + 30 * 60_000, + chatgptAccountId: "pool_acc", + }); + recordCodexUpstreamOutcome(config, "pool-a", 429, { + now, + resetAt: Math.floor((now + 4 * 24 * 60 * 60_000) / 1_000), + modelId: "gpt-5.3-codex-spark", + }); + Date.now = () => probeAt; + globalThis.fetch = ((_url: string | URL | Request, init?: RequestInit) => new Promise((_resolve, reject) => { + const signal = init?.signal; + if (!signal) throw new Error("expected compact request abort signal"); + signal.addEventListener("abort", () => reject(signal.reason), { once: true }); + markFetchStarted(); + })) as typeof fetch; + const pending = handleResponsesCompact( + compactionRequest(baseCompactionBody({ model: "gpt-5.3-codex-spark" }), abort.signal), + config, + { model: "", provider: "" }, + ); + await fetchStarted; + abort.abort(); + const cancelled = await pending; + expect(cancelled.status).toBe(499); + + Date.now = () => probeAt + CODEX_QUOTA_PROBE_INTERVAL_MS; + const nextProbe = await resolveCodexAuthContext( + new Headers({ authorization: "Bearer main-token" }), + config, + "pool", + { modelId: "gpt-5.3-codex-spark" }, + ); + expect(nextProbe).toMatchObject({ probeQuotaScope: "spark" }); + releaseCodexAuthContextProbeLease(nextProbe); + } finally { + Date.now = originalNow; + globalThis.fetch = originalFetch; + clearCodexUpstreamHealth(); + rmSync(testDir, { recursive: true, force: true }); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + if (previousCodexHome === undefined) delete process.env.CODEX_HOME; + else process.env.CODEX_HOME = previousCodexHome; + } + }); +}); + describe("routed compaction for key-mode openai-responses (#422)", () => { test("rewrites the wire: no trigger, no tools, summarizer prompt present", async () => { const bodies: Array> = []; diff --git a/tests/server-combo-failover-e2e.test.ts b/tests/server-combo-failover-e2e.test.ts index c464ed4a07..cd7fe3659b 100644 --- a/tests/server-combo-failover-e2e.test.ts +++ b/tests/server-combo-failover-e2e.test.ts @@ -850,6 +850,53 @@ describe("server combo failover 030 activation matrix", () => { expect(getCodexUpstreamHealth(rawAccountId)?.cooldownSource).toBe("retry-after"); }); + test("Spark reset cooldown fails over to the shared native quota on the same account (#590)", async () => { + const rawAccountId = "spark-scope-account"; + const config = comboConfig({ + openai: { + adapter: "openai-responses", + baseUrl: "https://chatgpt.com/backend-api/codex", + authMode: "forward", + codexAccountMode: "pool", + }, + }, [ + { provider: "openai", model: "gpt-5.3-codex-spark" }, + { provider: "openai", model: "gpt-5.6-terra" }, + ]); + config.codexAccounts = [{ + id: rawAccountId, + email: "pool@example.test", + isMain: false, + logLabel: "pspark1", + }]; + config.activeCodexAccountId = rawAccountId; + config.autoSwitchThreshold = 0; + saveCodexAccountCredential(rawAccountId, { + accessToken: "pool-access-token", + refreshToken: "pool-refresh-token", + expiresAt: Date.now() + 300_000, + chatgptAccountId: "acct-pool-spark", + }); + + const resetAt = Math.floor((Date.now() + 4 * 24 * 60 * 60_000) / 1000); + let upstreamCalls = 0; + customTransientResponse = async () => { + upstreamCalls += 1; + if (upstreamCalls === 1) { + return Response.json({ error: { message: "Spark quota exhausted" } }, { + status: 429, + headers: { "x-codex-primary-reset-at": String(resetAt) }, + }); + } + return Response.json(responsesSuccess("Terra fallback", "gpt-5.6-terra")); + }; + + const response = await post(config); + expect(response.status).toBe(200); + expect(upstreamCalls).toBe(2); + expect(await response.json()).toMatchObject({ model: "gpt-5.6-terra" }); + }); + test("keeps a failed estimate on A without overwriting B reported usage", async () => { customUsageEstimate = model => model === "m1" ? 41 : undefined; customFetchResponse = async request => {