From e875aa85e084f450c8be979fc7a8fcc7aaaa93f4 Mon Sep 17 00:00:00 2001
From: luvs01 <luvs01@hanmail.net>
Date: Mon, 31 Aug 2026 09:59:29 +0900
Subject: [PATCH] fix(codex): hand off affinity in refresh flight

---
 src/codex/account-store.ts        | 18 ++++++++++++++++++
 src/codex/routing.ts              | 10 +++++++++-
 tests/codex-account-store.test.ts | 13 ++++++++++++-
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/codex/account-store.ts b/src/codex/account-store.ts
index 011ac0692e..bccd1309bb 100644
--- a/src/codex/account-store.ts
+++ b/src/codex/account-store.ts
@@ -377,6 +377,15 @@ function withCredentialMutationLockSync<T>(fn: () => T): T {
 }
 
 type CodexTokenResult = { accessToken: string; chatgptAccountId: string; generation: number };
+type CodexRefreshGenerationHandoff = (accountId: string, fromGeneration: number, toGeneration: number) => void;
+const refreshGenerationHandoffs = new Set<CodexRefreshGenerationHandoff>();
+
+/** Register process-local state that must follow a credential refresh generation. */
+export function registerCodexRefreshGenerationHandoff(handoff: CodexRefreshGenerationHandoff): () => void {
+  refreshGenerationHandoffs.add(handoff);
+  return () => refreshGenerationHandoffs.delete(handoff);
+}
+
 type CodexRefreshResult = CodexTokenResult & {
   credential?: CodexAccountCredentials;
   /**
@@ -891,6 +900,15 @@ async function resolveCodexToken(
    * committed result, for every waiter, including none.
    */
   const refreshPromise = fetchPromise.then(async (result): Promise<CodexRefreshResult> => {
+    // Generation-dependent completion belongs to the flight, not to its initiating
+    // request. The owner may stop waiting after a disconnect while this detached work
+    // still commits G+1; advance process-local affinities before any waiter observes
+    // the result (and even when there are no surviving waiters).
+    if (result.selfRefreshed) {
+      for (const handoff of refreshGenerationHandoffs) {
+        handoff(id, result.generation - 1, result.generation);
+      }
+    }
     await notePlanFromRefreshedAccessToken(id, result.accessToken, result.generation);
     // One settlement path for the whole flight: the refreshing account, then any dormant alias that
     // adopted the same rotated JWT. An alias holds the identical access token, so a changed
diff --git a/src/codex/routing.ts b/src/codex/routing.ts
index 250aac9636..ed00d7061d 100644
--- a/src/codex/routing.ts
+++ b/src/codex/routing.ts
@@ -1,6 +1,10 @@
 import { randomUUID } from "node:crypto";
 import { saveConfigPreservingClaudeCode } from "../config";
-import { isCodexAccountGenerationLive, readCodexAccountRecord } from "./account-store";
+import {
+  isCodexAccountGenerationLive,
+  readCodexAccountRecord,
+  registerCodexRefreshGenerationHandoff,
+} from "./account-store";
 import { codexAccountLogLabel } from "./account-label";
 import { isCodexAccountPaused } from "./account-pause";
 import { clearCodexAccountPin, codexAccountPriorityLookup, pinnedCodexAccountId } from "./account-priority";
@@ -1013,6 +1017,10 @@ export function handOffThreadAffinityGeneration(
   return handedOff;
 }
 
+// A shared refresh can outlive the request that opened it. Register the affinity
+// handoff with the flight so a detached G -> G+1 commit cannot strand bindings at G.
+registerCodexRefreshGenerationHandoff(handOffThreadAffinityGeneration);
+
 function pruneExpiredThreadAffinities(now: number): void {
   for (const [threadId, affinities] of threadAccountMap) {
     for (const [scope, entry] of affinities) {
diff --git a/tests/codex-account-store.test.ts b/tests/codex-account-store.test.ts
index 3e2a9f3f77..f73ded646f 100644
--- a/tests/codex-account-store.test.ts
+++ b/tests/codex-account-store.test.ts
@@ -709,7 +709,12 @@ describe("codex-account-store CRUD", () => {
    * healthy account marked for reauthentication on behalf of a live request.
    */
   test("cancelling the caller that opened a refresh flight does not cancel a live joiner (#2892)", async () => {
-    const { forceRefreshCodexPoolToken, readCodexAccountRecord, saveCodexAccountCredential } =
+    const {
+      forceRefreshCodexPoolToken,
+      readCodexAccountRecord,
+      registerCodexRefreshGenerationHandoff,
+      saveCodexAccountCredential,
+    } =
       await import("../src/codex/account-store");
     saveCodexAccountCredential("cancel-owner", {
       accessToken: "rejected",
@@ -720,6 +725,8 @@ describe("codex-account-store CRUD", () => {
     const generation = readCodexAccountRecord("cancel-owner")!.generation;
 
     const originalFetch = globalThis.fetch;
+    const handoffs: Array<[string, number, number]> = [];
+    const unregisterHandoff = registerCodexRefreshGenerationHandoff((...handoff) => handoffs.push(handoff));
     let sawAbort = false;
     let calls = 0;
     let releaseFetch: (() => void) | undefined;
@@ -759,7 +766,11 @@ describe("codex-account-store CRUD", () => {
       expect(joined.accessToken).toBe("rotated");
       expect(calls).toBe(1);
       expect(readCodexAccountRecord("cancel-owner")!.credential!.accessToken).toBe("rotated");
+      // Completion is attached to the detached flight, not to either request's wait.
+      // The cancelled owner therefore cannot strand process-local affinities at G.
+      expect(handoffs).toEqual([["cancel-owner", generation, generation + 1]]);
     } finally {
+      unregisterHandoff();
       globalThis.fetch = originalFetch;
     }
   });