Skip to content

Commit 903b62c

Browse files
authored
fix(codex): same-request multi-account failover on quota 429 (lidge-jun#584) (lidge-jun#585)
* fix(codex): same-request pool failover on pre-stream 429/402 Retry exhausted primary accounts once on an eligible alternate within the same turn, and rebind over-threshold sticky threads immediately so Codex CLI does not stall when a secondary still has quota (lidge-jun#584). * fix(codex): attribute pool-retry transport failures to alternate account Also classify HTTP 402 as quota so same-request failover cools the depleted account consistently with 429. * fix(codex): keep subagent quota health scoped to the retried account After a successful same-request pool failover, update subagentFallbackAccountId so streamed terminal 429s attribute health to the account that actually served.
1 parent 94a83bc commit 903b62c

5 files changed

Lines changed: 281 additions & 87 deletions

File tree

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

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

106108
### anthropicAccountPool (experimental)

src/codex/routing.ts

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,9 @@ export function classifyCodexUpstreamOutcome(outcome: CodexUpstreamOutcome): Cod
158158
if (!Number.isFinite(outcome)) return "unknown";
159159
if (outcome >= 200 && outcome < 300) return "success";
160160
if (outcome === 401 || outcome === 403) return "credential";
161-
if (outcome === 429) return "quota";
161+
// 402 Payment Required is treated as quota exhaustion for pool cooldown/failover
162+
// (same-request alternate retry records this outcome for the depleted account).
163+
if (outcome === 429 || outcome === 402) return "quota";
162164
if (outcome >= 400 && outcome < 500) return "caller";
163165
if (outcome >= 500 && outcome < 600) return "transient";
164166
return "unknown";
@@ -637,21 +639,24 @@ export function resolveCodexAccountForThreadDetailed(
637639
// it crosses autoSwitchThreshold and a strictly-cooler account exists.
638640
// Without this the reuse branch returns before applyQuotaAutoSwitch and the
639641
// thread stays pinned for the full idle TTL (the WSL "never switches" report).
640-
if (now - entry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) {
642+
// Over-threshold pins re-eval immediately so a depleted primary does not keep
643+
// serving for up to 60s after a secondary with quota is available (#584).
644+
const threshold = config.autoSwitchThreshold ?? 80;
645+
const usage = threshold > 0
646+
? computeCodexUsageScore(
647+
getAccountQuota(entry.accountId),
648+
getPoolAccountPlan(config, entry.accountId),
649+
)
650+
: 0;
651+
const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold;
652+
if (overThreshold || now - entry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) {
641653
entry.lastReevalAt = now;
642-
const threshold = config.autoSwitchThreshold ?? 80;
643-
if (threshold > 0) {
644-
const usage = computeCodexUsageScore(
645-
getAccountQuota(entry.accountId),
646-
getPoolAccountPlan(config, entry.accountId),
647-
);
648-
if (!isUnknownUsage(usage) && usage >= threshold) {
649-
const best = pickLowerUsageAccount(config, entry.accountId, usage, now);
650-
if (best !== entry.accountId) {
651-
setActiveCodexAccount(config, best);
652-
bindThreadAffinity(threadId, best, now); // rebinds + resets clocks
653-
return { status: "selected", accountId: best };
654-
}
654+
if (overThreshold) {
655+
const best = pickLowerUsageAccount(config, entry.accountId, usage, now);
656+
if (best !== entry.accountId) {
657+
setActiveCodexAccount(config, best);
658+
bindThreadAffinity(threadId, best, now); // rebinds + resets clocks
659+
return { status: "selected", accountId: best };
655660
}
656661
}
657662
}

src/server/responses/core.ts

Lines changed: 168 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,140 @@ async function shouldRetryCodexPoolAccountModel400(
228228
}
229229
}
230230

231+
/** Pre-stream quota/billing rejections that warrant one alternate-account attempt (#584). */
232+
function shouldRetryCodexPoolAccountQuota(response: Response): boolean {
233+
return response.status === 429 || response.status === 402;
234+
}
235+
236+
interface CodexPoolAccountRetryArgs {
237+
req: Request;
238+
config: OcxConfig;
239+
route: { providerName: string; modelId: string; provider: OcxProviderConfig };
240+
parsed: OcxParsedRequest;
241+
logCtx: RequestLogContext;
242+
options: {
243+
abortSignal?: AbortSignal;
244+
onCodexAuthContextResolved?: (ctx: CodexAuthContext) => void;
245+
};
246+
firstAuthCtx: Extract<CodexAuthContext, { kind: "pool" | "main-pool" }>;
247+
firstResponse: Response;
248+
outcomeStatus: number;
249+
upstream: AbortController;
250+
connectMs: number;
251+
passthroughEstimate?: number;
252+
stream: boolean;
253+
}
254+
255+
type CodexPoolAccountRetryResult =
256+
| {
257+
kind: "retried";
258+
authCtx: Extract<CodexAuthContext, { kind: "pool" | "main-pool" }>;
259+
request: Awaited<ReturnType<ReturnType<typeof resolveAdapter>["buildRequest"]>>;
260+
upstreamResponse: Response;
261+
selectedForwardHeaders: Headers;
262+
}
263+
| { kind: "no-alternate" }
264+
| {
265+
kind: "transport";
266+
error: unknown;
267+
authCtx: Extract<CodexAuthContext, { kind: "pool" | "main-pool" }>;
268+
};
269+
270+
/**
271+
* One bounded alternate-account retry for Codex pool auth. Used for allow-listed
272+
* model-400 and for pre-stream 429/402 quota failures (#584).
273+
*/
274+
async function retryCodexPoolOnAlternateAccount(
275+
args: CodexPoolAccountRetryArgs,
276+
): Promise<CodexPoolAccountRetryResult> {
277+
const {
278+
req, config, route, parsed, logCtx, options, firstAuthCtx, firstResponse,
279+
outcomeStatus, upstream, connectMs, passthroughEstimate, stream,
280+
} = args;
281+
let retryAuthCtx: CodexAuthContext | undefined;
282+
try {
283+
retryAuthCtx = await resolveCodexAuthContext(
284+
req.headers,
285+
config,
286+
"pool",
287+
{ excludeAccountId: firstAuthCtx.accountId },
288+
);
289+
} catch (error) {
290+
if (
291+
!(error instanceof CodexPoolAuthenticationError)
292+
&& !(error instanceof CodexAuthContextError)
293+
&& !(error instanceof CodexAccountCooldownError)
294+
) throw error;
295+
}
296+
if (retryAuthCtx?.kind !== "pool" && retryAuthCtx?.kind !== "main-pool") {
297+
return { kind: "no-alternate" };
298+
}
299+
300+
const retryAfterRaw = firstResponse.headers.get("retry-after");
301+
if (outcomeStatus === 429 || outcomeStatus === 402) {
302+
const { applyAccountQuotaFromUpstreamHeaders } = await import("../../codex/auth-api");
303+
applyAccountQuotaFromUpstreamHeaders(firstAuthCtx.accountId, firstResponse.headers);
304+
}
305+
recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, outcomeStatus, {
306+
retryAfter: retryAfterRaw,
307+
resetAt: [
308+
firstResponse.headers.get("x-codex-primary-reset-at"),
309+
firstResponse.headers.get("x-codex-secondary-reset-at"),
310+
firstResponse.headers.get("x-codex-tertiary-reset-at"),
311+
].filter(Boolean),
312+
threadId: req.headers.get("x-codex-parent-thread-id"),
313+
probeLeaseId: codexProbeLeaseId(firstAuthCtx),
314+
});
315+
316+
const retryHeaders = headersForCodexAuthContext(req.headers, retryAuthCtx);
317+
const retryProvider = applyCodexAuthContextToProvider(
318+
stripCodexRuntimeProviderFields(route.provider),
319+
retryAuthCtx,
320+
"pool",
321+
);
322+
const retryAdapter = resolveAdapter(
323+
resolveWireProtocolOverride(route.providerName, route.modelId, retryProvider),
324+
config.cacheRetention,
325+
);
326+
const request = await retryAdapter.buildRequest(parsed, { headers: retryHeaders });
327+
recordAdapterReasoning(logCtx, request);
328+
329+
await firstResponse.body?.cancel().catch(() => undefined);
330+
options.onCodexAuthContextResolved?.(retryAuthCtx);
331+
route.provider = retryProvider;
332+
logCtx.provider = formatCodexProviderForLog(
333+
route.providerName,
334+
retryAuthCtx.accountId,
335+
config,
336+
);
337+
338+
noteAttemptSend(logCtx.activeAttempt, passthroughEstimate);
339+
try {
340+
const upstreamResponse = await fetchWithHeaderTimeout(
341+
request.url,
342+
{
343+
method: request.method,
344+
headers: request.headers,
345+
body: request.body,
346+
},
347+
upstream.signal,
348+
connectMs,
349+
stream,
350+
providerFetch(route.provider),
351+
);
352+
return {
353+
kind: "retried",
354+
authCtx: retryAuthCtx,
355+
request,
356+
upstreamResponse,
357+
selectedForwardHeaders: retryHeaders,
358+
};
359+
} catch (error) {
360+
// Attribute the transport failure to the alternate account (already selected).
361+
return { kind: "transport", error, authCtx: retryAuthCtx };
362+
}
363+
}
364+
231365

232366

233367
export function codexForwardTerminalOutcomeRecorder(
@@ -1273,77 +1407,46 @@ export async function handleResponses(
12731407
return transportFailureResponse(err);
12741408
}
12751409

1276-
if (
1277-
usesCodexForwardPoolAuth(authCtx, route.provider)
1278-
&& await shouldRetryCodexPoolAccountModel400(
1410+
if (usesCodexForwardPoolAuth(authCtx, route.provider)) {
1411+
let poolRetryOutcome: number | undefined;
1412+
if (await shouldRetryCodexPoolAccountModel400(
12791413
upstreamResponse,
12801414
route.modelId,
12811415
options.abortSignal,
1282-
)
1283-
) {
1284-
const firstAuthCtx = authCtx;
1285-
let retryAuthCtx: CodexAuthContext | undefined;
1286-
try {
1287-
retryAuthCtx = await resolveCodexAuthContext(
1288-
req.headers,
1289-
config,
1290-
"pool",
1291-
{ excludeAccountId: firstAuthCtx.accountId },
1292-
);
1293-
} catch (error) {
1294-
if (
1295-
!(error instanceof CodexPoolAuthenticationError)
1296-
&& !(error instanceof CodexAuthContextError)
1297-
&& !(error instanceof CodexAccountCooldownError)
1298-
) throw error;
1416+
)) {
1417+
poolRetryOutcome = 400;
1418+
} else if (shouldRetryCodexPoolAccountQuota(upstreamResponse)) {
1419+
// Pre-stream only: once SSE has begun, mid-stream quota stays terminal.
1420+
poolRetryOutcome = upstreamResponse.status;
12991421
}
13001422

1301-
if (retryAuthCtx?.kind === "pool" || retryAuthCtx?.kind === "main-pool") {
1302-
recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, 400, {
1303-
threadId: req.headers.get("x-codex-parent-thread-id"),
1304-
probeLeaseId: codexProbeLeaseId(firstAuthCtx),
1305-
});
1306-
1307-
const retryHeaders = headersForCodexAuthContext(req.headers, retryAuthCtx);
1308-
const retryProvider = applyCodexAuthContextToProvider(
1309-
stripCodexRuntimeProviderFields(route.provider),
1310-
retryAuthCtx,
1311-
"pool",
1312-
);
1313-
const retryAdapter = resolveAdapter(
1314-
resolveWireProtocolOverride(route.providerName, route.modelId, retryProvider),
1315-
config.cacheRetention,
1316-
);
1317-
request = await retryAdapter.buildRequest(parsed, { headers: retryHeaders });
1318-
recordAdapterReasoning(logCtx, request);
1319-
1320-
await upstreamResponse.body?.cancel().catch(() => undefined);
1321-
authCtx = retryAuthCtx;
1322-
options.onCodexAuthContextResolved?.(retryAuthCtx);
1323-
selectedForwardHeaders = retryHeaders;
1324-
route.provider = retryProvider;
1325-
logCtx.provider = formatCodexProviderForLog(
1326-
route.providerName,
1327-
retryAuthCtx.accountId,
1423+
if (poolRetryOutcome !== undefined) {
1424+
const retry = await retryCodexPoolOnAlternateAccount({
1425+
req,
13281426
config,
1329-
);
1330-
1331-
noteAttemptSend(logCtx.activeAttempt, passthroughEstimate);
1332-
try {
1333-
upstreamResponse = await fetchWithHeaderTimeout(
1334-
request.url,
1335-
{
1336-
method: request.method,
1337-
headers: request.headers,
1338-
body: request.body,
1339-
},
1340-
upstream.signal,
1341-
connectMs,
1342-
parsed.stream,
1343-
providerFetch(route.provider),
1344-
);
1345-
} catch (err) {
1346-
return transportFailureResponse(err);
1427+
route,
1428+
parsed,
1429+
logCtx,
1430+
options,
1431+
firstAuthCtx: authCtx,
1432+
firstResponse: upstreamResponse,
1433+
outcomeStatus: poolRetryOutcome,
1434+
upstream,
1435+
connectMs,
1436+
passthroughEstimate,
1437+
stream: parsed.stream,
1438+
});
1439+
if (retry.kind === "transport") {
1440+
authCtx = retry.authCtx;
1441+
return transportFailureResponse(retry.error);
1442+
}
1443+
if (retry.kind === "retried") {
1444+
authCtx = retry.authCtx;
1445+
request = retry.request;
1446+
upstreamResponse = retry.upstreamResponse;
1447+
selectedForwardHeaders = retry.selectedForwardHeaders;
1448+
// Keep subagent quota-failure health keyed to the account that actually served.
1449+
subagentFallbackAccountId = retry.authCtx.accountId;
13471450
}
13481451
}
13491452
}

tests/codex-routing.test.ts

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ describe("codex routing", () => {
197197
expect(classifyCodexUpstreamOutcome(401)).toBe("credential");
198198
expect(classifyCodexUpstreamOutcome(403)).toBe("credential");
199199
expect(classifyCodexUpstreamOutcome(429)).toBe("quota");
200+
expect(classifyCodexUpstreamOutcome(402)).toBe("quota");
200201
expect(classifyCodexUpstreamOutcome(422)).toBe("caller");
201202
expect(classifyCodexUpstreamOutcome(503)).toBe("transient");
202203
expect(classifyCodexUpstreamOutcome("connect_error")).toBe("transient");
@@ -1073,6 +1074,19 @@ describe("codex routing", () => {
10731074
expect(config.activeCodexAccountId).toBe("b");
10741075
});
10751076

1077+
test("bound thread over threshold switches immediately without waiting for re-eval (#584)", () => {
1078+
const config = makeConfig();
1079+
const now = 1_800_000_000_000;
1080+
updateAccountQuota("a", 10);
1081+
updateAccountQuota("b", 10);
1082+
expect(resolveCodexAccountForThread("t1", config, now)).toBe("a");
1083+
updateAccountQuota("a", 95);
1084+
updateAccountQuota("b", 5);
1085+
// Depleted primary must not stay pinned for up to 60s while a cooler account exists.
1086+
expect(resolveCodexAccountForThread("t1", config, now + 1_000)).toBe("b");
1087+
expect(config.activeCodexAccountId).toBe("b");
1088+
});
1089+
10761090
test("bound thread under threshold stays even if a lower account exists", () => {
10771091
const config = makeConfig();
10781092
const now = 1_800_000_000_000;
@@ -1087,21 +1101,31 @@ describe("codex routing", () => {
10871101
expect(config.activeCodexAccountId).toBe("a");
10881102
});
10891103

1090-
test("bound thread does not flap within the re-eval interval, then switches once", () => {
1104+
test("bound thread under threshold does not flap within the re-eval interval", () => {
10911105
const config = makeConfig();
10921106
const now = 1_800_000_000_000;
10931107
updateAccountQuota("a", 10);
10941108
updateAccountQuota("b", 10);
10951109
expect(resolveCodexAccountForThread("t1", config, now)).toBe("a");
1096-
updateAccountQuota("a", 95);
1110+
// Still under threshold — must not rebind on every reuse.
1111+
updateAccountQuota("a", 50);
10971112
updateAccountQuota("b", 5);
1098-
// Within the interval: no rebind yet.
10991113
expect(resolveCodexAccountForThread("t1", config, now + 1_000)).toBe("a");
1100-
// After the interval: switches once.
11011114
const later = now + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + 1;
1102-
expect(resolveCodexAccountForThread("t1", config, later)).toBe("b");
1115+
expect(resolveCodexAccountForThread("t1", config, later)).toBe("a");
1116+
});
1117+
1118+
test("bound thread over threshold switches once and does not ping-pong", () => {
1119+
const config = makeConfig();
1120+
const now = 1_800_000_000_000;
1121+
updateAccountQuota("a", 10);
1122+
updateAccountQuota("b", 10);
1123+
expect(resolveCodexAccountForThread("t1", config, now)).toBe("a");
1124+
updateAccountQuota("a", 95);
1125+
updateAccountQuota("b", 5);
1126+
expect(resolveCodexAccountForThread("t1", config, now + 1_000)).toBe("b");
11031127
// A subsequent interval does not ping-pong back: b is now the lowest.
1104-
const later2 = later + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + 1;
1128+
const later2 = now + 1_000 + CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + 1;
11051129
expect(resolveCodexAccountForThread("t1", config, later2)).toBe("b");
11061130
});
11071131

0 commit comments

Comments
 (0)