fix(codex): same-request multi-account failover on quota 429 (#584)

docs-site/src/content/docs/reference/configuration.md

@@ -100,7 +100,9 @@ also force the same native-provider recovery with `ocx recover-history --legacy-
 Use the dashboard's **Codex Auth** page to add pool accounts and refresh quotas. The config stores
 non-secret account metadata only; access and refresh tokens are kept in the hardened Codex account
 credential store. Existing thread ids keep account affinity, while new sessions can auto-route based
-on quota, cooldown, and health.
+on quota, cooldown, and health. A pre-stream upstream **429**/**402** on one pool account is retried
+once on an eligible alternate account in the same request (so Codex CLI does not stall on a depleted
+primary while another account still has quota).
 :::
 
 ### anthropicAccountPool (experimental)

src/codex/routing.ts

@@ -158,7 +158,9 @@ export function classifyCodexUpstreamOutcome(outcome: CodexUpstreamOutcome): Cod
   if (!Number.isFinite(outcome)) return "unknown";
   if (outcome >= 200 && outcome < 300) return "success";
   if (outcome === 401 || outcome === 403) return "credential";
-  if (outcome === 429) return "quota";
+  // 402 Payment Required is treated as quota exhaustion for pool cooldown/failover
+  // (same-request alternate retry records this outcome for the depleted account).
+  if (outcome === 429 || outcome === 402) return "quota";
   if (outcome >= 400 && outcome < 500) return "caller";
   if (outcome >= 500 && outcome < 600) return "transient";
   return "unknown";
@@ -637,21 +639,24 @@ export function resolveCodexAccountForThreadDetailed(
       // it crosses autoSwitchThreshold and a strictly-cooler account exists.
       // Without this the reuse branch returns before applyQuotaAutoSwitch and the
       // thread stays pinned for the full idle TTL (the WSL "never switches" report).
-      if (now - entry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) {
+      // Over-threshold pins re-eval immediately so a depleted primary does not keep
+      // serving for up to 60s after a secondary with quota is available (#584).
+      const threshold = config.autoSwitchThreshold ?? 80;
+      const usage = threshold > 0
+        ? computeCodexUsageScore(
+          getAccountQuota(entry.accountId),
+          getPoolAccountPlan(config, entry.accountId),
+        )
+        : 0;
+      const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold;
+      if (overThreshold || now - entry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) {
         entry.lastReevalAt = now;
-        const threshold = config.autoSwitchThreshold ?? 80;
-        if (threshold > 0) {
-          const usage = computeCodexUsageScore(
-            getAccountQuota(entry.accountId),
-            getPoolAccountPlan(config, entry.accountId),
-          );
-          if (!isUnknownUsage(usage) && usage >= threshold) {
-            const best = pickLowerUsageAccount(config, entry.accountId, usage, now);
-            if (best !== entry.accountId) {
-              setActiveCodexAccount(config, best);
-              bindThreadAffinity(threadId, best, now); // rebinds + resets clocks
-              return { status: "selected", accountId: best };
-            }
+        if (overThreshold) {
+          const best = pickLowerUsageAccount(config, entry.accountId, usage, now);
+          if (best !== entry.accountId) {
+            setActiveCodexAccount(config, best);
+            bindThreadAffinity(threadId, best, now); // rebinds + resets clocks
+            return { status: "selected", accountId: best };
           }
         }
       }

src/server/responses/core.ts

@@ -228,6 +228,140 @@ async function shouldRetryCodexPoolAccountModel400(
   }
 }
 
+/** Pre-stream quota/billing rejections that warrant one alternate-account attempt (#584). */
+function shouldRetryCodexPoolAccountQuota(response: Response): boolean {
+  return response.status === 429 || response.status === 402;
+}
+
+interface CodexPoolAccountRetryArgs {
+  req: Request;
+  config: OcxConfig;
+  route: { providerName: string; modelId: string; provider: OcxProviderConfig };
+  parsed: OcxParsedRequest;
+  logCtx: RequestLogContext;
+  options: {
+    abortSignal?: AbortSignal;
+    onCodexAuthContextResolved?: (ctx: CodexAuthContext) => void;
+  };
+  firstAuthCtx: Extract<CodexAuthContext, { kind: "pool" | "main-pool" }>;
+  firstResponse: Response;
+  outcomeStatus: number;
+  upstream: AbortController;
+  connectMs: number;
+  passthroughEstimate?: number;
+  stream: boolean;
+}
+
+type CodexPoolAccountRetryResult =
+  | {
+    kind: "retried";
+    authCtx: Extract<CodexAuthContext, { kind: "pool" | "main-pool" }>;
+    request: Awaited<ReturnType<ReturnType<typeof resolveAdapter>["buildRequest"]>>;
+    upstreamResponse: Response;
+    selectedForwardHeaders: Headers;
+  }
+  | { kind: "no-alternate" }
+  | {
+    kind: "transport";
+    error: unknown;
+    authCtx: Extract<CodexAuthContext, { kind: "pool" | "main-pool" }>;
+  };
+
+/**
+ * One bounded alternate-account retry for Codex pool auth. Used for allow-listed
+ * model-400 and for pre-stream 429/402 quota failures (#584).
+ */
+async function retryCodexPoolOnAlternateAccount(
+  args: CodexPoolAccountRetryArgs,
+): Promise<CodexPoolAccountRetryResult> {
+  const {
+    req, config, route, parsed, logCtx, options, firstAuthCtx, firstResponse,
+    outcomeStatus, upstream, connectMs, passthroughEstimate, stream,
+  } = args;
+  let retryAuthCtx: CodexAuthContext | undefined;
+  try {
+    retryAuthCtx = await resolveCodexAuthContext(
+      req.headers,
+      config,
+      "pool",
+      { excludeAccountId: firstAuthCtx.accountId },
+    );
+  } catch (error) {
+    if (
+      !(error instanceof CodexPoolAuthenticationError)
+      && !(error instanceof CodexAuthContextError)
+      && !(error instanceof CodexAccountCooldownError)
+    ) throw error;
+  }
+  if (retryAuthCtx?.kind !== "pool" && retryAuthCtx?.kind !== "main-pool") {
+    return { kind: "no-alternate" };
+  }
+
+  const retryAfterRaw = firstResponse.headers.get("retry-after");
+  if (outcomeStatus === 429 || outcomeStatus === 402) {
+    const { applyAccountQuotaFromUpstreamHeaders } = await import("../../codex/auth-api");
+    applyAccountQuotaFromUpstreamHeaders(firstAuthCtx.accountId, firstResponse.headers);
+  }
+  recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, outcomeStatus, {
+    retryAfter: retryAfterRaw,
+    resetAt: [
+      firstResponse.headers.get("x-codex-primary-reset-at"),
+      firstResponse.headers.get("x-codex-secondary-reset-at"),
+      firstResponse.headers.get("x-codex-tertiary-reset-at"),
+    ].filter(Boolean),
+    threadId: req.headers.get("x-codex-parent-thread-id"),
+    probeLeaseId: codexProbeLeaseId(firstAuthCtx),
+  });
+
+  const retryHeaders = headersForCodexAuthContext(req.headers, retryAuthCtx);
+  const retryProvider = applyCodexAuthContextToProvider(
+    stripCodexRuntimeProviderFields(route.provider),
+    retryAuthCtx,
+    "pool",
+  );
+  const retryAdapter = resolveAdapter(
+    resolveWireProtocolOverride(route.providerName, route.modelId, retryProvider),
+    config.cacheRetention,
+  );
+  const request = await retryAdapter.buildRequest(parsed, { headers: retryHeaders });
+  recordAdapterReasoning(logCtx, request);
+
+  await firstResponse.body?.cancel().catch(() => undefined);
+  options.onCodexAuthContextResolved?.(retryAuthCtx);
+  route.provider = retryProvider;
+  logCtx.provider = formatCodexProviderForLog(
+    route.providerName,
+    retryAuthCtx.accountId,
+    config,
+  );
+
+  noteAttemptSend(logCtx.activeAttempt, passthroughEstimate);
+  try {
+    const upstreamResponse = await fetchWithHeaderTimeout(
+      request.url,
+      {
+        method: request.method,
+        headers: request.headers,
+        body: request.body,
+      },
+      upstream.signal,
+      connectMs,
+      stream,
+      providerFetch(route.provider),
+    );
+    return {
+      kind: "retried",
+      authCtx: retryAuthCtx,
+      request,
+      upstreamResponse,
+      selectedForwardHeaders: retryHeaders,
+    };
+  } catch (error) {
+    // Attribute the transport failure to the alternate account (already selected).
+    return { kind: "transport", error, authCtx: retryAuthCtx };
+  }
+}
+
 
 
 export function codexForwardTerminalOutcomeRecorder(
@@ -1273,77 +1407,46 @@ export async function handleResponses(
       return transportFailureResponse(err);
     }
 
-    if (
-      usesCodexForwardPoolAuth(authCtx, route.provider)
-      && await shouldRetryCodexPoolAccountModel400(
+    if (usesCodexForwardPoolAuth(authCtx, route.provider)) {
+      let poolRetryOutcome: number | undefined;
+      if (await shouldRetryCodexPoolAccountModel400(
         upstreamResponse,
         route.modelId,
         options.abortSignal,
-      )
-    ) {
-      const firstAuthCtx = authCtx;
-      let retryAuthCtx: CodexAuthContext | undefined;
-      try {
-        retryAuthCtx = await resolveCodexAuthContext(
-          req.headers,
-          config,
-          "pool",
-          { excludeAccountId: firstAuthCtx.accountId },
-        );
-      } catch (error) {
-        if (
-          !(error instanceof CodexPoolAuthenticationError)
-          && !(error instanceof CodexAuthContextError)
-          && !(error instanceof CodexAccountCooldownError)
-        ) throw error;
+      )) {
+        poolRetryOutcome = 400;
+      } else if (shouldRetryCodexPoolAccountQuota(upstreamResponse)) {
+        // Pre-stream only: once SSE has begun, mid-stream quota stays terminal.
+        poolRetryOutcome = upstreamResponse.status;
       }
 
-      if (retryAuthCtx?.kind === "pool" || retryAuthCtx?.kind === "main-pool") {
-        recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, 400, {
-          threadId: req.headers.get("x-codex-parent-thread-id"),
-          probeLeaseId: codexProbeLeaseId(firstAuthCtx),
-        });
-
-        const retryHeaders = headersForCodexAuthContext(req.headers, retryAuthCtx);
-        const retryProvider = applyCodexAuthContextToProvider(
-          stripCodexRuntimeProviderFields(route.provider),
-          retryAuthCtx,
-          "pool",
-        );
-        const retryAdapter = resolveAdapter(
-          resolveWireProtocolOverride(route.providerName, route.modelId, retryProvider),
-          config.cacheRetention,
-        );
-        request = await retryAdapter.buildRequest(parsed, { headers: retryHeaders });
-        recordAdapterReasoning(logCtx, request);
-
-        await upstreamResponse.body?.cancel().catch(() => undefined);
-        authCtx = retryAuthCtx;
-        options.onCodexAuthContextResolved?.(retryAuthCtx);
-        selectedForwardHeaders = retryHeaders;
-        route.provider = retryProvider;
-        logCtx.provider = formatCodexProviderForLog(
-          route.providerName,
-          retryAuthCtx.accountId,
+      if (poolRetryOutcome !== undefined) {
+        const retry = await retryCodexPoolOnAlternateAccount({
+          req,
           config,
-        );
-
-        noteAttemptSend(logCtx.activeAttempt, passthroughEstimate);
-        try {
-          upstreamResponse = await fetchWithHeaderTimeout(
-            request.url,
-            {
-              method: request.method,
-              headers: request.headers,
-              body: request.body,
-            },
-            upstream.signal,
-            connectMs,
-            parsed.stream,
-            providerFetch(route.provider),
-          );
-        } catch (err) {
-          return transportFailureResponse(err);
+          route,
+          parsed,
+          logCtx,
+          options,
+          firstAuthCtx: authCtx,
+          firstResponse: upstreamResponse,
+          outcomeStatus: poolRetryOutcome,
+          upstream,
+          connectMs,
+          passthroughEstimate,
+          stream: parsed.stream,
+        });
+        if (retry.kind === "transport") {
+          authCtx = retry.authCtx;
+          return transportFailureResponse(retry.error);
+        }
+        if (retry.kind === "retried") {
+          authCtx = retry.authCtx;
+          request = retry.request;
+          upstreamResponse = retry.upstreamResponse;
+          selectedForwardHeaders = retry.selectedForwardHeaders;
+          // Keep subagent quota-failure health keyed to the account that actually served.
+          subagentFallbackAccountId = retry.authCtx.accountId;
         }
       }
     }

tests/codex-routing.test.ts

@@ -197,6 +197,7 @@ describe("codex routing", () => {
     expect(classifyCodexUpstreamOutcome(401)).toBe("credential");
     expect(classifyCodexUpstreamOutcome(403)).toBe("credential");
     expect(classifyCodexUpstreamOutcome(429)).toBe("quota");
+    expect(classifyCodexUpstreamOutcome(402)).toBe("quota");
     expect(classifyCodexUpstreamOutcome(422)).toBe("caller");
     expect(classifyCodexUpstreamOutcome(503)).toBe("transient");
     expect(classifyCodexUpstreamOutcome("connect_error")).toBe("transient");
@@ -1073,6 +1074,19 @@ describe("codex routing", () => {
     expect(config.activeCodexAccountId).toBe("b");
   });
 
+  test("bound thread over threshold switches immediately without waiting for re-eval (#584)", () => {
+    const config = makeConfig();
+    const now = 1_800_000_000_000;
+    updateAccountQuota("a", 10);
+    updateAccountQuota("b", 10);
+    expect(resolveCodexAccountForThread("t1", config, now)).toBe("a");
+    updateAccountQuota("a", 95);
+    updateAccountQuota("b", 5);
+    // Depleted primary must not stay pinned for up to 60s while a cooler account exists.
+    expect(resolveCodexAccountForThread("t1", config, now + 1_000)).toBe("b");
+    expect(config.activeCodexAccountId).toBe("b");
+  });
+
   test("bound thread under threshold stays even if a lower account exists", () => {
     const config = makeConfig();
     const now = 1_800_000_000_000;
@@ -1087,21 +1101,31 @@ describe("codex routing", () => {
     expect(config.activeCodexAccountId).toBe("a");
   });
 
-  test("bound thread does not flap within the re-eval interval, then switches once", () => {
+  test("bound thread under threshold does not flap within the re-eval interval", () => {
     const config = makeConfig();
     const now = 1_800_000_000_000;
     updateAccountQuota("a", 10);
     updateAccountQuota("b", 10);
     expect(resolveCodexAccountForThread("t1", config, now)).toBe("a");
-    updateAccountQuota("a", 95);
+    // Still under threshold — must not rebind on every reuse.
+    updateAccountQuota("a", 50);
     updateAccountQuota("b", 5);
-    // Within the interval: no rebind yet.
     expect(resolveCodexAccountForThread("t1", config, now + 1_000)).toBe("a");
-    // After the interval: switches once.
     const later = now + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + 1;
-    expect(resolveCodexAccountForThread("t1", config, later)).toBe("b");
+    expect(resolveCodexAccountForThread("t1", config, later)).toBe("a");
+  });
+
+  test("bound thread over threshold switches once and does not ping-pong", () => {
+    const config = makeConfig();
+    const now = 1_800_000_000_000;
+    updateAccountQuota("a", 10);
+    updateAccountQuota("b", 10);
+    expect(resolveCodexAccountForThread("t1", config, now)).toBe("a");
+    updateAccountQuota("a", 95);
+    updateAccountQuota("b", 5);
+    expect(resolveCodexAccountForThread("t1", config, now + 1_000)).toBe("b");
     // A subsequent interval does not ping-pong back: b is now the lowest.
-    const later2 = later + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + 1;
+    const later2 = now + 1_000 + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + 1;
     expect(resolveCodexAccountForThread("t1", config, later2)).toBe("b");
   });