Skip to content

Commit efbb34e

Browse files
ndycodendycode
committed
fix(routing): preserve failover identity bookkeeping
1 parent 9d81368 commit efbb34e

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
@@ -1542,7 +1542,9 @@ while (attempted.size < Math.max(1, accountCount)) {
15421542
if (!hadAccountId && tokenAccountId && accountId === tokenAccountId) {
15431543
account.accountIdSource = storedAccountIdSource ?? "token";
15441544
}
1545-
account.email = resolvedEmail;
1545+
if (resolvedEmail) {
1546+
account.email = resolvedEmail;
1547+
}
15461548

15471549
if (
15481550
accountCount > 1 &&
@@ -1601,6 +1603,7 @@ while (attempted.size < Math.max(1, accountCount)) {
16011603

16021604
let sameAccountRetryCount = 0;
16031605
let successAccountForResponse = account;
1606+
let successEntitlementAccountKey = entitlementAccountKey;
16041607
while (true) {
16051608
let response: Response;
16061609
const fetchStart = performance.now();
@@ -2233,6 +2236,7 @@ while (attempted.size < Math.max(1, accountCount)) {
22332236
}
22342237

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

‎test/index-retry.test.ts‎

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

‎test/index.test.ts‎

Lines changed: 148 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -425,12 +425,20 @@ vi.mock("../lib/accounts.js", async () => {
425425
accessToken?: string;
426426
idToken?: string;
427427
}) => {
428+
const tokenUtilsModule =
429+
tokenUtils as typeof import("../lib/auth/token-utils.js");
428430
const tokenAccountId = extractAccountIdMock(accessToken);
431+
const tokenEmail = tokenUtilsModule.sanitizeEmail(
432+
extractAccountEmailMock(accessToken, idToken),
433+
);
434+
const sanitizedStoredEmail = tokenUtilsModule.sanitizeEmail(storedEmail);
429435
return {
430-
accountId: (
431-
tokenUtils as typeof import("../lib/auth/token-utils.js")
432-
).resolveRequestAccountId(storedAccountId, source as never, tokenAccountId),
433-
email: extractAccountEmailMock(accessToken, idToken) ?? storedEmail,
436+
accountId: tokenUtilsModule.resolveRequestAccountId(
437+
storedAccountId,
438+
source as never,
439+
tokenAccountId,
440+
),
441+
email: tokenEmail ?? sanitizedStoredEmail,
434442
tokenAccountId,
435443
};
436444
},
@@ -440,7 +448,8 @@ vi.mock("../lib/accounts.js", async () => {
440448
formatAccountLabel: (_account: unknown, index: number) => `Account ${index + 1}`,
441449
formatCooldown: () => null,
442450
formatWaitTime: (ms: number) => `${Math.round(ms / 1000)}s`,
443-
sanitizeEmail: (email: string) => email,
451+
sanitizeEmail: (email: string | undefined) =>
452+
typeof email === "string" ? email.trim().toLowerCase() : undefined,
444453
shouldUpdateAccountIdFromToken: () => true,
445454
parseRateLimitReason: () => "unknown",
446455
lookupCodexCliTokensByEmail: vi.fn(async () => null),
@@ -1606,6 +1615,140 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
16061615
expect(fallbackAfterRequest?.email).toBe("beta-fresh@example.com");
16071616
});
16081617

1618+
it("preserves per-request fallback entitlement keys during concurrent fetch failover", async () => {
1619+
const { AccountManager } = await import("../lib/accounts.js");
1620+
const fetchHelpers = await import("../lib/request/fetch-helpers.js");
1621+
const entitlementModule = await import("../lib/entitlement-cache.js");
1622+
const accounts = [
1623+
{
1624+
index: 0,
1625+
accountId: "token-primary",
1626+
accountIdSource: "token",
1627+
email: "alpha@example.com",
1628+
refreshToken: "refresh-1",
1629+
accessToken: "access-alpha",
1630+
},
1631+
{
1632+
index: 1,
1633+
accountId: "workspace-fallback",
1634+
accountIdSource: "org",
1635+
email: "beta-stale@example.com",
1636+
refreshToken: "refresh-2",
1637+
accessToken: "access-beta",
1638+
idToken: "id-token-beta",
1639+
},
1640+
];
1641+
let selectionCount = 0;
1642+
const manager = {
1643+
...buildRoutingManager(accounts),
1644+
getCurrentOrNextForFamilyHybrid: () =>
1645+
selectionCount++ < 2 ? accounts[0] : accounts[1],
1646+
getCurrentOrNextForFamily: () =>
1647+
selectionCount++ < 2 ? accounts[0] : accounts[1],
1648+
};
1649+
vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(manager as never);
1650+
extractAccountIdMock
1651+
.mockReturnValueOnce("token-primary")
1652+
.mockReturnValueOnce("token-primary")
1653+
.mockReturnValueOnce("token-personal-a")
1654+
.mockReturnValueOnce("token-personal-b");
1655+
extractAccountEmailMock
1656+
.mockReturnValueOnce("alpha@example.com")
1657+
.mockReturnValueOnce("alpha@example.com")
1658+
.mockReturnValueOnce("beta-fresh-a@example.com")
1659+
.mockReturnValueOnce("beta-fresh-b@example.com");
1660+
vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation(
1661+
(_init, accountId, accessToken) =>
1662+
new Headers({
1663+
"x-test-account-id": String(accountId),
1664+
"x-test-access-token": String(accessToken),
1665+
}),
1666+
);
1667+
const blockedKeys: Array<[string, string | undefined]> = [];
1668+
const clearedKeys: Array<[string, string | undefined]> = [];
1669+
vi.spyOn(entitlementModule.EntitlementCache.prototype, "isBlocked").mockImplementation(
1670+
(accountKey, modelKey) => {
1671+
blockedKeys.push([accountKey, modelKey]);
1672+
return { blocked: false, waitMs: 0 };
1673+
},
1674+
);
1675+
vi.spyOn(entitlementModule.EntitlementCache.prototype, "clear").mockImplementation(
1676+
(accountKey, modelKey) => {
1677+
clearedKeys.push([accountKey, modelKey]);
1678+
},
1679+
);
1680+
1681+
let releasePrimaryFailures: (() => void) | undefined;
1682+
const primaryBarrier = new Promise<void>((resolve) => {
1683+
releasePrimaryFailures = resolve;
1684+
});
1685+
let primaryAttempts = 0;
1686+
const fallbackResponses: Array<{ accountId: string | null; accessToken: string | null }> =
1687+
[];
1688+
globalThis.fetch = vi.fn().mockImplementation(async (_url, init) => {
1689+
const headers = new Headers(init?.headers);
1690+
const accountId = headers.get("x-test-account-id");
1691+
const accessToken = headers.get("x-test-access-token");
1692+
if (accountId === "token-primary") {
1693+
primaryAttempts += 1;
1694+
if (primaryAttempts === 2) {
1695+
releasePrimaryFailures?.();
1696+
}
1697+
await primaryBarrier;
1698+
throw new Error("Network timeout");
1699+
}
1700+
1701+
fallbackResponses.push({ accountId, accessToken });
1702+
return new Response(JSON.stringify({ accountId, accessToken }), { status: 200 });
1703+
});
1704+
1705+
const { sdk } = await setupPlugin();
1706+
const [firstResponse, secondResponse] = await Promise.all([
1707+
sdk.fetch!("https://api.openai.com/v1/chat", {
1708+
method: "POST",
1709+
body: JSON.stringify({ model: "gpt-5.1" }),
1710+
}),
1711+
sdk.fetch!("https://api.openai.com/v1/chat", {
1712+
method: "POST",
1713+
body: JSON.stringify({ model: "gpt-5.1" }),
1714+
}),
1715+
]);
1716+
1717+
expect(firstResponse.status).toBe(200);
1718+
expect(secondResponse.status).toBe(200);
1719+
const routedFetchCalls = vi
1720+
.mocked(globalThis.fetch)
1721+
.mock.calls.map(([, init]) => {
1722+
const headers = new Headers(init?.headers);
1723+
return headers.get("x-test-account-id");
1724+
})
1725+
.filter((accountId): accountId is string => accountId !== null);
1726+
expect(primaryAttempts).toBeGreaterThanOrEqual(2);
1727+
expect(
1728+
routedFetchCalls.filter((accountId) => accountId === "workspace-fallback"),
1729+
).toHaveLength(2);
1730+
expect(fallbackResponses).toEqual([
1731+
{ accountId: "workspace-fallback", accessToken: "access-beta" },
1732+
{ accountId: "workspace-fallback", accessToken: "access-beta" },
1733+
]);
1734+
expect(blockedKeys).toEqual(
1735+
expect.arrayContaining([
1736+
["account:workspace-fallback::email:beta-fresh-a@example.com", "gpt-5.1"],
1737+
["account:workspace-fallback::email:beta-fresh-b@example.com", "gpt-5.1"],
1738+
]),
1739+
);
1740+
expect(clearedKeys).toEqual(
1741+
expect.arrayContaining([
1742+
["account:workspace-fallback::email:beta-fresh-a@example.com", "gpt-5.1"],
1743+
["account:workspace-fallback::email:beta-fresh-b@example.com", "gpt-5.1"],
1744+
]),
1745+
);
1746+
const fallbackAfterRequests = manager.getAccountsSnapshot()[1];
1747+
expect(fallbackAfterRequests?.accountId).toBe("workspace-fallback");
1748+
expect(fallbackAfterRequests?.accountIdSource).toBe("org");
1749+
expect(syncCodexCliSelectionMock).toHaveBeenCalledWith(1);
1750+
});
1751+
16091752
it("handles network errors and rotates to next account", async () => {
16101753
globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network timeout"));
16111754

0 commit comments

Comments
 (0)