Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions src/codex/account-store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
/**
Expand Down Expand Up @@ -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
Expand Down
10 changes: 9 additions & 1 deletion src/codex/routing.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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) {
Expand Down
13 changes: 12 additions & 1 deletion tests/codex-account-store.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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;
Expand Down Expand Up @@ -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;
}
});
Expand Down
Loading