fix(codex): refresh and retry once on a 401 (the biggest reliability gap) (#108)

The loader's proactive refresh only catches token EXPIRY. A Codex access
token can be revoked/invalidated server-side BEFORE its local `expires`
(password change, admin revoke, clock skew, or a lost cross-process
refresh race). That surfaced as a 401 returned verbatim — the request
just failed, and the only recovery was a manual `keys signin`.

Now a 401 on a Codex route forces one refresh and retries once. The
refresh re-reads the latest persisted auth first (a sibling process may
have already rotated the pair) and adopts a newer token if present,
otherwise spends the refresh token and persists the rotated pair. A dead
refresh token surfaces the reconnect message; a transient failure keeps
the original 401 rather than making things worse. Retries at most once,
so it can never loop.

Extracted sendWithCodex401Retry() as a pure, dependency-injected helper
so the send/refresh/retry control flow is unit-tested.

Author: aayambansal

backend/cli/src/plugin/codex.ts

@@ -71,6 +71,28 @@ export function classifyDevicePollStatus(status: number): "pending" | "transient
   return "fail"
 }
 
+/**
+ * Send a Codex request and, on a 401, refresh the access token once and retry
+ * once. The loader's proactive refresh only catches token EXPIRY; a token
+ * revoked/invalidated server-side before its local `expires` (password change,
+ * admin revoke, clock skew) surfaces as a 401 that was otherwise returned
+ * verbatim — the request just failed. `refresh` returns the new access token,
+ * `undefined` to give up (the original 401 is returned unchanged), or throws to
+ * surface a fatal error (e.g. the refresh token itself is dead). Retries at most
+ * once, so it can never loop.
+ */
+export async function sendWithCodex401Retry(
+  send: (accessToken: string) => Promise<Response>,
+  accessToken: string,
+  refresh: () => Promise<string | undefined>,
+): Promise<Response> {
+  const response = await send(accessToken)
+  if (response.status !== 401) return response
+  const refreshed = await refresh()
+  if (!refreshed) return response
+  return send(refreshed)
+}
+
 interface PkceCodes {
   verifier: string
   challenge: string
@@ -541,10 +563,9 @@ export async function CodexAuthPlugin(input: PluginInput): Promise<Hooks> {
               }
             }
 
-            // Set authorization header with access token
-            headers.set("authorization", `Bearer ${currentAuth.access}`)
-
-            // Set ChatGPT-Account-Id header for organization subscriptions
+            // Set ChatGPT-Account-Id header for organization subscriptions.
+            // (The authorization header is set per-attempt inside `send` below,
+            // so the 401-retry path can swap in a freshly-refreshed token.)
             if (authWithAccount.accountId) {
               headers.set("ChatGPT-Account-Id", authWithAccount.accountId)
             }
@@ -588,11 +609,55 @@ export async function CodexAuthPlugin(input: PluginInput): Promise<Hooks> {
               }
             }
 
-            const response = await fetch(url, {
-              ...init,
-              headers,
-              body: bodyForRequest,
-            })
+            // Send with the current bearer; re-called by the 401-retry path with
+            // a refreshed token. Each attempt re-sets the authorization header.
+            const send = (accessToken: string) => {
+              headers.set("authorization", `Bearer ${accessToken}`)
+              return fetch(url, { ...init, headers, body: bodyForRequest })
+            }
+
+            // Reactive (on-401) refresh: the token was rejected even though it
+            // isn't locally expired. Re-read the latest persisted auth first (a
+            // sibling process may have already rotated the pair) and adopt a
+            // newer token if one exists; otherwise spend the refresh token and
+            // persist the rotated pair. Returns the fresh access token, undefined
+            // to keep the original 401 (transient failure), or throws when the
+            // refresh token itself is dead.
+            const forceRefreshOn401 = async (): Promise<string | undefined> => {
+              try {
+                const latest = (await getAuth()) as typeof currentAuth & { accountId?: string }
+                if (latest.type !== "oauth") return undefined
+                if (latest.access && latest.access !== currentAuth.access && latest.expires > Date.now()) {
+                  currentAuth.access = latest.access
+                  authWithAccount.accountId = latest.accountId ?? authWithAccount.accountId
+                  return latest.access
+                }
+                const tokens = await refreshAccessTokenSingleFlight(latest.refresh)
+                const newAccountId = extractAccountId(tokens) || authWithAccount.accountId
+                await input.client.auth.set({
+                  path: { id: "openai-codex" },
+                  body: {
+                    type: "oauth",
+                    refresh: tokens.refresh_token,
+                    access: tokens.access_token,
+                    expires: Date.now() + (tokens.expires_in ?? 3600) * 1000,
+                    ...(newAccountId && { accountId: newAccountId }),
+                  },
+                })
+                currentAuth.access = tokens.access_token
+                authWithAccount.accountId = newAccountId
+                return tokens.access_token
+              } catch (e) {
+                if (e instanceof CodexRefreshInvalidError)
+                  throw new Error("Codex sign-in expired. Reconnect it with `openscience keys signin`.")
+                log.warn("codex 401-triggered refresh failed", { error: String(e) })
+                return undefined
+              }
+            }
+
+            const response = isCodexRoute
+              ? await sendWithCodex401Retry(send, currentAuth.access, forceRefreshOn401)
+              : await send(currentAuth.access)
 
             // Fallback to shared key if OAuth quota exceeded
             if (response.status === 429 || response.status === 403) {

backend/cli/test/plugin/codex-401-retry.test.ts

@@ -0,0 +1,84 @@
+import { test, expect } from "bun:test"
+import { sendWithCodex401Retry } from "../../src/plugin/codex"
+
+// A Codex token can be revoked/invalidated server-side before its local
+// `expires`, surfacing as a 401. That must trigger exactly one refresh + one
+// retry — never a verbatim failure, and never an infinite loop.
+
+test("a 2xx response returns immediately without refreshing", async () => {
+  let sends = 0
+  let refreshes = 0
+  const res = await sendWithCodex401Retry(
+    async () => {
+      sends++
+      return new Response("ok", { status: 200 })
+    },
+    "access-1",
+    async () => {
+      refreshes++
+      return "access-2"
+    },
+  )
+  expect(res.status).toBe(200)
+  expect(sends).toBe(1)
+  expect(refreshes).toBe(0)
+})
+
+test("a 401 refreshes once and retries once with the new token", async () => {
+  const tokens: string[] = []
+  let refreshes = 0
+  const res = await sendWithCodex401Retry(
+    async (token) => {
+      tokens.push(token)
+      return new Response(null, { status: tokens.length === 1 ? 401 : 200 })
+    },
+    "stale",
+    async () => {
+      refreshes++
+      return "fresh"
+    },
+  )
+  expect(res.status).toBe(200)
+  expect(refreshes).toBe(1)
+  expect(tokens).toEqual(["stale", "fresh"])
+})
+
+test("when refresh gives up (undefined), the original 401 is returned unchanged", async () => {
+  let sends = 0
+  const res = await sendWithCodex401Retry(
+    async () => {
+      sends++
+      return new Response("unauthorized", { status: 401 })
+    },
+    "stale",
+    async () => undefined,
+  )
+  expect(res.status).toBe(401)
+  expect(sends).toBe(1)
+})
+
+test("retries at most once — a persistent 401 does not loop", async () => {
+  let sends = 0
+  const res = await sendWithCodex401Retry(
+    async () => {
+      sends++
+      return new Response(null, { status: 401 })
+    },
+    "stale",
+    async () => "fresh",
+  )
+  expect(res.status).toBe(401)
+  expect(sends).toBe(2)
+})
+
+test("a fatal refresh error propagates to the caller", async () => {
+  await expect(
+    sendWithCodex401Retry(
+      async () => new Response(null, { status: 401 }),
+      "stale",
+      async () => {
+        throw new Error("Codex sign-in expired. Reconnect it with `openscience keys signin`.")
+      },
+    ),
+  ).rejects.toThrow("Reconnect it with")
+})