fix(codex): harden native main refresh publication

src/codex/main-account.ts

@@ -6,18 +6,23 @@ import {
   decodeJwtPayload,
   extractAccountId,
   refreshChatGPTToken,
+  ChatGPTTokenRefreshError,
 } from "../oauth/chatgpt";
 import type { OAuthCredentials } from "../oauth/types";
 import { extractChatgptPlanType } from "./plan";
 import { MAIN_CODEX_ACCOUNT_ID } from "./account-id";
-import {
-  refreshGrantFingerprintForToken,
-  withCodexRefreshFileLock,
-} from "./account-store";
-import { atomicWriteFile, resolveWriteTarget } from "../config/atomic-write";
+import { resolveWriteTarget } from "../config/atomic-write";
 import { resolveCodexHomeDir } from "./home";
 import { assertNotRealCodexHomeUnderTest } from "../lib/test-home-guard";
 import { clearAccountNeedsReauth } from "./account-runtime-state";
+import { withNativeMainExclusiveClaim } from "./native-main-claim";
+import { withNativeMainOwnerOperation } from "./native-main-owner";
+import { resolveNativeProfileContext, type NativeProfileContext } from "./native-profile-store";
+import {
+  NativeMainRefreshPublicationError,
+  publishNativeMainRefresh,
+  recoverNativeMainRefreshPublication,
+} from "./native-main-refresh-publication";
 
 export { MAIN_CODEX_ACCOUNT_ID } from "./account-id";
 
@@ -29,10 +34,13 @@ export { MAIN_CODEX_ACCOUNT_ID } from "./account-id";
 let mainAccountPlan: string | null = null;
 let jwtPlanAttempted = false;
 const MAIN_TOKEN_REFRESH_SKEW_MS = 60_000;
+const NATIVE_MAIN_REFRESH_WAIT_MS = 30_000;
+const MAX_NATIVE_MAIN_REFRESH_FLIGHTS = 32;
 let beforeMainAuthJsonRenameForTests: (() => void) | null = null;
 
 type MainAuthJsonCredential = {
   path: string;
+  raw: string;
   rawSha256: string;
   root: Record<string, unknown>;
   tokens: Record<string, unknown>;
@@ -46,6 +54,20 @@ export interface NativeMainRefreshDependencies {
   signal?: AbortSignal;
 }
 
+type NativeMainRefreshFlight = {
+  controller: AbortController;
+  deadline: ReturnType<typeof setTimeout>;
+  promise: Promise<{ accessToken: string; chatgptAccountId: string }>;
+};
+
+type NativeMainRefreshResolution = {
+  dependencies: NativeMainRefreshDependencies;
+  rejectedAccessToken: string | undefined;
+  replacementAttempted: boolean;
+};
+
+const nativeMainRefreshFlights = new Map<string, NativeMainRefreshFlight>();
+
 export class MainAuthJsonChangedDuringRefreshError extends Error {
   constructor() {
     super("Codex auth.json changed while its token was refreshing");
@@ -62,12 +84,15 @@ export class MainAccountTokenRefreshError extends Error {
   }
 }
 
-function nonEmptyString(value: unknown): string | undefined {
-  return typeof value === "string" && value.trim() ? value : undefined;
+export class MainAccountRefreshCancelledError extends Error {
+  constructor() {
+    super("Native credential refresh was cancelled.");
+    this.name = "MainAccountRefreshCancelledError";
+  }
 }
 
-function sha256(value: string): string {
-  return createHash("sha256").update(value).digest("hex");
+function nonEmptyString(value: unknown): string | undefined {
+  return typeof value === "string" && value.trim() ? value : undefined;
 }
 
 function readMainAuthJsonCredential(): MainAuthJsonCredential | null {
@@ -94,7 +119,8 @@ function readMainAuthJsonCredential(): MainAuthJsonCredential | null {
       ?? "";
     return {
       path,
-      rawSha256: sha256(raw),
+      raw,
+      rawSha256: createHash("sha256").update(raw).digest("hex"),
       root,
       tokens,
       ...(accessToken ? { accessToken } : {}),
@@ -131,6 +157,7 @@ function assertMainAuthJsonSnapshotUnchanged(expected: MainAuthJsonCredential):
 }
 
 function persistRefreshedMainAuthJson(
+  context: NativeProfileContext,
   expected: MainAuthJsonCredential,
   refreshed: OAuthCredentials,
 ): { accessToken: string; chatgptAccountId: string } {
@@ -146,20 +173,12 @@ function persistRefreshedMainAuthJson(
     refresh_token: refreshToken,
     account_id: chatgptAccountId,
   };
-  atomicWriteFile(
-    expected.path,
-    JSON.stringify({ ...expected.root, tokens }, null, 2) + "\n",
-    undefined,
-    {
-      beforeRename: () => {
-        assertMainAuthJsonSnapshotUnchanged(expected);
-        const hook = beforeMainAuthJsonRenameForTests;
-        beforeMainAuthJsonRenameForTests = null;
-        hook?.();
-      },
-      validateBeforeRename: () => assertMainAuthJsonSnapshotUnchanged(expected),
-    },
-  );
+  assertMainAuthJsonSnapshotUnchanged(expected);
+  const hook = beforeMainAuthJsonRenameForTests;
+  beforeMainAuthJsonRenameForTests = null;
+  hook?.();
+  assertMainAuthJsonSnapshotUnchanged(expected);
+  publishNativeMainRefresh(context, expected.path, expected.raw, JSON.stringify({ ...expected.root, tokens }, null, 2) + "\n");
   return { accessToken, chatgptAccountId };
 }
 
@@ -170,6 +189,7 @@ export function setMainAuthJsonBeforeRenameHookForTests(hook: (() => void) | nul
 async function resolveMainAccountToken(
   dependencies: NativeMainRefreshDependencies = {},
   rejectedAccessToken?: string,
+  replacementAttempted = false,
 ): Promise<{ accessToken: string; chatgptAccountId: string } | null> {
   const initial = readMainAuthJsonCredential();
   if (!initial) return null;
@@ -185,41 +205,99 @@ async function resolveMainAccountToken(
       : null;
   }
 
-  const signal = dependencies.signal
-    ? AbortSignal.any([dependencies.signal, AbortSignal.timeout(30_000)])
-    : AbortSignal.timeout(30_000);
-  const lockKey = refreshGrantFingerprintForToken(initial.refreshToken);
-  return withCodexRefreshFileLock(lockKey, signal, async () => {
-    const locked = readMainAuthJsonCredential();
-    if (!locked) throw new MainAuthJsonChangedDuringRefreshError();
-    if (!locked.refreshToken
-      || refreshGrantFingerprintForToken(locked.refreshToken) !== lockKey) {
-      if (locked.accessToken !== rejectedAccessToken
-        && mainAccessTokenFresh(locked.accessToken, Date.now(), 0)) {
-        return { accessToken: locked.accessToken!, chatgptAccountId: locked.chatgptAccountId };
-      }
-      throw new MainAuthJsonChangedDuringRefreshError();
-    }
-    if (locked.accessToken !== rejectedAccessToken
-      && mainAccessTokenFresh(locked.accessToken, Date.now(), MAIN_TOKEN_REFRESH_SKEW_MS)) {
-      return { accessToken: locked.accessToken!, chatgptAccountId: locked.chatgptAccountId };
-    }
-    const refresh = dependencies.refreshToken
-      ?? ((refreshToken: string, options: { signal: AbortSignal }) => refreshChatGPTToken(refreshToken, options));
-    let refreshed: OAuthCredentials;
-    try {
-      refreshed = await refresh(locked.refreshToken, { signal });
-    } catch (cause) {
-      const message = cause instanceof Error ? cause.message.toLowerCase() : "";
-      const reason = /invalid_grant|invalidated|revoked|expired/.test(message)
-        ? "reauth" as const
-        : "transient" as const;
-      throw new MainAccountTokenRefreshError(reason, { cause });
-    }
-    const result = persistRefreshedMainAuthJson(locked, refreshed);
-    clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID);
-    return result;
-  });
+  const context = resolveNativeProfileContext();
+  const current = nativeMainRefreshFlights.get(context.homeId);
+  const resolution = { dependencies, rejectedAccessToken, replacementAttempted };
+  if (current) return await resolveNativeMainRefreshFlight(context, current, resolution);
+  if (nativeMainRefreshFlights.size >= MAX_NATIVE_MAIN_REFRESH_FLIGHTS) {
+    throw new MainAccountTokenRefreshError("transient");
+  }
+  const controller = new AbortController();
+  const deadline = setTimeout(() => controller.abort(new Error("Native credential refresh timed out")), NATIVE_MAIN_REFRESH_WAIT_MS);
+  const flight: NativeMainRefreshFlight = {
+    controller,
+    deadline,
+    promise: runNativeMainRefreshFlight(context, dependencies, rejectedAccessToken, controller.signal),
+  };
+  nativeMainRefreshFlights.set(context.homeId, flight);
+  flight.promise.finally(() => {
+    clearTimeout(flight.deadline);
+    if (nativeMainRefreshFlights.get(context.homeId) === flight) nativeMainRefreshFlights.delete(context.homeId);
+  }).catch(() => undefined);
+  return await resolveNativeMainRefreshFlight(context, flight, resolution);
+}
+
+async function resolveNativeMainRefreshFlight(
+  context: NativeProfileContext,
+  flight: NativeMainRefreshFlight,
+  resolution: NativeMainRefreshResolution,
+): Promise<{ accessToken: string; chatgptAccountId: string }> {
+  const result = await waitForNativeMainRefresh(flight, resolution.dependencies.signal);
+  if (resolution.rejectedAccessToken === undefined || result.accessToken !== resolution.rejectedAccessToken) return result;
+  if (resolution.replacementAttempted) throw new MainAccountTokenRefreshError("transient");
+  if (nativeMainRefreshFlights.get(context.homeId) === flight) nativeMainRefreshFlights.delete(context.homeId);
+  const replacement = await resolveMainAccountToken(resolution.dependencies, resolution.rejectedAccessToken, true);
+  if (!replacement) throw new MainAccountTokenRefreshError("transient");
+  return replacement;
+}
+
+function abortError(_signal: AbortSignal): MainAccountRefreshCancelledError {
+  return new MainAccountRefreshCancelledError();
+}
+
+async function waitForNativeMainRefresh(
+  flight: NativeMainRefreshFlight,
+  signal: AbortSignal | undefined,
+): Promise<{ accessToken: string; chatgptAccountId: string }> {
+  if (!signal) return await flight.promise;
+  if (signal.aborted) throw abortError(signal);
+  return await Promise.race([
+    flight.promise,
+    new Promise<never>((_resolve, reject) => signal.addEventListener("abort", () => reject(abortError(signal)), { once: true })),
+  ]);
+}
+
+async function runNativeMainRefreshFlight(
+  context: NativeProfileContext,
+  dependencies: NativeMainRefreshDependencies,
+  rejectedAccessToken: string | undefined,
+  signal: AbortSignal,
+): Promise<{ accessToken: string; chatgptAccountId: string }> {
+  try {
+    return await withNativeMainOwnerOperation(context, async () => await withNativeMainExclusiveClaim(
+      context,
+      async () => {
+        recoverNativeMainRefreshPublication(context);
+        const locked = readMainAuthJsonCredential();
+        if (!locked) throw new MainAuthJsonChangedDuringRefreshError();
+        if (locked.accessToken !== rejectedAccessToken
+          && mainAccessTokenFresh(locked.accessToken, Date.now(), MAIN_TOKEN_REFRESH_SKEW_MS)) {
+          return { accessToken: locked.accessToken!, chatgptAccountId: locked.chatgptAccountId };
+        }
+        if (!locked.refreshToken) throw new MainAuthJsonChangedDuringRefreshError();
+        const refresh = dependencies.refreshToken
+          ?? ((token: string, options: { signal: AbortSignal }) => refreshChatGPTToken(token, options));
+        let refreshed: OAuthCredentials;
+        try {
+          refreshed = await refresh(locked.refreshToken, { signal });
+        } catch (cause) {
+          const terminal = cause instanceof ChatGPTTokenRefreshError
+            && cause.code === "invalid_grant"
+            && (cause.status === 400 || cause.status === 401);
+          throw new MainAccountTokenRefreshError(terminal ? "reauth" : "transient", { cause });
+        }
+        if (signal.aborted) throw new MainAccountTokenRefreshError("transient");
+        const result = persistRefreshedMainAuthJson(context, locked, refreshed);
+        clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID);
+        return result;
+      },
+      { waitMs: NATIVE_MAIN_REFRESH_WAIT_MS },
+    ));
+  } catch (cause) {
+    if (cause instanceof MainAccountTokenRefreshError || cause instanceof MainAuthJsonChangedDuringRefreshError) throw cause;
+    if (cause instanceof NativeMainRefreshPublicationError) throw new MainAccountTokenRefreshError("transient", { cause });
+    throw new MainAccountTokenRefreshError("transient", { cause });
+  }
 }
 
 /** Refresh the CLI-owned native credential before upstream I/O and publish it atomically. */