diff --git a/README.md b/README.md index bb134383f1..0bf33016e0 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,13 @@ account exclusion, affinity expiry, or 401/403 and 429 recovery can rebind them. selection order when one of them — usually your Codex Desktop login — should only be reached for once the others are drained. +For an opt-in soft threshold that can move a request away from a high-usage account +while preserving any usable remainder, see +[strict Codex pool quota admission](docs/codex-strict-quota.md). It reuses the +existing selectors, including fill-first; only a confirmed 100% window is blocked, +and recovery is verified from fresh quota metadata before an exhausted account is +re-enabled. + ### For agents ```bash diff --git a/docs-site/src/content/docs/ko/reference/cli/providers-accounts.md b/docs-site/src/content/docs/ko/reference/cli/providers-accounts.md index 31757f38d7..ef4c804884 100644 --- a/docs-site/src/content/docs/ko/reference/cli/providers-accounts.md +++ b/docs-site/src/content/docs/ko/reference/cli/providers-accounts.md @@ -215,6 +215,15 @@ OAuth 및 API 키 제공자에는 제공자의 할당량 보고 엔드포인트 { provider, autoSwitchThreshold: number, enabled: boolean } ``` +`codexAccountStrictQuota: true`이면 이 임계값은 선제 전환을 위한 소프트 선호입니다. 적격한 +계정 중 임계값 미만인 계정을 먼저 고르지만, 사용할 수 있는 계정에 100% 미만의 잔여량이 +있으면 계속 사용할 수 있습니다. 확인된 100% 창만 하드 차단되며 모든 적격 계정이 소진된 +경우에만 요청이 대기합니다. 읽기는 WHAM metadata를 병합해 짧게 캐시하고, 실패한 읽기는 +5분 backoff를 적용하며, reset 시간은 복구를 가정하지 않고 다음 읽기만 실행합니다. 독립적인 +`codexMainAccountHardLock` 스위치는 별도의 main 계정 제한으로 유지됩니다. 자세한 내용은 +[strict Codex pool quota guide](https://github.com/lidge-jun/opencodex/blob/main/docs/codex-strict-quota.md)를 +참조하세요. + ### `ocx account priority [<-100..100|first|earlier|normal|later|last|reset>] [--json]` Codex pool 계정 하나의 선택 순서를 읽거나 설정합니다. **값이 클수록 먼저** 쓰이고 기본값은 `0`, diff --git a/docs-site/src/content/docs/reference/cli/providers-accounts.md b/docs-site/src/content/docs/reference/cli/providers-accounts.md index cc2471a527..382d6022c4 100644 --- a/docs-site/src/content/docs/reference/cli/providers-accounts.md +++ b/docs-site/src/content/docs/reference/cli/providers-accounts.md @@ -339,6 +339,14 @@ exit 1. `--json` returns: { provider, autoSwitchThreshold: number, enabled: boolean } ``` +With `codexAccountStrictQuota: true`, this threshold is a soft switching preference: eligible +accounts below it are preferred, but a usable account may continue with any remaining quota below +100%. Only a confirmed 100% window is hard-blocked; requests wait only when every usable account +is exhausted. Reads use merged WHAM metadata with a short cache, failed reads back off for five +minutes, and a reset time triggers a read without implying recovery. The independent +`codexMainAccountHardLock` switch remains a separate main-account restriction. See the +[strict Codex pool quota guide](https://github.com/lidge-jun/opencodex/blob/main/docs/codex-strict-quota.md). + ### `ocx account priority [<-100..100|first|earlier|normal|later|last|reset>] [--json]` Reads or sets one Codex pool account's selection order: **higher is used earlier**, the default is diff --git a/docs/codex-strict-quota.md b/docs/codex-strict-quota.md new file mode 100644 index 0000000000..969a08452e --- /dev/null +++ b/docs/codex-strict-quota.md @@ -0,0 +1,74 @@ +# Strict Codex pool quota admission + +Strict quota admission is opt-in. It uses the existing account selector, threshold, +manual pin, credential store, and WHAM quota metadata endpoint. It does not introduce +another account selection strategy. + +To keep using a manually selected account until its threshold is reached, select +`fill-first` and enable strict admission in the configuration: + +```json +{ + "accountPoolStrategy": "fill-first", + "autoSwitchThreshold": 95, + "codexAccountStrictQuota": true +} +``` + +The existing `PUT /api/codex-auth/auto-switch` management endpoint also accepts +`{"threshold":95,"strictQuota":true}`. Omit `strictQuota` to keep its current value. +`GET /api/codex-auth/active` reports `codexAccountStrictQuota`. Existing management +authentication requirements apply. There is no new GUI control. + +## Selection and recovery + +- `autoSwitchThreshold` is a soft preference. When strict quota is enabled, selection + first filters for actually usable accounts (credentials, pause, reauthentication, + cooldown, model, and hard quota state), then prefers an account below the threshold. + If none is below it, the current usable account may continue with any remaining + quota below 100%. Only a confirmed 100% window is a hard quota block. When every + usable account is truly exhausted, the request waits for new evidence. +- A manual selection and thread affinity cannot make a paused, cooling, reauth-needed, + model-ineligible, or 100%-blocked account usable. With strict `fill-first`, an + eligible selected account stays active when no below-threshold replacement exists. + Higher-priority accounts are preferred when a switch is possible; this reuses the + existing priority-tier selector without creating a persistent manual pin for an + automatic selection. The independent `codexMainAccountHardLock` protection switch + remains separate and may still restrict the main account; strict quota does not + promise to override it. +- Quota reads use the WHAM metadata endpoint. Selection merges concurrent reads and + uses a short 10-second cache; a failed read earns a five-minute backoff. A reset + timestamp only makes the next metadata read due and never implies that quota has + recovered. Unknown or stale quota is never treated as zero usage. +- A measured block survives stale cache data, token refresh, and predicted reset + deadlines. A new valid quota reading must establish recovery. Partial or + credits-only responses cannot clear another window's known block. +- Only pending requests own recovery timers. Usage reads are shared and bounded; + with no pending request this feature performs no periodic work. Manual usage + refreshes wake pending requests. This feature sends no warmup model requests and + never redeems reset credits. +- Selecting main or enabling strict quota while main is active reads its identity-bound + usage in the management operation. If main usage is missing after startup or has + gone stale, a real waiting request requests metadata through a separate owned main + profile claim. Caller-owned authentication still does not read local credentials; + a failed metadata read keeps the request waiting with backoff. + +## Request boundaries + +A recognized pre-stream quota refusal may try each available account once. When +all candidate accounts are quota-blocked or unknown, a Responses request waits +for new evidence. Streaming requests emit `response.heartbeat` while waiting and +then forward the real upstream stream. Cancellation and service drain terminate +the wait and release its resources. Waiting does not synthesize a completed +response. + +Ordinary server errors, an uncertain WebSocket execution outcome, and a stream +that already produced output do not authorize this cross-account replay. The +existing stored-account 401 replay budget remains bounded across waiting cycles. +Client or network disconnects still terminate requests; this is not durable job +storage and does not promise recovery after the proxy process exits. + +Explicit Direct credentials and independent Spark/Reserve quota authorization +retain their own policies. An explicit account namespace does not silently switch +to another account. Authentication failures and operator-paused accounts remain +unavailable until their actual cause is repaired. diff --git a/readme/README.zh-CN.md b/readme/README.zh-CN.md index fa7cc35c3a..1491b793b2 100644 --- a/readme/README.zh-CN.md +++ b/readme/README.zh-CN.md @@ -128,7 +128,7 @@ npm 警告里给出的缩写命令缺少包名，会把当前目录重新安装 - **在 Codex 中使用任意 LLM。** 5 种协议 adapter 覆盖 Anthropic Messages、Google Gemini、Azure、OpenAI Responses 直通，以及所有 OpenAI 兼容 Chat Completions 端点 —— 即开箱即用的 **40+ provider**。 - **在 Claude 中也能使用任意 LLM。** `ocx claude` 可通过代理启动 Claude Code。Claude 仪表盘还提供独立的 Desktop 配置，可管理 Opus、Fable、Sonnet、Haiku 四个系列，并支持拖放、键盘操作和 JSON 导入/导出。 -- **安全地池化 ChatGPT 账户。** 现有 Codex 线程保持在一个账户上，而新会话可以从池中自动挑选使用量更低的账户，并带有配额刷新和非 PII 请求标签。 +- **安全地池化 ChatGPT 账户。** 现有 Codex 线程通常保持账户亲和，新会话按配置策略选取账户。可选的[严格额度门禁](../docs/codex-strict-quota.md)复用现有填满优先等策略，把阈值作为软切换偏好；只有确认达到 100% 才会硬拦，真实额度恢复后重新参与选择；带有配额刷新和非 PII 请求标签。 - **登录一次，免填 API key。** xAI、Anthropic、Kimi 支持 OAuth，可用现有账户认证，token 自动刷新。也可以转发 `codex login`、粘贴 API key，或使用 `${ENV_VAR}` 引用 —— 随你选择。 - **Codex 在哪里能用，它就在哪里能用。** 自动注入 Codex CLI、TUI、App 和 SDK。路由模型像原生模型一样出现在 Codex 的模型选择器里。 - **委派给合适的模型。** 在仪表盘或 config 中把最多 5 个路由/原生模型放进 Codex 的 subagent 选择器 —— 复杂任务交给 reasoning 模型，快速任务交给便宜模型。在 v2 多智能体表面（GPT-5.6 Sol/Terra）上，代理会注入精简的委派指引：首选子智能体模型与 effort（`injectionModel` / `injectionEffort`）、featured 模型清单及各自支持的 effort 阶梯，以及让跨模型 `spawn_agent` 覆盖得以应用的 `fork_turns` 规则。已知限制：原生父代理 spawn 路由子代理时，任务正文可能以后端加密形式到达而丢失（[#92](https://github.com/lidge-jun/opencodex/issues/92)）—— 需要可靠的跨 provider 委派请使用 v1 表面。想自定义文案，可在 `injectionPrompt` 中使用 `{{model}}` / `{{effort}}` / `{{roster}}` 占位符。 diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 0fbe7cf746..705287de67 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -462,6 +462,8 @@ "codex-restore-app-rewrite.test.ts": "codex-integration", "codex-retained-root-serialization.test.ts": "codex-integration", "codex-routing.test.ts": "codex-integration", + "codex-strict-quota.test.ts": "codex-integration", + "codex-strict-quota-refresh.test.ts": "codex-integration", "codex-runtime.test.ts": "codex-integration", "codex-service-manager-probe-hardening.test.ts": "codex-integration", "codex-service-manager-probe.test.ts": "codex-integration", @@ -1083,6 +1085,7 @@ "server-403-permission-e2e.test.ts": "server", "server-agent-task-recovery-replay.test.ts": "server", "server-auth.test.ts": "server", + "server-strict-quota-wait.test.ts": "server", "server-background-lifecycle.test.ts": "server", "server-clickjacking-headers.test.ts": "server", "server-combo-failover-e2e.test.ts": "server", diff --git a/src/codex/account-usability.ts b/src/codex/account-usability.ts index 3ce27475ee..5ed5e78436 100644 --- a/src/codex/account-usability.ts +++ b/src/codex/account-usability.ts @@ -12,12 +12,16 @@ import { isNativeMainTrafficBlocked } from "./native-profile-startup"; import { isMainAccountHardLocked } from "./main-account-hard-lock"; export interface CodexAccountUsabilityOptions { + /** Live admission policy when the routing config is a request-specific replay snapshot. */ + strictQuotaPolicy?: Readonly>; /** Route using cached runtime state only; the caller must reject selected main before auth. */ nativeMainSelectionOnly?: boolean; /** Test seam for proving whether routing attempted a physical native-token read. */ isMainAccountTokenLive?: typeof isMainAccountTokenLive; /** Confirmed account ids for an account-gated model; omitted for ordinary native models. */ modelEligibleAccountIds?: ReadonlySet; + /** Request-local retry exclusions are independent of the model's entitlement roster. */ + excludedAccountIds?: ReadonlySet; } export function isCodexAccountUsable( @@ -25,6 +29,7 @@ export function isCodexAccountUsable( accountId: string, options: CodexAccountUsabilityOptions = {}, ): boolean { + if (options.excludedAccountIds?.has(accountId)) return false; if (options.modelEligibleAccountIds && !options.modelEligibleAccountIds.has(accountId)) return false; if (accountId === MAIN_CODEX_ACCOUNT_ID) { if (isMainAccountHardLocked(config)) return false; diff --git a/src/codex/auth-api.ts b/src/codex/auth-api.ts index 51e3fed303..966fe77b28 100644 --- a/src/codex/auth-api.ts +++ b/src/codex/auth-api.ts @@ -6,6 +6,7 @@ import { withConfigMutationLockSync, } from "../config"; import { codexAccountLogLabel, withCodexAccountLogLabel } from "./account-label"; +import { notifyCodexQuotaChanges } from "./quota-events"; import { getCodexAccountCredential, getValidCodexToken, @@ -1463,6 +1464,25 @@ async function fetchPoolAccountQuota( } } +/** Explicitly read-only usage refresh for request-owned strict-quota recovery. */ +export async function refreshStrictCodexPoolQuotaSnapshots( + config: OcxConfig, accountIds: readonly string[], + policy: Pick = config, +): Promise { + if (policy.codexAccountStrictQuota !== true) return; + await mapWithConcurrency([...accountIds], POOL_QUOTA_REFRESH_CONCURRENCY, async id => { + if (isCodexAccountPaused(config, id)) return; + if (id === MAIN_CODEX_ACCOUNT_ID) { + // A background/request check cannot clear a reauth quarantine or redeem credits. + await fetchMainAccountInfoAttempt(true, 1, undefined, false, false); + return; + } + const account = configuredPoolAccount(config, id); + if (!account || isAccountNeedsReauth(id)) return; + await fetchPoolAccountQuota(id, true, account.plan); + }); +} + let primeInFlight: Promise | null = null; /** * Last prime attempt per pool account. A failed WHAM lookup stores no quota, so @@ -2142,6 +2162,14 @@ export async function handleCodexAuthAPI( else setCodexAccountPin(runtimeConfig, targetAccountId); resetCodexRoutingForManualSelection(targetAccountId); saveRuntimeConfig(config, runtimeConfig); + // Management owns the physical-main claim. Rebuild identity-bound usage here; + // caller-owned requests cannot read auth.json to recover an unknown main snapshot. + if (targetAccountId === MAIN_CODEX_ACCOUNT_ID && runtimeConfig.codexAccountStrictQuota === true + && (runtimeConfig.autoSwitchThreshold ?? 80) > 0) { + await fetchMainAccountInfoAttempt(true, 1, undefined, false, false); + } + // A pending strict-quota request must reconsider an operator-selected account now. + notifyCodexQuotaChanges(); return jsonResponse({ ok: true, activeCodexAccountId: body.accountId, appliesImmediately: true }); } @@ -2157,6 +2185,7 @@ export async function handleCodexAuthAPI( // lets a surface mark the account the operator actually chose. pinnedAccountId: pinnedCodexAccountId(runtimeConfig) ?? null, autoSwitchThreshold: runtimeConfig.autoSwitchThreshold ?? 80, + codexAccountStrictQuota: runtimeConfig.codexAccountStrictQuota === true, upstreamFailoverThreshold: runtimeConfig.upstreamFailoverThreshold ?? 3, accountPoolStrategy: normalizeAccountPoolStrategy(runtimeConfig.accountPoolStrategy), accountPoolStickyLimit: normalizeAccountPoolStickyLimit(runtimeConfig.accountPoolStickyLimit), @@ -2164,14 +2193,24 @@ export async function handleCodexAuthAPI( } if (url.pathname === "/api/codex-auth/auto-switch" && req.method === "PUT") { - let body: { threshold: number }; + let body: { threshold: number; strictQuota?: unknown }; try { body = (await req.json()) as typeof body; } catch { return jsonResponse({ error: "Invalid JSON" }, 400); } if (typeof body.threshold !== "number" || !Number.isInteger(body.threshold) || body.threshold < 0 || body.threshold > 100) { return jsonResponse({ error: "Threshold must be an integer 0-100" }, 400); } + if (body.strictQuota !== undefined && typeof body.strictQuota !== "boolean") { + return jsonResponse({ error: "strictQuota must be a boolean" }, 400); + } const runtimeConfig = getRuntimeConfig(config); runtimeConfig.autoSwitchThreshold = body.threshold; + if (typeof body.strictQuota === "boolean") runtimeConfig.codexAccountStrictQuota = body.strictQuota; saveRuntimeConfig(config, runtimeConfig); + if (runtimeConfig.codexAccountStrictQuota === true && body.threshold > 0 + && (getEffectiveActiveCodexAccountId(runtimeConfig) ?? MAIN_CODEX_ACCOUNT_ID) === MAIN_CODEX_ACCOUNT_ID + && !isCodexAccountPaused(runtimeConfig, MAIN_CODEX_ACCOUNT_ID)) { + await fetchMainAccountInfoAttempt(true, 1, undefined, false, false); + } + notifyCodexQuotaChanges(); return jsonResponse({ ok: true }); } diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 2f319b3144..42a02a00a6 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -27,6 +27,7 @@ import { codexQuotaScopeForModel, computeCodexUsageScore, getCodexQuotaHealthSnapshot, + getEffectiveActiveCodexAccountId, isEffectiveCodexAccountPinned, releaseCodexQuotaProbeLease, releaseCodexQuotaScopeProbeLease, @@ -34,6 +35,7 @@ import { tryAcquireCodexQuotaScopeProbeLease, pickAlternateCodexAccount, resolveCodexAccountForThreadDetailed, + strictQuotaReplacement, } from "./routing"; import { entitledCodexAccountIdsForModel, @@ -64,6 +66,8 @@ import { CODEX_RESERVE_HELPER_UNSUPPORTED_MESSAGE, isCodexReserveHelperUnsupport import type { DataPlaneAdmission } from "../server/auth-cors"; import { getMainReserveAuthorization, isMainReserveAuthorizationLive, type MainReserveAuthorization } from "./reserve-availability"; import { UpstreamRetryEvidenceError } from "../lib/upstream-retry"; +import { getCodexStrictQuotaStatus, isCodexStrictQuotaEnabled, isCodexStrictQuotaEligible } from "./strict-quota"; +import { refreshStrictCodexQuotasOnDemand } from "./strict-quota-refresh"; const CODEX_AFFINITY_COMPONENT_MAX_BYTES = 512; const CODEX_APP_AFFINITY_KEY = randomBytes(32); @@ -76,7 +80,8 @@ const CODEX_APP_AFFINITY_KEY = randomBytes(32); * path. This keeps the keyring boundary intact instead of reading auth.json just to classify a * request that already brought its own credential (#3157). */ -function requestOwnedMainPinHasQuotaHeadroom(config: OcxConfig): boolean { +function requestOwnedMainPinHasQuotaHeadroom(config: CodexAuthPolicyConfig, quotaScope?: CodexQuotaScope): boolean { + if (isCodexStrictQuotaEnabled(config, quotaScope)) return isCodexStrictQuotaEligible(config, MAIN_CODEX_ACCOUNT_ID, quotaScope); const threshold = config.autoSwitchThreshold ?? 80; if (threshold <= 0) return true; const usage = computeCodexUsageScore(getAccountQuota(MAIN_CODEX_ACCOUNT_ID)); @@ -112,7 +117,9 @@ export function codexPoolAffinityKey(headers: Headers): string | undefined { } export type CodexAuthContext = - | { kind: "main"; accountId: null; reserveAuthorization?: MainReserveAuthorization } + | { kind: "main"; accountId: null; reserveAuthorization?: MainReserveAuthorization; + /** Caller-owned credential chosen by Pool; Direct contexts deliberately omit this marker. */ + poolQuotaScope?: CodexQuotaScope; fixedAccount?: boolean } | { kind: "pool"; accountId: string; @@ -311,6 +318,15 @@ export class CodexMainAccountHardLockError extends CodexAccountCooldownError { } } +/** A local refusal before dispatch. Only ordinary Pool requests may wait and reselect. */ +export class CodexStrictQuotaUnavailableError extends CodexAccountCooldownError { + constructor(readonly waitable = true) { + super(MAIN_CODEX_ACCOUNT_ID, 0); + this.name = "CodexStrictQuotaUnavailableError"; + this.message = "No Codex account has confirmed remaining quota"; + } +} + export class CodexReserveUnavailableError extends CodexAccountCooldownError { constructor() { super(MAIN_CODEX_ACCOUNT_ID, 0); @@ -332,6 +348,7 @@ export class CodexReserveHelperUnsupportedError extends CodexReserveUnavailableE export type CodexAuthPolicyConfig = Readonly>; interface CodexAuthMaterializationOptions { @@ -407,7 +424,10 @@ function assertMaterializedReserve(headers: Headers, ctx: CodexAuthContext, opti } } -/** A dispatch never renews permission: the next request may obtain a fresh bounded proof. */ +/** + * Final Pool quota and Reserve admission after pacing/retry/WS setup. The legacy export name + * stays stable for transports. A dispatch never performs a refresh or renews permission. + */ export function createCodexReserveDispatchGuard( ctx: CodexAuthContext, config: CodexAuthPolicyConfig, @@ -418,15 +438,25 @@ export function createCodexReserveDispatchGuard( // Snapshot the resolved source value, not the caller's mutable admission object. Config stays // live so policy changes remain visible after pacing and retry backoff. const source = admission?.source; - if (modelId !== NATIVE_RESERVE_MODEL || source !== "loopback") return undefined; + const poolSelection = ctx.kind !== "main" || ctx.poolQuotaScope !== undefined; + const selectedId = ctx.kind === "main" ? MAIN_CODEX_ACCOUNT_ID : ctx.accountId; + const fixedAccount = ctx.fixedAccount === true; + const scope = codexQuotaScopeForModel(modelId); + const reserveDispatch = modelId === NATIVE_RESERVE_MODEL && source === "loopback"; + if (!poolSelection && !reserveDispatch) return undefined; // Only immutable request facts decide whether to install the callback. Flag/role eligibility // is checked inside it, including an opt-in enabled while a send waits for pacing or WS open. - const ingress = Object.freeze({ source }); + const ingress = source === undefined ? undefined : Object.freeze({ source }); return headers => { - if (isCodexReserveHelperUnsupported(config, modelId, ingress, terminalHelper)) { - throw new CodexReserveHelperUnsupportedError(); + if (poolSelection && !isCodexStrictQuotaEligible(config, selectedId, scope)) { + throw new CodexStrictQuotaUnavailableError(!fixedAccount); + } + if (reserveDispatch) { + if (isCodexReserveHelperUnsupported(config, modelId, ingress, terminalHelper)) { + throw new CodexReserveHelperUnsupportedError(); + } + assertMaterializedReserve(headers, ctx, { config, modelId, admission: ingress }); } - assertMaterializedReserve(headers, ctx, { config, modelId, admission: ingress }); }; } @@ -492,7 +522,8 @@ export function cooldownAccountLabel(accountId: string): string { * injected `openai_base_url` in config.toml. */ export function cooldownErrorMessage(err: CodexAccountCooldownError, accountSelector?: string): string { - if (err instanceof CodexMainAccountHardLockError || err instanceof CodexReserveUnavailableError) return err.message; + if (err instanceof CodexStrictQuotaUnavailableError || err instanceof CodexMainAccountHardLockError + || err instanceof CodexReserveUnavailableError) return err.message; const until = new Date(err.cooldownUntil).toISOString(); const scopeLabels: Record = { spark: "Spark quota", shared: "shared native quota", reserve: "Reserve quota", @@ -517,7 +548,7 @@ export function cooldownErrorResponse( ): Response { const res = formatErrorResponse(429, "rate_limit_error", cooldownErrorMessage(err, accountSelector)); const headers = new Headers(res.headers); - if (!(err instanceof CodexReserveUnavailableError) + if (!(err instanceof CodexStrictQuotaUnavailableError) && !(err instanceof CodexReserveUnavailableError) && (!(err instanceof CodexMainAccountHardLockError) || err.resetAt !== undefined)) { headers.set("Retry-After", String(Math.max(1, Math.ceil((err.cooldownUntil - now) / 1000)))); } @@ -535,7 +566,8 @@ export class CodexThreadAffinityExpiredError extends Error { } export function shouldMarkAccountNeedsReauthForCodexAuthFailure(cause: unknown): boolean { - return !(cause instanceof CodexMainAccountHardLockError) + return !(cause instanceof CodexStrictQuotaUnavailableError) + && !(cause instanceof CodexMainAccountHardLockError) && !(cause instanceof CodexReserveUnavailableError) && !(cause instanceof CodexCredentialGenerationConflictError) && !(cause instanceof CodexCredentialRefreshLockTimeoutError) @@ -553,6 +585,8 @@ export interface ResolveCodexAuthContextOptions { /** Live policy owner when the routing config is a caller-specific replay snapshot. */ codexAuthPolicy?: CodexAuthPolicyConfig; excludeAccountId?: string; + /** Accounts already rejected in this logical request; never visit one twice. */ + excludeAccountIds?: ReadonlySet; /** Resolve exactly this account without consulting or mutating Pool selection. */ accountId?: string; /** Final native model selected for this request, used to select its quota group. */ @@ -598,16 +632,28 @@ export async function resolveCodexAuthContext( throw new CodexReserveUnavailableError(); } const fixedAccountId = reserve ? MAIN_CODEX_ACCOUNT_ID : options.accountId; - const preserveRequestOwnedMainPin = requestScopedMainCredential + const quotaScope = codexQuotaScopeForModel(options.modelId); + const canPreserveRequestOwnedMainPin = () => mode === "pool" && requestScopedMainCredential && fixedAccountId === undefined + && options.excludeAccountId !== MAIN_CODEX_ACCOUNT_ID + && !options.excludeAccountIds?.has(MAIN_CODEX_ACCOUNT_ID) && config.activeCodexAccountPinned === MAIN_CODEX_ACCOUNT_ID && isEffectiveCodexAccountPinned(config) && !policy.pausedCodexAccountIds?.includes(MAIN_CODEX_ACCOUNT_ID) && !(callerMatchesObservedMain(headers) && isMainAccountHardLocked(policy)) - && requestOwnedMainPinHasQuotaHeadroom(config); - if (fixedAccountId !== undefined && options.excludeAccountId !== undefined) { + && requestOwnedMainPinHasQuotaHeadroom(policy, quotaScope); + const preserveRequestOwnedMainPin = canPreserveRequestOwnedMainPin(); + if (fixedAccountId !== undefined && (options.excludeAccountId !== undefined || options.excludeAccountIds?.size)) { throw new Error("Codex auth context cannot select and exclude an account simultaneously"); } + const callerOwnedContext = (): Extract => { + if (mode === "pool" && !isCodexStrictQuotaEligible(policy, MAIN_CODEX_ACCOUNT_ID, quotaScope)) { + throw new CodexStrictQuotaUnavailableError(fixedAccountId === undefined); + } + return { kind: "main", accountId: null, + ...(mode === "pool" ? { poolQuotaScope: quotaScope ?? "shared", + ...(fixedAccountId !== undefined ? { fixedAccount: true } : {}) } : {}) }; + }; const resolveCallerOwnedMainContext = async (): Promise => { if (!hasCallerCodexBearer(headers)) throw new CodexDirectAuthenticationError(); const substituteStoredMain = options.substituteMainCredentialForDirect === true; @@ -618,7 +664,7 @@ export async function resolveCodexAuthContext( const token = selectedCodexToken(selected); const reserveAuthorization = await authorizeReserveCredential(token, captureMainQuotaWriter(token.chatgptAccountId), policy, options.signal, undefined, writerGeneration); - return { kind: "main", accountId: null, reserveAuthorization }; + return { ...callerOwnedContext(), reserveAuthorization }; } if (options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId)) { const entitled = await ( @@ -629,7 +675,7 @@ export async function resolveCodexAuthContext( } } if (callerMatchesObservedMain(headers)) assertMainAccountPolicy(policy); - return { kind: "main", accountId: null }; + return callerOwnedContext(); } // Admission-bearer Direct requests later replace the proxy secret with the stored @@ -663,7 +709,7 @@ export async function resolveCodexAuthContext( } } assertMainAccountPolicy(policy); - return { kind: "main", accountId: null }; + return callerOwnedContext(); } finally { // The short selector reservation ends here. A successful claim remains owned by // the enclosing turn lease until the request or transferred stream settles. @@ -675,14 +721,38 @@ export async function resolveCodexAuthContext( // is the one exception where that exclusion is selection evidence in the opposite direction. // Validate the caller's own gated-model roster before using it, and fall through to a Pool model // detour when it lacks the grant. This branch performs no physical-main credential read. - if (preserveRequestOwnedMainPin) { + let preferStoredQuotaReplacement = false; + if (preserveRequestOwnedMainPin && isCodexStrictQuotaEnabled(policy, quotaScope)) { + const mainStatus = getCodexStrictQuotaStatus(policy, MAIN_CODEX_ACCOUNT_ID, quotaScope); + if (mainStatus.state === "ready" && mainStatus.usedPercent! >= mainStatus.threshold!) { + // A caller-owned main pin may use its remainder, but a usable below-threshold + // stored account still takes precedence. Exclude physical main from both reads. + const excluded = new Set([MAIN_CODEX_ACCOUNT_ID, ...(options.excludeAccountIds ?? [])]); + if (options.excludeAccountId) excluded.add(options.excludeAccountId); + const entitlement = options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId) + ? await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config, { + excludeAccountIds: excluded, signal: options.signal, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, + }) : undefined; + const modelEligibleAccountIds = entitlement + ? entitledCodexAccountIdsForModel(entitlement, options.modelId) : undefined; + const candidateOptions = { strictQuotaPolicy: policy, excludedAccountIds: excluded, modelEligibleAccountIds }; + const candidates = new Set((config.codexAccounts ?? []).map(account => account.id) + .filter(id => !policy.pausedCodexAccountIds?.includes(id) && isCodexAccountUsable(config, id, candidateOptions))); + await refreshStrictCodexQuotasOnDemand(config, candidates, { policy, signal: options.signal, forSelection: true }); + options.signal?.throwIfAborted(); + preferStoredQuotaReplacement = strictQuotaReplacement(config, MAIN_CODEX_ACCOUNT_ID, Date.now(), quotaScope, candidateOptions) !== null; + } + } + if (preserveRequestOwnedMainPin && !preferStoredQuotaReplacement) { const callerEntitled = !options.modelId || !ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId) || await ( options.isDirectCallerEntitledToCodexModel ?? isDirectCallerEntitledToCodexModel )(headers, options.modelId); - if (callerEntitled && !(callerMatchesObservedMain(headers) && isMainAccountHardLocked(policy))) { - return { kind: "main", accountId: null }; + if (callerEntitled && canPreserveRequestOwnedMainPin()) { + options.signal?.throwIfAborted(); + return callerOwnedContext(); } } // An explicit namespace binding is stronger than the provider's default mode. It must use the @@ -712,7 +782,6 @@ export async function resolveCodexAuthContext( const nativeMainSelectionOnly = !nativeMainTrafficBlocked && selectionAdmission?.mainProfileDraining === true; let accountId: string; - const quotaScope = codexQuotaScopeForModel(options.modelId); try { const excludeAccountIds = nativeMainReadsForbidden ? new Set([MAIN_CODEX_ACCOUNT_ID]) @@ -732,6 +801,7 @@ export async function resolveCodexAuthContext( ? new Set([...entitledAccountIds].filter(candidate => !excludeAccountIds?.has(candidate))) : undefined; const selectionOptions = { + strictQuotaPolicy: policy, // Temporary switch drain keeps the candidate until the atomic claim rejects // it. Retained recovery makes main wholly ineligible so pool routing continues. nativeMainSelectionOnly, @@ -742,13 +812,14 @@ export async function resolveCodexAuthContext( ? () => preserveRequestOwnedMainPin : options.isMainAccountTokenLive, modelEligibleAccountIds, + excludedAccountIds: options.excludeAccountIds, }; // A pre-drain selector reserves the native identity while reconciliation and // routing inspect it. Selectors arriving after the fence skip reconciliation // and may still route to non-main pool accounts without touching switch state. if (reserve && !nativeMainReadsForbidden && !selectionAdmission) throw new CodexMainProfileDrainingError(); if (!nativeMainReadsForbidden) reconcileMainCodexAccountRuntimeState(); - const resolution = fixedAccountId !== undefined + const resolveSelection = () => fixedAccountId !== undefined ? { status: "selected" as const, accountId: fixedAccountId } : options.excludeAccountId ? (() => { @@ -771,6 +842,37 @@ export async function resolveCodexAuthContext( selectionOptions, options.modelId, ); + const strict = isCodexStrictQuotaEnabled(policy, quotaScope); + const potentialIds = () => [MAIN_CODEX_ACCOUNT_ID, ...(config.codexAccounts ?? []).map(account => account.id)] + .filter(id => !policy.pausedCodexAccountIds?.includes(id) + && !(id === MAIN_CODEX_ACCOUNT_ID && nativeMainReadsForbidden) + && isCodexAccountUsable(config, id, { ...selectionOptions, excludedAccountIds: undefined })); + // Refresh an unknown/stale chosen account before letting another one steal its work. + // The existing fill-first/manual selection remains the owner of that preference. + if (strict && !options.excludeAccountId && !options.excludeAccountIds?.size) { + const preferred = fixedAccountId ?? getEffectiveActiveCodexAccountId(config) ?? MAIN_CODEX_ACCOUNT_ID; + if (potentialIds().includes(preferred) + && getCodexStrictQuotaStatus(policy, preferred, quotaScope).state === "unknown") { + await refreshStrictCodexQuotasOnDemand(config, new Set([preferred]), { policy, signal: options.signal }); + options.signal?.throwIfAborted(); + } + } + // A switch re-reads candidate metadata, including previously exhausted accounts + // that may have reset or received quota since the last request. + const preferredStatus = getCodexStrictQuotaStatus(policy, + fixedAccountId ?? getEffectiveActiveCodexAccountId(config) ?? MAIN_CODEX_ACCOUNT_ID, quotaScope); + if (strict && (options.excludeAccountId || options.excludeAccountIds?.size + || preferredStatus.state !== "ready" || preferredStatus.usedPercent! >= preferredStatus.threshold!)) { + const ids = fixedAccountId === undefined ? potentialIds() : potentialIds().filter(id => id === fixedAccountId); + await refreshStrictCodexQuotasOnDemand(config, new Set(ids), { policy, signal: options.signal, forSelection: true }); + options.signal?.throwIfAborted(); + } + let resolution = resolveSelection(); + if (strict && resolution.status === "none" && fixedAccountId === undefined) { + await refreshStrictCodexQuotasOnDemand(config, new Set(potentialIds()), { policy, signal: options.signal }); + options.signal?.throwIfAborted(); + resolution = resolveSelection(); + } if (resolution.status === "expired") throw new CodexThreadAffinityExpiredError(resolution.accountId); const selected = resolution.status === "selected" ? resolution.accountId : null; if (!selected) { @@ -782,9 +884,28 @@ export async function resolveCodexAuthContext( requestScopedMainCredential && fixedAccountId === undefined && options.excludeAccountId !== MAIN_CODEX_ACCOUNT_ID + && !options.excludeAccountIds?.has(MAIN_CODEX_ACCOUNT_ID) + && !policy.pausedCodexAccountIds?.includes(MAIN_CODEX_ACCOUNT_ID) + && (!strict || isCodexStrictQuotaEligible(policy, MAIN_CODEX_ACCOUNT_ID, quotaScope)) ) { return await resolveCallerOwnedMainContext(); } + // A request-owned main cannot inspect auth.json here. Let the pending-request + // owner obtain metadata under its own native claim instead of reporting a + // profile drain when only main quota evidence is missing or expired. + if (strict && requestScopedMainCredential && fixedAccountId === undefined + && options.excludeAccountId !== MAIN_CODEX_ACCOUNT_ID + && !options.excludeAccountIds?.has(MAIN_CODEX_ACCOUNT_ID) + && !policy.pausedCodexAccountIds?.includes(MAIN_CODEX_ACCOUNT_ID) + && !(callerMatchesObservedMain(headers) && isMainAccountHardLocked(policy)) + && !isCodexStrictQuotaEligible(policy, MAIN_CODEX_ACCOUNT_ID, quotaScope)) { + throw new CodexStrictQuotaUnavailableError(); + } + if (strict && fixedAccountId === undefined && potentialIds().some(id => + !isCodexStrictQuotaEligible(policy, id, quotaScope) + || getCodexQuotaHealthSnapshot(id, quotaScope) !== null)) { + throw new CodexStrictQuotaUnavailableError(); + } if (fixedAccountId !== undefined) { throw new CodexPoolAuthenticationError( modelEligibleAccountIds && !modelEligibleAccountIds.has(fixedAccountId) @@ -815,6 +936,10 @@ export async function resolveCodexAuthContext( ); } accountId = selected; + if (options.excludeAccountIds?.has(accountId)) throw new CodexPoolAuthenticationError(); + if (strict && !isCodexStrictQuotaEligible(policy, accountId, quotaScope)) { + throw new CodexStrictQuotaUnavailableError(fixedAccountId === undefined); + } if (accountId === MAIN_CODEX_ACCOUNT_ID) assertMainAccountPolicy(policy); if (accountId === MAIN_CODEX_ACCOUNT_ID && nativeMainTrafficBlocked) { throw new CodexMainProfileDrainingError(); @@ -899,11 +1024,14 @@ export async function resolveCodexAuthContext( ...(options.nativeMainRefreshDependencies ?? {}), }); if (token) mainQuotaWriter = observeSelectedMainCredential(token, mainQuotaWriter); + if (!isCodexStrictQuotaEligible(policy, accountId, quotaScope)) { + throw new CodexStrictQuotaUnavailableError(fixedAccountId === undefined); + } assertMainAccountPolicy(policy); } catch (cause) { if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); - if (cause instanceof CodexMainAccountHardLockError) throw cause; + if (cause instanceof CodexMainAccountHardLockError || cause instanceof CodexStrictQuotaUnavailableError) throw cause; if (!options.signal?.aborted && shouldMarkAccountNeedsReauthForCodexAuthFailure(cause)) { markAccountNeedsReauth(accountId, writerGeneration); } @@ -938,6 +1066,9 @@ export async function resolveCodexAuthContext( try { const token = await getValidCodexToken(accountId); + if (!isCodexStrictQuotaEligible(policy, accountId, quotaScope)) { + throw new CodexStrictQuotaUnavailableError(fixedAccountId === undefined); + } return { kind: "pool", accountId, @@ -954,6 +1085,7 @@ export async function resolveCodexAuthContext( } catch (cause) { if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); + if (cause instanceof CodexStrictQuotaUnavailableError) throw cause; if (!options.signal?.aborted && shouldMarkAccountNeedsReauthForCodexAuthFailure(cause)) { markAccountNeedsReauth(accountId, writerGeneration); } @@ -1013,12 +1145,19 @@ export function materializeCodexUpstreamAuth( ctx: CodexAuthContext, options: CodexAuthMaterializationOptions = {}, ): Headers { + if (ctx.kind === "main" && ctx.poolQuotaScope !== undefined && options.config + && !isCodexStrictQuotaEligible(options.config, MAIN_CODEX_ACCOUNT_ID, ctx.poolQuotaScope)) { + throw new CodexStrictQuotaUnavailableError(ctx.fixedAccount !== true); + } const selected = new Headers(); for (const name of FORWARD_HEADERS) { const value = headers.get(name); if (value) selected.set(name, value); } if (ctx.kind === "pool" || ctx.kind === "main-pool") { + if (options.config && !isCodexStrictQuotaEligible(options.config, ctx.accountId, ctx.quotaScope)) { + throw new CodexStrictQuotaUnavailableError(ctx.fixedAccount !== true); + } selected.set("authorization", `Bearer ${ctx.accessToken}`); selected.set("chatgpt-account-id", ctx.chatgptAccountId); if (ctx.kind === "main-pool") { @@ -1101,6 +1240,10 @@ export async function materializeCodexUpstreamAuthAsync( ctx: CodexAuthContext, options: CodexAuthMaterializationOptions = {}, ): Promise { + if (ctx.kind === "main" && ctx.poolQuotaScope !== undefined && options.config + && !isCodexStrictQuotaEligible(options.config, MAIN_CODEX_ACCOUNT_ID, ctx.poolQuotaScope)) { + throw new CodexStrictQuotaUnavailableError(ctx.fixedAccount !== true); + } if (requiresReserveAuthorization(options.config, options.modelId, options.admission)) { return materializeReserveUpstreamAuth(headers, ctx, options); } @@ -1122,6 +1265,10 @@ export async function materializeCodexUpstreamAuthAsync( selected.set("authorization", `Bearer ${stored.accessToken}`); if (stored.chatgptAccountId) selected.set("chatgpt-account-id", stored.chatgptAccountId); observeSelectedMainCredential(stored, writer); + if (ctx.poolQuotaScope !== undefined && options.config + && !isCodexStrictQuotaEligible(options.config, MAIN_CODEX_ACCOUNT_ID, ctx.poolQuotaScope)) { + throw new CodexStrictQuotaUnavailableError(ctx.fixedAccount !== true); + } assertMainAccountPolicy(options.config); // An opt-in enabled during token refresh must not turn a proof-less context into Reserve. assertMaterializedReserve(selected, ctx, options); diff --git a/src/codex/quota-events.ts b/src/codex/quota-events.ts new file mode 100644 index 0000000000..ac42598cd6 --- /dev/null +++ b/src/codex/quota-events.ts @@ -0,0 +1,15 @@ +/** Dependency-free wakeups for request-owned quota waiters. No timers or network work. */ +const listeners = new Set<() => void>(); +let revision = 0; +export function getCodexQuotaRevision(): number { return revision; } +export function subscribeCodexQuotaChanges(listener: () => void): () => void { + listeners.add(listener); + return () => { listeners.delete(listener); }; +} +export function notifyCodexQuotaChanges(): void { + revision++; + for (const listener of listeners) { + try { listener(); } + catch { console.warn("[codex] quota change listener failed"); } + } +} diff --git a/src/codex/quota-types.ts b/src/codex/quota-types.ts index 6c06de6ae9..837295f5be 100644 --- a/src/codex/quota-types.ts +++ b/src/codex/quota-types.ts @@ -49,3 +49,14 @@ export type WhamUsageResponse = { rate_limit_reset_credits?: { available_count: number } | null; additional_rate_limits?: WhamAdditionalRateLimit[] | null; }; + +/** Durable usage evidence; each window owns its clock independently of partial updates. */ +export type StrictAccountQuota = { + windows: Array<{ + scope: "shared"; + key: "weekly" | "monthly" | "short"; + usedPercent: number; + observedAt: number; + resetAt?: number; + }>; +}; diff --git a/src/codex/quota.ts b/src/codex/quota.ts index 0648e4a717..b0b867ee0c 100644 --- a/src/codex/quota.ts +++ b/src/codex/quota.ts @@ -1,3 +1,7 @@ +import { createHash } from "node:crypto"; +import { readCodexAccountRecord } from "./account-store"; +import { notifyCodexQuotaChanges } from "./quota-events"; +import type { StrictAccountQuota } from "./quota-types"; import { existsSync, readFileSync, unlinkSync } from "node:fs"; import { join } from "node:path"; import { atomicWriteFile, getConfigDir } from "../config"; @@ -19,6 +23,7 @@ type QuotaDiskFile = { version: 1; quotas: Record; mainPolicyQuota?: MainPolicyQuota; + strictQuotas?: Record; }; type MainPolicyQuota = { identityKey: string; quota: StoredAccountQuota }; @@ -46,6 +51,47 @@ const MONTHLY_WINDOW_MIN_MINUTES = MONTHLY_WINDOW_MIN_SECONDS / 60; const WEEKLY_WINDOW_MIN_MINUTES = WEEKLY_WINDOW_MIN_SECONDS / 60; const accountQuota = new Map(); +const strictQuotas = new Map(); +// Parser provenance is private and cannot alter legacy dashboard DTOs. +const invalidStrictSnapshots = new WeakSet(); +// Only a structurally complete WHAM response can retire a previously observed window. +// Header subsets and caller-built parsed snapshots never receive this private provenance. +const retiredStrictWindows = new WeakMap>(); +function strictQuotaIdentity(accountId: string): string | null { + if (accountId === MAIN_CODEX_ACCOUNT_ID) return getObservedMainQuotaIdentityKey() ?? null; + const record = readCodexAccountRecord(accountId); + if (!record?.credential || record.deletedAt != null) return null; + // Token refresh changes generation without changing the upstream quota owner. Retain + // blocked windows across that refresh; a different account identity starts unknown. + const owner = record.credential.chatgptAccountId; + return owner ? createHash("sha256").update(owner).digest("hex") : `generation:${record.generation}`; +} +function observeStrictQuota(accountId: string, quota: Omit | null, now: number): void { + const identity = strictQuotaIdentity(accountId); + if (!identity) { strictQuotas.delete(accountId); return; } + if (!quota || invalidStrictSnapshots.has(quota)) return; + const previous = strictQuotas.get(accountId); + const retired = retiredStrictWindows.get(quota); + const windows = previous?.identity === identity + ? previous.quota.windows.filter(window => !retired?.has(window.key)) : []; + for (const key of ["weekly", "monthly", "short"] as const) { + const usedPercent = quota[`${key}Percent`]; + if (typeof usedPercent !== "number" || !Number.isFinite(usedPercent) || usedPercent < 0 || usedPercent > 100) continue; + const resetAt = quota[`${key}ResetAt`]; + const next = { scope: "shared" as const, key, usedPercent, observedAt: now, + ...(typeof resetAt === "number" && Number.isFinite(resetAt) && resetAt >= 0 ? { resetAt } : {}) }; + const index = windows.findIndex(window => window.key === key); + if (index < 0) windows.push(next); else windows[index] = next; + } + if (windows.length) strictQuotas.set(accountId, { identity, quota: { windows } }); +} +/** Reads only added-account store or already-observed main identity, never native-main ownership. */ +export function getStrictAccountQuota(accountId: string): StrictAccountQuota | null { + hydrateAccountQuotasFromDisk(); + const cached = strictQuotas.get(accountId); + return cached && cached.identity === strictQuotaIdentity(accountId) ? structuredClone(cached.quota) : null; +} + let lastReconciledGeneration = 0; let liveAccountIds = new Set(); @@ -259,7 +305,11 @@ export function setAccountQuotaFromParsed( } : null; } + // Strict pool admission evaluates all reported shared windows, independently of the + // narrower legacy main hard-lock window policy. The writer still proves main ownership. + observeStrictQuota(accountId, isMain ? (mainWriter ? quota : null) : quota, updatedAt); schedulePersistAccountQuotas(); + notifyCodexQuotaChanges(); // Credits carry the previous usage tuple; they must not refresh its observation clock. if (!(quota.resetCredits !== undefined && !snapshotHasUsage(quota))) { notifyCodexQuotaSnapshot(accountId, next); @@ -469,6 +519,7 @@ export function parseUpstreamQuotaHeaders(headers: Headers): Omit window && window.scope === "shared" + && ["weekly", "monthly", "short"].includes(window.key) + && typeof window.usedPercent === "number" && Number.isFinite(window.usedPercent) + && window.usedPercent >= 0 && window.usedPercent <= 100 + && typeof window.observedAt === "number" && Number.isFinite(window.observedAt) && window.observedAt >= 0) + .map(window => ({ scope: "shared" as const, key: window.key, usedPercent: window.usedPercent, + observedAt: window.observedAt, + ...(typeof window.resetAt === "number" && Number.isFinite(window.resetAt) && window.resetAt >= 0 + ? { resetAt: window.resetAt } : {}) })); + if (windows.length) strictQuotas.set(id, { identity: entry.identity, quota: { windows } }); + } const now = Date.now(); for (const [accountId, quota] of Object.entries(parsed.quotas)) { if (!quota || typeof quota !== "object" || typeof quota.updatedAt !== "number") continue; @@ -608,6 +677,7 @@ function schedulePersistAccountQuotas(): void { const body: QuotaDiskFile = { version: 1, quotas, + strictQuotas: Object.fromEntries(strictQuotas), ...(mainPolicyQuota ? { mainPolicyQuota } : {}), }; atomicWriteFile(join(getConfigDir(), QUOTA_CACHE_FILENAME), `${JSON.stringify(body)}\n`); @@ -660,12 +730,16 @@ export function clearAccountQuota(accountId?: string): void { if (accountId) { hydrateAccountQuotasFromDisk(); accountQuota.delete(accountId); + strictQuotas.delete(accountId); + notifyCodexQuotaChanges(); if (accountId === MAIN_CODEX_ACCOUNT_ID) mainPolicyQuota = null; schedulePersistAccountQuotas(); forgetCodexQuotaBaseline(accountId); return; } accountQuota.clear(); + strictQuotas.clear(); + notifyCodexQuotaChanges(); forgetCodexQuotaBaseline(); mainPolicyQuota = null; diskHydrated = false; @@ -688,6 +762,7 @@ export function reconcileCodexQuotaAccounts(context: GenerationContext): number for (const accountId of accountQuota.keys()) { if (context.codexAccountIds.has(accountId)) continue; accountQuota.delete(accountId); + strictQuotas.delete(accountId); removed += 1; } liveAccountIds = new Set(context.codexAccountIds); @@ -703,6 +778,8 @@ function filterMainPolicyMonthlyQuota( ): Omit | null { if (!quota || monthlyOnlyPlan || quota.monthlyIsPrimaryWindow === true) return quota; const filtered = { ...quota }; + const retired = retiredStrictWindows.get(quota); + if (retired) retiredStrictWindows.set(filtered, retired); delete filtered.monthlyPercent; delete filtered.monthlyResetAt; delete filtered.monthlyIsPrimaryWindow; @@ -717,6 +794,30 @@ export function parseMainPolicyUsageQuota(data: WhamUsageResponse): Omit + !!window && typeof window.used_percent === "number" && Number.isFinite(window.used_percent) + && window.used_percent >= 0 && window.used_percent <= 100 + && typeof window.limit_window_seconds === "number" && Number.isFinite(window.limit_window_seconds) + && window.limit_window_seconds > 0; + if (!rate || !complete(rate.primary_window) || !Object.hasOwn(rate, "secondary_window")) return; + if (rate.secondary_window !== null && !complete(rate.secondary_window)) return; + if (rate.tertiary_window != null && !complete(rate.tertiary_window)) return; + const declared = new Set<"weekly" | "monthly" | "short">(); + for (const window of [rate.primary_window, rate.secondary_window]) { + if (!window) continue; + declared.add(isExplicitShortWindow(window) ? "short" : isExplicitMonthlyWindow(window) ? "monthly" : "weekly"); + } + // The established parser treats tertiary as supplementary monthly. Its omission is not + // proof that a prior monthly window disappeared, so only explicit null retires that bar. + if (rate.tertiary_window) declared.add("monthly"); + const retired = new Set<"weekly" | "monthly" | "short">(); + for (const key of ["weekly", "short"] as const) if (!declared.has(key)) retired.add(key); + if (rate.tertiary_window === null && !declared.has("monthly")) retired.add("monthly"); + retiredStrictWindows.set(quota, retired); +} + export function parseUsageQuota(data: WhamUsageResponse): Omit | null { const resetCredits = typeof data.rate_limit_reset_credits?.available_count === "number" ? data.rate_limit_reset_credits.available_count @@ -815,5 +916,10 @@ export function parseUsageQuota(data: WhamUsageResponse): Omit isInvalidPolicyUsagePercent(window?.used_percent))) { + invalidStrictSnapshots.add(quota); + } else { + markCompleteStrictWindowAuthority(data, quota); + } return hasKnownQuotaValue(quota) || resetCredits !== undefined ? quota : null; } diff --git a/src/codex/routing.ts b/src/codex/routing.ts index dbf9cab086..ab7af2e9c3 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -1,3 +1,4 @@ +import { getCodexStrictQuotaStatus, isCodexStrictQuotaEligible, isCodexStrictQuotaEnabled } from "./strict-quota"; import { randomUUID } from "node:crypto"; import { saveConfigPreservingClaudeCode } from "../config"; import { isCodexAccountGenerationLive, readCodexAccountRecord } from "./account-store"; @@ -930,7 +931,8 @@ function isCodexAccountSelectable( quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, ): boolean { - return !isCodexAccountPaused(config, accountId) + return isCodexStrictQuotaEligible(selectionOptions?.strictQuotaPolicy ?? config, accountId, quotaScope, now) + && !isCodexAccountPaused(config, accountId) && getCodexQuotaHealthSnapshot(accountId, quotaScope, now) === null && !isCodexAccountSoftAvoided(accountId, now) && isCodexAccountUsable(config, accountId, selectionOptions); @@ -1167,6 +1169,7 @@ function getEligiblePoolAccounts( && !isCodexAccountPaused(config, account.id) && !isAccountNeedsReauth(account.id) && (!skipFailoverReadyCandidates || !shouldFailover(config, account.id, now))) + .filter(account => isCodexStrictQuotaEligible(selectionOptions?.strictQuotaPolicy ?? config, account.id, quotaScope, now)) .filter(account => getCodexQuotaHealthSnapshot(account.id, quotaScope, now) === null) .filter(account => !isCodexAccountSoftAvoided(account.id, now)) .filter(account => isCodexAccountUsable(config, account.id, selectionOptions)) @@ -1175,6 +1178,7 @@ function getEligiblePoolAccounts( // first-class rotation candidate when its read-only token is usable (Option A). if ( excludeId !== MAIN_CODEX_ACCOUNT_ID + && isCodexStrictQuotaEligible(selectionOptions?.strictQuotaPolicy ?? config, MAIN_CODEX_ACCOUNT_ID, quotaScope, now) && !isCodexAccountPaused(config, MAIN_CODEX_ACCOUNT_ID) && (!isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID) || hasMainAccountRefreshGrant()) && getCodexQuotaHealthSnapshot(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) === null @@ -1184,14 +1188,26 @@ function getEligiblePoolAccounts( ) { ids.unshift(MAIN_CODEX_ACCOUNT_ID); } + const policy = selectionOptions?.strictQuotaPolicy ?? config; + const preferred = isCodexStrictQuotaEnabled(policy, quotaScope) + ? ids.filter(id => hasCodexQuotaHeadroom(config, id, quotaScope, selectionOptions, now)) : []; + // Prefer fresh below-threshold capacity only after request eligibility has been applied. + // When none exists, retain every usable remainder instead of entering quota wait. + const candidates = preferred.length ? preferred : ids; // 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, + candidates, codexAccountPriorityLookup(config), - id => hasCodexQuotaHeadroom(config, id, selectionOptions, now), - pinnedCodexAccountId(config), + id => hasCodexQuotaHeadroom(config, id, quotaScope, selectionOptions, now), + pinnedCodexAccountId(config) ?? ( + isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) + && normalizeAccountPoolStrategy(config.accountPoolStrategy) === "fill-first" + // Finish the current usable account before reopening higher tiers. This ceiling + // is runtime selection only; it never creates a persisted manual pin. + ? getEffectiveActiveCodexAccountId(config) : undefined + ), ); } @@ -1221,10 +1237,16 @@ function stickyLimitForConfig(config: OcxConfig): number { function hasCodexQuotaHeadroom( config: OcxConfig, accountId: string, + quotaScope: CodexQuotaScope | undefined, selectionOptions?: CodexAccountUsabilityOptions, now: number = Date.now(), ): boolean { - const threshold = config.autoSwitchThreshold ?? 80; + const policy = selectionOptions?.strictQuotaPolicy ?? config; + if (isCodexStrictQuotaEnabled(policy, quotaScope)) { + const status = getCodexStrictQuotaStatus(policy, accountId, quotaScope, now); + return status.state === "ready" && status.usedPercent! < status.threshold!; + } + const threshold = policy.autoSwitchThreshold ?? 80; if (threshold <= 0) return true; const usage = computeCodexUsageScore( getAccountQuota(accountId), @@ -1249,18 +1271,22 @@ function pickFillFirstCodexAccount( if (eligible.length === 0) return null; const active = getEffectiveActiveCodexAccountId(config); - if (active && eligible.includes(active) && hasCodexQuotaHeadroom(config, active, selectionOptions, now)) { + if (active && eligible.includes(active) + && (hasCodexQuotaHeadroom(config, active, quotaScope, selectionOptions, now) + || (isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) + && !eligible.some(id => hasCodexQuotaHeadroom(config, id, quotaScope, selectionOptions, now))))) { return active; } - return pickNextFillFirstCodexAccount(config, active ?? null, eligible, now, selectionOptions); + return pickNextFillFirstCodexAccount(config, active ?? null, quotaScope, 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()), + quotaScope: CodexQuotaScope | undefined, + eligible: readonly string[] = listEligibleCodexAccountIds(config, Date.now(), quotaScope), now = Date.now(), selectionOptions?: CodexAccountUsabilityOptions, ): string | null { @@ -1269,7 +1295,7 @@ function pickNextFillFirstCodexAccount( 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; + if (hasCodexQuotaHeadroom(config, id, quotaScope, selectionOptions, now)) return id; } return ordered[0] ?? null; } @@ -1284,7 +1310,7 @@ function pickNextFillFirstCodexAccount( const startIdx = stableAll.indexOf(afterId); if (startIdx < 0) { for (const id of ordered) { - if (hasCodexQuotaHeadroom(config, id, selectionOptions, now)) return id; + if (hasCodexQuotaHeadroom(config, id, quotaScope, selectionOptions, now)) return id; } return ordered[0] ?? null; } @@ -1295,7 +1321,7 @@ function pickNextFillFirstCodexAccount( const candidate = stableAll[(startIdx + step) % stableAll.length]!; if (!eligible.includes(candidate)) continue; if (!fallback) fallback = candidate; - if (hasCodexQuotaHeadroom(config, candidate, selectionOptions, now)) return candidate; + if (hasCodexQuotaHeadroom(config, candidate, quotaScope, selectionOptions, now)) return candidate; } return fallback ?? ordered[0] ?? null; } @@ -1382,10 +1408,11 @@ function sharedStateSelectionOptions( selectionOptions?: CodexAccountUsabilityOptions, ): Pick< CodexAccountUsabilityOptions, - "nativeMainSelectionOnly" | "isMainAccountTokenLive" + "nativeMainSelectionOnly" | "isMainAccountTokenLive" | "strictQuotaPolicy" > | undefined { if (!selectionOptions) return undefined; return { + ...(selectionOptions.strictQuotaPolicy ? { strictQuotaPolicy: selectionOptions.strictQuotaPolicy } : {}), ...(selectionOptions.nativeMainSelectionOnly !== undefined ? { nativeMainSelectionOnly: selectionOptions.nativeMainSelectionOnly } : {}), @@ -1488,7 +1515,7 @@ export function pickAlternateCodexAccount( } if (strategy === "fill-first") { const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); - return pickNextFillFirstCodexAccount(config, excludeId, eligible, now, selectionOptions); + return pickNextFillFirstCodexAccount(config, excludeId, quotaScope, eligible, now, selectionOptions); } return pickLowestUsageCodexAccount(config, excludeId, now, quotaScope, selectionOptions); } @@ -1605,7 +1632,7 @@ function pickPriorityPreemption( if ( pinned !== undefined && eligible.includes(pinned) - && hasCodexQuotaHeadroom(config, pinned, selectionOptions, now) + && hasCodexQuotaHeadroom(config, pinned, quotaScope, selectionOptions, now) ) return null; const priorityOf = codexAccountPriorityLookup(config); if (priorityOf(eligible[0]!) <= priorityOf(active)) return null; @@ -1613,7 +1640,7 @@ function pickPriorityPreemption( // picking one would hand the request straight back to a drained account. return pickLowestUsageAmong( config, - eligible.filter(id => hasCodexQuotaHeadroom(config, id, selectionOptions, now)), + eligible.filter(id => hasCodexQuotaHeadroom(config, id, quotaScope, selectionOptions, now)), selectionOptions, now, ); @@ -1630,13 +1657,16 @@ function releaseDrainedCodexAccountPin( config: OcxConfig, selectionOptions?: Pick< CodexAccountUsabilityOptions, - "nativeMainSelectionOnly" | "isMainAccountTokenLive" + "nativeMainSelectionOnly" | "isMainAccountTokenLive" | "strictQuotaPolicy" >, now: number = Date.now(), ): void { const pinned = pinnedCodexAccountId(config); if (pinned === undefined) return; - const knownUnavailable = isAccountNeedsReauth(pinned) || isCodexAccountPaused(config, pinned); + // Strict admission can prove a drain the legacy scorer calls unknown (a short-only + // reading below 100, for example). Retire that pin now so recovery cannot revive it. + const strictDrained = getCodexStrictQuotaStatus(selectionOptions?.strictQuotaPolicy ?? config, pinned, "shared", now).state === "blocked"; + const knownUnavailable = strictDrained || isAccountNeedsReauth(pinned) || isCodexAccountPaused(config, pinned); if (knownUnavailable) { clearCodexAccountPin(config); saveConfigPreservingClaudeCode(config); @@ -1647,7 +1677,7 @@ function releaseDrainedCodexAccountPin( // is readable. Cached reauth and configured pause state were handled above. if (pinned === MAIN_CODEX_ACCOUNT_ID && selectionOptions?.nativeMainSelectionOnly === true) return; const drained = !isCodexAccountUsable(config, pinned, selectionOptions) - || !hasCodexQuotaHeadroom(config, pinned, selectionOptions, now); + || !hasCodexQuotaHeadroom(config, pinned, "shared", selectionOptions, now); if (!drained) return; clearCodexAccountPin(config); saveConfigPreservingClaudeCode(config); @@ -1661,7 +1691,14 @@ function applyQuotaAutoSwitch( selectionOptions?: CodexAccountUsabilityOptions, commitSharedSelection = true, ): string { - const threshold = config.autoSwitchThreshold ?? 80; + if (isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope)) { + const replacement = strictQuotaReplacement(config, active, now, quotaScope, selectionOptions); + if (replacement && commitSharedSelection && !isIndependentCodexQuotaScope(quotaScope)) { + setActiveCodexAccount(config, replacement); + } + return replacement ?? active; + } + const threshold = (selectionOptions?.strictQuotaPolicy ?? config).autoSwitchThreshold ?? 80; if (threshold <= 0) return active; const quota = getAccountQuota(active); const activeUsage = computeCodexUsageScore( @@ -1701,7 +1738,7 @@ function isHealthySharedCodexSelection( selectionOptions: CodexAccountUsabilityOptions | undefined, ): boolean { return isCodexAccountSelectable(config, accountId, now, quotaScope, selectionOptions) - && hasCodexQuotaHeadroom(config, accountId, selectionOptions, now) + && hasCodexQuotaHeadroom(config, accountId, quotaScope, selectionOptions, now) && !shouldFailover(config, accountId, now); } @@ -1764,6 +1801,18 @@ export function resolveCodexAccountForThread( return resolution.status === "selected" ? resolution.accountId : null; } +/** Soft-threshold rotation for every strict strategy, with no churn among remainders. */ +export function strictQuotaReplacement( + config: OcxConfig, active: string, now: number, quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + if (!isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) + || hasCodexQuotaHeadroom(config, active, quotaScope, selectionOptions, now)) return null; + const preferred = getEligiblePoolAccounts(config, active, now, quotaScope, selectionOptions, true) + .filter(id => hasCodexQuotaHeadroom(config, id, quotaScope, selectionOptions, now)); + return preferred[0] ?? null; +} + function previewReusableAffinityAccount( entry: ThreadAffinityEntry | undefined, config: OcxConfig, @@ -1780,10 +1829,13 @@ function previewReusableAffinityAccount( ) { return null; } + if (isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope)) { + return strictQuotaReplacement(config, entry.accountId, now, quotaScope, selectionOptions) ?? entry.accountId; + } // Quota strategy only: non-quota strategies keep affinity for ongoing threads // (new-session-only rotation — docs / affinity policy A). if (normalizeAccountPoolStrategy(config.accountPoolStrategy) === "quota") { - const threshold = config.autoSwitchThreshold ?? 80; + const threshold = (selectionOptions?.strictQuotaPolicy ?? config).autoSwitchThreshold ?? 80; if (threshold > 0) { const usage = computeCodexUsageScore( getAccountQuota(entry.accountId), @@ -1818,8 +1870,11 @@ function reevaluateAffinityQuota( quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, ): string | null { + if (isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope)) { + return strictQuotaReplacement(config, entry.accountId, now, quotaScope, selectionOptions); + } if (normalizeAccountPoolStrategy(config.accountPoolStrategy) !== "quota") return null; - const threshold = config.autoSwitchThreshold ?? 80; + const threshold = (selectionOptions?.strictQuotaPolicy ?? config).autoSwitchThreshold ?? 80; const usage = threshold > 0 ? computeCodexUsageScore( getAccountQuota(entry.accountId), @@ -1906,14 +1961,18 @@ export function previewCodexAccountForRequest( const fallback = pickLowestUsageCodexAccount(config, active, now, quotaScope, selectionOptions); if (fallback) active = fallback; else if ( - hasConfiguredPoolAccount(config, active, selectionOptions) + !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) + && hasConfiguredPoolAccount(config, active, selectionOptions) && !isCodexAccountPaused(config, active) ) return active; else return null; } active = pickPriorityPreemption(config, active, now, quotaScope, selectionOptions) ?? active; - const threshold = config.autoSwitchThreshold ?? 80; + if (isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope)) { + active = strictQuotaReplacement(config, active, now, quotaScope, selectionOptions) ?? active; + } else { + const threshold = (selectionOptions?.strictQuotaPolicy ?? config).autoSwitchThreshold ?? 80; if (threshold > 0) { const usage = computeCodexUsageScore( getAccountQuota(active), @@ -1924,16 +1983,17 @@ export function previewCodexAccountForRequest( active = pickLowerUsageAccount(config, active, usage, now, quotaScope, selectionOptions); } } + } if (shouldFailover(config, active, now)) { const best = pickLowestUsageCodexAccount(config, active, now, quotaScope, selectionOptions); if (best) active = best; } if (!isCodexAccountUsable(config, active, selectionOptions)) { - return hasConfiguredPoolAccount(config, active, selectionOptions) ? active : null; + return !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) && hasConfiguredPoolAccount(config, active, selectionOptions) ? active : null; } if (isCodexAccountPaused(config, active)) return null; if (getCodexQuotaHealthSnapshot(active, quotaScope, now)) { - return hasConfiguredPoolAccount(config, active, selectionOptions) ? active : null; + return !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) && hasConfiguredPoolAccount(config, active, selectionOptions) ? active : null; } return active; } @@ -2016,7 +2076,7 @@ export function resolveCodexAccountForThreadDetailed( && isCodexAccountSelectable(config, entry.accountId, now, quotaScope, selectionOptions); const failoverReady = shouldFailover(config, entry.accountId, now); const healthyForSharedAffinity = selectableForSharedState - && hasCodexQuotaHeadroom(config, entry.accountId, sharedSelectionOptions, now) + && hasCodexQuotaHeadroom(config, entry.accountId, quotaScope, sharedSelectionOptions, now) && !failoverReady; if ( selectableForRequest @@ -2092,7 +2152,8 @@ export function resolveCodexAccountForThreadDetailed( const selected = pickLowestUsageCodexAccount(config, undefined, now, quotaScope, selectionOptions); if (!selected) { if ( - selectionOptions?.nativeMainSelectionOnly === true + !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) + && selectionOptions?.nativeMainSelectionOnly === true && selectionOptions.modelEligibleAccountIds !== undefined ) { return { status: "selected", accountId: MAIN_CODEX_ACCOUNT_ID }; @@ -2112,7 +2173,7 @@ export function resolveCodexAccountForThreadDetailed( sharedSelectionOptions, ); const activeHealthyForSharedSelection = activeSelectableForSharedState - && hasCodexQuotaHeadroom(config, active, sharedSelectionOptions, now) + && hasCodexQuotaHeadroom(config, active, quotaScope, sharedSelectionOptions, now) && !shouldFailover(config, active, now); if (!isCodexAccountSelectable(config, active, now, quotaScope, selectionOptions)) { const fallback = pickLowestUsageCodexAccount(config, active, now, quotaScope, selectionOptions); @@ -2125,7 +2186,8 @@ export function resolveCodexAccountForThreadDetailed( } active = fallback; } else if ( - selectionOptions?.nativeMainSelectionOnly === true + !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) + && selectionOptions?.nativeMainSelectionOnly === true && selectionOptions.modelEligibleAccountIds !== undefined ) { // Entitlement discovery intentionally excludes main while a temporary drain @@ -2135,7 +2197,8 @@ export function resolveCodexAccountForThreadDetailed( // active account or persist/bind this synthetic selection. return { status: "selected", accountId: MAIN_CODEX_ACCOUNT_ID }; } else if ( - hasConfiguredPoolAccount(config, active, selectionOptions) + !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) + && hasConfiguredPoolAccount(config, active, selectionOptions) && !isCodexAccountPaused(config, active) ) { return { status: "selected", accountId: active }; @@ -2176,13 +2239,13 @@ export function resolveCodexAccountForThreadDetailed( !preserveSharedSelectionForModelDetour, ); if (!isCodexAccountUsable(config, active, selectionOptions)) { - return hasConfiguredPoolAccount(config, active, selectionOptions) + return !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) && hasConfiguredPoolAccount(config, active, selectionOptions) ? { status: "selected", accountId: active } : { status: "none" }; } if (isCodexAccountPaused(config, active)) return { status: "none" }; if (getCodexQuotaHealthSnapshot(active, quotaScope, now)) { - return hasConfiguredPoolAccount(config, active, selectionOptions) + return !isCodexStrictQuotaEnabled(selectionOptions?.strictQuotaPolicy ?? config, quotaScope) && hasConfiguredPoolAccount(config, active, selectionOptions) ? { status: "selected", accountId: active } : { status: "none" }; } diff --git a/src/codex/strict-quota-refresh.ts b/src/codex/strict-quota-refresh.ts new file mode 100644 index 0000000000..61fdc99a53 --- /dev/null +++ b/src/codex/strict-quota-refresh.ts @@ -0,0 +1,213 @@ +import type { OcxConfig } from "../types"; +import { MAIN_CODEX_ACCOUNT_ID, isSelectableCodexPoolAccount } from "./account-id"; +import { readCodexAccountRecord } from "./account-store"; +import { captureMainAccountIdentityGeneration, getObservedMainQuotaIdentityKey } from "./main-account-cache"; +import { getCodexQuotaHealthSnapshot } from "./routing"; +import { getStrictAccountQuota } from "./quota"; +import { getCodexQuotaRevision, subscribeCodexQuotaChanges } from "./quota-events"; +import { CODEX_STRICT_QUOTA_FRESHNESS_MS, getCodexStrictQuotaStatus, isCodexStrictQuotaEnabled, type CodexStrictQuotaConfig } from "./strict-quota"; + +type Refresh = (config: OcxConfig, accountIds: readonly string[], policy?: CodexStrictQuotaConfig) => Promise; +type Probe = { attemptedAt: number; credentialKey: string; failed?: boolean }; +const SELECTION_QUOTA_FRESHNESS_MS = 10_000; +export type StrictCodexQuotaRefreshResult = { + /** Attempted does not claim a successful quota read; eligibility still comes from the cache. */ + status: "off" | "idle" | "attempted" | "failed"; + accountIds: readonly string[]; +}; +const probes = new Map(); +const groups = new Map(); +// One batch at a time preserves the auth API's bounded concurrency across overlapping requests. +let flight: Promise | undefined; +let refresh: Refresh = async (config, ids, policy) => { + const { refreshStrictCodexPoolQuotaSnapshots } = await import("./auth-api"); + await refreshStrictCodexPoolQuotaSnapshots(config, ids, policy); +}; +let clock = () => Date.now(); +let schedule = (fn: () => void, delay: number) => setTimeout(fn, delay); +let cancel = (timer: ReturnType) => clearTimeout(timer); + +function accountIds(config: OcxConfig): string[] { + return [...new Set([MAIN_CODEX_ACCOUNT_ID, ...(config.codexAccounts ?? []) + .filter(isSelectableCodexPoolAccount).map(account => account.id)])] + .filter(id => !config.pausedCodexAccountIds?.includes(id)); +} +function credentialKey(id: string): string { + if (id === MAIN_CODEX_ACCOUNT_ID) { + return `${getObservedMainQuotaIdentityKey() ?? "unknown"}:${captureMainAccountIdentityGeneration()}`; + } + const record = readCodexAccountRecord(id); + return record && record.deletedAt == null ? String(record.generation) : "absent"; +} + +/** A predicted reset schedules a read; it never changes eligibility by itself. */ +function probeDueAt(config: CodexStrictQuotaConfig, id: string, now: number): number { + const state = getCodexStrictQuotaStatus(config, id, "shared", now); + const prior = probes.get(id); + // Authentication repair must not inherit backoff earned by the old credential. + if (prior && prior.credentialKey !== credentialKey(id)) return now; + const attemptedAt = prior?.attemptedAt; + const observedAt = state.updatedAt; + if (attemptedAt === undefined && observedAt === undefined) return now; + let due = Math.max(attemptedAt ?? 0, observedAt ?? 0) + CODEX_STRICT_QUOTA_FRESHNESS_MS; + for (const window of getStrictAccountQuota(id)?.windows ?? []) { + const raw = window.resetAt; + if (typeof raw !== "number" || !Number.isFinite(raw) || raw <= 0) continue; + const resetMs = raw < 1_000_000_000_000 ? raw * 1000 : raw; + // One early probe per prediction. A failed post-reset read earns normal backoff. + if (resetMs > Math.max(attemptedAt ?? 0, window.observedAt)) due = Math.min(due, resetMs + 1000); + } + return due; +} + +/** Cancelling one caller does not abort the shared metadata read needed by other callers. */ +function waitForRefresh(work: Promise, signal?: AbortSignal): Promise { + if (!signal) return work; + if (signal.aborted) return Promise.reject(signal.reason ?? new DOMException("Aborted", "AbortError")); + return new Promise((resolve, reject) => { + const abort = () => { signal.removeEventListener("abort", abort); reject(signal.reason ?? new DOMException("Aborted", "AbortError")); }; + signal.addEventListener("abort", abort, { once: true }); + work.then(value => { signal.removeEventListener("abort", abort); resolve(value); }, + error => { signal.removeEventListener("abort", abort); reject(error); }); + }); +} + +/** Request-triggered usage reads. No idle timer and no inference/reset-credit calls. */ +export async function refreshStrictCodexQuotasOnDemand( + config: OcxConfig, requestedIds?: ReadonlySet, + options: { policy?: CodexStrictQuotaConfig; signal?: AbortSignal; forSelection?: boolean } = {}, +): Promise { + options.signal?.throwIfAborted(); + const policy = options.policy ?? config; + if (!isCodexStrictQuotaEnabled(policy)) return { status: "off", accountIds: [] }; + let joined: StrictCodexQuotaRefreshResult | undefined; + // A different batch may cover only part of this request. Re-evaluate after it settles; + // per-account attempt markers suppress duplicate reads, including failed reads. + while (flight) joined = await waitForRefresh(flight, options.signal); + const configured = accountIds(config); + const now = clock(); + const ids = configured.filter(id => { + if (requestedIds && !requestedIds.has(id)) return false; + const state = getCodexStrictQuotaStatus(policy, id, "shared", now); + const prior = probes.get(id); + // A real selection consults recent metadata, including accounts formerly exhausted + // but topped up since. Collapse bursts; an unsuccessful read still earns normal backoff. + const selectionDue = options.forSelection === true && !prior?.failed + && now >= Math.max(prior?.attemptedAt ?? 0, state.updatedAt ?? 0) + SELECTION_QUOTA_FRESHNESS_MS; + return selectionDue || (state.state !== "ready" && probeDueAt(policy, id, now) <= now); + }); + if (!ids.length) { + const joinedIds = joined?.accountIds.filter(id => configured.includes(id) && (!requestedIds || requestedIds.has(id))) ?? []; + return joinedIds.length ? { status: joined!.status, accountIds: joinedIds } : { status: "idle", accountIds: [] }; + } + const observations = new Map(ids.map(id => [id, JSON.stringify(getStrictAccountQuota(id))])); + for (const id of ids) probes.set(id, { attemptedAt: now, credentialKey: credentialKey(id) }); + // Start in a microtask so the shared flight exists before any synchronous injected work. + const markFailed = (id: string) => { + const prior = probes.get(id); + if (prior?.attemptedAt === now) prior.failed = true; + }; + const batch = Promise.resolve().then(() => refresh(config, ids, policy)).then( + (): StrictCodexQuotaRefreshResult => { + for (const id of ids) { + // Some metadata APIs report a failed read without throwing. No new observation + // must not become a successful short-backoff refresh just because the promise resolved. + if (JSON.stringify(getStrictAccountQuota(id)) === observations.get(id)) markFailed(id); + } + return { status: "attempted", accountIds: ids }; + }, + (): StrictCodexQuotaRefreshResult => { + for (const id of ids) markFailed(id); + return { status: "failed", accountIds: ids }; + }, + ); + const owned = batch.finally(() => { if (flight === owned) flight = undefined; }); + flight = owned; + return waitForRefresh(owned, options.signal); +} + +type WaitGroup = { + listeners: Set<() => void>; + timer?: ReturnType; + unsubscribe: () => void; +}; + +/** Wait only while a real request is pending. The last cancellation removes the sole timer. */ +export function waitForStrictCodexQuotaChange( + config: OcxConfig, signal?: AbortSignal, observedRevision?: number, +): Promise { + if (signal?.aborted) return Promise.reject(signal.reason ?? new DOMException("Aborted", "AbortError")); + return new Promise((resolve, reject) => { + let group = groups.get(config); + if (!group) { + const created: WaitGroup = { listeners: new Set(), unsubscribe: () => {} }; + const wake = () => { for (const listener of [...created.listeners]) listener(); }; + created.unsubscribe = subscribeCodexQuotaChanges(wake); + groups.set(config, created); + group = created; + } + const owner = group; + const cleanup = () => { + owner.listeners.delete(done); + signal?.removeEventListener("abort", aborted); + if (owner.listeners.size === 0) { + if (owner.timer !== undefined) { cancel(owner.timer); delete owner.timer; } + owner.unsubscribe(); + if (groups.get(config) === owner) groups.delete(config); + } + }; + const done = () => { cleanup(); resolve(); }; + const aborted = () => { cleanup(); reject(signal?.reason ?? new DOMException("Aborted", "AbortError")); }; + owner.listeners.add(done); + signal?.addEventListener("abort", aborted, { once: true }); + // Subscribe first, then compare: a manual reset between refusal and registration + // must cause a retry now, rather than being lost until the next metadata probe. + if (observedRevision !== undefined && getCodexQuotaRevision() !== observedRevision) { + done(); return; + } + if (owner.timer === undefined) { + const now = clock(); + // A request has already probed its real candidates before waiting. Unobserved rows + // (notably a missing or fenced native main) must not turn the group into a 1s poll. + const candidates = accountIds(config).filter(id => probes.has(id) || getStrictAccountQuota(id) !== null); + const cooldowns = accountIds(config).flatMap(id => { + const until = getCodexQuotaHealthSnapshot(id, "shared", now)?.cooldownUntil; + return until !== undefined && until > now ? [until] : []; + }); + // Real candidates were attempted before entering this wait. A past deadline can + // belong to a stale native-main snapshot this request cannot use; it must not + // repeatedly wake the whole group. Events still wake immediately after a new read. + const probeDeadlines = candidates.map(id => probeDueAt(config, id, now)).filter(due => due > now); + const due = Math.min(now + CODEX_STRICT_QUOTA_FRESHNESS_MS, ...probeDeadlines, ...cooldowns); + owner.timer = schedule(() => { + delete owner.timer; + // The pending requests perform the next coalesced refresh before selecting. + for (const listener of [...owner.listeners]) listener(); + }, Math.max(1000, due - now)); + owner.timer.unref?.(); + } + }); +} + +export function strictCodexQuotaWaiterCount(): number { + let count = 0; + for (const group of groups.values()) count += group.listeners.size; + return count; +} + +/** Test injection keeps request/wakeup tests off user credentials and real timers. */ +export function setStrictCodexQuotaRefreshForTests(fn: Refresh, runtime?: { + now?: () => number; + setTimeout?: typeof schedule; + clearTimeout?: typeof cancel; +}): () => void { + if (flight || groups.size) throw new Error("Cannot replace quota runtime while requests are active"); + const previous = { refresh, clock, schedule, cancel }; + refresh = fn; clock = runtime?.now ?? clock; + schedule = runtime?.setTimeout ?? schedule; cancel = runtime?.clearTimeout ?? cancel; + probes.clear(); + return () => { + if (flight || groups.size) throw new Error("Quota test left active requests"); + ({ refresh, clock, schedule, cancel } = previous); probes.clear(); + }; +} diff --git a/src/codex/strict-quota.ts b/src/codex/strict-quota.ts new file mode 100644 index 0000000000..8a1aeae8d2 --- /dev/null +++ b/src/codex/strict-quota.ts @@ -0,0 +1,49 @@ +import type { OcxConfig } from "../types"; +import { getStrictAccountQuota } from "./quota"; + +export const CODEX_STRICT_QUOTA_FRESHNESS_MS = 5 * 60_000; +export type CodexStrictQuotaConfig = Pick; +type QuotaScope = "shared" | "spark" | "reserve"; +export type CodexStrictQuotaStatus = { + state: "off" | "unknown" | "ready" | "blocked"; + threshold?: number; + usedPercent?: number; + resetAt?: number; + updatedAt?: number; +}; +export function isCodexStrictQuotaEnabled(config: CodexStrictQuotaConfig, quotaScope?: QuotaScope): boolean { + const threshold = config.autoSwitchThreshold ?? 80; + // Independent scopes have separate admission authority and no ordinary-quota evidence. + return config.codexAccountStrictQuota === true && Number.isFinite(threshold) && threshold > 0 + && (quotaScope === undefined || quotaScope === "shared"); +} +export function getCodexStrictQuotaStatus( + config: CodexStrictQuotaConfig, accountId: string, quotaScope?: QuotaScope, now = Date.now(), +): CodexStrictQuotaStatus { + if (!isCodexStrictQuotaEnabled(config, quotaScope)) return { state: "off" }; + const threshold = Math.min(config.autoSwitchThreshold ?? 80, 100); + const windows = getStrictAccountQuota(accountId)?.windows ?? []; + if (!windows.length) return { state: "unknown", threshold }; + const hottest = windows.reduce((a, b) => a.usedPercent >= b.usedPercent ? a : b); + const details = { threshold, usedPercent: hottest.usedPercent, resetAt: hottest.resetAt, + updatedAt: Math.min(...windows.map(window => window.observedAt)) }; + // A deadline is a prediction. Only a new valid reading can release a measured block. + // The switch threshold is a preference, not lost capacity. Only observed exhaustion + // blocks admission; routing prefers below-threshold candidates when one is usable. + if (hottest.usedPercent >= 100) return { state: "blocked", ...details }; + const fresh = windows.every(window => { + const rawReset = window.resetAt; + const resetMs = typeof rawReset === "number" && Number.isFinite(rawReset) && rawReset > 0 + ? (rawReset < 1_000_000_000_000 ? rawReset * 1000 : rawReset) : undefined; + // Reaching a predicted reset requests new metadata; it never invents fresh quota. + return now >= window.observedAt && now - window.observedAt <= CODEX_STRICT_QUOTA_FRESHNESS_MS + && !(resetMs !== undefined && resetMs > window.observedAt && now >= resetMs); + }); + return { state: fresh ? "ready" : "unknown", ...details }; +} +export function isCodexStrictQuotaEligible( + config: CodexStrictQuotaConfig, accountId: string, quotaScope?: QuotaScope, now = Date.now(), +): boolean { + const { state } = getCodexStrictQuotaStatus(config, accountId, quotaScope, now); + return state === "off" || state === "ready"; +} diff --git a/src/config.ts b/src/config.ts index 728b3969ea..497e4693fb 100644 --- a/src/config.ts +++ b/src/config.ts @@ -1129,6 +1129,7 @@ const configSchema = z.object({ // edit turns the tier off rather than rejecting the config that carries it. ultraFastTier: z.boolean().optional().catch(false), codexMainAccountHardLock: z.boolean().optional().catch(false), + codexAccountStrictQuota: z.boolean().optional().catch(false), // Future versions remain opaque through passthrough-compatible whole-config saves. // Only version 1 grants deletion authority in the rebase path. configRebaseProvenance: z.unknown().optional(), diff --git a/src/server/responses/codex-auth-error.ts b/src/server/responses/codex-auth-error.ts index 8d84c54392..ec88140f58 100644 --- a/src/server/responses/codex-auth-error.ts +++ b/src/server/responses/codex-auth-error.ts @@ -9,6 +9,7 @@ import { CodexMainSubstitutionUnavailableError, CodexPoolAuthenticationError, CodexThreadAffinityExpiredError, + CodexStrictQuotaUnavailableError, } from "../../codex/auth-context"; import { MAIN_CODEX_ACCOUNT_ID, @@ -16,6 +17,7 @@ import { MainAuthJsonChangedDuringRefreshError, } from "../../codex/main-account"; import { NativeProfileError } from "../../codex/native-profile-types"; +import { markStrictQuotaWaitResponse } from "./strict-quota-response"; export interface CodexAuthContextErrorResponseOptions { accountSelector?: string; @@ -46,6 +48,12 @@ export function mapCodexAuthContextErrorToResponse( error: unknown, options: CodexAuthContextErrorResponseOptions, ): Response | undefined { + if (error instanceof CodexStrictQuotaUnavailableError) { + const response = formatErrorResponse(429, "codex_quota_unavailable", error.message); + if (error.waitable) markStrictQuotaWaitResponse(response); + response.headers.set("Retry-After", "30"); + return response; + } if (error instanceof CodexAccountCooldownError) { return cooldownErrorResponse(error, options.now, options.accountSelector); } diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index d1c3e8f1e0..72194ef653 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -1,3 +1,5 @@ +import { isCodexStrictQuotaEnabled } from "../../codex/strict-quota"; +import { isStrictQuotaWaitResponse, markStrictQuotaWaitResponse } from "./strict-quota-response"; import type { Server } from "bun"; import { randomUUID } from "node:crypto"; import { bridgeToResponsesSSE, buildResponseJSON, formatErrorResponse, type ResponsesTerminalStatus } from "../../bridge"; @@ -152,6 +154,7 @@ import { unwrapUpstreamRetryEvidenceError, codexPoolAffinityKey, CodexAccountCooldownError, + CodexStrictQuotaUnavailableError, CodexAuthContextError, CodexMainProfileDrainingError, CodexPoolAuthenticationError, @@ -1018,6 +1021,8 @@ interface CodexPoolAccountRetryArgs { * out of budget. */ sameAccountOnly?: boolean; + /** Request-owned exclusion/budget for strict quota traversal only. */ + attemptedAccountIds?: Set; upstream: AbortController; connectMs: number; passthroughEstimate?: number; @@ -1037,7 +1042,7 @@ type CodexPoolAccountRetryResult = upstreamResponse: Response; selectedForwardHeaders: Headers; } - | { kind: "no-alternate" } + | { kind: "no-alternate"; quotaWaitable?: boolean } | { kind: "transport"; error: unknown; @@ -1149,6 +1154,63 @@ function shouldDeferCodexResetDerivedCooldown(response: Response, enabled?: bool */ async function retryCodexPoolOnAlternateAccount( args: CodexPoolAccountRetryArgs, +): Promise { + if (!isCodexStrictQuotaEnabled(args.config, args.firstAuthCtx.quotaScope) + || (args.outcomeStatus !== 429 && args.outcomeStatus !== 402) + || args.firstAuthCtx.fixedAccount || args.sameAccountOnly) { + return retryCodexPoolOnOneAlternateAccount(args); + } + // Only authoritative pre-stream quota refusals enter this loop. A sent WS frame, + // a body failure, or an ambiguous transport error never authorizes another send. + const attemptedAccountIds = new Set([args.firstAuthCtx.accountId]); + const accountBudget = new Set([ + MAIN_CODEX_ACCOUNT_ID, ...(args.config.codexAccounts ?? []).map(account => account.id), + ]).size; + let current = args; + let last: CodexPoolAccountRetryResult = { kind: "no-alternate" }; + while (attemptedAccountIds.size < accountBudget) { + if (args.upstream.signal.aborted || args.options.abortSignal?.aborted) { + return { kind: "transport", error: args.upstream.signal.reason + ?? args.options.abortSignal?.reason, authCtx: current.firstAuthCtx }; + } + const retry = await retryCodexPoolOnOneAlternateAccount({ ...current, attemptedAccountIds }); + if (retry.kind === "no-alternate") { + if (!retry.quotaWaitable) return last; + if (last.kind === "retried") { + last.upstreamResponse = codexQuotaWaitResponse(last.upstreamResponse); + return last; + } + return retry; + } + if (retry.kind === "transport") return retry; + last = retry; + if (retry.authCtx.kind === "main" + || !await shouldRetryCodexPoolAccountQuota(retry.upstreamResponse, args.options.abortSignal)) return retry; + current = { + ...args, + firstAuthCtx: retry.authCtx, + firstResponse: retry.upstreamResponse, + outcomeStatus: retry.upstreamResponse.status >= 500 ? 429 : retry.upstreamResponse.status, + }; + } + // Every available account in the frozen budget has explicitly rejected this turn. + if (last.kind === "retried") { + last.upstreamResponse = codexQuotaWaitResponse(last.upstreamResponse); + return last; + } + return { kind: "no-alternate", quotaWaitable: true }; +} + +/** Internal policy signal; only attach to an authoritative pre-stream quota rejection. */ +function codexQuotaWaitResponse(response: Response): Response { + const headers = new Headers(response.headers); + headers.set("x-opencodex-quota-wait", "1"); + const marked = new Response(response.body, { status: response.status, statusText: response.statusText, headers }); + return markStrictQuotaWaitResponse(marked); +} + +async function retryCodexPoolOnOneAlternateAccount( + args: CodexPoolAccountRetryArgs, ): Promise { const { req, config, route, parsed, logCtx, options, firstAuthCtx, firstResponse, @@ -1190,6 +1252,8 @@ async function retryCodexPoolOnAlternateAccount( "pool", { excludeAccountId: firstAuthCtx.accountId, + ...(args.attemptedAccountIds ? { excludeAccountIds: args.attemptedAccountIds } : {}), + signal: options.abortSignal ?? upstream.signal, admission: options.admission, codexAuthPolicy: options.codexAuthPolicy, modelId: route.modelId, @@ -1199,6 +1263,9 @@ async function retryCodexPoolOnAlternateAccount( }, ); } catch (error) { + if (args.attemptedAccountIds && error instanceof CodexStrictQuotaUnavailableError && error.waitable) { + return { kind: "no-alternate", quotaWaitable: true }; + } const unexpectedRetryError = !(error instanceof CodexPoolAuthenticationError) && !(error instanceof CodexAuthContextError) @@ -1234,6 +1301,16 @@ async function retryCodexPoolOnAlternateAccount( return { kind: "no-alternate" }; } + if (args.attemptedAccountIds) { + const selectedId = retryAuthCtx.accountId ?? MAIN_CODEX_ACCOUNT_ID; + // Defence in depth against a selector racing a config/credential update. + if (args.attemptedAccountIds.has(selectedId)) { + releaseCodexAuthContextProbeLease(retryAuthCtx); + return { kind: "no-alternate" }; + } + args.attemptedAccountIds.add(selectedId); + } + const quotaMeta = { ...codexQuotaOutcomeMeta(firstResponse), ...(await codexDenialOutcomeMeta(firstResponse)) }; if (outcomeStatus === 429 || outcomeStatus === 402) { const { applyAccountQuotaFromUpstreamHeaders } = await import("../../codex/auth-api"); @@ -1289,6 +1366,16 @@ async function retryCodexPoolOnAlternateAccount( recordAdapterTier(logCtx, request); await firstResponse.body?.cancel().catch(() => undefined); + if (args.attemptedAccountIds) { + if (logCtx.activeAttempt) finishRequestAttempt(logCtx.activeAttempt, firstResponse.status, + Date.now() - (logCtx.activeAttemptStartedAt ?? Date.now())); + const attempt = beginRequestAttempt((logCtx.attempts?.length ?? 0) + 1, + formatCodexProviderForLog(route.providerName, retryAuthCtx.accountId, config), + route.modelId, retryAdapter.name); + logCtx.activeAttempt = attempt; + logCtx.activeAttemptStartedAt = Date.now(); + (logCtx.attempts ??= []).push(attempt); + } options.onCodexAuthContextResolved?.(retryAuthCtx); route.provider = retryProvider; logCtx.provider = formatCodexProviderForLog( @@ -1317,6 +1404,9 @@ async function retryCodexPoolOnAlternateAccount( let upstreamResponse: Response; try { while (true) { + if (upstream.signal.aborted || options.abortSignal?.aborted) { + return { kind: "transport", error: upstream.signal.reason ?? options.abortSignal?.reason, authCtx: retryAuthCtx }; + } noteAttemptSend(logCtx.activeAttempt, passthroughEstimate); try { upstreamResponse = await fetchWithHeaderTimeout( @@ -1512,6 +1602,13 @@ export interface ConsumedComboFailure { +export interface ResponsesReplaySnapshot { + sourceBody: unknown; + previousResponseInputExpanded: boolean; + providerContinuation: OcxProviderContinuationState | undefined; + recoveredPlaintext: boolean; +} + export interface HandleResponsesOptions { /** Original live policy owner; separate from caller-specific routing/sidecar snapshots. */ codexAuthPolicy?: CodexAuthPolicyConfig; @@ -1568,12 +1665,11 @@ export interface HandleResponsesOptions { /** Internal recursion guard; callers outside this module must not set it. */ comboAttempt?: boolean; /** Internal combo handoff for one parent-validated continuation snapshot. */ - comboReplaySnapshot?: { - sourceBody: unknown; - previousResponseInputExpanded: boolean; - providerContinuation: OcxProviderContinuationState | undefined; - recoveredPlaintext: boolean; - }; + comboReplaySnapshot?: ResponsesReplaySnapshot; + /** Internal same-request quota handoff; never populated from client JSON. */ + quotaReplaySnapshot?: ResponsesReplaySnapshot; + /** Lazy capture: materialize only when a trusted quota refusal enters waiting. */ + onQuotaReplaySnapshot?: (capture: () => ResponsesReplaySnapshot) => void; /** Internal combo handoff: allow a later same-provider model after a reset-derived 429/402. */ deferCodexResetDerivedCooldown?: boolean; /** 030-owned handoff when a child consumed the original failure under bounds. */ @@ -2762,6 +2858,7 @@ function finalizeOwnedTranslatorBudget(response: Response, budget: TranslatorBud statusText: response.statusText, headers: response.headers, }); + if (isStrictQuotaWaitResponse(response)) markStrictQuotaWaitResponse(finalizedResponse); if (isNativePassthroughSseResponse(response)) { markNativePassthroughSseResponse(finalizedResponse); } @@ -2917,8 +3014,9 @@ async function handleResponsesInner( const inboundClientThreadId = req.headers.get("x-codex-parent-thread-id")?.trim() || undefined; const cursorClientThreadId = codexPoolAffinityKey(req.headers); const originalBody = body; - if (options.comboReplaySnapshot) { - copyPreviousResponseReplayProvenance(options.comboReplaySnapshot.sourceBody, body); + const replaySnapshot = options.comboReplaySnapshot ?? options.quotaReplaySnapshot; + if (replaySnapshot) { + copyPreviousResponseReplayProvenance(replaySnapshot.sourceBody, body); } else { body = expandPreviousResponseInput(body, inboundClientThreadId); if (previousResponseScopeMismatch(body)) { @@ -2932,7 +3030,7 @@ async function handleResponsesInner( ); } } - const previousResponseInputExpanded = options.comboReplaySnapshot?.previousResponseInputExpanded + const previousResponseInputExpanded = replaySnapshot?.previousResponseInputExpanded ?? (body !== originalBody && typeof (body as { previous_response_id?: unknown }).previous_response_id === "string"); @@ -2978,13 +3076,13 @@ async function handleResponsesInner( effort: effortRow.effort, }; } - if (options.comboReplaySnapshot?.recoveredPlaintext) { + if (replaySnapshot?.recoveredPlaintext) { markBodyNonPersistable(parsed._rawBody); } toolBridgeMaps = buildToolBridgeMaps(parsed, translatorBudget); if (previousResponseInputExpanded) parsed._previousResponseInputExpanded = true; - const providerContinuationCandidate = options.comboReplaySnapshot - ? options.comboReplaySnapshot.providerContinuation + const providerContinuationCandidate = replaySnapshot + ? replaySnapshot.providerContinuation : previousResponseProviderState(parsed.previousResponseId); if (providerContinuationCandidate) parsed._providerContinuationCandidate = providerContinuationCandidate; if (inboundClientThreadId) { @@ -3020,6 +3118,21 @@ async function handleResponsesInner( } return formatErrorResponse(400, "invalid_request_error", err instanceof Error ? err.message : String(err)); } + // Preserve the validated expanded input before route-specific top-level rewrites. The + // shallow envelope retains no extra copy of large content; the full clone happens only + // when the caller actually enters quota waiting. Scope/provenance remain proxy-private. + let quotaRecoveredPlaintext = replaySnapshot?.recoveredPlaintext ?? false; + if (options.onQuotaReplaySnapshot) { + const replayBody = { ...(body as Record) }; + copyPreviousResponseReplayProvenance(body, replayBody); + const providerContinuation = parsed._providerContinuationCandidate; + options.onQuotaReplaySnapshot(() => { + const sourceBody = structuredClone(replayBody); + copyPreviousResponseReplayProvenance(replayBody, sourceBody); + return { sourceBody, previousResponseInputExpanded, providerContinuation, + recoveredPlaintext: quotaRecoveredPlaintext }; + }); + } options.onRequestBodyRead?.(); const responseStateOptions = (force = false): { force?: boolean; clientThreadId?: string } => ({ ...(force ? { force: true } : {}), @@ -3290,6 +3403,7 @@ async function handleResponsesInner( // text. Bar it from the continuation cache before any recording path can reach it — // that cache is persisted to disk, which would defeat the recovery cache's TTL. markBodyNonPersistable(parsed._rawBody); + quotaRecoveredPlaintext = true; // The ciphertext-only pass intentionally excludes routed candidates. Once recovery // makes the assignment readable, run selection again with the full configured chain @@ -4855,6 +4969,9 @@ async function handleResponsesInner( ); }, }); + if (retry.kind === "no-alternate" && retry.quotaWaitable) { + upstreamResponse = codexQuotaWaitResponse(upstreamResponse); + } if (retry.kind === "transport") { authCtx = retry.authCtx; return transportFailureResponse(retry.error); @@ -4891,6 +5008,9 @@ async function handleResponsesInner( break; } const headers = sanitizePassthroughHeaders(upstreamResponse.headers); + // An upstream-provided header is not authority to replay the request. + headers.delete("x-opencodex-quota-wait"); + if (isStrictQuotaWaitResponse(upstreamResponse)) headers.set("x-opencodex-quota-wait", "1"); const resolvedModel = headers.get("openai-model")?.trim(); if (resolvedModel && !logCtx.preserveResolvedModelFromRoute) logCtx.resolvedModel = resolvedModel; if (isUsageDebugEnabled()) { @@ -4980,7 +5100,7 @@ async function handleResponsesInner( return new Response(upstreamResponse.body, { status: upstreamResponse.status, statusText: upstreamResponse.statusText, - headers: sanitizePassthroughHeaders(upstreamResponse.headers), + headers, }); } if (!upstreamResponse.ok) { @@ -5004,10 +5124,13 @@ async function handleResponsesInner( translatorBudget, ); } - return formatPassthroughUpstreamError(upstreamResponse.status, errorText, { + const formattedError = formatPassthroughUpstreamError(upstreamResponse.status, errorText, { statusText: upstreamResponse.statusText, headers, }); + return isStrictQuotaWaitResponse(upstreamResponse) + ? markStrictQuotaWaitResponse(formattedError) + : formattedError; } // Bun#32111 workaround: passthrough SSE uses tee()+native relay to avoid the diff --git a/src/server/responses/policy-fallback.ts b/src/server/responses/policy-fallback.ts index a4f06d0fa6..6dcfee107d 100644 --- a/src/server/responses/policy-fallback.ts +++ b/src/server/responses/policy-fallback.ts @@ -1,10 +1,13 @@ +import { getCodexQuotaRevision } from "../../codex/quota-events"; +import { isStrictQuotaWaitResponse } from "./strict-quota-response"; +import { waitForStrictQuotaResponse, type StrictQuotaWaitOptions } from "./strict-quota-wait"; import { comboFailureDecision } from "../../combos/failover"; import { readBoundedResponseBody } from "../../lib/bounded-body"; import { readJsonRequestBody } from "../request-decompress"; import { finishRequestAttempt, type RequestLogContext } from "../request-log"; import type { OcxConfig } from "../../types"; import type { RouteCandidateTrace, RouteDecisionTraceV1 } from "../../routing/trace"; -import { handleResponses as handleResponsesCore } from "./core"; +import { handleResponses as handleResponsesCore, type ResponsesReplaySnapshot } from "./core"; import { requestPacingOverloadResponse } from "./pacing-overload"; type CoreHandler = typeof handleResponsesCore; @@ -12,6 +15,7 @@ type CoreOptions = Parameters[3]; export interface PolicyFallbackDeps { runCore?: CoreHandler; + quotaWait?: Pick; } function candidateKey(candidate: Pick): string { @@ -41,6 +45,14 @@ export function rankPolicyFallbackCandidates( .map(({ candidate }) => candidate); } +function requestWithBody(req: Request, rawBody: Record, signal = req.signal): Request { + const headers = new Headers(req.headers); + headers.delete("content-encoding"); + headers.delete("content-length"); + headers.set("content-type", "application/json"); + return new Request(req.url, { method: req.method, headers, body: JSON.stringify(rawBody), signal }); +} + function requestWithCandidate( req: Request, rawBody: Record, @@ -117,8 +129,10 @@ export async function handleResponsesWithPolicyFallback( const runCore = deps.runCore ?? handleResponsesCore; let requestBodyReadNotified = false; let storedPool401ReplayDispatched = false; + let captureQuotaReplay: (() => ResponsesReplaySnapshot) | undefined; const coreOptions: CoreOptions = { ...options, + onQuotaReplaySnapshot: capture => { captureQuotaReplay = capture; }, ...(options.onRequestBodyRead ? { onRequestBodyRead: () => { if (requestBodyReadNotified) return; @@ -140,6 +154,7 @@ export async function handleResponsesWithPolicyFallback( } let response: Response; + let quotaRevisionBeforeCore = getCodexQuotaRevision(); try { response = await runCore(req, config, logCtx, coreOptions); } catch (error) { @@ -149,7 +164,46 @@ export async function handleResponsesWithPolicyFallback( } const initialTrace = logCtx.routeDecision; const initialRequestedModel = logCtx.requestedModel; - if (!rawBody || !isPolicyDecision(initialTrace)) return response; + const settleQuotaWait = (first: Response, raw: Record): Promise => { + const snapshot = captureQuotaReplay?.(); + captureQuotaReplay = undefined; + const body = (snapshot?.sourceBody ?? raw) as Record; + return waitForStrictQuotaResponse({ + config, quotaPolicy: options.codexAuthPolicy, initial: first, stream: body.stream === true, + signals: [req.signal, options.abortSignal], lease: options.turnAdmissionLease, + canReplay: () => !storedPool401ReplayDispatched, + finishAttempt: status => finishFailedPolicyAttempt(logCtx, status), + observedRevision: () => quotaRevisionBeforeCore, + onFailure: status => { logCtx.terminalHttpStatus = status; }, + ...deps.quotaWait, + resume: async signal => { + quotaRevisionBeforeCore = getCodexQuotaRevision(); + try { + return await runCore(requestWithBody(req, body, signal), config, logCtx, { + ...coreOptions, abortSignal: signal, quotaReplaySnapshot: snapshot, + onQuotaReplaySnapshot: undefined, + // The heartbeat response is inspected by the ordinary outer SSE logger. Native + // callbacks would otherwise finalize the same logical request a second time. + ...(body.stream === true ? { + onNativePassthroughTerminal: undefined, onNativePassthroughCancel: undefined, + } : {}), + }); + } catch (error) { + const overload = requestPacingOverloadResponse(error); + if (overload) return overload; + throw error; + } finally { + logCtx.requestedModel = initialRequestedModel; + logCtx.routeDecision = initialTrace; + } + }, + }); + }; + if (!rawBody || !isPolicyDecision(initialTrace)) { + return rawBody && !storedPool401ReplayDispatched && isStrictQuotaWaitResponse(response) + ? settleQuotaWait(response, rawBody) : response; + } + let quotaReplayBody = rawBody; const tried = new Set([ candidateKey({ provider: initialTrace.selected.provider, model: initialTrace.selected.model }), @@ -158,13 +212,16 @@ export async function handleResponsesWithPolicyFallback( while (!storedPool401ReplayDispatched && await shouldHopPolicyCandidate(response, req.signal)) { if (req.signal.aborted) return response; const next = rankPolicyFallbackCandidates(initialTrace, tried)[0]; - if (!next) return response; + if (!next) break; tried.add(candidateKey(next)); finishFailedPolicyAttempt(logCtx, response.status); const retryRequest = requestWithCandidate(req, rawBody, next); + quotaReplayBody = { ...rawBody, model: `${next.provider}/${next.model}` }; try { try { + quotaRevisionBeforeCore = getCodexQuotaRevision(); + captureQuotaReplay = undefined; response = await runCore(retryRequest, config, logCtx, coreOptions); } catch (error) { const overload = requestPacingOverloadResponse(error); @@ -177,7 +234,9 @@ export async function handleResponsesWithPolicyFallback( } } - return response; + // Exhaust the operator's existing policy candidates before waiting on the final quota lane. + return !storedPool401ReplayDispatched && isStrictQuotaWaitResponse(response) + ? settleQuotaWait(response, quotaReplayBody) : response; } export const handleResponses = handleResponsesWithPolicyFallback; diff --git a/src/server/responses/strict-quota-response.ts b/src/server/responses/strict-quota-response.ts new file mode 100644 index 0000000000..bee8585bbb --- /dev/null +++ b/src/server/responses/strict-quota-response.ts @@ -0,0 +1,10 @@ +/** Process-local proof: an upstream header alone never authorizes replaying a request. */ +const waitableResponses = new WeakSet(); +export function markStrictQuotaWaitResponse(response: Response): Response { + response.headers.set("x-opencodex-quota-wait", "1"); + waitableResponses.add(response); + return response; +} +export function isStrictQuotaWaitResponse(response: Response): boolean { + return waitableResponses.has(response); +} diff --git a/src/server/responses/strict-quota-wait.ts b/src/server/responses/strict-quota-wait.ts new file mode 100644 index 0000000000..59fa00ca35 --- /dev/null +++ b/src/server/responses/strict-quota-wait.ts @@ -0,0 +1,200 @@ +import type { OcxConfig } from "../../types"; +import type { CodexAuthPolicyConfig } from "../../codex/auth-context"; +import type { AdmissionLease } from "../../lib/admission"; +import { readBoundedResponseBody } from "../../lib/bounded-body"; +import { MAIN_CODEX_ACCOUNT_ID } from "../../codex/account-id"; +import { refreshStrictCodexQuotasOnDemand, waitForStrictCodexQuotaChange } from "../../codex/strict-quota-refresh"; +import { registerTurn, unregisterTurn } from "../lifecycle"; +import { isStrictQuotaWaitResponse } from "./strict-quota-response"; + +const HEARTBEAT = new TextEncoder().encode( + 'event: response.heartbeat\ndata: {"type":"response.heartbeat"}\n\n', +); + +export interface StrictQuotaWaitOptions { + config: OcxConfig; + quotaPolicy?: CodexAuthPolicyConfig; + initial: Response; + stream: boolean; + signals: readonly (AbortSignal | undefined)[]; + lease?: AdmissionLease; + canReplay: () => boolean; + resume: (signal: AbortSignal) => Promise; + finishAttempt: (status: number) => void; + observedRevision?: () => number; + onFailure?: (status: number) => void; + /** Test seam: production uses the coalesced, event-driven quota waiter. */ + waitForChange?: typeof waitForStrictCodexQuotaChange; + heartbeatMs?: number; +} + +/** A real rejected replay needs a failed terminal once the heartbeat committed SSE headers. */ +async function rejectedResponseFrame(response: Response, signal: AbortSignal): Promise { + let message = response.ok + ? "The resumed request returned an unexpected non-streaming response" + : `The resumed request failed (HTTP ${response.status})`; + let code = response.ok ? "protocol_error" : "upstream_error"; + try { + const body = await readBoundedResponseBody(response, { signal }); + if (body.displaySafe && !body.truncated) { + const parsed = JSON.parse(body.text) as { error?: { message?: unknown; code?: unknown; type?: unknown } }; + if (typeof parsed.error?.message === "string") message = parsed.error.message.slice(0, 2048); + if (typeof parsed.error?.code === "string") code = parsed.error.code.slice(0, 128); + else if (typeof parsed.error?.type === "string") code = parsed.error.type.slice(0, 128); + } + } catch { + // An unreadable error body keeps the truthful HTTP-status terminal, never a success. + } + return new TextEncoder().encode(`event: response.failed\ndata: ${JSON.stringify({ + type: "response.failed", response: { status: "failed", error: { code, message } }, + })}\n\n`); +} + +/** + * Keep one admitted request alive until a fresh quota observation permits another safe attempt. + * Only process-local refusal evidence enters this owner; it never retries an accepted stream. + */ +export async function waitForStrictQuotaResponse(options: StrictQuotaWaitOptions): Promise { + const ac = new AbortController(); + const parentListeners: Array<() => void> = []; + for (const signal of new Set(options.signals.filter((value): value is AbortSignal => !!value))) { + const abort = () => ac.abort(signal.reason); + if (signal.aborted) abort(); + else { + signal.addEventListener("abort", abort, { once: true }); + parentListeners.push(() => signal.removeEventListener("abort", abort)); + } + } + const ownsTurn = !!options.lease && "bindAbortController" in options.lease; + if (ownsTurn) registerTurn(ac, options.lease); + let timer: ReturnType | undefined; + let reader: ReadableStreamDefaultReader | undefined; + let controller: ReadableStreamDefaultController | undefined; + let closed = false; + let waiting = true; + const cleanup = () => { + if (closed) return; + closed = true; + if (timer !== undefined) clearInterval(timer); + timer = undefined; + for (const remove of parentListeners) remove(); + ac.signal.removeEventListener("abort", aborted); + if (ownsTurn) unregisterTurn(ac); + }; + const aborted = () => { + const reason = ac.signal.reason ?? new DOMException("Aborted", "AbortError"); + void reader?.cancel(reason).catch(() => {}); + try { controller?.error(reason); } catch { /* already closed */ } + cleanup(); + }; + ac.signal.addEventListener("abort", aborted, { once: true }); + + const resume = async (): Promise => { + let response = options.initial; + try { + while (isStrictQuotaWaitResponse(response) && options.canReplay()) { + ac.signal.throwIfAborted(); + // Subscribe before releasing the old response, so quota updates during disposal wake us. + const changed = (options.waitForChange ?? waitForStrictCodexQuotaChange)(options.config, ac.signal, options.observedRevision?.()); + void changed.catch(() => {}); + await response.body?.cancel().catch(() => {}); + options.finishAttempt(response.status); + // Caller-owned routing cannot inspect the physical main profile. This pending + // request uses the metadata owner's separate claim to recover missing/stale usage; + // no caller credential crosses that boundary and no work runs while idle. + const refreshConfig = options.quotaPolicy ? { ...options.config, ...options.quotaPolicy } : options.config; + await refreshStrictCodexQuotasOnDemand(refreshConfig, new Set([MAIN_CODEX_ACCOUNT_ID]), { signal: ac.signal }); + await changed; + ac.signal.throwIfAborted(); + response = await options.resume(ac.signal); + } + ac.signal.throwIfAborted(); + return response; + } catch (error) { + await response.body?.cancel(error).catch(() => {}); + throw error; + } + }; + + if (!options.stream) { + try { + const response = await resume(); + waiting = false; + if (!response.body) { cleanup(); return response; } + reader = response.body.getReader(); + const body = new ReadableStream({ + start(value) { controller = value; }, + async pull(value) { + try { + ac.signal.throwIfAborted(); + const next = await reader!.read(); + if (closed) return; + if (next.done) { cleanup(); value.close(); } + else value.enqueue(next.value); + } catch (error) { cleanup(); if (!ac.signal.aborted) value.error(error); } + }, + async cancel(reason) { ac.abort(reason); await reader?.cancel(reason).catch(() => {}); cleanup(); }, + }); + const headers = new Headers(response.headers); + headers.delete("x-opencodex-quota-wait"); + return new Response(body, { status: response.status, statusText: response.statusText, headers }); + } catch (error) { cleanup(); throw error; } + } + + // Start no model work while backpressured: resume waits on metadata, and the final body is + // read only by pull(). Heartbeats occupy at most one queued chunk while no client is reading. + const pending = resume(); + void pending.then(response => { + waiting = false; + if (timer !== undefined) clearInterval(timer); + timer = undefined; + if (closed) void response.body?.cancel().catch(() => {}); + }, () => { + waiting = false; + if (timer !== undefined) clearInterval(timer); + timer = undefined; + }); + const body = new ReadableStream({ + start(value) { + controller = value; + if (ac.signal.aborted) { aborted(); return; } + value.enqueue(HEARTBEAT); + timer = setInterval(() => { + if (!closed && waiting && (value.desiredSize ?? 0) > 0) value.enqueue(HEARTBEAT); + }, options.heartbeatMs ?? 2_000); + timer.unref?.(); + }, + async pull(value) { + try { + if (!reader) { + const response = await pending; + if (closed) { await response.body?.cancel(); return; } + waiting = false; + if (timer !== undefined) clearInterval(timer); + timer = undefined; + if (!response.ok || !response.headers.get("content-type")?.includes("text/event-stream") || !response.body) { + options.onFailure?.(response.ok ? 502 : response.status); + value.enqueue(await rejectedResponseFrame(response, ac.signal)); + cleanup(); value.close(); return; + } + reader = response.body.getReader(); + } + const next = await reader.read(); + if (closed) return; + if (next.done) { cleanup(); value.close(); } + else value.enqueue(next.value); + } catch (error) { + if (closed) return; + if (ac.signal.aborted) { aborted(); return; } + options.onFailure?.(502); + value.enqueue(new TextEncoder().encode('event: response.failed\ndata: {"type":"response.failed","response":{"status":"failed","error":{"code":"proxy_error","message":"The quota wait could not resume this request"}}}\n\n')); + cleanup(); value.close(); + } + }, + async cancel(reason) { ac.abort(reason); await reader?.cancel(reason).catch(() => {}); cleanup(); }, + }); + return new Response(body, { headers: { + "content-type": "text/event-stream; charset=utf-8", "cache-control": "no-cache", + "x-accel-buffering": "no", connection: "keep-alive", + } }); +} diff --git a/src/types/config.ts b/src/types/config.ts index 27b1a8ec81..bd91c0cc37 100644 --- a/src/types/config.ts +++ b/src/types/config.ts @@ -764,6 +764,8 @@ export interface OcxConfig { activeCodexAccountId?: string; /** Auto-switch threshold (0-100). Default 80. 0 = disabled. */ autoSwitchThreshold?: number; + /** Default off. Require fresh unexhausted quota; prefer accounts below autoSwitchThreshold. */ + codexAccountStrictQuota?: boolean; /** New-session account rotation strategy for the Codex pool. Default quota (today's behaviour). */ accountPoolStrategy?: OcxAccountPoolRotationStrategy; /** Successful new-session binds retained on one round-robin selection. Default 1; range 1..100. */ diff --git a/tests/codex-integration/codex-auth-api.test.ts b/tests/codex-integration/codex-auth-api.test.ts index 04e19d49fa..988039a16b 100644 --- a/tests/codex-integration/codex-auth-api.test.ts +++ b/tests/codex-integration/codex-auth-api.test.ts @@ -9,6 +9,9 @@ import { getNativeMainProfileRequestCount, resetLifecycleDrainStateForTests, } from "../../src/server/lifecycle"; +import { getCodexStrictQuotaStatus } from "../../src/codex/strict-quota"; +import { resolveCodexAuthContext } from "../../src/codex/auth-context"; +import { getCodexQuotaRevision } from "../../src/codex/quota-events"; import { fallbackCodexAccountLogLabel } from "../../src/codex/account-label"; import { handleCodexAuthAPI, updateAccountQuota, getAccountQuota, @@ -1486,6 +1489,7 @@ describe("codex-auth API", () => { pinnedAccountId: null, autoSwitchThreshold: 55, upstreamFailoverThreshold: 3, + codexAccountStrictQuota: false, accountPoolStrategy: "quota", accountPoolStickyLimit: 1, }); @@ -3478,6 +3482,79 @@ describe("codex-auth API", () => { } }); + test("strict quota opt-in uses the existing threshold API and validates before mutation", async () => { + const config = makeConfig({ autoSwitchThreshold: 80 }); + const put = async (body: unknown) => { + const req = new Request("http://localhost/api/codex-auth/auto-switch", { + method: "PUT", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body), + }); + return handleCodexAuthAPI(req, new URL(req.url), config); + }; + expect((await put({ threshold: 95, strictQuota: "true" }))!.status).toBe(400); + expect(config.autoSwitchThreshold).toBe(80); + expect(config.codexAccountStrictQuota).toBeUndefined(); + expect((await put({ threshold: 95, strictQuota: true }))!.status).toBe(200); + expect(config.codexAccountStrictQuota).toBe(true); + expect((await put({ threshold: 90 }))!.status).toBe(200); + expect(config.codexAccountStrictQuota).toBe(true); + const req = new Request("http://localhost/api/codex-auth/active"); + const state = await (await handleCodexAuthAPI(req, new URL(req.url), config))!.json(); + expect(state).toMatchObject({ autoSwitchThreshold: 90, codexAccountStrictQuota: true }); + expect((await put({ threshold: 90, strictQuota: false }))!.status).toBe(200); + expect(config.codexAccountStrictQuota).toBe(false); + }); + + for (const action of ["select-main", "enable-strict"] as const) { + test(`${action} primes unknown main usage under management ownership before waking requests`, async () => { + const config = makeConfig({ codexAccountStrictQuota: true, autoSwitchThreshold: action === "enable-strict" ? 0 : 95, + activeCodexAccountId: MAIN_CODEX_ACCOUNT_ID, activeCodexAccountPinned: MAIN_CODEX_ACCOUNT_ID }); + writeFileSync(join(TEST_CODEX_HOME, "auth.json"), JSON.stringify({ + tokens: { access_token: "main-strict-owned", account_id: "acct-main-strict-owned" }, + })); + expect(getCodexStrictQuotaStatus({ ...config, autoSwitchThreshold: 95 }, MAIN_CODEX_ACCOUNT_ID).state).toBe("unknown"); + const urls: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + urls.push(String(input)); + if (String(input) !== "https://chatgpt.com/backend-api/wham/usage") throw new Error("unexpected non-metadata request"); + return Response.json({ plan_type: "plus", rate_limit: { secondary_window: { used_percent: 0 } } }); + }) as typeof fetch; + const req = new Request(`http://localhost/api/codex-auth/${action === "select-main" ? "active" : "auto-switch"}`, { + method: "PUT", headers: { "Content-Type": "application/json" }, + body: JSON.stringify(action === "select-main" ? { accountId: MAIN_CODEX_ACCOUNT_ID } : { threshold: 95, strictQuota: true }), + }); + const revision = getCodexQuotaRevision(); + expect((await handleCodexAuthAPI(req, new URL(req.url), config))!.status).toBe(200); + expect(urls).toEqual(["https://chatgpt.com/backend-api/wham/usage"]); + expect(getCodexQuotaRevision()).toBeGreaterThan(revision); + expect(getCodexStrictQuotaStatus(config, MAIN_CODEX_ACCOUNT_ID)).toMatchObject({ state: "ready", usedPercent: 0 }); + const context = await resolveCodexAuthContext(new Headers({ authorization: "Bearer main-strict-owned", "chatgpt-account-id": "acct-main-strict-owned" }), config, "pool", { + requestScopedMainCredential: true, + isMainAccountTokenLive: () => { throw new Error("request must not inspect native main"); }, + }); + expect(context.kind).toBe("main"); + }); + } + test("failed strict main metadata priming preserves unknown quota and never redeems credits", async () => { + const config = makeConfig({ codexAccountStrictQuota: true, autoSwitchThreshold: 95, + activeCodexAccountId: MAIN_CODEX_ACCOUNT_ID, activeCodexAccountPinned: MAIN_CODEX_ACCOUNT_ID }); + writeFileSync(join(TEST_CODEX_HOME, "auth.json"), JSON.stringify({ + tokens: { access_token: "main-strict-failed", account_id: "acct-main-strict-failed" }, + })); + const urls: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL) => { + urls.push(String(input)); + return new Response("unavailable", { status: 503 }); + }) as typeof fetch; + const req = new Request("http://localhost/api/codex-auth/active", { method: "PUT", + headers: { "Content-Type": "application/json" }, body: JSON.stringify({ accountId: MAIN_CODEX_ACCOUNT_ID }) }); + expect((await handleCodexAuthAPI(req, new URL(req.url), config))!.status).toBe(200); + expect(urls).toEqual(["https://chatgpt.com/backend-api/wham/usage"]); + expect(getCodexStrictQuotaStatus(config, MAIN_CODEX_ACCOUNT_ID).state).toBe("unknown"); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + await expect(resolveCodexAuthContext(new Headers({ authorization: "Bearer main-strict-failed", "chatgpt-account-id": "acct-main-strict-failed" }), config, "pool", { + requestScopedMainCredential: true, accountId: MAIN_CODEX_ACCOUNT_ID, + })).rejects.toMatchObject({ name: "CodexStrictQuotaUnavailableError" }); + }); test("PUT /api/codex-auth/active mutates live runtime config", async () => { const config = makeConfig({ codexAccounts: [{ id: "pool-next", email: "pool-next@example.test", isMain: false }], @@ -3487,9 +3564,11 @@ describe("codex-auth API", () => { headers: { "Content-Type": "application/json" }, body: JSON.stringify({ accountId: "pool-next" }), }); + const revision = getCodexQuotaRevision(); const resp = await handleCodexAuthAPI(req, new URL(req.url), config); expect(resp!.status).toBe(200); expect(await resp!.json()).toMatchObject({ activeCodexAccountId: "pool-next", appliesImmediately: true }); + expect(getCodexQuotaRevision()).toBeGreaterThan(revision); expect(config.activeCodexAccountId).toBe("pool-next"); }); diff --git a/tests/codex-integration/codex-auth-context.test.ts b/tests/codex-integration/codex-auth-context.test.ts index bb4f24a098..46b985c6f7 100644 --- a/tests/codex-integration/codex-auth-context.test.ts +++ b/tests/codex-integration/codex-auth-context.test.ts @@ -7,6 +7,8 @@ import { assertCodexAuthContextNotCooled, CODEX_MAIN_PROFILE_MAINTENANCE_MESSAGE, CodexAccountCooldownError, + CodexStrictQuotaUnavailableError, + createCodexReserveDispatchGuard, CodexAuthContextError, CodexDirectAuthenticationError, CodexMainProfileDrainingError, @@ -73,6 +75,8 @@ import { } from "../../src/server/lifecycle"; import type { CodexModelEntitlementSnapshot } from "../../src/codex/model-entitlements"; import { hasForwardableCodexBearer } from "../../src/server/auth-cors"; +import { refreshStrictCodexQuotasOnDemand, setStrictCodexQuotaRefreshForTests } from "../../src/codex/strict-quota-refresh"; +import { captureMainQuotaWriter, observeMainQuotaIdentity } from "../../src/codex/main-account-cache"; import { removeTreeWithRetry } from "../helpers/remove-tree"; let testDir: string; @@ -2228,3 +2232,206 @@ describe("native-main fence names its gate reason", () => { } }); }); + + +describe("strict quota Pool auth admission", () => { + function strictFixture() { + const cfg = config(); + cfg.codexAccountStrictQuota = true; cfg.autoSwitchThreshold = 95; + cfg.activeCodexAccountId = MAIN_CODEX_ACCOUNT_ID; cfg.activeCodexAccountPinned = MAIN_CODEX_ACCOUNT_ID; + resetCodexRoutingForManualSelection(MAIN_CODEX_ACCOUNT_ID); + observeMainQuotaIdentity("strict-main-owner"); + const writer = captureMainQuotaWriter("strict-main-owner")!; + const mainQuota = (weeklyPercent: number) => setAccountQuotaFromParsed(MAIN_CODEX_ACCOUNT_ID, { weeklyPercent }, undefined, writer); + mainQuota(20); + saveCodexAccountCredential("pool-a", { accessToken: "strict-pool-access", refreshToken: "strict-pool-refresh", + chatgptAccountId: "strict-pool-owner", expiresAt: Date.now() + 3600000 }); + setAccountQuotaFromParsed("pool-a", { weeklyPercent: 10 }); + const headers = new Headers({ authorization: "Bearer access-token-strict-test", "chatgpt-account-id": "strict-main-owner" }); + return { cfg, mainQuota, headers }; + } + const poolEntitlements = async (): Promise => ({ + modelsByAccount: new Map([["pool-a", new Set(["gpt-daybreak-blue-latest"])]]), + confirmedAccountIds: new Set(["pool-a"]), credentialIdentities: new Map(), + }); + for (const evidence of ["missing", "stale"] as const) { + test(`${evidence} caller-main evidence requests quota recovery instead of a native-profile drain`, async () => { + const { cfg, mainQuota, headers } = strictFixture(); + cfg.codexAccounts = []; + if (evidence === "missing") clearAccountQuota(MAIN_CODEX_ACCOUNT_ID); + else { + const now = Date.now(); + const clock = spyOn(Date, "now").mockReturnValue(now - 6 * 60_000); + try { mainQuota(0); } finally { clock.mockRestore(); } + } + await expect(resolveCodexAuthContext(headers, cfg, "pool", { + requestScopedMainCredential: true, + isMainAccountTokenLive: () => { throw new Error("must not inspect native main"); }, + getMainAccountToken: () => { throw new Error("must not inspect native main"); }, + })).rejects.toMatchObject({ name: "CodexStrictQuotaUnavailableError", waitable: true }); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { + requestScopedMainCredential: true, accountId: MAIN_CODEX_ACCOUNT_ID, + })).rejects.toMatchObject({ name: "CodexStrictQuotaUnavailableError", waitable: false }); + }); + } + test("a manually pinned caller-owned main with fresh zero usage remains available at 95", async () => { + const { cfg, mainQuota, headers } = strictFixture(); mainQuota(0); + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", { + requestScopedMainCredential: true, + isMainAccountTokenLive: () => { throw new Error("must not inspect native main"); }, + }); + expect(ctx).toMatchObject({ kind: "main", poolQuotaScope: "shared" }); + }); + for (const replacement of ["below", "remainder", "paused", "model-ineligible"] as const) { + test(`caller-owned main remainder considers ${replacement} replacement without native reads`, async () => { + const { cfg, mainQuota, headers } = strictFixture(); mainQuota(99.9); + if (replacement === "remainder") setAccountQuotaFromParsed("pool-a", { weeklyPercent: 98 }); + if (replacement === "paused") cfg.pausedCodexAccountIds = ["pool-a"]; + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", { + requestScopedMainCredential: true, + ...(replacement === "model-ineligible" ? { + modelId: "gpt-daybreak-blue-latest", + isDirectCallerEntitledToCodexModel: async () => true, + resolveCodexModelEntitlements: async (): Promise => ({ + modelsByAccount: new Map([["pool-a", new Set()]]), + confirmedAccountIds: new Set(["pool-a"]), credentialIdentities: new Map(), + }), + } : {}), + isMainAccountTokenLive: () => { throw new Error("must not inspect native main"); }, + getMainAccountToken: () => { throw new Error("must not inspect native main"); }, + }); + expect(ctx.kind).toBe(replacement === "below" ? "pool" : "main"); + if (ctx.kind === "main") expect(() => materializeCodexUpstreamAuth(headers, ctx, { config: cfg })).not.toThrow(); + }); + } + test("a replacement exhausted during model discovery does not strand a usable main remainder", async () => { + const { cfg, mainQuota, headers } = strictFixture(); mainQuota(99); + let discoveries = 0; + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", { + requestScopedMainCredential: true, modelId: "gpt-daybreak-blue-latest", + isDirectCallerEntitledToCodexModel: async () => true, + resolveCodexModelEntitlements: async () => { + if (++discoveries === 2) setAccountQuotaFromParsed("pool-a", { weeklyPercent: 100 }); + return poolEntitlements(); + }, + isMainAccountTokenLive: () => { throw new Error("must not inspect native main"); }, + }); + expect(ctx.kind).toBe("main"); + }); + test("an exhausted caller main can select another account's remainder", async () => { + const { cfg, mainQuota, headers } = strictFixture(); mainQuota(100); + setAccountQuotaFromParsed("pool-a", { weeklyPercent: 99 }); + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", { requestScopedMainCredential: true }); + expect(ctx).toMatchObject({ kind: "pool", accountId: "pool-a" }); + }); + test("caller-owned main selected by Pool is rechecked before materialization and dispatch", async () => { + const { cfg, mainQuota, headers } = strictFixture(); + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", { requestScopedMainCredential: true }); + expect(ctx).toMatchObject({ kind: "main", poolQuotaScope: "shared" }); + const guard = createCodexReserveDispatchGuard(ctx, cfg, "gpt-6-astra"); + expect(guard).toBeDefined(); + mainQuota(100); + expect(() => materializeCodexUpstreamAuth(headers, ctx, { config: cfg })).toThrow(CodexStrictQuotaUnavailableError); + expect(() => guard!(headers)).toThrow(CodexStrictQuotaUnavailableError); + }); + for (const change of ["quota", "pause", "exclude"] as const) { + test(`main pin is rechecked after async entitlement ${change} change`, async () => { + const { cfg, mainQuota, headers } = strictFixture(); + const excluded = new Set(); + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", { + requestScopedMainCredential: true, modelId: "gpt-daybreak-blue-latest", excludeAccountIds: excluded, + isDirectCallerEntitledToCodexModel: async () => { + if (change === "quota") mainQuota(100); + if (change === "pause") cfg.pausedCodexAccountIds = [MAIN_CODEX_ACCOUNT_ID]; + if (change === "exclude") excluded.add(MAIN_CODEX_ACCOUNT_ID); + return true; + }, + resolveCodexModelEntitlements: poolEntitlements, + isMainAccountTokenLive: () => { throw new Error("must not inspect native main"); }, + getMainAccountToken: () => { throw new Error("must not inspect native main"); }, + getValidMainAccountToken: async () => { throw new Error("must not inspect native main"); }, + }); + expect(ctx).toMatchObject({ kind: "pool", accountId: "pool-a" }); + }); + } + test("explicit Pool main caller credential remains fixed at resolve and late dispatch", async () => { + const { cfg, mainQuota, headers } = strictFixture(); + const options = { requestScopedMainCredential: true, accountId: MAIN_CODEX_ACCOUNT_ID }; + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", options); + expect(ctx).toMatchObject({ kind: "main", poolQuotaScope: "shared", fixedAccount: true }); + const guard = createCodexReserveDispatchGuard(ctx, cfg, "gpt-6-astra")!; + mainQuota(100); + await expect(resolveCodexAuthContext(headers, cfg, "pool", options)).rejects.toMatchObject({ + name: "CodexStrictQuotaUnavailableError", waitable: false, + }); + for (const dispatch of [() => guard(headers), () => materializeCodexUpstreamAuth(headers, ctx, { config: cfg })]) { + expect(dispatch).toThrow(CodexStrictQuotaUnavailableError); + try { dispatch(); } catch (error) { expect((error as CodexStrictQuotaUnavailableError).waitable).toBe(false); } + } + }); + test("dispatch callback remains installed while policy is off, then sees live enablement", () => { + const { cfg, headers } = strictFixture(); cfg.codexAccountStrictQuota = false; + const ctx = { kind: "pool" as const, accountId: "pool-a", writerGeneration: 0, generation: 1, + accessToken: "strict-pool-access", chatgptAccountId: "strict-pool-owner" }; + const guard = createCodexReserveDispatchGuard(ctx, cfg, "gpt-5.6-luna"); + expect(guard).toBeDefined(); + setAccountQuotaFromParsed("pool-a", { weeklyPercent: 100 }); + expect(() => guard!(headers)).not.toThrow(); + cfg.codexAccountStrictQuota = true; + expect(() => guard!(headers)).toThrow(CodexStrictQuotaUnavailableError); + }); + test("explicit Direct and independent scopes keep their own admission", async () => { + const { cfg, mainQuota, headers } = strictFixture(); mainQuota(100); + const ctx = await resolveCodexAuthContext(headers, cfg, "direct", { requestScopedMainCredential: true }); + expect(ctx).toEqual({ kind: "main", accountId: null }); + expect(createCodexReserveDispatchGuard(ctx, cfg, "gpt-6-astra")).toBeUndefined(); + expect(materializeCodexUpstreamAuth(headers, ctx, { config: cfg }).get("authorization")).toBe(headers.get("authorization")); + const pool = { kind: "pool" as const, accountId: "pool-a", writerGeneration: 0, generation: 1, + accessToken: "strict-pool-access", chatgptAccountId: "strict-pool-owner" }; + setAccountQuotaFromParsed("pool-a", { weeklyPercent: 100 }); + expect(() => createCodexReserveDispatchGuard(pool, cfg, "gpt-5.3-codex-spark")!(headers)).not.toThrow(); + }); + test("live strict policy selects an alternative without replacing the replay config owner", async () => { + const { cfg, headers } = strictFixture(); + cfg.codexAccountStrictQuota = false; cfg.autoSwitchThreshold = 0; cfg.accountPoolStrategy = "fill-first"; + cfg.pausedCodexAccountIds = [MAIN_CODEX_ACCOUNT_ID]; cfg.activeCodexAccountId = "pool-a"; + delete cfg.activeCodexAccountPinned; resetCodexRoutingForManualSelection("pool-a"); + cfg.codexAccounts!.push({ id: "pool-b", isMain: false }); + saveCodexAccountCredential("pool-b", { accessToken: "pool-b-access", refreshToken: "pool-b-refresh", + chatgptAccountId: "pool-b-owner", expiresAt: Date.now() + 3600000 }); + setAccountQuotaFromParsed("pool-a", { weeklyPercent: 100 }); + setAccountQuotaFromParsed("pool-b", { weeklyPercent: 10 }); + const ctx = await resolveCodexAuthContext(headers, cfg, "pool", { codexAuthPolicy: { + codexAccountStrictQuota: true, autoSwitchThreshold: 95, pausedCodexAccountIds: [MAIN_CODEX_ACCOUNT_ID], + } }); + expect(ctx).toMatchObject({ kind: "pool", accountId: "pool-b" }); + expect(cfg.codexAccountStrictQuota).toBe(false); expect(cfg.autoSwitchThreshold).toBe(0); + }); + test("client cancellation stops waiting for shared metadata without aborting other callers", async () => { + const { cfg, headers } = strictFixture(); + cfg.pausedCodexAccountIds = [MAIN_CODEX_ACCOUNT_ID]; cfg.activeCodexAccountId = "pool-a"; + delete cfg.activeCodexAccountPinned; resetCodexRoutingForManualSelection("pool-a"); clearAccountQuota("pool-a"); + let started!: () => void; let release!: () => void; + const entered = new Promise(resolve => { started = resolve; }); + const restore = setStrictCodexQuotaRefreshForTests(async () => { + started(); await new Promise(resolve => { release = resolve; }); + }); + const controller = new AbortController(); + const work = resolveCodexAuthContext(headers, cfg, "pool", { signal: controller.signal }); + const stopped = work.catch(error => error); + try { + await entered; controller.abort(); + expect(await stopped).toBeInstanceOf(DOMException); + } finally { + release(); await refreshStrictCodexQuotasOnDemand(cfg, new Set(["pool-a"])); restore(); + } + }); + test("strict refusal is neither reauth nor a fabricated reset deadline", async () => { + const err = new CodexStrictQuotaUnavailableError(false); + expect(shouldMarkAccountNeedsReauthForCodexAuthFailure(err)).toBe(false); + expect(cooldownErrorMessage(err)).toBe(err.message); + const response = cooldownErrorResponse(err); + expect(response.headers.get("retry-after")).toBeNull(); + expect(await response.text()).not.toContain("1970"); + }); +}); diff --git a/tests/codex-integration/codex-strict-quota-refresh.test.ts b/tests/codex-integration/codex-strict-quota-refresh.test.ts new file mode 100644 index 0000000000..3cdd4aa42d --- /dev/null +++ b/tests/codex-integration/codex-strict-quota-refresh.test.ts @@ -0,0 +1,238 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { OcxConfig } from "../../src/types"; +import { saveCodexAccountCredential } from "../../src/codex/account-store"; +import { clearAccountQuota, setAccountQuotaFromParsed, updateAccountQuota } from "../../src/codex/quota"; +import { clearCodexUpstreamHealth, recordCodexUpstreamOutcome } from "../../src/codex/routing"; +import { captureMainQuotaWriter, observeMainQuotaIdentity } from "../../src/codex/main-account-cache"; +import { getCodexQuotaRevision } from "../../src/codex/quota-events"; +import { getCodexStrictQuotaStatus } from "../../src/codex/strict-quota"; +import { refreshStrictCodexQuotasOnDemand, setStrictCodexQuotaRefreshForTests, + strictCodexQuotaWaiterCount, waitForStrictCodexQuotaChange } from "../../src/codex/strict-quota-refresh"; + +let dir: string; +let previousHome: string | undefined; +let previousCodexHome: string | undefined; +let restore: (() => void) | undefined; +let now: number; +let calls: string[][]; +const timers = new Map, { fn: () => void; delay: number }>(); +let cancellations: number; +const config = (): OcxConfig => ({ providers: {}, codexAccountStrictQuota: true, autoSwitchThreshold: 95, + pausedCodexAccountIds: ["__main__"], codexAccounts: [{ id: "a" }, { id: "b" }, { id: "c" }] } as OcxConfig); +function credential(id: string) { + saveCodexAccountCredential(id, { accessToken: `test-access-${id}`, refreshToken: `test-refresh-${id}`, + chatgptAccountId: `test-owner-${id}`, expiresAt: Date.now() + 3600000 }); +} +function runtime(refresh: (cfg: OcxConfig, ids: readonly string[]) => Promise = async () => {}) { + restore = setStrictCodexQuotaRefreshForTests(async (cfg, ids) => { + calls.push([...ids]); await refresh(cfg, ids); + }, { now: () => now, + setTimeout: (fn, delay) => { + const timer = { unref() {} } as ReturnType; + timers.set(timer, { fn, delay }); return timer; + }, + clearTimeout: timer => { cancellations++; timers.delete(timer); }, + }); +} +function fireTimer() { + const [timer, { fn, delay }] = [...timers][0]!; + timers.delete(timer); now += delay; fn(); +} +beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), "ocx-strict-refresh-")); + previousHome = process.env.OPENCODEX_HOME; previousCodexHome = process.env.CODEX_HOME; + process.env.OPENCODEX_HOME = dir; process.env.CODEX_HOME = dir; + clearAccountQuota(); clearCodexUpstreamHealth(); for (const id of ["a", "b", "c"]) credential(id); + now = Date.now(); calls = []; cancellations = 0; +}); +afterEach(() => { + expect(strictCodexQuotaWaiterCount()).toBe(0); expect(timers.size).toBe(0); + restore?.(); restore = undefined; clearAccountQuota(); clearCodexUpstreamHealth(); + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousHome; + if (previousCodexHome === undefined) delete process.env.CODEX_HOME; else process.env.CODEX_HOME = previousCodexHome; + rmSync(dir, { recursive: true, force: true }); +}); +describe("strict quota request refresh", () => { + test("one batch for simultaneous callers, with no idle timers", async () => { + let release!: () => void; + runtime(async () => await new Promise(resolve => { release = resolve; })); + const cfg = config(); + const work = Array.from({ length: 20 }, () => refreshStrictCodexQuotasOnDemand(cfg)); + await Promise.resolve(); expect(calls).toEqual([["a", "b", "c"]]); + release(); const results = await Promise.all(work); + expect(results.every(result => result.status === "attempted")).toBe(true); + expect(calls).toHaveLength(1); expect(timers.size).toBe(0); + }); + test("overlapping batches serialize across config instances instead of multiplying API concurrency", async () => { + let firstRelease!: () => void; + let active = 0; let maximum = 0; + runtime(async () => { + active++; maximum = Math.max(maximum, active); + if (calls.length === 1) await new Promise(resolve => { firstRelease = resolve; }); + active--; + }); + const first = refreshStrictCodexQuotasOnDemand(config(), new Set(["a", "b"])); + const second = refreshStrictCodexQuotasOnDemand(config(), new Set(["b", "c"])); + await Promise.resolve(); firstRelease(); await Promise.all([first, second]); + expect(calls).toEqual([["a", "b"], ["c"]]); expect(maximum).toBe(1); + }); + test("failure is observable and earns five-minute backoff without changing eligibility", async () => { + runtime(async () => { throw new Error("sensitive upstream failure detail"); }); + const cfg = config(); + expect(await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]))).toEqual({ status: "failed", accountIds: ["a"] }); + now += 299999; + expect((await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]))).status).toBe("idle"); + expect(getCodexStrictQuotaStatus(cfg, "a").state).toBe("unknown"); + now++; + expect((await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]))).status).toBe("failed"); + expect(calls).toHaveLength(2); + }); + test("unknown after credential repair does not inherit old-attempt backoff", async () => { + runtime(); const cfg = config(); + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); + credential("a"); + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); + expect(calls).toEqual([["a"], ["a"]]); + }); + for (const unit of ["seconds", "milliseconds"] as const) { + test(`${unit} reset permits one early read, never automatic recovery`, async () => { + runtime(); const cfg = config(); + const reset = Date.now() + 10000; + setAccountQuotaFromParsed("a", { weeklyPercent: 100, weeklyResetAt: unit === "seconds" ? reset / 1000 : reset }); + now = reset - 1; + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); expect(calls).toHaveLength(0); + now = reset + 1000; + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); expect(calls).toEqual([["a"]]); + expect(getCodexStrictQuotaStatus(cfg, "a", "shared", now).state).toBe("blocked"); + now += 1000; + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); expect(calls).toHaveLength(1); + }); + } + test("requested ready account does not trigger reads of other blocked accounts", async () => { + runtime(); updateAccountQuota("a", 1); updateAccountQuota("b", 100); + now = Date.now(); + await refreshStrictCodexQuotasOnDemand(config(), new Set(["a"])); + expect(calls).toHaveLength(0); + }); +}); +describe("strict quota live waiters", () => { + test("manual quota update wakes all waiters and removes their sole timer", async () => { + runtime(); const cfg = config(); + const one = waitForStrictCodexQuotaChange(cfg); const two = waitForStrictCodexQuotaChange(cfg); + expect(timers.size).toBe(1); expect(strictCodexQuotaWaiterCount()).toBe(2); + updateAccountQuota("a", 0); await Promise.all([one, two]); + expect(cancellations).toBe(1); expect(calls).toHaveLength(0); + }); + test("last cancellation removes timer and subscription with no idle poll", async () => { + runtime(); const cfg = config(); const a = new AbortController(); const b = new AbortController(); + const first = waitForStrictCodexQuotaChange(cfg, a.signal).catch(error => error.name); + const second = waitForStrictCodexQuotaChange(cfg, b.signal).catch(error => error.name); + a.abort(); expect(timers.size).toBe(1); b.abort(); + expect(await first).toBe("AbortError"); expect(await second).toBe("AbortError"); + expect(timers.size).toBe(0); expect(cancellations).toBe(1); + updateAccountQuota("a", 1); expect(calls).toHaveLength(0); + }); + test("timer only wakes requests; it never refreshes by itself", async () => { + runtime(); const cfg = config(); await refreshStrictCodexQuotasOnDemand(cfg); + const waiting = waitForStrictCodexQuotaChange(cfg); + expect([...timers.values()][0]!.delay).toBe(300000); + fireTimer(); await waiting; + expect(calls).toHaveLength(1); expect(strictCodexQuotaWaiterCount()).toBe(0); + await refreshStrictCodexQuotasOnDemand(cfg); expect(calls).toHaveLength(2); + }); + test("unobserved native main and unprobed accounts do not cause one-second wakeups", async () => { + runtime(); const cfg = config(); cfg.pausedCodexAccountIds = []; + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); + const waiting = waitForStrictCodexQuotaChange(cfg); + expect([...timers.values()][0]!.delay).toBe(300000); + fireTimer(); await waiting; + expect(calls).toEqual([["a"]]); + }); + test("stale observed but ineligible native main does not shorten a real pool attempt's backoff", async () => { + runtime(); const cfg = config(); cfg.pausedCodexAccountIds = []; + observeMainQuotaIdentity("strict-wait-ineligible-main"); + const writer = captureMainQuotaWriter("strict-wait-ineligible-main")!; + setAccountQuotaFromParsed("__main__", { weeklyPercent: 100 }, undefined, writer); + now = Date.now() + 300001; + // The resolver cannot use physical main for this request, so it probes only its real pool candidate. + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); + const waiting = waitForStrictCodexQuotaChange(cfg); + const delay = [...timers.values()][0]!.delay; + fireTimer(); await waiting; + expect(delay).toBe(300000); + expect(calls).toEqual([["a"]]); + expect(getCodexStrictQuotaStatus(cfg, "__main__", "shared", now).state).toBe("blocked"); + }); + test("a wholly unobserved pool uses the default waiting interval", async () => { + runtime(); const waiting = waitForStrictCodexQuotaChange(config()); + expect([...timers.values()][0]!.delay).toBe(300000); + fireTimer(); await waiting; + expect(calls).toHaveLength(0); + }); + test("quota update between refusal and waiter registration is not lost", async () => { + runtime(); const revision = getCodexQuotaRevision(); + updateAccountQuota("a", 0); + await waitForStrictCodexQuotaChange(config(), undefined, revision); + expect(timers.size).toBe(0); expect(strictCodexQuotaWaiterCount()).toBe(0); + expect(calls).toHaveLength(0); + }); + test("a live upstream cooldown wakes at expiry rather than the quota freshness interval", async () => { + runtime(); const cfg = config(); updateAccountQuota("a", 20); + now = Date.now(); + recordCodexUpstreamOutcome(cfg, "a", 429, { retryAfter: "30", modelId: "gpt-5.6-luna", now, fixedAccount: true }); + const waiting = waitForStrictCodexQuotaChange(cfg); + expect([...timers.values()][0]!.delay).toBe(30000); + fireTimer(); await waiting; + const afterExpiry = waitForStrictCodexQuotaChange(cfg); + expect([...timers.values()][0]!.delay).toBeGreaterThan(1000); + fireTimer(); await afterExpiry; + expect(calls).toHaveLength(0); + }); + test("already cancelled requests do not subscribe or schedule", async () => { + runtime(); const controller = new AbortController(); controller.abort(); + await expect(waitForStrictCodexQuotaChange(config(), controller.signal)).rejects.toThrow(); + }); +}); + +describe("selection-time quota metadata", () => { + test("a top-up is discovered while switching even inside the ordinary five-minute cache", async () => { + const cfg = config(); updateAccountQuota("a", 100); + now = Date.now() + 10_001; + runtime(async () => updateAccountQuota("a", 0)); + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]), { forSelection: true }); + expect(calls).toEqual([["a"]]); + expect(getCodexStrictQuotaStatus(cfg, "a").state).toBe("ready"); + expect(timers.size).toBe(0); + }); + test("healthy capacity crossing its predicted reset is re-read without assuming recovery", async () => { + const cfg = config(); const reset = Date.now() + 10000; + setAccountQuotaFromParsed("a", { weeklyPercent: 50, weeklyResetAt: reset / 1000 }); + runtime(); now = reset + 1000; + expect(getCodexStrictQuotaStatus(cfg, "a", "shared", now).state).toBe("unknown"); + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"])); + expect(calls).toEqual([["a"]]); + expect(getCodexStrictQuotaStatus(cfg, "a", "shared", now).state).toBe("unknown"); + }); + test("failed switch-time metadata reads retain backoff instead of probing every turn", async () => { + const cfg = config(); updateAccountQuota("a", 100); + now = Date.now() + 10_001; + runtime(async () => { throw new Error("metadata unavailable"); }); + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]), { forSelection: true }); + now += 10_001; + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]), { forSelection: true }); + expect(calls).toHaveLength(1); + expect(getCodexStrictQuotaStatus(cfg, "a").state).toBe("blocked"); + }); +}); + +test("a resolved metadata failure with no new observation keeps selection backoff", async () => { + const cfg = config(); updateAccountQuota("a", 100); + now = Date.now() + 10_001; runtime(); + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]), { forSelection: true }); + now += 10_001; + await refreshStrictCodexQuotasOnDemand(cfg, new Set(["a"]), { forSelection: true }); + expect(calls).toEqual([["a"]]); +}); diff --git a/tests/codex-integration/codex-strict-quota.test.ts b/tests/codex-integration/codex-strict-quota.test.ts new file mode 100644 index 0000000000..8245fcd740 --- /dev/null +++ b/tests/codex-integration/codex-strict-quota.test.ts @@ -0,0 +1,301 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { createHash } from "node:crypto"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { OcxConfig } from "../../src/types"; +import { saveCodexAccountCredential } from "../../src/codex/account-store"; +import { clearAccountQuota, getStrictAccountQuota, parseMainPolicyUsageQuota, parseUsageQuota, setAccountQuotaFromParsed, updateAccountQuota } from "../../src/codex/quota"; +import { getCodexStrictQuotaStatus } from "../../src/codex/strict-quota"; +import { observeMainQuotaIdentity, captureMainQuotaWriter } from "../../src/codex/main-account-cache"; +import { subscribeCodexQuotaChanges } from "../../src/codex/quota-events"; +import { clearThreadAccountMap, clearCodexUpstreamHealth, resetCodexRoutingForManualSelection, previewCodexAccountForRequest, resolveCodexAccountForThreadDetailed } from "../../src/codex/routing"; + +let dir: string; +let priorHome: string | undefined; +let priorCodex: string | undefined; +const config = (extra: Partial = {}): OcxConfig => ({ + providers: {}, codexAccounts: [{ id: "a" }, { id: "b" }], activeCodexAccountId: "a", + autoSwitchThreshold: 95, codexAccountStrictQuota: true, ...extra, +} as OcxConfig); +function credential(id: string, owner = `test-account-${id}`) { + saveCodexAccountCredential(id, { accessToken: `test-access-${id}`, refreshToken: `test-refresh-${id}`, + chatgptAccountId: owner, expiresAt: Date.now() + 3600000 }); +} +beforeEach(() => { + dir = mkdtempSync(join(tmpdir(), "ocx-strict-quota-")); + priorHome = process.env.OPENCODEX_HOME; priorCodex = process.env.CODEX_HOME; + process.env.OPENCODEX_HOME = dir; process.env.CODEX_HOME = dir; + clearAccountQuota(); clearThreadAccountMap(); clearCodexUpstreamHealth(); + credential("a"); credential("b"); +}); +afterEach(() => { + clearAccountQuota(); clearThreadAccountMap(); clearCodexUpstreamHealth(); + if (priorHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = priorHome; + if (priorCodex === undefined) delete process.env.CODEX_HOME; else process.env.CODEX_HOME = priorCodex; + rmSync(dir, { recursive: true, force: true }); +}); +describe("strict quota policy", () => { + test("off by default, unknown fails closed, independent scopes keep their own policy", () => { + expect(getCodexStrictQuotaStatus(config({ codexAccountStrictQuota: false }), "a").state).toBe("off"); + expect(getCodexStrictQuotaStatus(config({ autoSwitchThreshold: 0 }), "a").state).toBe("off"); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("unknown"); + expect(getCodexStrictQuotaStatus(config(), "a", "reserve").state).toBe("off"); + expect(getCodexStrictQuotaStatus(config(), "a", "spark").state).toBe("off"); + }); + test("switch threshold preserves remaining quota until actual exhaustion", () => { + for (const threshold of [95, 99, 100]) { + for (const used of [95, 99, 99.9]) { + updateAccountQuota("a", used); + expect(getCodexStrictQuotaStatus(config({ autoSwitchThreshold: threshold }), "a").state).toBe("ready"); + } + updateAccountQuota("a", 100); + expect(getCodexStrictQuotaStatus(config({ autoSwitchThreshold: threshold }), "a").state).toBe("blocked"); + } + }); + test("reset passage and unrelated partial/credits updates cannot recover an observed short block", () => { + setAccountQuotaFromParsed("a", { shortPercent: 100, shortResetAt: Date.now() / 1000 - 1, weeklyPercent: 10 }); + setAccountQuotaFromParsed("a", { weeklyPercent: 0, resetCredits: 4 }); + expect(getCodexStrictQuotaStatus(config(), "a", "shared", Date.now() + 86400000).state).toBe("blocked"); + setAccountQuotaFromParsed("a", { shortPercent: 0 }); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("ready"); + }); + test("each window must be fresh; credit and other-window writes do not refresh old evidence", () => { + updateAccountQuota("a", 20); + const observed = getStrictAccountQuota("a")!.windows[0]!.observedAt; + setAccountQuotaFromParsed("a", { resetCredits: 3 }); + expect(getStrictAccountQuota("a")!.windows[0]!.observedAt).toBe(observed); + expect(getCodexStrictQuotaStatus(config(), "a", "shared", observed + 300001).state).toBe("unknown"); + setAccountQuotaFromParsed("a", { monthlyPercent: 1 }); + expect(getStrictAccountQuota("a")!.windows.find(w => w.key === "weekly")!.observedAt).toBe(observed); + }); + test("invalid upstream evidence is not clamped into recovery", () => { + updateAccountQuota("a", 100); + for (const used_percent of [-1, NaN, Infinity]) { + setAccountQuotaFromParsed("a", parseUsageQuota({ rate_limit: { primary_window: { used_percent } } })); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("blocked"); + } + updateAccountQuota("a", -1); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("blocked"); + }); + test("persisted blocks outlive the legacy six-hour cache and predicted reset", () => { + const old = Date.now() - 86400000; + writeFileSync(join(dir, "codex-quota-cache.json"), JSON.stringify({ version: 1, quotas: {}, + strictQuotas: { a: { identity: createHash("sha256").update("test-account-a").digest("hex"), + quota: { windows: [{ scope: "shared", key: "weekly", usedPercent: 100, observedAt: old, resetAt: old / 1000 }] } } } })); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("blocked"); + updateAccountQuota("a", 0); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("ready"); + }); + test("credential replacement invalidates old evidence and requires a fresh read", () => { + updateAccountQuota("a", 10); credential("a", "replacement-account"); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("unknown"); + updateAccountQuota("a", 1); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("ready"); + }); + test("token refresh retains blocked windows until those windows are freshly observed", () => { + setAccountQuotaFromParsed("a", { shortPercent: 100, weeklyPercent: 1 }); + credential("a"); updateAccountQuota("a", 0); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("blocked"); + }); + test("complete monthly-only WHAM retires old weekly and short windows", () => { + setAccountQuotaFromParsed("a", { weeklyPercent: 100, shortPercent: 100 }); + setAccountQuotaFromParsed("a", parseUsageQuota({ rate_limit: { + primary_window: { used_percent: 0, limit_window_seconds: 30 * 86400 }, secondary_window: null, + } })); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("ready"); + expect(getStrictAccountQuota("a")!.windows.map(w => w.key)).toEqual(["monthly"]); + }); + test("partial or invalid monthly response cannot retire a blocked weekly window", () => { + setAccountQuotaFromParsed("a", { weeklyPercent: 100 }); + for (const rate_limit of [ + { primary_window: { used_percent: 0, limit_window_seconds: 30 * 86400 } }, + { primary_window: { used_percent: -1, limit_window_seconds: 30 * 86400 }, secondary_window: null }, + { primary_window: { used_percent: 0, limit_window_seconds: 30 * 86400 }, secondary_window: { used_percent: NaN, limit_window_seconds: 7 * 86400 } }, + ]) { + setAccountQuotaFromParsed("a", parseUsageQuota({ rate_limit })); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("blocked"); + } + setAccountQuotaFromParsed("a", { monthlyPercent: 0, monthlyIsPrimaryWindow: true }); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("blocked"); + }); + test("main policy retains complete WHAM authority through validated parsing", () => { + observeMainQuotaIdentity("strict-monthly-migration"); + const writer = captureMainQuotaWriter("strict-monthly-migration")!; + setAccountQuotaFromParsed("__main__", { weeklyPercent: 100 }, undefined, writer); + const data = { rate_limit: { primary_window: { used_percent: 0, limit_window_seconds: 30 * 86400 }, secondary_window: null } }; + setAccountQuotaFromParsed("__main__", parseUsageQuota(data), undefined, writer, parseMainPolicyUsageQuota(data)); + expect(getCodexStrictQuotaStatus(config(), "__main__").state).toBe("ready"); + setAccountQuotaFromParsed("__main__", { monthlyPercent: 100 }, undefined, writer); + const weekly = { rate_limit: { primary_window: { used_percent: 0, limit_window_seconds: 7 * 86400 }, + secondary_window: null, tertiary_window: null } }; + setAccountQuotaFromParsed("__main__", parseUsageQuota(weekly), undefined, writer, parseMainPolicyUsageQuota(weekly)); + expect(getCodexStrictQuotaStatus(config(), "__main__").state).toBe("ready"); + }); + test("main and added accounts apply the same strict maximum across reported windows", () => { + observeMainQuotaIdentity("strict-shared-max"); + const writer = captureMainQuotaWriter("strict-shared-max")!; + const data = { rate_limit: { primary_window: { used_percent: 20, limit_window_seconds: 7 * 86400 }, + tertiary_window: { used_percent: 100, limit_window_seconds: 30 * 86400 } } }; + setAccountQuotaFromParsed("a", parseUsageQuota(data)); + setAccountQuotaFromParsed("__main__", parseUsageQuota(data), undefined, writer, parseMainPolicyUsageQuota(data)); + expect(getCodexStrictQuotaStatus(config(), "a").state).toBe("blocked"); + expect(getCodexStrictQuotaStatus(config(), "__main__").state).toBe("blocked"); + }); + test("quota notification subscription is removable and has no background polling", () => { + let changes = 0; const stop = subscribeCodexQuotaChanges(() => changes++); + updateAccountQuota("a", 20); expect(changes).toBe(1); + stop(); updateAccountQuota("a", 30); expect(changes).toBe(1); + }); +}); +describe("strict pool routing", () => { + for (const accountPoolStrategy of ["quota", "fill-first"] as const) { + test(`${accountPoolStrategy}: current stays until threshold, affinity and preview agree`, () => { + const cfg = config({ accountPoolStrategy }); + updateAccountQuota("a", 90); updateAccountQuota("b", 1); + expect(previewCodexAccountForRequest("thread", cfg)).toBe("a"); + expect(resolveCodexAccountForThreadDetailed("thread", cfg)).toEqual({ status: "selected", accountId: "a" }); + updateAccountQuota("a", 95); + expect(previewCodexAccountForRequest("thread", cfg)).toBe("b"); + expect(resolveCodexAccountForThreadDetailed("thread", cfg)).toEqual({ status: "selected", accountId: "b" }); + updateAccountQuota("a", 0); + expect(previewCodexAccountForRequest(null, cfg)).toBe("b"); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "b" }); + }); + } + test("manual fill-first B remains selected when higher-priority A recovers", () => { + const cfg = config({ accountPoolStrategy: "fill-first", activeCodexAccountId: "b", + activeCodexAccountPinned: "b", codexAccountPriorities: { a: 100, b: 0 } }); + resetCodexRoutingForManualSelection("b"); + updateAccountQuota("a", 99); updateAccountQuota("b", 60); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "b" }); + updateAccountQuota("a", 0); + expect(previewCodexAccountForRequest(null, cfg)).toBe("b"); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "b" }); + updateAccountQuota("b", 95); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "a" }); + }); + test("a strict-only drained pin is retired and cannot reclaim the active account after recovery", () => { + const cfg = config({ accountPoolStrategy: "fill-first", activeCodexAccountId: "b", + activeCodexAccountPinned: "b", codexAccountPriorities: { a: 100, b: 0 } }); + resetCodexRoutingForManualSelection("b"); + updateAccountQuota("a", 20); + // A short-only reading is enough for strict admission; the legacy scorer keeps it unknown. + setAccountQuotaFromParsed("b", { shortPercent: 0 }); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "b" }); + setAccountQuotaFromParsed("b", { shortPercent: 95 }); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "a" }); + const pinAfterDrain = cfg.activeCodexAccountPinned; + setAccountQuotaFromParsed("b", { shortPercent: 0 }); + const afterRecovery = resolveCodexAccountForThreadDetailed(null, cfg); + expect(pinAfterDrain).toBeUndefined(); + expect(afterRecovery).toEqual({ status: "selected", accountId: "a" }); + }); + test("strict fill-first finishes an automatically selected account before a higher tier returns", () => { + const cfg = config({ accountPoolStrategy: "fill-first", codexAccountPriorities: { a: 100, b: 0 } }); + updateAccountQuota("a", 95); updateAccountQuota("b", 20); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "b" }); + updateAccountQuota("a", 0); + const previewAfterRecovery = previewCodexAccountForRequest(null, cfg); + const afterRecovery = resolveCodexAccountForThreadDetailed(null, cfg); + updateAccountQuota("b", 95); + expect(previewAfterRecovery).toBe("b"); + expect(afterRecovery).toEqual({ status: "selected", accountId: "b" }); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "a" }); + expect(cfg.activeCodexAccountPinned).toBeUndefined(); + }); + test("round-robin still rotates new threads but excludes over-threshold candidates", () => { + const cfg = config({ accountPoolStrategy: "round-robin" }); + updateAccountQuota("a", 20); updateAccountQuota("b", 30); + const first = resolveCodexAccountForThreadDetailed("rr-1", cfg); + const second = resolveCodexAccountForThreadDetailed("rr-2", cfg); + expect(first.status).toBe("selected"); expect(second.status).toBe("selected"); + expect(first).not.toEqual(second); + updateAccountQuota("a", 95); + expect(previewCodexAccountForRequest("rr-3", cfg)).toBe("b"); + expect(resolveCodexAccountForThreadDetailed("rr-3", cfg)).toEqual({ status: "selected", accountId: "b" }); + }); + test("all blocked or unknown never falls back to active", () => { + updateAccountQuota("a", 100); + expect(previewCodexAccountForRequest(null, config())).toBeNull(); + expect(resolveCodexAccountForThreadDetailed(null, config())).toEqual({ status: "none" }); + }); + test("manual pause survives new healthy quota", () => { + updateAccountQuota("a", 1); updateAccountQuota("b", 1); + const cfg = config({ pausedCodexAccountIds: ["a", "b"] }); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "none" }); + expect(cfg.pausedCodexAccountIds).toEqual(["a", "b"]); + }); +}); + +describe("independent scope selection", () => { + for (const quotaScope of ["spark", "reserve"] as const) { + for (const evidence of ["unknown", "stale"] as const) { + test(`${quotaScope}: ${evidence} shared quota does not rotate fill-first`, () => { + const cfg = config({ accountPoolStrategy: "fill-first", pausedCodexAccountIds: ["__main__"] }); + resetCodexRoutingForManualSelection("a"); + if (evidence === "stale") { + updateAccountQuota("a", 10); updateAccountQuota("b", 20); + } + const now = Date.now() + (evidence === "stale" ? 300_001 : 0); + expect(getCodexStrictQuotaStatus(cfg, "a", "shared", now).state).toBe("unknown"); + // Independent scopes retain their own admission policy even when strict + // shared evidence would reject both accounts. + for (const threadId of [null, "independent-scope", "independent-scope"]) { + expect(previewCodexAccountForRequest(threadId, cfg, now, quotaScope)).toBe("a"); + expect(resolveCodexAccountForThreadDetailed(threadId, cfg, now, quotaScope)) + .toEqual({ status: "selected", accountId: "a" }); + } + expect(cfg.activeCodexAccountId).toBe("a"); + }); + } + } +}); + +describe("soft threshold remainder routing", () => { + for (const accountPoolStrategy of ["quota", "fill-first"] as const) { + test(`${accountPoolStrategy}: one exhausted account does not strand the last five percent`, () => { + const cfg = config({ accountPoolStrategy, activeCodexAccountId: "b" }); + resetCodexRoutingForManualSelection("b"); + updateAccountQuota("a", 100); updateAccountQuota("b", 94); + expect(resolveCodexAccountForThreadDetailed("remainder", cfg)).toEqual({ status: "selected", accountId: "b" }); + for (const used of [95, 99, 99.9]) { + updateAccountQuota("b", used); + expect(previewCodexAccountForRequest("remainder", cfg)).toBe("b"); + expect(resolveCodexAccountForThreadDetailed("remainder", cfg)).toEqual({ status: "selected", accountId: "b" }); + expect(resolveCodexAccountForThreadDetailed(null, cfg)).toEqual({ status: "selected", accountId: "b" }); + } + updateAccountQuota("b", 100); + expect(previewCodexAccountForRequest("remainder", cfg)).toBeNull(); + expect(resolveCodexAccountForThreadDetailed("remainder", cfg)).toEqual({ status: "none" }); + }); + test(`${accountPoolStrategy}: no churn between remainders, but restored headroom wins`, () => { + const cfg = config({ accountPoolStrategy }); + updateAccountQuota("a", 96); updateAccountQuota("b", 95); + for (let i = 0; i < 3; i++) { + expect(previewCodexAccountForRequest("stable", cfg)).toBe("a"); + expect(resolveCodexAccountForThreadDetailed("stable", cfg)).toEqual({ status: "selected", accountId: "a" }); + } + updateAccountQuota("b", 0); + expect(previewCodexAccountForRequest("stable", cfg)).toBe("b"); + expect(resolveCodexAccountForThreadDetailed("stable", cfg)).toEqual({ status: "selected", accountId: "b" }); + }); + test(`${accountPoolStrategy}: request-ineligible headroom does not hide usable remainder`, () => { + const cfg = config({ accountPoolStrategy }); + updateAccountQuota("a", 99.5); updateAccountQuota("b", 0); + const selection = { modelEligibleAccountIds: new Set(["a"]) }; + expect(previewCodexAccountForRequest("scoped", cfg, Date.now(), "shared", selection)).toBe("a"); + expect(resolveCodexAccountForThreadDetailed("scoped", cfg, Date.now(), "shared", selection)) + .toEqual({ status: "selected", accountId: "a" }); + cfg.pausedCodexAccountIds = ["a"]; + expect(previewCodexAccountForRequest("scoped", cfg, Date.now(), "shared", selection)).toBeNull(); + }); + } + test("exhaustion selects another remainder and a fresh reset restores eligibility", () => { + const cfg = config({ accountPoolStrategy: "fill-first" }); + updateAccountQuota("a", 100); updateAccountQuota("b", 99.5); + expect(resolveCodexAccountForThreadDetailed("fallback", cfg)).toEqual({ status: "selected", accountId: "b" }); + updateAccountQuota("a", 0); + expect(previewCodexAccountForRequest("fallback", cfg)).toBe("a"); + }); +}); diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index db2583b00b..6b4e394339 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -297,6 +297,8 @@ "codex-restore-app-rewrite.test.ts": "codex-integration", "codex-retained-root-serialization.test.ts": "codex-integration", "codex-routing.test.ts": "codex-integration", + "codex-strict-quota.test.ts": "codex-integration", + "codex-strict-quota-refresh.test.ts": "codex-integration", "codex-runtime.test.ts": "codex-integration", "codex-service-manager-probe-hardening.test.ts": "codex-integration", "codex-service-manager-probe.test.ts": "codex-integration", @@ -918,6 +920,7 @@ "server-403-permission-e2e.test.ts": "server", "server-agent-task-recovery-replay.test.ts": "server", "server-auth.test.ts": "server", + "server-strict-quota-wait.test.ts": "server", "server-background-lifecycle.test.ts": "server", "server-clickjacking-headers.test.ts": "server", "server-combo-failover-e2e.test.ts": "server", diff --git a/tests/server/server-auth.test.ts b/tests/server/server-auth.test.ts index 11fddb9772..f8026373e4 100644 --- a/tests/server/server-auth.test.ts +++ b/tests/server/server-auth.test.ts @@ -1,3 +1,6 @@ +import { handleResponsesWithPolicyFallback } from "../../src/server/responses/policy-fallback"; +import { clearResponseStateForTests, rememberResponseState } from "../../src/responses/state"; +import { handleResponses as handleResponsesCore } from "../../src/server/responses/core"; import { waitForNativeMainStartupGate } from "../../src/codex/native-profile-startup"; import { afterEach, beforeEach, describe, expect, spyOn, test } from "bun:test"; import { createHash } from "node:crypto"; @@ -178,6 +181,7 @@ afterEach(() => { clearThreadAccountMap(); clearAccountNeedsReauth("pool-a"); clearAccountNeedsReauth("pool-b"); + clearAccountNeedsReauth("pool-c"); clearAccountQuota(); resetCodexModelEntitlementCacheForTests(); resetDebugSettingsForTests(); @@ -231,6 +235,8 @@ async function startPoolRetryHarness( reply: (accountId: string, request: Request) => Response | Promise, options: { secondAccount?: boolean; + thirdAccount?: boolean; + strictQuota?: boolean; streamMode?: "legacy-tee" | "eager-relay"; accountMode?: "direct" | "pool"; activeAccountId?: string; @@ -305,6 +311,7 @@ async function startPoolRetryHarness( : []), ], activeCodexAccountId: options.activeAccountId ?? "pool-a", + ...(options.strictQuota !== undefined ? { codexAccountStrictQuota: options.strictQuota } : {}), ...(options.accountNamespaces ? { codexAccountNamespaces: options.accountNamespaces } : {}), ...(options.pausedAccountIds ? { pausedCodexAccountIds: options.pausedAccountIds } : {}), ...(options.visionSidecarModel ? { visionSidecar: { model: options.visionSidecarModel } } : {}), @@ -312,7 +319,17 @@ async function startPoolRetryHarness( ...(options.streamMode ? { streamMode: options.streamMode } : {}), ...(options.combos ? { combos: options.combos } : {}), } as OcxConfig; + if (options.thirdAccount) config.codexAccounts!.push({ + id: "pool-c", email: "pool-c@example.test", isMain: false, chatgptAccountId: "acct-pool-c", + }); saveConfig(config); + if (options.thirdAccount) { + saveCodexAccountCredential("pool-c", { + accessToken: "pool-c-token", refreshToken: "pool-c-refresh", + expiresAt: Date.now() + 10 * 60_000, chatgptAccountId: "acct-pool-c", + }); + updateAccountQuota("pool-c", 30); + } if (!options.omitCredentialAccountIds?.includes("pool-a")) { saveCodexAccountCredential("pool-a", { accessToken: "pool-a-token", @@ -3101,6 +3118,129 @@ describe("server local API auth", () => { } }, { timeout: SERVER_BUDGET_MS }); + test.each([false, true])("strict quota traversal=%s respects the account retry budget", async strictQuota => { + const harness = await startPoolRetryHarness(accountId => accountId === "acct-pool-c" + ? Response.json({ id: "third-account-success", status: "completed", output: [] }) + : Response.json({ error: { message: "quota exceeded" } }, { status: 429, headers: { "retry-after": "60" } }), + { thirdAccount: true, strictQuota }); + try { + const response = await harness.request(); + expect(response.status).toBe(strictQuota ? 200 : 429); + await response.text(); + expect(harness.dispatches).toEqual(strictQuota + ? ["acct-pool-a", "acct-pool-b", "acct-pool-c"] : ["acct-pool-a", "acct-pool-b"]); + if (strictQuota) { + const requestId = response.headers.get("x-opencodex-request-id"); + const logs = getRequestLogEntries().filter(entry => entry.requestId === requestId); + expect(logs).toHaveLength(1); + expect(logs[0]?.attempts?.map(attempt => attempt.status)).toEqual([429, 429, 200]); + expect(logs[0]?.attempts?.map(attempt => attempt.sendCount)).toEqual([1, 1, 1]); + } + } finally { await stopPoolRetryHarness(harness); } + }, { timeout: SERVER_BUDGET_MS }); + + test("strict quota traversal stops after every configured credential explicitly refuses", async () => { + const harness = await startPoolRetryHarness(() => Response.json({ error: { message: "quota exceeded" } }, + { status: 429, headers: { "retry-after": "60" } }), { thirdAccount: true, strictQuota: true }); + try { + // Exercise one bounded core cycle; the outer server owns waiting for recovery. + const response = await handleResponsesCore(new Request("http://localhost/v1/responses", { + method: "POST", headers: { "content-type": "application/json", authorization: "Bearer inbound-token" }, + body: JSON.stringify({ model: POOL_RETRY_MODEL, input: "hello", stream: false }), + }), harness.config, { model: "unknown", provider: "unknown" }); + expect(response.status).toBe(429); + expect(response.headers.get("x-opencodex-quota-wait")).toBe("1"); + expect(await response.text()).toContain("quota exceeded"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b", "acct-pool-c"]); + } finally { await stopPoolRetryHarness(harness); } + }, { timeout: SERVER_BUDGET_MS }); + + test.each(["owner-task", "foreign-task"])("strict quota wait preserves validated context after cache expiration (%s)", async clientTask => { + const bodies: Record[] = []; + let recovered = false; + const harness = await startPoolRetryHarness(async (_accountId, req) => { + bodies.push(await req.json() as Record); + return recovered + ? Response.json({ id: "quota-restored", status: "completed", output: [] }) + : Response.json({ error: { message: "quota exceeded" } }, { status: 429 }); + }, { strictQuota: true }); + try { + clearResponseStateForTests(); + rememberResponseState({ input: [{ role: "user", content: "prior user context" }] }, + { id: "resp_quota_snapshot", status: "completed", output: [{ type: "message", role: "assistant", content: [{ type: "output_text", text: "prior answer" }] }] }, + undefined, { clientThreadId: "owner-task" }); + const response = await handleResponsesWithPolicyFallback(new Request("http://localhost/v1/responses", { + method: "POST", headers: { "content-type": "application/json", authorization: "Bearer inbound-token", "x-codex-parent-thread-id": clientTask }, + body: JSON.stringify({ model: POOL_RETRY_MODEL, previous_response_id: "resp_quota_snapshot", + input: [{ role: "user", content: "new delta" }], stream: false }), + }), harness.config, { model: "unknown", provider: "unknown" }, {}, { + quotaWait: { waitForChange: async () => { + expect(recovered).toBe(false); + clearResponseStateForTests(); + clearCodexUpstreamHealth(); + updateAccountQuota("pool-a", 0); + updateAccountQuota("pool-b", 0); + recovered = true; + } }, + }); + expect(response.status).toBe(200); + expect(await response.text()).toContain("quota-restored"); + expect(bodies).toHaveLength(3); + const finalInput = JSON.stringify(bodies.at(-1)?.input); + expect(finalInput).toContain("new delta"); + if (clientTask === "owner-task") { + expect(finalInput).toContain("prior user context"); + expect(finalInput).toContain("prior answer"); + } else { + expect(finalInput).not.toContain("prior user context"); + expect(finalInput).not.toContain("prior answer"); + expect(bodies.at(-1)?.previous_response_id).toBeUndefined(); + } + } finally { clearResponseStateForTests(); await stopPoolRetryHarness(harness); } + }, { timeout: SERVER_BUDGET_MS }); + + test("strict quota traversal honors cancellation before a third send", async () => { + const controller = new AbortController(); + const harness = await startPoolRetryHarness(async accountId => { + if (accountId === "acct-pool-b") { + controller.abort(); + // Give the server the client cancellation before returning the rejected response. + await Bun.sleep(20); + } + return Response.json({ error: { message: "quota exceeded" } }, { status: 429 }); + }, { thirdAccount: true, strictQuota: true }); + try { + await expect(harness.request({ signal: controller.signal })).rejects.toThrow(); + await Bun.sleep(30); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b"]); + } finally { await stopPoolRetryHarness(harness); } + }, { timeout: SERVER_BUDGET_MS }); + + test("strict quota traversal stops at a non-quota alternate failure", async () => { + const harness = await startPoolRetryHarness(accountId => Response.json({ error: { message: "unavailable" } }, + { status: accountId === "acct-pool-a" ? 402 : 503, headers: { "x-opencodex-quota-wait": "1" } }), { thirdAccount: true, strictQuota: true }); + try { + const response = await harness.request(); + expect(response.status).toBe(503); + expect(response.headers.get("x-opencodex-quota-wait")).toBeNull(); + await response.text(); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b"]); + } finally { await stopPoolRetryHarness(harness); } + }, { timeout: SERVER_BUDGET_MS }); + + test("strict quota traversal never retries after an alternate starts SSE output", async () => { + const harness = await startPoolRetryHarness(accountId => accountId === "acct-pool-a" + ? Response.json({ error: { message: "quota exceeded" } }, { status: 429 }) + : new Response('event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"partial"}\n\nevent: response.failed\ndata: {"type":"response.failed","response":{"status":"failed","error":{"code":"rate_limit_exceeded","message":"quota exceeded"}}}\n\n', + { headers: { "content-type": "text/event-stream" } }), { thirdAccount: true, strictQuota: true }); + try { + const response = await harness.request({ stream: true }); + const body = await response.text(); + expect(body).toContain("partial"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b"]); + } finally { await stopPoolRetryHarness(harness); } + }, { timeout: SERVER_BUDGET_MS }); + test("#584: pre-stream 429 retries once on another eligible pool account", async () => { const harness = await startPoolRetryHarness(accountId => accountId === "acct-pool-a" ? new Response(JSON.stringify({ error: { message: "rate limited" } }), { diff --git a/tests/server/server-strict-quota-wait.test.ts b/tests/server/server-strict-quota-wait.test.ts new file mode 100644 index 0000000000..b9d8156257 --- /dev/null +++ b/tests/server/server-strict-quota-wait.test.ts @@ -0,0 +1,308 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import type { ServerWebSocket } from "bun"; +import type { OcxConfig } from "../../src/types"; +import { handleResponsesWithPolicyFallback, type PolicyFallbackDeps } from "../../src/server/responses/policy-fallback"; +import { markStrictQuotaWaitResponse } from "../../src/server/responses/strict-quota-response"; +import { beginRequestAttempt, type RequestLogContext } from "../../src/server/request-log"; +import { abortAndReleaseAllTurns, getActiveTurnCount, tryAdmitTurn } from "../../src/server/lifecycle"; +import { notifyCodexQuotaChanges } from "../../src/codex/quota-events"; +import { setStrictCodexQuotaRefreshForTests, strictCodexQuotaWaiterCount } from "../../src/codex/strict-quota-refresh"; +import { responseWithDeferredRequestLog } from "../../src/server/relay"; +import { sendResponseToWebSocket, type WsData } from "../../src/server/ws-bridge"; + +const config = { providers: {}, codexAccounts: [], codexAccountStrictQuota: true, autoSwitchThreshold: 95 } as OcxConfig; +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); +const completed = 'event: response.completed\ndata: {"type":"response.completed","response":{"id":"real-response","status":"completed","output":[]}}\n\n'; +function rejection(): Response { + return markStrictQuotaWaitResponse(Response.json({ error: { message: "quota unavailable" } }, { status: 429 })); +} +function request(stream = false, signal?: AbortSignal, extra: Record = {}): Request { + return new Request("http://localhost/v1/responses", { + method: "POST", signal, headers: { "content-type": "application/json", "session-id": "root-session" }, + body: JSON.stringify({ model: "openai/gpt-test", input: [{ role: "user", content: "original" }], stream, ...extra }), + }); +} +function log(): RequestLogContext { return { model: "unknown", provider: "unknown" }; } +function attempt(ctx: RequestLogContext): void { + const row = beginRequestAttempt((ctx.attempts?.length ?? 0) + 1, "openai/pool", "gpt-test", "test"); + row.sendCount = 1; + (ctx.attempts ??= []).push(row); + ctx.activeAttempt = row; + ctx.activeAttemptStartedAt = Date.now(); +} +function waitGate() { + const listeners = new Set<() => void>(); + return { + count: () => listeners.size, + wake: () => { for (const wake of [...listeners]) wake(); }, + waitForChange: (_config: OcxConfig, signal?: AbortSignal): Promise => new Promise((resolve, reject) => { + const cleanup = () => { listeners.delete(done); signal?.removeEventListener("abort", abort); }; + const done = () => { cleanup(); resolve(); }; + const abort = () => { cleanup(); reject(signal?.reason ?? new DOMException("Aborted", "AbortError")); }; + if (signal?.aborted) { abort(); return; } + listeners.add(done); signal?.addEventListener("abort", abort, { once: true }); + }), + }; +} +async function until(predicate: () => boolean): Promise { + for (let i = 0; i < 100 && !predicate(); i++) await Bun.sleep(2); + expect(predicate()).toBe(true); +} + +let precedingTurns = 0; +let restoreRefresh: () => void; +let refreshOnWait: (ids: readonly string[]) => Promise; +beforeEach(() => { + precedingTurns = getActiveTurnCount(); + refreshOnWait = async () => {}; + restoreRefresh = setStrictCodexQuotaRefreshForTests(async (_config, ids) => refreshOnWait(ids)); +}); +afterEach(async () => { + // An aborted SSE consumer may finish before the shared metadata microtask settles. + await Bun.sleep(0); + // Other server suites share this process; assert this test returns its own admission. + expect(getActiveTurnCount()).toBe(precedingTurns); + expect(strictCodexQuotaWaiterCount()).toBe(0); + restoreRefresh(); +}); + +describe("strict quota request wait", () => { + test("a successful response at the 95% soft threshold is sent immediately", async () => { + let sends = 0; + const response = await handleResponsesWithPolicyFallback(request(), config, log(), {}, { + runCore: async () => { sends++; return Response.json({ id: "soft-threshold-response", status: "completed" }); }, + }); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ id: "soft-threshold-response", status: "completed" }); + expect(sends).toBe(1); + }); + + test("a pending request refreshes unknown main metadata and resumes without a management action", async () => { + let known = false; let sends = 0; const reads: string[][] = []; + refreshOnWait = async ids => { + reads.push([...ids]); + expect(strictCodexQuotaWaiterCount()).toBe(1); + known = true; notifyCodexQuotaChanges(); + }; + const response = await handleResponsesWithPolicyFallback(request(), config, log(), {}, { + runCore: async () => { sends++; return known ? Response.json({ id: "fresh-main" }) : rejection(); }, + }); + expect(await response.json()).toEqual({ id: "fresh-main" }); + expect(reads).toEqual([["__main__"]]); expect(sends).toBe(2); + }); + + test("main metadata uses live policy while replay retains its captured config", async () => { + const captured = { ...config, codexAccountStrictQuota: false, pausedCodexAccountIds: ["__main__"] }; + const live = { ...config, pausedCodexAccountIds: [] }; + let known = false; const reads: string[][] = []; + refreshOnWait = async ids => { reads.push([...ids]); known = true; notifyCodexQuotaChanges(); }; + const response = await handleResponsesWithPolicyFallback(request(), captured, log(), { codexAuthPolicy: live }, { + runCore: async (_req, owner, _log, options) => { + expect(owner).toBe(captured); expect(options?.codexAuthPolicy).toBe(live); + return known ? Response.json({ id: "live-main" }) : rejection(); + }, + }); + expect(await response.json()).toEqual({ id: "live-main" }); expect(reads).toEqual([["__main__"]]); + }); + + test("failed main metadata keeps waiting and is not retried on every wake", async () => { + const ac = new AbortController(); let reads = 0; let sends = 0; + refreshOnWait = async () => { reads++; throw new Error("metadata unavailable"); }; + const pending = handleResponsesWithPolicyFallback(request(false, ac.signal), config, log(), {}, { + runCore: async () => { sends++; return rejection(); }, + }); + await until(() => reads === 1); + notifyCodexQuotaChanges(); + await until(() => sends === 2 && strictCodexQuotaWaiterCount() === 1); + expect(reads).toBe(1); + ac.abort(new Error("stop waiting")); + await expect(pending).rejects.toThrow("stop waiting"); + }); + + test("non-streaming resumes the exact input after quota wakeup and preserves attempts", async () => { + const gate = waitGate(); const ctx = log(); const seen: unknown[] = []; + const runCore: NonNullable = async (req, _config, logCtx, options) => { + seen.push(await req.json()); + expect(req.headers.get("session-id")).toBe("root-session"); + attempt(logCtx); options?.onRequestBodyRead?.(); + return seen.length === 1 ? rejection() : Response.json({ id: "real-response", status: "completed" }); + }; + let bodyReads = 0; + const pending = handleResponsesWithPolicyFallback(request(false, undefined, { previous_response_id: "prior-local-id" }), config, ctx, + { onRequestBodyRead: () => { bodyReads++; } }, { runCore, quotaWait: gate }); + await until(() => gate.count() === 1); + expect(seen).toHaveLength(1); + gate.wake(); + const response = await pending; + expect(await response.json()).toEqual({ id: "real-response", status: "completed" }); + expect(seen).toHaveLength(2); expect(seen[1]).toEqual(seen[0]); + expect(bodyReads).toBe(1); expect(gate.count()).toBe(0); + expect(ctx.attempts).toHaveLength(2); expect(ctx.attempts?.[0]?.status).toBe(429); + }); + + test("non-streaming cancellation releases the admitted pending request", async () => { + const ac = new AbortController(); const gate = waitGate(); const lease = tryAdmitTurn()!; + const pending = handleResponsesWithPolicyFallback(request(false), config, log(), { abortSignal: ac.signal, turnAdmissionLease: lease }, { + runCore: async () => rejection(), quotaWait: gate, + }); + await until(() => gate.count() === 1); + ac.abort(new Error("request cancelled")); + await expect(pending).rejects.toThrow("request cancelled"); expect(gate.count()).toBe(0); + }); + + test("forged upstream wait headers cannot authorize replay", async () => { + let sends = 0; const gate = waitGate(); + const response = await handleResponsesWithPolicyFallback(request(), config, log(), {}, { + runCore: async () => { sends++; return Response.json({}, { status: 429, headers: { "x-opencodex-quota-wait": "1" } }); }, quotaWait: gate, + }); + expect(response.status).toBe(429); expect(sends).toBe(1); expect(gate.count()).toBe(0); + }); + + test("stored 401 replay budget forbids waiting or another account cycle", async () => { + let sends = 0; const gate = waitGate(); + const response = await handleResponsesWithPolicyFallback(request(), config, log(), {}, { + runCore: async (_req, _config, _log, options) => { sends++; options?.onStoredPool401ReplayDispatched?.(); return rejection(); }, quotaWait: gate, + }); + expect(response.status).toBe(429); expect(sends).toBe(1); expect(gate.count()).toBe(0); + }); + + test("typed SSE heartbeats keep waiting alive and only real resumed events complete the turn", async () => { + const gate = waitGate(); let sends = 0; + const response = await handleResponsesWithPolicyFallback(request(true), config, log(), {}, { + runCore: async () => ++sends === 1 ? rejection() : new Response(completed, { headers: { "content-type": "text/event-stream" } }), + quotaWait: { ...gate, heartbeatMs: 5 }, + }); + expect(response.status).toBe(200); + const reader = response.body!.getReader(); + expect(decoder.decode((await reader.read()).value)).toContain('"type":"response.heartbeat"'); + expect(decoder.decode((await reader.read()).value)).toContain('"type":"response.heartbeat"'); + expect(sends).toBe(1); + gate.wake(); + let output = ""; + while (true) { const next = await reader.read(); if (next.done) break; output += decoder.decode(next.value); } + expect(output).toContain('"id":"real-response"'); expect(output.match(/event: response.completed/g)).toHaveLength(1); + expect(sends).toBe(2); expect(gate.count()).toBe(0); + }); + + test.each(["request", "websocket"] as const)("%s abort cancels a pending wait without dispatch", async kind => { + const ac = new AbortController(); const gate = waitGate(); let sends = 0; + const lease = tryAdmitTurn()!; + const response = await handleResponsesWithPolicyFallback(request(true, kind === "request" ? ac.signal : undefined), config, log(), + { turnAdmissionLease: lease, ...(kind === "websocket" ? { abortSignal: ac.signal } : {}) }, { + runCore: async () => { sends++; return rejection(); }, quotaWait: gate, + }); + const reader = response.body!.getReader(); await reader.read(); + expect(getActiveTurnCount()).toBe(precedingTurns + 1); + ac.abort(new Error("client cancelled")); + await expect(reader.read()).rejects.toThrow("client cancelled"); + expect(sends).toBe(1); expect(gate.count()).toBe(0); + }); + + test("service drain aborts the admitted wait and releases its subscription", async () => { + const gate = waitGate(); const lease = tryAdmitTurn()!; + const response = await handleResponsesWithPolicyFallback(request(true), config, log(), { turnAdmissionLease: lease }, { + runCore: async () => rejection(), quotaWait: gate, + }); + const reader = response.body!.getReader(); await reader.read(); + abortAndReleaseAllTurns(new Error("server shutdown")); + precedingTurns = 0; + await expect(reader.read()).rejects.toThrow("server shutdown"); expect(gate.count()).toBe(0); + }); + + test("response cancellation cleans up before any client pull", async () => { + const gate = waitGate(); const lease = tryAdmitTurn()!; + const response = await handleResponsesWithPolicyFallback(request(true), config, log(), { turnAdmissionLease: lease }, { + runCore: async () => rejection(), quotaWait: { ...gate, heartbeatMs: 5 }, + }); + await response.body!.cancel(); expect(gate.count()).toBe(0); + }); + + test("quota updates wake the production waiter with no idle subscription left", async () => { + let sends = 0; + const pending = handleResponsesWithPolicyFallback(request(), config, log(), {}, { + runCore: async () => ++sends === 1 ? rejection() : Response.json({ id: "real-response" }), + }); + await until(() => strictCodexQuotaWaiterCount() === 1); + notifyCodexQuotaChanges(); + expect(await (await pending).json()).toEqual({ id: "real-response" }); expect(sends).toBe(2); + }); + + test("a quota reset between core admission and wait subscription is not lost", async () => { + let sends = 0; + const response = await handleResponsesWithPolicyFallback(request(), config, log(), {}, { + runCore: async () => { + sends++; + if (sends === 1) { notifyCodexQuotaChanges(); return rejection(); } + return Response.json({ id: "already-reset" }); + }, + }); + expect(await response.json()).toEqual({ id: "already-reset" }); expect(sends).toBe(2); + }); + + test("a terminal non-quota rejection after waiting becomes a real failed SSE terminal", async () => { + let sends = 0; const gate = waitGate(); const ctx = log(); + const response = await handleResponsesWithPolicyFallback(request(true), config, ctx, {}, { + runCore: async () => ++sends === 1 ? rejection() : Response.json({ error: { code: "upstream_error", message: "provider unavailable" } }, { status: 503 }), + quotaWait: gate, + }); + gate.wake(); const body = await response.text(); + expect(body).toContain("response.failed"); expect(body).toContain("provider unavailable"); expect(body).not.toContain("response.completed"); + expect(ctx.terminalHttpStatus).toBe(503); expect(sends).toBe(2); + }); + + test("a resumed stored 401 consumes the same logical budget and cannot wait again", async () => { + const gate = waitGate(); let sends = 0; + const response = await handleResponsesWithPolicyFallback(request(true), config, log(), {}, { + runCore: async (_req, _config, _log, options) => { + if (++sends > 1) options?.onStoredPool401ReplayDispatched?.(); + return rejection(); + }, quotaWait: gate, + }); + gate.wake(); const body = await response.text(); + expect(body).toContain("response.failed"); expect(sends).toBe(2); expect(gate.count()).toBe(0); + }); + + test("resumed native SSE uses exactly one outer request-log owner", async () => { + const gate = waitGate(); const ctx = log(); let sends = 0; let nativeFinalized = 0; let logged = 0; + const response = await handleResponsesWithPolicyFallback(request(true), config, ctx, { + onNativePassthroughTerminal: () => { nativeFinalized++; }, + }, { + runCore: async (_req, _config, logCtx, options) => { + attempt(logCtx); + if (++sends === 1) return rejection(); + options?.onNativePassthroughTerminal?.("completed"); + return new Response(completed, { headers: { "content-type": "text/event-stream" } }); + }, quotaWait: gate, + }); + const loggedResponse = responseWithDeferredRequestLog(response, "quota-wait-log", Date.now(), ctx, () => { logged++; }); + gate.wake(); await loggedResponse.text(); + expect(nativeFinalized).toBe(0); expect(logged).toBe(1); expect(ctx.attempts).toHaveLength(2); + }); + + test("WS bridge relays typed heartbeat and the real recovered terminal", async () => { + const gate = waitGate(); let sends = 0; const frames: Array<{ type: string }> = []; + const response = await handleResponsesWithPolicyFallback(request(true), config, log(), {}, { + runCore: async () => ++sends === 1 ? rejection() : new Response(completed, { headers: { "content-type": "text/event-stream" } }), quotaWait: gate, + }); + const ws = { readyState: 1, data: {}, send: (text: string) => { frames.push(JSON.parse(text)); return text.length; } } as unknown as ServerWebSocket; + const pump = sendResponseToWebSocket(ws, response, () => true); + await until(() => frames.some(frame => frame.type === "response.heartbeat")); + gate.wake(); await pump; + expect(frames.filter(frame => frame.type === "response.completed")).toHaveLength(1); expect(gate.count()).toBe(0); + }); + + test("slow readers do not accumulate heartbeat frames or eagerly drain recovered output", async () => { + const gate = waitGate(); let sends = 0; let pulls = 0; + const response = await handleResponsesWithPolicyFallback(request(true), config, log(), {}, { + runCore: async () => ++sends === 1 ? rejection() : new Response(new ReadableStream({ + pull(controller) { pulls++; controller.enqueue(encoder.encode(completed)); if (pulls === 5) controller.close(); }, + }), { headers: { "content-type": "text/event-stream" } }), quotaWait: { ...gate, heartbeatMs: 2 }, + }); + await Bun.sleep(25); gate.wake(); await Bun.sleep(25); + expect(pulls).toBeLessThanOrEqual(1); + const text = await response.text(); + expect(text.match(/event: response.heartbeat/g)?.length).toBeLessThanOrEqual(2); + expect(pulls).toBe(5); + }); +});