fix(codex): reset main runtime state after account switch

src/codex/account-lifecycle.ts

@@ -1,15 +1,46 @@
 import { removeCodexAccountCredential } from "./account-store";
 import { clearAccountNeedsReauth } from "./account-runtime-state";
+import { getMainChatgptAccountId } from "./auth-collision";
+import { MAIN_CODEX_ACCOUNT_ID, setMainAccountPlan } from "./main-account";
 import { clearAccountQuota } from "./quota";
 import { clearCodexUpstreamHealthForAccount, clearThreadAccountMapForAccount } from "./routing";
 import { invalidateCodexWebSocketsForAccount } from "./websocket-registry";
+import { clearMainAccountInfoCache } from "./main-account-cache";
 import type { OcxConfig } from "../types";
 
+let observedMainChatgptAccountId: string | undefined;
+
 export function purgeCodexAccountRuntimeState(accountId: string): void {
   clearAccountNeedsReauth(accountId);
   clearAccountQuota(accountId);
   clearThreadAccountMapForAccount(accountId);
   clearCodexUpstreamHealthForAccount(accountId);
+  if (accountId === MAIN_CODEX_ACCOUNT_ID) clearMainAccountInfoCache();
+}
+
+/**
+ * The main Codex login is stored under the stable `__main__` alias, while
+ * `~/.codex/auth.json` can be replaced with credentials for another physical
+ * ChatGPT account. Drop alias-keyed runtime state when that identity changes so
+ * cooldown, quota, reauth, and thread affinity do not leak across accounts.
+ */
+export function reconcileMainCodexAccountRuntimeState(): boolean {
+  const currentAccountId = getMainChatgptAccountId();
+  // A missing/malformed auth.json is an unknown identity, not a confirmed account switch. Keep the
+  // prior observation and its safety state until a real account id can be read again.
+  if (currentAccountId === null) return false;
+  const previousAccountId = observedMainChatgptAccountId;
+  observedMainChatgptAccountId = currentAccountId;
+  if (previousAccountId === undefined || previousAccountId === currentAccountId) return false;
+
+  purgeCodexAccountRuntimeState(MAIN_CODEX_ACCOUNT_ID);
+  setMainAccountPlan(null);
+  invalidateCodexWebSocketsForAccount(MAIN_CODEX_ACCOUNT_ID);
+  return true;
+}
+
+export function resetMainCodexAccountIdentityTrackingForTests(): void {
+  observedMainChatgptAccountId = undefined;
 }
 
 export function deleteCodexAccount(runtimeConfig: OcxConfig, accountId: string): void {

src/codex/auth-api.ts

@@ -9,8 +9,8 @@ import {
   CodexCredentialRefreshLockTimeoutError,
   TokenRefreshError,
 } from "./account-store";
-import { deleteCodexAccount } from "./account-lifecycle";
-import { checkAccountIdCollision, readCodexTokens } from "./auth-collision";
+import { deleteCodexAccount, reconcileMainCodexAccountRuntimeState } from "./account-lifecycle";
+import { checkAccountIdCollision, getMainChatgptAccountId, readCodexTokens } from "./auth-collision";
 export { checkAccountIdCollision, getMainChatgptAccountId } from "./auth-collision";
 export { clearAccountNeedsReauth, isAccountNeedsReauth, markAccountNeedsReauth } from "./account-runtime-state";
 import { clearAccountNeedsReauth, isAccountNeedsReauth, markAccountNeedsReauth } from "./account-runtime-state";
@@ -26,6 +26,13 @@ import {
 export { clearAccountQuota, getAccountQuota, parseUsageQuota, updateAccountQuota } from "./quota";
 import { extractAccountId, decodeJwtPayload } from "../oauth/chatgpt";
 import { MAIN_CODEX_ACCOUNT_ID, setMainAccountPlan } from "./main-account";
+import {
+  clearMainAccountInfoCache,
+  getMainAccountInfoCache,
+  setMainAccountInfoCache,
+  type MainAccountInfo,
+} from "./main-account-cache";
+export { clearMainAccountInfoCache } from "./main-account-cache";
 import { maskEmail } from "../lib/privacy";
 import { CodexWarmupError, codexWarmupFailureReason, warmCodexAccount } from "./warmup";
 export { maskEmail } from "../lib/privacy";
@@ -179,7 +186,6 @@ function expireCodexAuthFlow(flowId: string | null, error = "Login cancelled"):
   }
 }
 
-let mainAccountCache: { email: string | null; plan: string | null; quota: Omit<StoredAccountQuota, "updatedAt"> | null; ts: number } | null = null;
 const MAIN_CACHE_TTL = 5 * 60_000;
 const POOL_CACHE_TTL = 5 * 60_000;
 const POOL_QUOTA_REFRESH_CONCURRENCY = 4;
@@ -245,36 +251,62 @@ async function isTerminalMainAuthResponse(resp: Response): Promise<boolean> {
   }
 }
 
-export async function fetchMainAccountInfo(forceRefresh = false): Promise<{ email: string | null; plan: string | null; quota: Omit<StoredAccountQuota, "updatedAt"> | null }> {
+export async function fetchMainAccountInfo(forceRefresh = false): Promise<MainAccountInfo> {
+  return fetchMainAccountInfoAttempt(forceRefresh, 1);
+}
+
+const EMPTY_MAIN_ACCOUNT_INFO: MainAccountInfo = { email: null, plan: null, quota: null };
+
+async function retryMainAccountInfoIfIdentityChanged(
+  requestAccountId: string | null,
+  retriesRemaining: number,
+): Promise<MainAccountInfo | null> {
+  const currentAccountId = getMainChatgptAccountId();
+  if (currentAccountId === null || currentAccountId === requestAccountId) return null;
+  reconcileMainCodexAccountRuntimeState();
+  return retriesRemaining > 0
+    ? fetchMainAccountInfoAttempt(true, retriesRemaining - 1)
+    : EMPTY_MAIN_ACCOUNT_INFO;
+}
+
+async function fetchMainAccountInfoAttempt(forceRefresh: boolean, retriesRemaining: number): Promise<MainAccountInfo> {
+  reconcileMainCodexAccountRuntimeState();
   const tokens = readCodexTokens();
   if (!tokens) {
-    mainAccountCache = null;
+    clearMainAccountInfoCache();
     markAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID);
-    return { email: null, plan: null, quota: null };
+    return EMPTY_MAIN_ACCOUNT_INFO;
   }
-  if (!forceRefresh && mainAccountCache && Date.now() - mainAccountCache.ts < MAIN_CACHE_TTL) {
-    return mainAccountCache;
+  const requestAccountId = extractAccountId(tokens.id_token, tokens.access_token) ?? (tokens.account_id || null);
+  const cached = getMainAccountInfoCache();
+  if (!forceRefresh && cached && Date.now() - cached.ts < MAIN_CACHE_TTL) {
+    return cached;
   }
   try {
     const resp = await fetch("https://chatgpt.com/backend-api/wham/usage", {
       headers: { Authorization: `Bearer ${tokens.access_token}`, "ChatGPT-Account-Id": tokens.account_id },
       signal: AbortSignal.timeout(8000),
     });
     if (!resp.ok) {
-      if (await isTerminalMainAuthResponse(resp)) {
-        mainAccountCache = null;
+      const terminalAuthFailure = await isTerminalMainAuthResponse(resp);
+      const retried = await retryMainAccountInfoIfIdentityChanged(requestAccountId, retriesRemaining);
+      if (retried) return retried;
+      if (terminalAuthFailure) {
+        clearMainAccountInfoCache();
         markAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID);
       }
-      return { email: null, plan: null, quota: null };
+      return EMPTY_MAIN_ACCOUNT_INFO;
     }
     const data = (await resp.json()) as WhamUsageResponse;
+    const retried = await retryMainAccountInfoIfIdentityChanged(requestAccountId, retriesRemaining);
+    if (retried) return retried;
     const result = {
       email: data.email ?? null,
       plan: data.plan_type ?? null,
       quota: parseUsageQuota(data),
       ts: Date.now(),
     };
-    mainAccountCache = result;
+    setMainAccountInfoCache(result);
     clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID);
     // Mirror main quota + plan into the shared stores so the rotation engine can
     // score and auto-switch the main account exactly like a pool account (Option A).
@@ -291,7 +323,8 @@ export async function fetchMainAccountInfo(forceRefresh = false): Promise<{ emai
     }
     return result;
   } catch {
-    return { email: null, plan: null, quota: null };
+    const retried = await retryMainAccountInfoIfIdentityChanged(requestAccountId, retriesRemaining);
+    return retried ?? EMPTY_MAIN_ACCOUNT_INFO;
   }
 }
 
@@ -369,6 +402,10 @@ export async function primeCodexPoolQuotas(config: OcxConfig, reason: string): P
     || providerCodexAccountMode(OPENAI_CODEX_PROVIDER_ID, openai) !== "pool"
   ) return;
   if (primeInFlight) return primeInFlight;
+  // Seed the observed physical main identity before startup/lazy priming can populate quota or
+  // plan state. Otherwise the first post-startup account switch sees no previous identity and
+  // skips the purge that protects the stable __main__ alias.
+  reconcileMainCodexAccountRuntimeState();
   primeInFlight = (async () => {
     const runtimeConfig = getRuntimeConfig(config);
     const pool = (runtimeConfig.codexAccounts ?? []).filter(a => !a.isMain);

src/codex/auth-context.ts

@@ -6,6 +6,7 @@ import {
 } from "./account-store";
 import { markAccountNeedsReauth } from "./account-runtime-state";
 import { isCodexAccountUsable } from "./account-usability";
+import { reconcileMainCodexAccountRuntimeState } from "./account-lifecycle";
 import { MAIN_CODEX_ACCOUNT_ID, getMainAccountToken } from "./main-account";
 import {
   getCodexAccountCooldownUntil,
@@ -107,6 +108,7 @@ export async function resolveCodexAuthContext(
     if (!hasCallerCodexBearer(headers)) throw new CodexDirectAuthenticationError();
     return { kind: "main", accountId: null };
   }
+  reconcileMainCodexAccountRuntimeState();
   const threadId = headers.get("x-codex-parent-thread-id");
   const resolution = options.excludeAccountId
     ? (() => {

src/codex/main-account-cache.ts

@@ -0,0 +1,25 @@
+import type { StoredAccountQuota } from "./quota";
+
+export interface MainAccountInfo {
+  email: string | null;
+  plan: string | null;
+  quota: Omit<StoredAccountQuota, "updatedAt"> | null;
+}
+
+export interface CachedMainAccountInfo extends MainAccountInfo {
+  ts: number;
+}
+
+let cachedMainAccountInfo: CachedMainAccountInfo | null = null;
+
+export function getMainAccountInfoCache(): CachedMainAccountInfo | null {
+  return cachedMainAccountInfo;
+}
+
+export function setMainAccountInfoCache(value: CachedMainAccountInfo): void {
+  cachedMainAccountInfo = value;
+}
+
+export function clearMainAccountInfoCache(): void {
+  cachedMainAccountInfo = null;
+}