@bitkyc08/opencodex 2.37.0 → 2.39.0-preview.20260901

package/src/codex/quota-401-recovery.ts

@@ -0,0 +1,190 @@
+import { randomUUID } from "node:crypto";
+import { isCodexAccountGenerationLive, type CodexRefreshProvenance } from "./account-store";
+import { QUOTA_RECOVERY_LEASE_MS } from "./quota-recovery-timing";
+
+/**
+ * One refresh-and-replay per credential lineage, for a WHAM 401 (#3019).
+ *
+ * A bare 401 from `backend-api/wham/usage` is what a stale-but-refreshable bearer produces
+ * after a plan change, so quarantining on it tells the operator to re-authenticate an
+ * account that was fine. The fix is to refresh and replay once — and `once` is the whole
+ * problem, because an unbounded retry against an upstream 401 is a self-inflicted
+ * credential-stuffing loop.
+ *
+ * ## Why a claim id rather than the lineage
+ *
+ * Keying the budget on the lineage cannot separate an old claimant from a later retry on
+ * the same lineage: claim L, a transient failure releases L, another poll claims L, and
+ * the first caller's late settlement spends the second caller's claim. The claim id makes
+ * every mutation a compare-and-set on `(accountId, lineage, claimId)`, so a stale
+ * completion is a no-op instead of somebody else's budget.
+ *
+ * ## Why `spent` never expires and `claimed` does
+ *
+ * Expiring a spent record would hand the same lineage another refresh, which is exactly
+ * the property this module exists to hold. But a caller can be cancelled or die between
+ * claim and settle, so `claimed` carries a bounded lease: an expired lease is reclaimable,
+ * and is never promoted to `spent` because the refresh may not have happened.
+ */
+
+export type RecoveryRecord =
+  | { state: "claimed"; lineage: number; claimId: string; expiresAt: number }
+  /**
+   * `terminal` marks a lineage whose refresh proved the grant is dead. It is spent AND the
+   * account must keep reporting needs-reauth: without it, the next bare 401 finds the
+   * budget used, reports transient, and a dead credential looks healthy.
+   */
+  | { state: "spent"; lineage: number; terminal?: true }
+  | { state: "backoff"; lineage: number; nextAttemptAt: number };
+
+export type ClaimResult =
+  | { granted: true; claimId: string }
+  | { granted: false; reason: "spent" | "in-flight" | "backoff" };
+
+const records = new Map<string, RecoveryRecord>();
+
+function leaseExpired(record: RecoveryRecord, now: number): boolean {
+  return record.state === "claimed" && record.expiresAt <= now;
+}
+
+/**
+ * Claim the one refresh this lineage is entitled to.
+ *
+ * A live claim blocks EVERY lineage for the account, not just its own. The refresh it
+ * fences commits `G -> G+1` before the claimant can settle, so during that window a
+ * `G+1` claim would otherwise be granted and the late settlement would land on nothing.
+ */
+export function claimQuotaRecovery(
+  accountId: string,
+  lineage: number,
+  now: number = Date.now(),
+): ClaimResult {
+  const existing = records.get(accountId);
+  if (existing && !leaseExpired(existing, now)) {
+    if (existing.state === "claimed") return { granted: false, reason: "in-flight" };
+    if (existing.state === "spent") {
+      // Only this lineage is fenced. A different one is a new grant with its own budget.
+      if (existing.lineage === lineage) return { granted: false, reason: "spent" };
+    } else if (existing.lineage === lineage && existing.nextAttemptAt > now) {
+      return { granted: false, reason: "backoff" };
+    }
+  }
+  const claimId = randomUUID();
+  records.set(accountId, { state: "claimed", lineage, claimId, expiresAt: now + QUOTA_RECOVERY_LEASE_MS });
+  return { granted: true, claimId };
+}
+
+function heldClaim(accountId: string, claimId: string): Extract<RecoveryRecord, { state: "claimed" }> | null {
+  const existing = records.get(accountId);
+  if (!existing || existing.state !== "claimed" || existing.claimId !== claimId) return null;
+  return existing;
+}
+
+/**
+ * Record the outcome of a claimed refresh.
+ *
+ * Which lineage gets fenced depends on the outcome, and getting this backwards is how a
+ * fresh credential loses the recovery it is entitled to:
+ *
+ * - `self-refresh` and `joined-lineage` are this lineage's own attempt, so the RETURNED
+ *   generation is spent. The returned one, never the rejected one: a successful token
+ *   response can rotate only the refresh grant, leaving the access token byte-identical,
+ *   and the credential has already moved by then.
+ * - `external-replacement` means somebody else's credential is now stored. The claimed
+ *   OLD lineage is spent — this caller did use its attempt — and the returned generation
+ *   is left untouched so it can claim on its own next 401.
+ */
+export function settleQuotaRecovery(
+  accountId: string,
+  claimId: string,
+  outcome: { provenance: CodexRefreshProvenance; generation: number },
+): void {
+  const held = heldClaim(accountId, claimId);
+  if (!held) return; // superseded or already settled: never disturb a newer claimant
+  const fenced = outcome.provenance === "external-replacement" ? held.lineage : outcome.generation;
+  records.set(accountId, { state: "spent", lineage: fenced });
+}
+
+/**
+ * Record a refresh that failed with proof the grant itself is dead.
+ *
+ * Distinct from {@link releaseQuotaRecovery}: a revoked or expired grant will not become
+ * valid on the next poll, so backing off would let a later 401 report the account healthy
+ * while it is not. The lineage is fenced durably and the caller quarantines.
+ */
+export function settleQuotaRecoveryTerminal(accountId: string, claimId: string): void {
+  const held = heldClaim(accountId, claimId);
+  if (!held) return;
+  records.set(accountId, { state: "spent", lineage: held.lineage, terminal: true });
+}
+
+/** Did this lineage's one refresh prove the grant dead? */
+export function quotaRecoveryTerminalFor(accountId: string, lineage: number): boolean {
+  const record = records.get(accountId);
+  return record?.state === "spent" && record.lineage === lineage && record.terminal === true;
+}
+
+/**
+ * Release a claim whose refresh failed without proving anything about the credential.
+ *
+ * Into BACKOFF, not into eligibility: a failed quota request does not refresh the quota
+ * timestamp, so successive dashboard and background polls would each issue another token
+ * refresh — the loop, reopened from the other side.
+ */
+export function releaseQuotaRecovery(
+  accountId: string,
+  claimId: string,
+  backoffMs: number,
+  now: number = Date.now(),
+): void {
+  const held = heldClaim(accountId, claimId);
+  if (!held) return;
+  records.set(accountId, { state: "backoff", lineage: held.lineage, nextAttemptAt: now + backoffMs });
+}
+
+/**
+ * Drop records for accounts that no longer exist, and fences whose credential has moved.
+ *
+ * A live claim is exempt. Its refresh commits the generation forward, so "the fenced
+ * generation is no longer stored" is the EXPECTED state mid-flight, and removing it there
+ * would let a second claim in before the first settles.
+ */
+export function reconcileQuotaRecovery(liveAccountIds: ReadonlySet<string>, now: number = Date.now()): number {
+  let removed = 0;
+  for (const [accountId, record] of [...records]) {
+    if (!liveAccountIds.has(accountId)) {
+      records.delete(accountId);
+      removed += 1;
+      continue;
+    }
+    if (record.state === "claimed" && !leaseExpired(record, now)) continue;
+    if (!isCodexAccountGenerationLive(accountId, record.lineage)) {
+      records.delete(accountId);
+      removed += 1;
+    }
+  }
+  return removed;
+}
+
+/** Drop expired backoff windows and abandoned leases. */
+export function sweepExpiredQuotaRecovery(now: number = Date.now()): number {
+  let removed = 0;
+  for (const [accountId, record] of [...records]) {
+    // A spent record is durable: expiring it would grant the same lineage another refresh.
+    if (record.state === "spent") continue;
+    const stale = record.state === "backoff" ? record.nextAttemptAt <= now : leaseExpired(record, now);
+    if (stale) {
+      records.delete(accountId);
+      removed += 1;
+    }
+  }
+  return removed;
+}
+
+export function quotaRecoveryRecordForTests(accountId: string): RecoveryRecord | undefined {
+  return records.get(accountId);
+}
+
+export function resetQuotaRecoveryForTests(): void {
+  records.clear();
+}

package/src/codex/quota-recovery-timing.ts

@@ -0,0 +1,28 @@
+/**
+ * Timing contract shared by the credential refresh, the WHAM quota request, and the
+ * 401-recovery budget (#3019).
+ *
+ * A leaf on purpose: `auth-api.ts` imports the recovery store to claim and settle, so the
+ * recovery store cannot import `auth-api.ts` back for the WHAM timeout. Both of them, and
+ * `account-store.ts`, import this instead.
+ *
+ * These were three inline literals in three files. The lease has to outlast the operations
+ * it fences — a lease that expires mid-refresh admits a second claim for a lineage that is
+ * already being refreshed — so it is derived here rather than restated as a round number
+ * that drifts away from what it covers.
+ */
+
+/** The refresh flight's own `AbortSignal.timeout` ceiling. */
+export const CODEX_REFRESH_FLIGHT_CEILING_MS = 30_000;
+
+/** Deadline for one `backend-api/wham/usage` request. */
+export const WHAM_REQUEST_TIMEOUT_MS = 8_000;
+
+/**
+ * How long a recovery claim stays valid without settlement.
+ *
+ * The sequence a claim covers is: WHAM request → refresh → WHAM replay. Both quota legs sit
+ * inside the lease, which is why the request timeout is counted twice.
+ */
+export const QUOTA_RECOVERY_LEASE_MS =
+  CODEX_REFRESH_FLIGHT_CEILING_MS + WHAM_REQUEST_TIMEOUT_MS * 2;

package/src/codex/quota.ts

@@ -110,6 +110,12 @@ function mayCommitAccountQuota(accountId: string, writerGeneration: number): boo
 // Valid upstream percentages are normalized to 0..100. Keep "unknown" outside that domain so an
 // actually exhausted account is still eligible for threshold rotation.
 export const CODEX_UNKNOWN_USAGE_SCORE = 101;
+/**
+ * A window reading at or above this is a measured refusal, not a position on a scale.
+ *
+ * Separate from `CODEX_UNKNOWN_USAGE_SCORE` because they mean opposite things: unknown is
+ * "we have not observed this account", 100 is "we observed it and it is full".
+ */
 export const CODEX_EXHAUSTED_USAGE_PERCENT = 100;
 
 export function isCodexQuotaExhausted(