stack 5/7: probe reset-derived cooldowns without waiting to be selected (#915)

src/codex/auth-api.ts

@@ -16,11 +16,13 @@ import {
 import { deleteCodexAccount, reconcileMainCodexAccountRuntimeState } from "./account-lifecycle";
 import { isCodexAccountPaused, setCodexAccountPaused } from "./account-pause";
 import {
+  claimDueCodexQuotaRecoveryProbes,
   clearCodexAccountCooldown,
   clearThreadAccountMapForAccount,
   getEffectiveActiveCodexAccountId,
   reconcileCodexActiveAfterExclusion,
   resetCodexRoutingForManualSelection,
+  settleCodexQuotaRecoveryProbe,
 } from "./routing";
 import {
   normalizeAccountPoolStickyLimit,
@@ -35,6 +37,7 @@ import { clearAccountNeedsReauth, isAccountNeedsReauth, markAccountNeedsReauth }
 import {
   clearAccountQuota,
   getAccountQuota,
+  isCompleteCodexQuotaRecoverySnapshot,
   isCodexQuotaExhausted,
   listAccountQuotas,
   parseUsageQuota,
@@ -53,7 +56,7 @@ export {
 } from "./quota";
 import { extractAccountId, decodeJwtPayload } from "../oauth/chatgpt";
 import { getMainAccountPlan, MAIN_CODEX_ACCOUNT_ID, setMainAccountPlan } from "./main-account";
-import { captureConfigGeneration } from "../lib/state-store-sweeper";
+import { captureConfigGeneration, registerStateSweepAfterTick } from "../lib/state-store-sweeper";
 import { reconcileLiveStateStores } from "../lib/state-store-registrations";
 import {
   captureMainAccountIdentityGeneration,
@@ -816,6 +819,49 @@ async function fetchPoolAccountQuota(accountId: string, forceRefresh = false, co
 }
 
 let primeInFlight: Promise<void> | null = null;
+let cooldownRecoveryInFlight: Promise<void> | null = null;
+
+export async function runCodexCooldownRecoveryProbes(config: OcxConfig, now = Date.now()): Promise<void> {
+  const openai = config.providers[OPENAI_CODEX_PROVIDER_ID];
+  if (!openai
+    || openai.disabled === true
+    || !isCanonicalOpenAiForwardProvider(openai)
+    || providerCodexAccountMode(OPENAI_CODEX_PROVIDER_ID, openai) !== "pool") return;
+  if (cooldownRecoveryInFlight) return cooldownRecoveryInFlight;
+  cooldownRecoveryInFlight = (async () => {
+    const claims = claimDueCodexQuotaRecoveryProbes(config, POOL_QUOTA_REFRESH_CONCURRENCY, now);
+    await mapWithConcurrency(claims, POOL_QUOTA_REFRESH_CONCURRENCY, async claim => {
+      const account = configuredPoolAccount(config, claim.accountId);
+      if (!account) {
+        settleCodexQuotaRecoveryProbe(claim, false, {}, now);
+        return;
+      }
+      try {
+        const result = await fetchPoolAccountQuota(claim.accountId, true, account.plan);
+        // Defence in depth: `spark` is already excluded at the claim site, since generic WHAM
+        // cannot prove a spark recovery. Keep the settle-side guard so a future claim change
+        // cannot silently start clearing spark on generic evidence.
+        const recovered = claim.scope !== "spark"
+          && isCompleteCodexQuotaRecoverySnapshot(result.freshQuota ?? null, result.freshPlan ?? account.plan);
+        settleCodexQuotaRecoveryProbe(claim, recovered, {
+          credentialGeneration: result.freshCredentialGeneration,
+        }, now);
+      } catch {
+        settleCodexQuotaRecoveryProbe(claim, false, {}, now);
+      }
+    });
+  })().catch(() => {
+    // Background recovery is best-effort; routing keeps the cooldown on failure.
+  }).finally(() => { cooldownRecoveryInFlight = null; });
+  return cooldownRecoveryInFlight;
+}
+
+export function registerCodexCooldownRecoveryProbeWorker(config: OcxConfig): void {
+  registerStateSweepAfterTick({
+    name: "codex-cooldown-recovery",
+    afterTick: () => { void runCodexCooldownRecoveryProbes(config); },
+  });
+}
 
 export interface PrimeCodexPoolQuotasOptions {
   /** Test seams for proving fenced/recovery priming performs no native-main work. */
@@ -908,6 +954,11 @@ export function clearCodexQuotaPrimeState(): void {
   primeInFlight = null;
 }
 
+/** Test-only reset for the worker-level single-flight. */
+export function clearCodexCooldownRecoveryProbeState(): void {
+  cooldownRecoveryInFlight = null;
+}
+
 export function effectiveCodexAuthAccountId(config: OcxConfig): string {
   return getEffectiveActiveCodexAccountId(config) ?? MAIN_CODEX_ACCOUNT_ID;
 }

src/codex/quota.ts

@@ -67,15 +67,48 @@ export function isCodexQuotaExhausted(
   plan?: string | null,
 ): boolean {
   if (!quota) return false;
-  const normalizedPlan = plan?.trim().toLowerCase();
-  const values = normalizedPlan === "go" || normalizedPlan === "free"
+  const values = codexQuotaWindowForPlan(plan) === "monthly"
     ? [quota.monthlyPercent]
     : [quota.weeklyPercent, quota.monthlyPercent];
   return values.some(value => typeof value === "number"
     && Number.isFinite(value)
     && value >= CODEX_EXHAUSTED_USAGE_PERCENT);
 }
 
+/**
+ * Which usage window a plan reports in. This is the SINGLE rule shared by quota
+ * parsing, exhaustion, and recovery — they must not diverge.
+ *
+ * An allowlist of "known" plans was tried here and was wrong: the upstream model
+ * snapshot alone carries 21 distinct plan strings (`edu_plus`, `finserv`, `k12`,
+ * `quorum`, `self_serve_business_usage_based`, ...), and `CodexAccount.plan` is an
+ * unrestricted string, so any list is a list of the plans someone remembered.
+ * Twelve real plans would have been refused recovery and stayed cooled forever —
+ * the very defect this unit exists to fix, reintroduced as a typo-shaped hole.
+ *
+ * The honest rule is the parser's own: Go and Free report a 30-day window,
+ * everything else (including an absent plan) reports weekly. Recovery reads the
+ * window the parser actually wrote rather than second-guessing it.
+ */
+export function codexQuotaWindowForPlan(plan?: string | null): "monthly" | "weekly" {
+  const normalized = plan?.trim().toLowerCase();
+  return normalized === "go" || normalized === "free" ? "monthly" : "weekly";
+}
+
+export function isCompleteCodexQuotaRecoverySnapshot(
+  quota: Pick<StoredAccountQuota, "weeklyPercent" | "monthlyPercent"> | null,
+  plan?: string | null,
+): boolean {
+  if (!quota || isCodexQuotaExhausted(quota, plan)) return false;
+  // Recovery still fails closed on MISSING EVIDENCE — a credits-only or windowless payload
+  // carries no usage reading at all and must never clear a cooldown. What it does not do is
+  // fail closed on an unfamiliar plan NAME, which only ever meant "cooled forever".
+  const required = codexQuotaWindowForPlan(plan) === "monthly"
+    ? quota.monthlyPercent
+    : quota.weeklyPercent;
+  return typeof required === "number" && Number.isFinite(required);
+}
+
 export function normalizeUsagePercent(value: unknown): number | undefined {
   const numeric = typeof value === "number"
     ? value
@@ -367,7 +400,7 @@ export function parseUsageQuota(data: WhamUsageResponse): Omit<StoredAccountQuot
   }
 
   const quota: Omit<StoredAccountQuota, "updatedAt"> = {};
-  const thirtyDayOnly = data.plan_type?.trim().toLowerCase() === "go" || data.plan_type?.trim().toLowerCase() === "free";
+  const thirtyDayOnly = codexQuotaWindowForPlan(data.plan_type) === "monthly";
   const primaryWindow = data.rate_limit.primary_window;
   const secondaryWindow = data.rate_limit.secondary_window;
   const tertiaryWindow = data.rate_limit.tertiary_window;

src/codex/routing.ts

@@ -133,6 +133,18 @@ export type CodexCooldownSource = "retry-after" | "reset-derived" | "default";
  */
 export type CodexQuotaScope = "shared" | "spark";
 
+export type CodexQuotaRecoveryProbeClaim = {
+  accountId: string;
+  scope?: CodexQuotaScope;
+  leaseId: string;
+  cooldownGeneration: number;
+  credentialGeneration: number;
+};
+
+export type CodexQuotaRecoveryProbeProof = {
+  credentialGeneration?: number;
+};
+
 /**
  * Requests without a resolved native model retain the historic one-account-per-
  * thread behavior. Requests with a known quota scope get an independent
@@ -417,6 +429,119 @@ function canAcquireQuotaProbeLease(health: CodexUpstreamHealth | undefined, now:
   return now - origin >= CODEX_QUOTA_PROBE_INTERVAL_MS;
 }
 
+/**
+ * Claim due reset-derived cooldown probes without consulting account selection.
+ * Pool credentials only: the main account has no quota-refresh single-flight.
+ */
+export function claimDueCodexQuotaRecoveryProbes(
+  config: OcxConfig,
+  limit: number,
+  now = Date.now(),
+): CodexQuotaRecoveryProbeClaim[] {
+  const boundedLimit = Math.max(0, Math.floor(limit));
+  if (boundedLimit === 0) return [];
+  const candidates: Array<{
+    accountId: string;
+    scope?: CodexQuotaScope;
+    health: CodexUpstreamHealth;
+    credentialGeneration: number;
+    order: number;
+  }> = [];
+  for (const [order, account] of (config.codexAccounts ?? []).entries()) {
+    if (!isSelectableCodexPoolAccount(account)
+      || isCodexAccountPaused(config, account.id)
+      || isAccountNeedsReauth(account.id)) continue;
+    const record = readCodexAccountRecord(account.id);
+    if (!record?.credential || record.deletedAt != null) continue;
+    const due = [
+      { scope: undefined, health: upstreamHealth.get(account.id) },
+      ...[...(quotaScopedHealth.get(account.id) ?? [])].map(([scope, health]) => ({ scope, health })),
+    ].filter((entry): entry is { scope?: CodexQuotaScope; health: CodexUpstreamHealth } =>
+      // `spark` is deliberately never claimed. `GET /backend-api/wham/usage` takes no scope
+      // parameter and returns generic weekly/monthly windows, so its result can never prove a
+      // spark recovery — a claim here would spend an upstream call to settle `false` every
+      // time, and (with one claim per account per pass) delay the shared scope that CAN recover.
+      entry.scope !== "spark"
+      && entry.health?.cooldownSource === "reset-derived"
+      && canAcquireQuotaProbeLease(entry.health, now))
+      .sort((a, b) =>
+        (a.health.lastProbeAt ?? a.health.cooldownSince ?? 0)
+        - (b.health.lastProbeAt ?? b.health.cooldownSince ?? 0));
+    const candidate = due[0];
+    if (candidate) candidates.push({
+      accountId: account.id,
+      ...(candidate.scope ? { scope: candidate.scope } : {}),
+      health: candidate.health,
+      credentialGeneration: record.generation,
+      order,
+    });
+  }
+  candidates.sort((a, b) => {
+    const age = (a.health.lastProbeAt ?? a.health.cooldownSince ?? 0)
+      - (b.health.lastProbeAt ?? b.health.cooldownSince ?? 0);
+    return age || a.order - b.order;
+  });
+  return candidates.slice(0, boundedLimit).map(candidate => {
+    const leaseId = randomUUID();
+    const next = {
+      ...candidate.health,
+      probeLeaseId: leaseId,
+      probeLeaseGeneration: candidate.health.cooldownGeneration ?? 0,
+      lastProbeAt: now,
+    };
+    if (candidate.scope) setScopedHealth(candidate.accountId, candidate.scope, next);
+    else upstreamHealth.set(candidate.accountId, next);
+    return {
+      accountId: candidate.accountId,
+      ...(candidate.scope ? { scope: candidate.scope } : {}),
+      leaseId,
+      cooldownGeneration: candidate.health.cooldownGeneration ?? 0,
+      credentialGeneration: candidate.credentialGeneration,
+    };
+  });
+}
+
+/** Settle one background recovery claim without mutating account-wide outcome state. */
+export function settleCodexQuotaRecoveryProbe(
+  claim: CodexQuotaRecoveryProbeClaim,
+  recovered: boolean,
+  proof: CodexQuotaRecoveryProbeProof,
+  now = Date.now(),
+): boolean {
+  const health = claim.scope
+    ? scopedHealthFor(claim.accountId, claim.scope)
+    : upstreamHealth.get(claim.accountId);
+  if (!health || health.probeLeaseId !== claim.leaseId) return false;
+  const fenced = (health.cooldownGeneration ?? 0) === claim.cooldownGeneration
+    && (health.probeLeaseGeneration ?? 0) === claim.cooldownGeneration
+    && claim.credentialGeneration === proof.credentialGeneration
+    && isCodexAccountGenerationLive(claim.accountId, claim.credentialGeneration);
+  if (!recovered || !fenced) {
+    const released = withProbeLeaseReleased(health, now);
+    if (claim.scope) setScopedHealth(claim.accountId, claim.scope, released);
+    else upstreamHealth.set(claim.accountId, released);
+    return false;
+  }
+  if (claim.scope) {
+    deleteScopedHealth(claim.accountId, claim.scope);
+  } else {
+    const {
+      cooldownUntil: _until,
+      cooldownSince: _since,
+      cooldownSource: _source,
+      probeLeaseId: _leaseId,
+      probeLeaseGeneration: _leaseGeneration,
+      ...rest
+    } = health;
+    upstreamHealth.set(claim.accountId, {
+      ...rest,
+      cooldownGeneration: claim.cooldownGeneration + 1,
+      lastProbeAt: now,
+    });
+  }
+  return true;
+}
+
 /** Acquire the recovery probe for one confirmed model-specific quota group. */
 export function tryAcquireCodexQuotaScopeProbeLease(
   accountId: string,

src/server/index.ts

@@ -20,6 +20,7 @@ import {
 } from "../config";
 import { reconcileOAuthProviders } from "../oauth";
 import { invalidateCodexModelsCache } from "../codex/catalog";
+import { registerCodexCooldownRecoveryProbeWorker } from "../codex/auth-api";
 import { startMemoryWatchdog } from "./memory-watchdog";
 import {
   reconcileLiveStateStores,
@@ -416,6 +417,7 @@ export function startServer(port?: number, deps: StartServerDeps = {}) {
   registerAppOwnedMemorySweepFallback();
   configureAppOwnedMemoryBudget(resolveAppOwnedMemoryBudgetBytes(config.appOwnedMemoryBudgetMb));
   enforceAppOwnedMemoryBudget();
+  registerCodexCooldownRecoveryProbeWorker(config);
   startStateStoreSweeper();
   // Issue #42 Phase 3: opt-in archived auto-cleanup (default OFF). Unref'd hourly
   // tick for daily/weekly; startup evaluation is fire-and-forget after listen.