Skip to content

Commit 71a0d96

Browse files
ndycodendycode
committed
fix(runtime): preserve fallback account source on snapshot updates
1 parent 5576232 commit 71a0d96

2 files changed

Lines changed: 155 additions & 21 deletions

File tree

‎index.ts‎

Lines changed: 28 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1496,11 +1496,14 @@ while (attempted.size < Math.max(1, accountCount)) {
14961496
continue;
14971497
}
14981498

1499-
const hadAccountId = !!account.accountId;
1499+
const storedAccountId = account.accountId;
1500+
const storedAccountIdSource = account.accountIdSource;
1501+
const storedEmail = account.email;
1502+
const hadAccountId = !!storedAccountId;
15001503
const runtimeIdentity = resolveRuntimeRequestIdentity({
1501-
storedAccountId: account.accountId,
1502-
source: account.accountIdSource,
1503-
storedEmail: account.email,
1504+
storedAccountId,
1505+
source: storedAccountIdSource,
1506+
storedEmail,
15041507
accessToken: accountAuth.access,
15051508
idToken: accountAuth.idToken,
15061509
});
@@ -1517,16 +1520,11 @@ while (attempted.size < Math.max(1, accountCount)) {
15171520
}
15181521
const resolvedEmail = runtimeIdentity.email;
15191522
const entitlementAccountKey = resolveEntitlementAccountKey({
1520-
accountId: account.accountId ?? accountId,
1523+
accountId: storedAccountId ?? accountId,
15211524
email: resolvedEmail,
15221525
refreshToken: account.refreshToken,
15231526
index: account.index,
15241527
});
1525-
account.accountId = accountId;
1526-
if (!hadAccountId && tokenAccountId && accountId === tokenAccountId) {
1527-
account.accountIdSource = account.accountIdSource ?? "token";
1528-
}
1529-
account.email = resolvedEmail;
15301528
const entitlementBlock = entitlementCache.isBlocked(
15311529
entitlementAccountKey,
15321530
model ?? modelFamily,
@@ -1539,6 +1537,11 @@ while (attempted.size < Math.max(1, accountCount)) {
15391537
);
15401538
continue;
15411539
}
1540+
account.accountId = accountId;
1541+
if (!hadAccountId && tokenAccountId && accountId === tokenAccountId) {
1542+
account.accountIdSource = storedAccountIdSource ?? "token";
1543+
}
1544+
account.email = resolvedEmail;
15421545

15431546
if (
15441547
accountCount > 1 &&
@@ -2104,10 +2107,14 @@ while (attempted.size < Math.max(1, accountCount)) {
21042107
continue;
21052108
}
21062109

2110+
const fallbackStoredAccountId = fallbackAccount.accountId;
2111+
const fallbackStoredAccountIdSource = fallbackAccount.accountIdSource;
2112+
const fallbackStoredEmail = fallbackAccount.email;
2113+
const hadFallbackAccountId = !!fallbackStoredAccountId;
21072114
const fallbackRuntimeIdentity = resolveRuntimeRequestIdentity({
2108-
storedAccountId: fallbackAccount.accountId,
2109-
source: fallbackAccount.accountIdSource,
2110-
storedEmail: fallbackAccount.email,
2115+
storedAccountId: fallbackStoredAccountId,
2116+
source: fallbackStoredAccountIdSource,
2117+
storedEmail: fallbackStoredEmail,
21112118
accessToken: fallbackAuth.access,
21122119
idToken: fallbackAuth.idToken,
21132120
});
@@ -2118,27 +2125,28 @@ while (attempted.size < Math.max(1, accountCount)) {
21182125
}
21192126
const fallbackResolvedEmail = fallbackRuntimeIdentity.email;
21202127
const fallbackEntitlementAccountKey = resolveEntitlementAccountKey({
2121-
accountId: fallbackAccount.accountId ?? fallbackAccountId,
2128+
accountId: fallbackStoredAccountId ?? fallbackAccountId,
21222129
email: fallbackResolvedEmail,
21232130
refreshToken: fallbackAccount.refreshToken,
21242131
index: fallbackAccount.index,
21252132
});
2133+
2134+
if (!accountManager.consumeToken(fallbackAccount, modelFamily, model)) {
2135+
continue;
2136+
}
21262137
fallbackAccount.accountId = fallbackAccountId;
21272138
if (
2128-
!fallbackAccount.accountIdSource &&
2139+
!hadFallbackAccountId &&
21292140
fallbackTokenAccountId &&
21302141
fallbackAccountId === fallbackTokenAccountId
21312142
) {
2132-
fallbackAccount.accountIdSource = "token";
2143+
fallbackAccount.accountIdSource =
2144+
fallbackStoredAccountIdSource ?? "token";
21332145
}
21342146
if (fallbackResolvedEmail) {
21352147
fallbackAccount.email = fallbackResolvedEmail;
21362148
}
21372149

2138-
if (!accountManager.consumeToken(fallbackAccount, modelFamily, model)) {
2139-
continue;
2140-
}
2141-
21422150
const fallbackHeaders = createCodexHeaders(
21432151
requestInit,
21442152
fallbackAccountId,

‎test/index.test.ts‎

Lines changed: 127 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1248,7 +1248,10 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
12481248
const entitlementModule = await import("../lib/entitlement-cache.js");
12491249
const isBlockedSpy = vi
12501250
.spyOn(entitlementModule.EntitlementCache.prototype, "isBlocked")
1251-
.mockReturnValue({ blocked: false, waitMs: 0 });
1251+
.mockImplementation(() => {
1252+
expect(mockStorage.accounts[0]?.email).toBe("stale@example.com");
1253+
return { blocked: false, waitMs: 0 };
1254+
});
12521255
globalThis.fetch = vi.fn().mockResolvedValue(
12531256
new Response(JSON.stringify({ content: "test" }), { status: 200 }),
12541257
);
@@ -1405,6 +1408,129 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
14051408
expect(syncCodexCliSelectionMock).toHaveBeenCalledWith(1);
14061409
});
14071410

1411+
it("locks fallback entitlement checks to a stable account snapshot before mutating the shared account", async () => {
1412+
const { AccountManager } = await import("../lib/accounts.js");
1413+
const fetchHelpers = await import("../lib/request/fetch-helpers.js");
1414+
const entitlementModule = await import("../lib/entitlement-cache.js");
1415+
const manager = buildRoutingManager([
1416+
{
1417+
index: 0,
1418+
accountId: "token-primary",
1419+
accountIdSource: "token",
1420+
email: "alpha@example.com",
1421+
refreshToken: "refresh-1",
1422+
accessToken: "access-alpha",
1423+
},
1424+
{
1425+
index: 1,
1426+
accountId: "workspace-fallback",
1427+
accountIdSource: "org",
1428+
email: "beta-stale@example.com",
1429+
refreshToken: "refresh-2",
1430+
accessToken: "access-beta",
1431+
idToken: "id-token-beta",
1432+
},
1433+
]);
1434+
vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(manager as never);
1435+
extractAccountIdMock
1436+
.mockReturnValueOnce("token-primary")
1437+
.mockReturnValueOnce("token-personal-fallback");
1438+
extractAccountEmailMock
1439+
.mockReturnValueOnce("alpha@example.com")
1440+
.mockReturnValueOnce("beta-fresh@example.com");
1441+
vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation(
1442+
(_init, accountId, accessToken) =>
1443+
new Headers({
1444+
"x-test-account-id": String(accountId),
1445+
"x-test-access-token": String(accessToken),
1446+
}),
1447+
);
1448+
const isBlockedSpy = vi
1449+
.spyOn(entitlementModule.EntitlementCache.prototype, "isBlocked")
1450+
.mockImplementation((key) => {
1451+
if (key === "account:workspace-fallback::email:beta-fresh@example.com") {
1452+
const fallbackSnapshot = manager.getAccountsSnapshot()[1];
1453+
expect(fallbackSnapshot?.accountId).toBe("workspace-fallback");
1454+
expect(fallbackSnapshot?.email).toBe("beta-stale@example.com");
1455+
}
1456+
return { blocked: false, waitMs: 0 };
1457+
});
1458+
globalThis.fetch = vi
1459+
.fn()
1460+
.mockRejectedValueOnce(new Error("Network timeout"))
1461+
.mockRejectedValueOnce(new Error("Network timeout"))
1462+
.mockResolvedValueOnce(new Response(JSON.stringify({ content: "ok" }), { status: 200 }));
1463+
1464+
const { sdk } = await setupPlugin();
1465+
const response = await sdk.fetch!("https://api.openai.com/v1/chat", {
1466+
method: "POST",
1467+
body: JSON.stringify({ model: "gpt-5.1" }),
1468+
});
1469+
1470+
expect(response.status).toBe(200);
1471+
expect(isBlockedSpy).toHaveBeenCalledWith(
1472+
"account:workspace-fallback::email:beta-fresh@example.com",
1473+
"gpt-5.1",
1474+
);
1475+
const fallbackAfterRequest = manager.getAccountsSnapshot()[1];
1476+
expect(fallbackAfterRequest?.accountId).toBe("workspace-fallback");
1477+
expect(fallbackAfterRequest?.email).toBe("beta-fresh@example.com");
1478+
});
1479+
1480+
it("preserves a non-token fallback accountIdSource when the stored accountId was originally missing", async () => {
1481+
const { AccountManager } = await import("../lib/accounts.js");
1482+
const fetchHelpers = await import("../lib/request/fetch-helpers.js");
1483+
const manager = buildRoutingManager([
1484+
{
1485+
index: 0,
1486+
accountId: "token-primary",
1487+
accountIdSource: "token",
1488+
email: "alpha@example.com",
1489+
refreshToken: "refresh-1",
1490+
accessToken: "access-alpha",
1491+
},
1492+
{
1493+
index: 1,
1494+
accountIdSource: "org",
1495+
email: "beta-stale@example.com",
1496+
refreshToken: "refresh-2",
1497+
accessToken: "access-beta",
1498+
idToken: "id-token-beta",
1499+
},
1500+
]);
1501+
vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(manager as never);
1502+
extractAccountIdMock
1503+
.mockReturnValueOnce("token-primary")
1504+
.mockReturnValueOnce("workspace-fallback");
1505+
extractAccountEmailMock
1506+
.mockReturnValueOnce("alpha@example.com")
1507+
.mockReturnValueOnce("beta-fresh@example.com");
1508+
vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation(
1509+
(_init, accountId, accessToken) =>
1510+
new Headers({
1511+
"x-test-account-id": String(accountId),
1512+
"x-test-access-token": String(accessToken),
1513+
}),
1514+
);
1515+
globalThis.fetch = vi
1516+
.fn()
1517+
.mockRejectedValueOnce(new Error("Network timeout"))
1518+
.mockRejectedValueOnce(new Error("Network timeout"))
1519+
.mockResolvedValueOnce(new Response(JSON.stringify({ content: "ok" }), { status: 200 }));
1520+
1521+
const { sdk } = await setupPlugin();
1522+
const response = await sdk.fetch!("https://api.openai.com/v1/chat", {
1523+
method: "POST",
1524+
body: JSON.stringify({ model: "gpt-5.1" }),
1525+
});
1526+
1527+
expect(response.status).toBe(200);
1528+
const fallbackAfterRequest = manager.getAccountsSnapshot()[1];
1529+
expect(fallbackAfterRequest?.accountId).toBe("workspace-fallback");
1530+
expect(fallbackAfterRequest?.accountIdSource).toBe("org");
1531+
expect(fallbackAfterRequest?.email).toBe("beta-fresh@example.com");
1532+
});
1533+
14081534
it("handles network errors and rotates to next account", async () => {
14091535
globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network timeout"));
14101536

0 commit comments

Comments
 (0)