fix(codex): persist refreshed pool account plans

Author: luvs01

src/codex/auth-api.ts

@@ -1,10 +1,11 @@
-import { loadConfig, saveConfigPreservingClaudeCode } from "../config";
+import { loadConfig, readConfigDiagnostics, saveConfigPreservingClaudeCode } from "../config";
 import { withCodexAccountLogLabel } from "./account-label";
 import {
   getCodexAccountCredential,
   getValidCodexToken,
   isCodexAccountGenerationLive,
   markCodexAccountValidated,
+  readCodexAccountRecord,
   saveCodexAccountCredential,
   CodexCredentialGenerationConflictError,
   CodexCredentialRefreshLockTimeoutError,
@@ -413,6 +414,8 @@ async function fetchMainAccountInfoAttempt(forceRefresh: boolean, retriesRemaini
 interface PoolQuotaResult {
   quota: StoredAccountQuota | null;
   needsReauth: boolean;
+  /** Credential generation whose cache or network result this DTO state belongs to. */
+  credentialGeneration?: number;
   /** Present only when this call freshly parsed a WHAM usage response. */
   freshQuota?: Omit<StoredAccountQuota, "updatedAt">;
   /** Present only when this call's WHAM response included a non-empty `plan_type`. */
@@ -423,6 +426,16 @@ interface PoolQuotaResult {
   freshResetCredits?: number;
 }
 
+interface PoolQuotaRefreshFlight {
+  state: {
+    startCredentialGeneration?: number;
+    resolvedCredentialGeneration?: number;
+  };
+  promise: Promise<PoolQuotaResult>;
+}
+
+const poolQuotaRefreshInFlight = new Map<string, Set<PoolQuotaRefreshFlight>>();
+
 export interface CodexAuthAccountDto {
   id: string;
   alias?: string;
@@ -440,39 +453,144 @@ export interface CodexAuthAccountDto {
   healthAction?: string;
 }
 
-async function fetchPoolAccountQuota(accountId: string, forceRefresh = false, configuredPlan?: string): Promise<PoolQuotaResult> {
-  const existing = getAccountQuota(accountId);
-  if (!forceRefresh && existing && Date.now() - existing.updatedAt < POOL_CACHE_TTL) {
-    return { quota: existing, needsReauth: false };
+interface FreshPoolPlanUpdate {
+  accountId: string;
+  plan: string;
+  credentialGeneration: number;
+}
+
+/**
+ * Persist only validated plan leaves against the latest disk snapshot. A quota GET must not save
+ * the long-lived runtime object wholesale: unrelated manual/provider writes may have landed while
+ * WHAM requests were in flight. Missing or malformed files fail closed: a read path must not
+ * recreate a deleted config from the server's older in-memory snapshot.
+ */
+function reconcileFreshPoolAccountPlans(runtimeConfig: OcxConfig, updates: FreshPoolPlanUpdate[]): void {
+  if (updates.length === 0) return;
+  const diagnostics = readConfigDiagnostics();
+  if (diagnostics.source !== "file") return;
+  const persistedConfig = diagnostics.config;
+  const accepted: FreshPoolPlanUpdate[] = [];
+  let persistedChanged = false;
+
+  for (const update of updates) {
+    if (!isCodexAccountGenerationLive(update.accountId, update.credentialGeneration)) continue;
+    const liveAccount = configuredPoolAccount(runtimeConfig, update.accountId);
+    const persistedAccount = configuredPoolAccount(persistedConfig, update.accountId);
+    if (!liveAccount || !persistedAccount) continue;
+    accepted.push(update);
+    if (persistedAccount.plan !== update.plan) {
+      persistedAccount.plan = update.plan;
+      persistedChanged = true;
+    }
   }
+
+  if (persistedChanged) saveConfigPreservingClaudeCode(persistedConfig);
+  for (const update of accepted) {
+    if (!isCodexAccountGenerationLive(update.accountId, update.credentialGeneration)) continue;
+    const liveAccount = configuredPoolAccount(runtimeConfig, update.accountId);
+    if (liveAccount) liveAccount.plan = update.plan;
+  }
+}
+
+async function fetchFreshPoolAccountQuota(
+  accountId: string,
+  existing: StoredAccountQuota | null,
+  configuredPlan?: string,
+  onCredentialGeneration?: (generation: number) => void,
+): Promise<PoolQuotaResult> {
+  let requestCredentialGeneration = readCodexAccountRecord(accountId)?.generation;
   try {
     const { accessToken, chatgptAccountId, generation } = await getValidCodexToken(accountId);
+    requestCredentialGeneration = generation;
+    onCredentialGeneration?.(generation);
     const resp = await fetch("https://chatgpt.com/backend-api/wham/usage", {
       headers: { Authorization: `Bearer ${accessToken}`, "ChatGPT-Account-Id": chatgptAccountId },
       signal: AbortSignal.timeout(8000),
     });
-    if (!resp.ok) return { quota: existing ?? null, needsReauth: resp.status === 401 };
+    if (!resp.ok) {
+      return {
+        quota: existing ?? null,
+        needsReauth: resp.status === 401,
+        credentialGeneration: generation,
+      };
+    }
     const data = (await resp.json()) as WhamUsageResponse;
     const freshPlan = nonEmptyPlan(data.plan_type) ?? undefined;
     const quota = parseUsageQuota({ ...data, plan_type: freshPlan ?? configuredPlan });
     const freshResetCredits = quota?.resetCredits;
-    if (!quota) return { quota: existing ?? null, needsReauth: false };
+    if (!quota) {
+      return {
+        quota: isCodexAccountGenerationLive(accountId, generation) ? existing ?? null : getAccountQuota(accountId),
+        needsReauth: false,
+        credentialGeneration: generation,
+      };
+    }
     if (!isCodexAccountGenerationLive(accountId, generation)) {
-      return { quota: getAccountQuota(accountId), needsReauth: false };
+      return { quota: null, needsReauth: false, credentialGeneration: generation };
     }
     setAccountQuotaFromParsed(accountId, quota);
     return {
       quota: getAccountQuota(accountId),
       needsReauth: false,
+      credentialGeneration: generation,
       freshQuota: quota,
       freshCredentialGeneration: generation,
       ...(freshPlan !== undefined ? { freshPlan } : {}),
       ...(freshResetCredits !== undefined ? { freshResetCredits } : {}),
     };
   } catch (e) {
-    if (e instanceof CodexCredentialGenerationConflictError || e instanceof CodexCredentialRefreshLockTimeoutError) return { quota: existing ?? null, needsReauth: false };
-    if (e instanceof TokenRefreshError) return { quota: existing ?? null, needsReauth: true };
-    return { quota: existing ?? null, needsReauth: false };
+    if (e instanceof CodexCredentialGenerationConflictError || e instanceof CodexCredentialRefreshLockTimeoutError) {
+      return { quota: null, needsReauth: false, credentialGeneration: requestCredentialGeneration };
+    }
+    if (e instanceof TokenRefreshError) {
+      return { quota: existing ?? null, needsReauth: true, credentialGeneration: requestCredentialGeneration };
+    }
+    return { quota: existing ?? null, needsReauth: false, credentialGeneration: requestCredentialGeneration };
+  }
+}
+
+async function fetchPoolAccountQuota(accountId: string, forceRefresh = false, configuredPlan?: string): Promise<PoolQuotaResult> {
+  const existing = getAccountQuota(accountId);
+  if (!forceRefresh && existing && Date.now() - existing.updatedAt < POOL_CACHE_TTL) {
+    return {
+      quota: existing,
+      needsReauth: false,
+      credentialGeneration: readCodexAccountRecord(accountId)?.generation,
+    };
+  }
+  // A token refresh may increment the generation (and rotate the refresh token) before WHAM
+  // completes. Join a flight whose starting or resolved generation is still current, but let a
+  // replacement credential with the same pool id start its own request.
+  const record = readCodexAccountRecord(accountId);
+  const flights = poolQuotaRefreshInFlight.get(accountId);
+  const current = flights && [...flights].find(flight => {
+    const generation = flight.state.resolvedCredentialGeneration
+      ?? flight.state.startCredentialGeneration;
+    return generation !== undefined && isCodexAccountGenerationLive(accountId, generation);
+  });
+  if (current) return current.promise;
+
+  const state: PoolQuotaRefreshFlight["state"] = {
+    startCredentialGeneration: record?.generation,
+  };
+  const refresh = fetchFreshPoolAccountQuota(
+    accountId,
+    existing,
+    configuredPlan,
+    generation => { state.resolvedCredentialGeneration = generation; },
+  );
+  const flight: PoolQuotaRefreshFlight = { state, promise: refresh };
+  const activeFlights = flights ?? new Set<PoolQuotaRefreshFlight>();
+  activeFlights.add(flight);
+  if (!flights) poolQuotaRefreshInFlight.set(accountId, activeFlights);
+  try {
+    return await refresh;
+  } finally {
+    activeFlights.delete(flight);
+    if (activeFlights.size === 0 && poolQuotaRefreshInFlight.get(accountId) === activeFlights) {
+      poolQuotaRefreshInFlight.delete(accountId);
+    }
   }
 }
 
@@ -542,12 +660,48 @@ export async function listCodexAuthAccounts(config: OcxConfig, forceRefresh = fa
   const runtimeConfig = getRuntimeConfig(config);
   const poolAccounts = (runtimeConfig.codexAccounts ?? []).filter(isSelectableCodexPoolAccount);
   const mainInfo = await fetchMainAccountInfo(forceRefresh);
-  const withQuota = await mapWithConcurrency(poolAccounts, POOL_QUOTA_REFRESH_CONCURRENCY, async a => {
-    const cred = getCodexAccountCredential(a.id);
+  const refreshedPool = await mapWithConcurrency(poolAccounts, POOL_QUOTA_REFRESH_CONCURRENCY, async account => {
+    const cred = getCodexAccountCredential(account.id);
     const quotaResult = cred
-      ? await fetchPoolAccountQuota(a.id, forceRefresh, a.plan)
+      ? await fetchPoolAccountQuota(account.id, forceRefresh, account.plan)
       : { quota: null, needsReauth: true };
-    return poolAccountDto(a, quotaResult, !!cred, isCodexAccountPaused(runtimeConfig, a.id));
+    return { accountId: account.id, quotaResult };
+  });
+
+  // WHAM plan_type is authoritative only for the credential generation that fetched it. Collect
+  // changes after every parallel read settles, then apply one narrow disk patch for the batch.
+  const planUpdates = refreshedPool.flatMap(({ accountId, quotaResult }): FreshPoolPlanUpdate[] => {
+    const plan = quotaResult.freshPlan;
+    const credentialGeneration = quotaResult.freshCredentialGeneration;
+    return plan && credentialGeneration !== undefined
+      ? [{ accountId, plan, credentialGeneration }]
+      : [];
+  });
+  reconcileFreshPoolAccountPlans(runtimeConfig, planUpdates);
+
+  const withQuota = refreshedPool.flatMap(({ accountId, quotaResult }) => {
+    const currentAccount = configuredPoolAccount(runtimeConfig, accountId);
+    if (!currentAccount) return [];
+    const currentCredential = getCodexAccountCredential(accountId);
+    if (!currentCredential) {
+      return [poolAccountDto(
+        currentAccount,
+        { quota: null, needsReauth: true },
+        false,
+        isCodexAccountPaused(runtimeConfig, accountId),
+      )];
+    }
+    const resultGeneration = quotaResult.credentialGeneration ?? quotaResult.freshCredentialGeneration;
+    const effectiveQuotaResult = resultGeneration !== undefined
+      && !isCodexAccountGenerationLive(accountId, resultGeneration)
+      ? { quota: null, needsReauth: false }
+      : quotaResult;
+    return [poolAccountDto(
+      currentAccount,
+      effectiveQuotaResult,
+      true,
+      isCodexAccountPaused(runtimeConfig, accountId),
+    )];
   });
   const hasMainCredential = readCodexTokens() !== null;
   const mainNeedsReauth = !hasMainCredential || isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID);
@@ -1128,6 +1282,9 @@ export async function handleCodexAuthAPI(
                 expiresAt: cred.expires,
                 chatgptAccountId: oauthAccountId,
               });
+              // A successful reauthentication replaces the credential generation. Do not let a
+              // failed optional WHAM probe make the replacement inherit quota from the old record.
+              if (reauth) clearAccountQuota(accountId);
               markCodexAccountValidated(accountId, warmup.validatedAt);
               clearAccountNeedsReauth(accountId);
               if (quota) {