Fix native main token refresh and replay

src/codex/account-store.ts

@@ -11,13 +11,15 @@ import {
   withConfigMutationLockSync,
 } from "../config";
 import { assertNotRealHomeUnderTest } from "../lib/test-home-guard";
+import { CHATGPT_CLIENT_ID, CHATGPT_TOKEN_URL } from "../oauth/chatgpt";
 import type { CodexAccountCredentialRecord, CodexAccountCredentials } from "../types";
 
 type LegacyCodexAccountStore = Record<string, CodexAccountCredentials>;
 type CodexAccountStore = Record<string, CodexAccountCredentialRecord>;
 type RawCodexAccountStore = Record<string, CodexAccountCredentials | CodexAccountCredentialRecord>;
 
-const REFRESH_SKEW_MS = 60_000;
+export const CODEX_REFRESH_SKEW_MS = 60_000;
+const REFRESH_SKEW_MS = CODEX_REFRESH_SKEW_MS;
 const REFRESH_LOCK_STALE_MS = 60_000;
 const REFRESH_LOCK_WAIT_MS = REFRESH_LOCK_STALE_MS + 5_000;
 const REFRESH_LOCK_POLL_MS = 50;
@@ -47,6 +49,23 @@ function isCredential(value: unknown): value is CodexAccountCredentials {
     && typeof value.chatgptAccountId === "string";
 }
 
+function requiredTokenResponseString(data: Record<string, unknown>, field: string): string {
+  const value = data[field];
+  if (typeof value !== "string" || value.trim().length === 0) {
+    throw new TokenRefreshError("unknown", `Codex token refresh returned a malformed ${field}.`);
+  }
+  return value;
+}
+
+function optionalTokenResponseString(data: Record<string, unknown>, field: string): string | undefined {
+  const value = data[field];
+  if (value === undefined || value === null) return undefined;
+  if (typeof value !== "string" || value.trim().length === 0) {
+    throw new TokenRefreshError("unknown", `Codex token refresh returned a malformed ${field}.`);
+  }
+  return value;
+}
+
 function isCredentialRecord(value: unknown): value is CodexAccountCredentialRecord {
   return isObject(value)
     && typeof value.generation === "number"
@@ -229,13 +248,10 @@ export function tombstoneCodexAccount(id: string): number {
   });
 }
 
-const CHATGPT_TOKEN_URL = "https://auth.openai.com/oauth/token";
-const CHATGPT_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann";
-
 export class TokenRefreshError extends Error {
   reason: "expired" | "revoked" | "unknown";
-  constructor(reason: "expired" | "revoked" | "unknown", message: string) {
-    super(message);
+  constructor(reason: "expired" | "revoked" | "unknown", message: string, options?: ErrorOptions) {
+    super(message, options);
     this.name = "TokenRefreshError";
     this.reason = reason;
   }
@@ -331,7 +347,7 @@ function isRefreshLockStale(path: string): boolean {
   }
 }
 
-async function withCodexRefreshFileLock<T>(lockKey: string, signal: AbortSignal, fn: () => Promise<T>): Promise<T> {
+export async function withCodexRefreshFileLock<T>(lockKey: string, signal: AbortSignal, fn: () => Promise<T>): Promise<T> {
   hardenConfigDir();
   const dir = getConfigDir();
   if (!existsSync(dir)) mkdirSync(dir, { recursive: true, mode: 0o700 });
@@ -372,20 +388,68 @@ async function withCodexRefreshFileLock<T>(lockKey: string, signal: AbortSignal,
   }
 }
 
-function findFreshCredentialForGrant(
-  refreshGrantFingerprint: string,
-  excludeId: string,
-): CodexAccountCredentials | null {
-  const now = Date.now();
+export function findFreshCredentialForGrant(query: {
+  refreshGrantFingerprint: string;
+  excludeId?: string;
+  now?: number;
+}): CodexAccountCredentials | null {
+  const now = query.now ?? Date.now();
   const records = loadCodexAccountRecordStore();
   for (const [candidateId, candidate] of Object.entries(records)) {
-    if (candidateId === excludeId || candidate.deletedAt != null || !candidate.credential) continue;
-    if (recordGrantFingerprint(candidate) !== refreshGrantFingerprint) continue;
+    if (candidateId === query.excludeId || candidate.deletedAt != null || !candidate.credential) continue;
+    if (recordGrantFingerprint(candidate) !== query.refreshGrantFingerprint) continue;
     if (candidate.credential.expiresAt > now + REFRESH_SKEW_MS) return candidate.credential;
   }
   return null;
 }
 
+export function findUniqueFreshCredentialForChatgptAccount(query: {
+  chatgptAccountId: string;
+  excludeId?: string;
+  now?: number;
+}): CodexAccountCredentials | null {
+  const chatgptAccountId = query.chatgptAccountId.trim();
+  if (!chatgptAccountId) return null;
+  const now = query.now ?? Date.now();
+  const records = loadCodexAccountRecordStore();
+  let match: CodexAccountCredentials | null = null;
+  for (const [candidateId, candidate] of Object.entries(records)) {
+    if (candidateId === query.excludeId || candidate.deletedAt != null || !candidate.credential) continue;
+    if (candidate.credential.chatgptAccountId !== chatgptAccountId) continue;
+    if (candidate.credential.expiresAt <= now + REFRESH_SKEW_MS) continue;
+    if (match) return null;
+    match = candidate.credential;
+  }
+  return match;
+}
+
+export function publishFreshCredentialForGrant(query: {
+  refreshGrantFingerprint: string;
+  credential: CodexAccountCredentials;
+  excludeId?: string;
+}): void {
+  withCredentialMutationLockSync(() => {
+    const store = loadCodexAccountRecordStore();
+    let changed = false;
+    for (const [candidateId, candidate] of Object.entries(store)) {
+      if (candidateId === query.excludeId || candidate.deletedAt != null || !candidate.credential) continue;
+      if (recordGrantFingerprint(candidate) !== query.refreshGrantFingerprint) continue;
+      store[candidateId] = {
+        credential: {
+          ...query.credential,
+          chatgptAccountId: candidate.credential.chatgptAccountId || query.credential.chatgptAccountId,
+        },
+        generation: candidate.generation + 1,
+        refreshGrantFingerprint: refreshGrantFingerprintForToken(query.credential.refreshToken),
+        replacedAt: candidate.replacedAt,
+        ...preservedValidationMetadata(candidate),
+      };
+      changed = true;
+    }
+    if (changed) persist(store);
+  });
+}
+
 async function notePlanFromRefreshedAccessToken(
   id: string,
   accessToken: string,
@@ -471,7 +535,7 @@ export async function getValidCodexToken(id: string): Promise<CodexTokenResult>
         credential: lockedCred,
       };
     }
-    const sameGrantFreshCredential = findFreshCredentialForGrant(refreshGrantFingerprint, id);
+    const sameGrantFreshCredential = findFreshCredentialForGrant({ refreshGrantFingerprint, excludeId: id });
     if (sameGrantFreshCredential) {
       if (!saveCodexAccountCredentialIfGeneration(id, startGeneration, sameGrantFreshCredential)) {
         throw new CodexCredentialGenerationConflictError();
@@ -505,7 +569,9 @@ export async function getValidCodexToken(id: string): Promise<CodexTokenResult>
         : "unknown" as const;
       throw new TokenRefreshError(reason, `Codex token refresh failed (${reason}); reauthenticate the account.`);
     }
-    const data = (await res.json()) as { access_token: string; refresh_token?: string; expires_in: number };
+    const data = (await res.json()) as Record<string, unknown>;
+    const accessToken = requiredTokenResponseString(data, "access_token");
+    const refreshToken = optionalTokenResponseString(data, "refresh_token");
     // Guard against a missing/non-finite/negative expires_in (malformed upstream
     // response): a NaN expiry would never compare as expired, and a negative
     // duration would stamp an already-past expiry — both block refresh semantics.
@@ -519,8 +585,8 @@ export async function getValidCodexToken(id: string): Promise<CodexTokenResult>
     const safeExpiresAt = Number.isFinite(expiresAt) ? expiresAt : Date.now() + 3600 * 1000;
 
     const updated: CodexAccountCredentials = {
-      accessToken: data.access_token,
-      refreshToken: data.refresh_token ?? lockedCred.refreshToken,
+      accessToken,
+      refreshToken: refreshToken ?? lockedCred.refreshToken,
       expiresAt: safeExpiresAt,
       chatgptAccountId: lockedCred.chatgptAccountId,
     };

src/codex/account-usability.ts

@@ -1,6 +1,6 @@
 import { getCodexAccountCredential } from "./account-store";
 import { isAccountNeedsReauth } from "./account-runtime-state";
-import { MAIN_CODEX_ACCOUNT_ID, isMainAccountTokenLive } from "./main-account";
+import { MAIN_CODEX_ACCOUNT_ID, isMainAccountCredentialUsable } from "./main-account";
 import { hasLegacyMainCodexPoolAccount, isSelectableCodexPoolAccount } from "./account-id";
 import type { OcxConfig } from "../types";
 import { isNativeMainTrafficBlocked } from "./native-profile-startup";
@@ -9,7 +9,7 @@ export interface CodexAccountUsabilityOptions {
   /** 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;
+  isMainAccountTokenLive?: typeof isMainAccountCredentialUsable;
   /** Confirmed account ids for an account-gated model; omitted for ordinary native models. */
   modelEligibleAccountIds?: ReadonlySet<string>;
 }
@@ -32,8 +32,8 @@ export function isCodexAccountUsable(
     // before reservation or token materialization. Treat cached main as a routing
     // candidate without touching the credential file so affinity is not rebound.
     if (options.nativeMainSelectionOnly) return true;
-    // Main account: credential is the read-only ~/.codex/auth.json token (Option A).
-    return (options.isMainAccountTokenLive ?? isMainAccountTokenLive)();
+    // Main account: a refresh grant is enough to route; materialization refreshes before I/O.
+    return (options.isMainAccountTokenLive ?? isMainAccountCredentialUsable)();
   }
   const exists = (config.codexAccounts ?? [])
     .some(account => isSelectableCodexPoolAccount(account) && account.id === accountId);

src/codex/auth-context.ts

@@ -4,6 +4,7 @@ import {
   CodexCredentialRefreshLockTimeoutError,
   CodexCredentialRefreshBusyError,
   CodexCredentialRefreshStaleError,
+  TokenRefreshError,
   getValidCodexToken,
   isCodexAccountGenerationLive,
 } from "./account-store";
@@ -12,7 +13,14 @@ import { isCodexAccountPaused } from "./account-pause";
 import { ConfigMutationLockError } from "../config";
 import { isCodexAccountUsable } from "./account-usability";
 import { reconcileMainCodexAccountRuntimeState } from "./account-lifecycle";
-import { MAIN_CODEX_ACCOUNT_ID, getMainAccountToken, isMainAccountTokenLive } from "./main-account";
+import {
+  MAIN_CODEX_ACCOUNT_ID,
+  MainAuthJsonChangedDuringRefreshError,
+  getMainAccountToken,
+  getValidMainAccountToken,
+  isMainAccountTokenLive,
+  type NativeMainRefreshDependencies,
+} from "./main-account";
 import { isNativeMainTrafficBlocked, nativeMainStartupGateSnapshot } from "./native-profile-startup";
 import type { NativeMainStartupBlockReason } from "./native-profile-startup";
 import {
@@ -29,6 +37,7 @@ import {
   entitledCodexAccountIdsForModel,
   isDirectCallerEntitledToCodexModel,
   resolveCodexModelEntitlements,
+  type CodexModelEntitlementResolveOptions,
   type CodexModelEntitlementSnapshot,
 } from "./model-entitlements";
 import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "./catalog/native-models";
@@ -318,9 +327,16 @@ export function shouldMarkAccountNeedsReauthForCodexAuthFailure(cause: unknown):
     && !(cause instanceof CodexCredentialRefreshLockTimeoutError)
     && !(cause instanceof CodexCredentialRefreshBusyError)
     && !(cause instanceof CodexCredentialRefreshStaleError)
+    && !(cause instanceof MainAuthJsonChangedDuringRefreshError)
+    && !(cause instanceof TokenRefreshError && cause.reason === "unknown")
     && !(cause instanceof ConfigMutationLockError);
 }
 
+type ResolveCodexModelEntitlementsForAuthContext = (
+  config: Pick<OcxConfig, "codexAccounts">,
+  options?: Pick<CodexModelEntitlementResolveOptions, "nativeMainRefreshDependencies" | "signal">,
+) => Promise<CodexModelEntitlementSnapshot>;
+
 export interface ResolveCodexAuthContextOptions {
   excludeAccountId?: string;
   /** Resolve exactly this account without consulting or mutating Pool selection. */
@@ -332,17 +348,27 @@ export interface ResolveCodexAuthContextOptions {
   /** Test-only native credential read seams. */
   isMainAccountTokenLive?: () => boolean;
   getMainAccountToken?: typeof getMainAccountToken;
+  getValidMainAccountToken?: typeof getValidMainAccountToken;
+  nativeMainRefreshDependencies?: NativeMainRefreshDependencies;
+  signal?: AbortSignal;
   primeCodexPoolQuotas?: (config: OcxConfig, reason: string) => Promise<void>;
   /** Test seam for account-gated native model discovery. */
-  resolveCodexModelEntitlements?: (
-    config: Pick<OcxConfig, "codexAccounts">,
-  ) => Promise<CodexModelEntitlementSnapshot>;
+  resolveCodexModelEntitlements?: ResolveCodexModelEntitlementsForAuthContext;
   /** Direct requests admitted with a proxy bearer substitute the stored native-main credential. */
   substituteMainCredentialForDirect?: boolean;
   /** Test seam for a Direct request's own forwarded ChatGPT credential. */
   isDirectCallerEntitledToCodexModel?: (headers: Headers, modelId: string) => Promise<boolean>;
 }
 
+function entitlementOptionsForAuthContext(
+  options: Pick<ResolveCodexAuthContextOptions, "nativeMainRefreshDependencies" | "signal">,
+): Pick<CodexModelEntitlementResolveOptions, "nativeMainRefreshDependencies" | "signal"> {
+  return {
+    ...(options.nativeMainRefreshDependencies ? { nativeMainRefreshDependencies: options.nativeMainRefreshDependencies } : {}),
+    ...(options.signal ? { signal: options.signal } : {}),
+  };
+}
+
 export interface CodexAccountSelectionAdmission {
   readonly mainProfileDraining: boolean;
   claimMainProfile(): boolean;
@@ -367,7 +393,10 @@ export async function resolveCodexAuthContext(
     if (options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId)) {
       const entitled = options.substituteMainCredentialForDirect
         ? entitledCodexAccountIdsForModel(
-            await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config),
+            await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(
+              config,
+              entitlementOptionsForAuthContext(options),
+            ),
             options.modelId,
           )?.has(MAIN_CODEX_ACCOUNT_ID) === true
         : await (options.isDirectCallerEntitledToCodexModel ?? isDirectCallerEntitledToCodexModel)(
@@ -382,7 +411,10 @@ export async function resolveCodexAuthContext(
   }
   const affinityKey = fixedAccountId === undefined ? codexPoolAffinityKey(headers) : undefined;
   const entitlementSnapshot = options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId)
-    ? await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config)
+    ? await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(
+      config,
+      entitlementOptionsForAuthContext(options),
+    )
     : undefined;
   const modelEligibleAccountIds = entitlementSnapshot
     ? entitledCodexAccountIdsForModel(entitlementSnapshot, options.modelId)
@@ -486,7 +518,12 @@ export async function resolveCodexAuthContext(
   // best-effort prime so the NEXT routing decision has real scores. This never
   // blocks the current request, and the helper's single-flight guard collapses
   // repeated triggers into one pass.
-  if (fixedAccountId === undefined && !nativeMainReadsForbidden && !getAccountQuota(accountId)) {
+  if (
+    accountId !== MAIN_CODEX_ACCOUNT_ID
+    && fixedAccountId === undefined
+    && !nativeMainReadsForbidden
+    && !getAccountQuota(accountId)
+  ) {
     if (options.primeCodexPoolQuotas) {
       void options.primeCodexPoolQuotas(config, "pre-route").catch(() => {});
     } else {
@@ -520,8 +557,21 @@ export async function resolveCodexAuthContext(
   }
 
   if (accountId === MAIN_CODEX_ACCOUNT_ID) {
-    // Main account in rotation: inject the read-only auth.json token and fail closed if it vanished.
-    const token = (options.getMainAccountToken ?? getMainAccountToken)();
+    // Main account in rotation: refresh auth.json before upstream I/O and fail closed if it vanished.
+    let token: { accessToken: string; chatgptAccountId: string } | null;
+    try {
+      token = await (options.getValidMainAccountToken ?? getValidMainAccountToken)({
+        signal: options.signal,
+        dependencies: options.nativeMainRefreshDependencies,
+      });
+    } catch (cause) {
+      if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId);
+      else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId);
+      if (shouldMarkAccountNeedsReauthForCodexAuthFailure(cause)) {
+        markAccountNeedsReauth(accountId, writerGeneration);
+      }
+      throw new CodexAuthContextError(accountId, cause);
+    }
     if (!token) {
       // Nothing will reach upstream, so give the probe back instead of burning it.
       if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId);
@@ -530,6 +580,15 @@ export async function resolveCodexAuthContext(
         fixedAccountId !== undefined ? "Selected Codex account is unavailable" : undefined,
       );
     }
+    if (fixedAccountId === undefined && !nativeMainReadsForbidden && !getAccountQuota(accountId)) {
+      if (options.primeCodexPoolQuotas) {
+        void options.primeCodexPoolQuotas(config, "pre-route").catch(() => {});
+      } else {
+        import("./auth-api")
+          .then(({ primeCodexPoolQuotas }) => primeCodexPoolQuotas(config, "pre-route"))
+          .catch(() => {});
+      }
+    }
     return {
       kind: "main-pool",
       accountId,
@@ -644,6 +703,36 @@ export function materializeCodexUpstreamAuth(
   return selected;
 }
 
+export async function materializeCodexUpstreamAuthAsync(
+  headers: Headers,
+  ctx: CodexAuthContext,
+  options: {
+    substituteMainCredential?: boolean;
+    signal?: AbortSignal;
+    getValidMainAccountToken?: typeof getValidMainAccountToken;
+    nativeMainRefreshDependencies?: NativeMainRefreshDependencies;
+  } = {},
+): Promise<Headers> {
+  if (ctx.kind !== "main" || options.substituteMainCredential !== true) {
+    return materializeCodexUpstreamAuth(headers, ctx, options);
+  }
+  const selected = new Headers();
+  for (const name of FORWARD_HEADERS) {
+    const value = headers.get(name);
+    if (value) selected.set(name, value);
+  }
+  const stored = await (options.getValidMainAccountToken ?? getValidMainAccountToken)({
+    signal: options.signal,
+    dependencies: options.nativeMainRefreshDependencies,
+  });
+  if (!stored?.accessToken) {
+    throw new CodexMainSubstitutionUnavailableError();
+  }
+  selected.set("authorization", `Bearer ${stored.accessToken}`);
+  if (stored.chatgptAccountId) selected.set("chatgpt-account-id", stored.chatgptAccountId);
+  return selected;
+}
+
 /** @deprecated Prefer materializeCodexUpstreamAuth; kept for call sites without admission context. */
 export function headersForCodexAuthContext(headers: Headers, ctx: CodexAuthContext): Headers {
   return materializeCodexUpstreamAuth(headers, ctx);