Skip to content

Commit b208746

Browse files
ndycodendycode
committed
fix(routing): preserve failover identity bookkeeping
1 parent 06340e6 commit b208746

3 files changed

Lines changed: 155 additions & 11 deletions

File tree

‎index.ts‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1541,7 +1541,9 @@ while (attempted.size < Math.max(1, accountCount)) {
15411541
if (!hadAccountId && tokenAccountId && accountId === tokenAccountId) {
15421542
account.accountIdSource = storedAccountIdSource ?? "token";
15431543
}
1544-
account.email = resolvedEmail;
1544+
if (resolvedEmail) {
1545+
account.email = resolvedEmail;
1546+
}
15451547

15461548
if (
15471549
accountCount > 1 &&
@@ -1600,6 +1602,7 @@ while (attempted.size < Math.max(1, accountCount)) {
16001602

16011603
let sameAccountRetryCount = 0;
16021604
let successAccountForResponse = account;
1605+
let successEntitlementAccountKey = entitlementAccountKey;
16031606
while (true) {
16041607
let response: Response;
16051608
const fetchStart = performance.now();
@@ -2232,6 +2235,7 @@ while (attempted.size < Math.max(1, accountCount)) {
22322235
}
22332236

22342237
successAccountForResponse = fallbackAccount;
2238+
successEntitlementAccountKey = fallbackEntitlementAccountKey;
22352239
runtimeMetrics.streamFailoverRecoveries += 1;
22362240
if (fallbackAccount.index !== account.index) {
22372241
runtimeMetrics.streamFailoverCrossAccountRecoveries += 1;
@@ -2341,10 +2345,7 @@ while (attempted.size < Math.max(1, accountCount)) {
23412345
if (successAccountForResponse.index !== account.index) {
23422346
accountManager.markSwitched(successAccountForResponse, "rotation", modelFamily);
23432347
}
2344-
const successAccountKey =
2345-
successAccountForResponse.index === account.index
2346-
? entitlementAccountKey
2347-
: resolveEntitlementAccountKey(successAccountForResponse);
2348+
const successAccountKey = successEntitlementAccountKey;
23482349
accountManager.recordSuccess(successAccountForResponse, modelFamily, model);
23492350
capabilityPolicyStore.recordSuccess(
23502351
successAccountKey,

‎test/index-retry.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ vi.mock("../lib/accounts.js", async () => {
140140
source as never,
141141
tokenAccountId,
142142
),
143-
email: tokenEmail ?? sanitizedStoredEmail ?? "user@example.com",
143+
email: tokenEmail ?? sanitizedStoredEmail,
144144
tokenAccountId,
145145
};
146146
},

‎test/index.test.ts‎

Lines changed: 148 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -424,12 +424,20 @@ vi.mock("../lib/accounts.js", async () => {
424424
accessToken?: string;
425425
idToken?: string;
426426
}) => {
427+
const tokenUtilsModule =
428+
tokenUtils as typeof import("../lib/auth/token-utils.js");
427429
const tokenAccountId = extractAccountIdMock(accessToken);
430+
const tokenEmail = tokenUtilsModule.sanitizeEmail(
431+
extractAccountEmailMock(accessToken, idToken),
432+
);
433+
const sanitizedStoredEmail = tokenUtilsModule.sanitizeEmail(storedEmail);
428434
return {
429-
accountId: (
430-
tokenUtils as typeof import("../lib/auth/token-utils.js")
431-
).resolveRequestAccountId(storedAccountId, source as never, tokenAccountId),
432-
email: extractAccountEmailMock(accessToken, idToken) ?? storedEmail,
435+
accountId: tokenUtilsModule.resolveRequestAccountId(
436+
storedAccountId,
437+
source as never,
438+
tokenAccountId,
439+
),
440+
email: tokenEmail ?? sanitizedStoredEmail,
433441
tokenAccountId,
434442
};
435443
},
@@ -439,7 +447,8 @@ vi.mock("../lib/accounts.js", async () => {
439447
formatAccountLabel: (_account: unknown, index: number) => `Account ${index + 1}`,
440448
formatCooldown: () => null,
441449
formatWaitTime: (ms: number) => `${Math.round(ms / 1000)}s`,
442-
sanitizeEmail: (email: string) => email,
450+
sanitizeEmail: (email: string | undefined) =>
451+
typeof email === "string" ? email.trim().toLowerCase() : undefined,
443452
shouldUpdateAccountIdFromToken: () => true,
444453
parseRateLimitReason: () => "unknown",
445454
lookupCodexCliTokensByEmail: vi.fn(async () => null),
@@ -1535,6 +1544,140 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
15351544
expect(fallbackAfterRequest?.email).toBe("beta-fresh@example.com");
15361545
});
15371546

1547+
it("preserves per-request fallback entitlement keys during concurrent fetch failover", async () => {
1548+
const { AccountManager } = await import("../lib/accounts.js");
1549+
const fetchHelpers = await import("../lib/request/fetch-helpers.js");
1550+
const entitlementModule = await import("../lib/entitlement-cache.js");
1551+
const accounts = [
1552+
{
1553+
index: 0,
1554+
accountId: "token-primary",
1555+
accountIdSource: "token",
1556+
email: "alpha@example.com",
1557+
refreshToken: "refresh-1",
1558+
accessToken: "access-alpha",
1559+
},
1560+
{
1561+
index: 1,
1562+
accountId: "workspace-fallback",
1563+
accountIdSource: "org",
1564+
email: "beta-stale@example.com",
1565+
refreshToken: "refresh-2",
1566+
accessToken: "access-beta",
1567+
idToken: "id-token-beta",
1568+
},
1569+
];
1570+
let selectionCount = 0;
1571+
const manager = {
1572+
...buildRoutingManager(accounts),
1573+
getCurrentOrNextForFamilyHybrid: () =>
1574+
selectionCount++ < 2 ? accounts[0] : accounts[1],
1575+
getCurrentOrNextForFamily: () =>
1576+
selectionCount++ < 2 ? accounts[0] : accounts[1],
1577+
};
1578+
vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(manager as never);
1579+
extractAccountIdMock
1580+
.mockReturnValueOnce("token-primary")
1581+
.mockReturnValueOnce("token-primary")
1582+
.mockReturnValueOnce("token-personal-a")
1583+
.mockReturnValueOnce("token-personal-b");
1584+
extractAccountEmailMock
1585+
.mockReturnValueOnce("alpha@example.com")
1586+
.mockReturnValueOnce("alpha@example.com")
1587+
.mockReturnValueOnce("beta-fresh-a@example.com")
1588+
.mockReturnValueOnce("beta-fresh-b@example.com");
1589+
vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation(
1590+
(_init, accountId, accessToken) =>
1591+
new Headers({
1592+
"x-test-account-id": String(accountId),
1593+
"x-test-access-token": String(accessToken),
1594+
}),
1595+
);
1596+
const blockedKeys: Array<[string, string | undefined]> = [];
1597+
const clearedKeys: Array<[string, string | undefined]> = [];
1598+
vi.spyOn(entitlementModule.EntitlementCache.prototype, "isBlocked").mockImplementation(
1599+
(accountKey, modelKey) => {
1600+
blockedKeys.push([accountKey, modelKey]);
1601+
return { blocked: false, waitMs: 0 };
1602+
},
1603+
);
1604+
vi.spyOn(entitlementModule.EntitlementCache.prototype, "clear").mockImplementation(
1605+
(accountKey, modelKey) => {
1606+
clearedKeys.push([accountKey, modelKey]);
1607+
},
1608+
);
1609+
1610+
let releasePrimaryFailures: (() => void) | undefined;
1611+
const primaryBarrier = new Promise<void>((resolve) => {
1612+
releasePrimaryFailures = resolve;
1613+
});
1614+
let primaryAttempts = 0;
1615+
const fallbackResponses: Array<{ accountId: string | null; accessToken: string | null }> =
1616+
[];
1617+
globalThis.fetch = vi.fn().mockImplementation(async (_url, init) => {
1618+
const headers = new Headers(init?.headers);
1619+
const accountId = headers.get("x-test-account-id");
1620+
const accessToken = headers.get("x-test-access-token");
1621+
if (accountId === "token-primary") {
1622+
primaryAttempts += 1;
1623+
if (primaryAttempts === 2) {
1624+
releasePrimaryFailures?.();
1625+
}
1626+
await primaryBarrier;
1627+
throw new Error("Network timeout");
1628+
}
1629+
1630+
fallbackResponses.push({ accountId, accessToken });
1631+
return new Response(JSON.stringify({ accountId, accessToken }), { status: 200 });
1632+
});
1633+
1634+
const { sdk } = await setupPlugin();
1635+
const [firstResponse, secondResponse] = await Promise.all([
1636+
sdk.fetch!("https://api.openai.com/v1/chat", {
1637+
method: "POST",
1638+
body: JSON.stringify({ model: "gpt-5.1" }),
1639+
}),
1640+
sdk.fetch!("https://api.openai.com/v1/chat", {
1641+
method: "POST",
1642+
body: JSON.stringify({ model: "gpt-5.1" }),
1643+
}),
1644+
]);
1645+
1646+
expect(firstResponse.status).toBe(200);
1647+
expect(secondResponse.status).toBe(200);
1648+
const routedFetchCalls = vi
1649+
.mocked(globalThis.fetch)
1650+
.mock.calls.map(([, init]) => {
1651+
const headers = new Headers(init?.headers);
1652+
return headers.get("x-test-account-id");
1653+
})
1654+
.filter((accountId): accountId is string => accountId !== null);
1655+
expect(primaryAttempts).toBeGreaterThanOrEqual(2);
1656+
expect(
1657+
routedFetchCalls.filter((accountId) => accountId === "workspace-fallback"),
1658+
).toHaveLength(2);
1659+
expect(fallbackResponses).toEqual([
1660+
{ accountId: "workspace-fallback", accessToken: "access-beta" },
1661+
{ accountId: "workspace-fallback", accessToken: "access-beta" },
1662+
]);
1663+
expect(blockedKeys).toEqual(
1664+
expect.arrayContaining([
1665+
["account:workspace-fallback::email:beta-fresh-a@example.com", "gpt-5.1"],
1666+
["account:workspace-fallback::email:beta-fresh-b@example.com", "gpt-5.1"],
1667+
]),
1668+
);
1669+
expect(clearedKeys).toEqual(
1670+
expect.arrayContaining([
1671+
["account:workspace-fallback::email:beta-fresh-a@example.com", "gpt-5.1"],
1672+
["account:workspace-fallback::email:beta-fresh-b@example.com", "gpt-5.1"],
1673+
]),
1674+
);
1675+
const fallbackAfterRequests = manager.getAccountsSnapshot()[1];
1676+
expect(fallbackAfterRequests?.accountId).toBe("workspace-fallback");
1677+
expect(fallbackAfterRequests?.accountIdSource).toBe("org");
1678+
expect(syncCodexCliSelectionMock).toHaveBeenCalledWith(1);
1679+
});
1680+
15381681
it("handles network errors and rotates to next account", async () => {
15391682
globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network timeout"));
15401683

0 commit comments

Comments
 (0)