Skip to content

Commit a0caf1f

Browse files
ndycodendycode
committed
Merge branch 'pr-133-head' into release-integration-20260320
2 parents baba870 + a8aec23 commit a0caf1f

8 files changed

Lines changed: 783 additions & 38 deletions

File tree

‎index.ts‎

Lines changed: 74 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ import {
109109
formatAccountLabel,
110110
formatCooldown,
111111
formatWaitTime,
112+
resolveRuntimeRequestIdentity,
112113
sanitizeEmail,
113114
selectBestAccountCandidate,
114115
shouldUpdateAccountIdFromToken,
@@ -1495,13 +1496,19 @@ while (attempted.size < Math.max(1, accountCount)) {
14951496
continue;
14961497
}
14971498

1498-
const hadAccountId = !!account.accountId;
1499-
const tokenAccountId = extractAccountId(accountAuth.access);
1500-
const accountId = resolveRequestAccountId(
1501-
account.accountId,
1502-
account.accountIdSource,
1503-
tokenAccountId,
1504-
);
1499+
const storedAccountId = account.accountId;
1500+
const storedAccountIdSource = account.accountIdSource;
1501+
const storedEmail = account.email;
1502+
const hadAccountId = !!storedAccountId;
1503+
const runtimeIdentity = resolveRuntimeRequestIdentity({
1504+
storedAccountId,
1505+
source: storedAccountIdSource,
1506+
storedEmail,
1507+
accessToken: accountAuth.access,
1508+
idToken: accountAuth.idToken,
1509+
});
1510+
const tokenAccountId = runtimeIdentity.tokenAccountId;
1511+
const accountId = runtimeIdentity.accountId;
15051512
if (!accountId) {
15061513
accountManager.markAccountCoolingDown(
15071514
account,
@@ -1511,19 +1518,13 @@ while (attempted.size < Math.max(1, accountCount)) {
15111518
accountManager.saveToDiskDebounced();
15121519
continue;
15131520
}
1514-
const resolvedEmail =
1515-
extractAccountEmail(accountAuth.access) ?? account.email;
1521+
const resolvedEmail = runtimeIdentity.email;
15161522
const entitlementAccountKey = resolveEntitlementAccountKey({
1517-
accountId: account.accountId ?? accountId,
1523+
accountId: storedAccountId ?? accountId,
15181524
email: resolvedEmail,
15191525
refreshToken: account.refreshToken,
15201526
index: account.index,
15211527
});
1522-
account.accountId = accountId;
1523-
if (!hadAccountId && tokenAccountId && accountId === tokenAccountId) {
1524-
account.accountIdSource = account.accountIdSource ?? "token";
1525-
}
1526-
account.email = resolvedEmail;
15271528
const entitlementBlock = entitlementCache.isBlocked(
15281529
entitlementAccountKey,
15291530
model ?? modelFamily,
@@ -1536,6 +1537,13 @@ while (attempted.size < Math.max(1, accountCount)) {
15361537
);
15371538
continue;
15381539
}
1540+
account.accountId = accountId;
1541+
if (!hadAccountId && tokenAccountId && accountId === tokenAccountId) {
1542+
account.accountIdSource = storedAccountIdSource ?? "token";
1543+
}
1544+
if (resolvedEmail) {
1545+
account.email = resolvedEmail;
1546+
}
15391547

15401548
if (
15411549
accountCount > 1 &&
@@ -1594,6 +1602,7 @@ while (attempted.size < Math.max(1, accountCount)) {
15941602

15951603
let sameAccountRetryCount = 0;
15961604
let successAccountForResponse = account;
1605+
let successEntitlementAccountKey = entitlementAccountKey;
15971606
while (true) {
15981607
let response: Response;
15991608
const fetchStart = performance.now();
@@ -2101,19 +2110,58 @@ while (attempted.size < Math.max(1, accountCount)) {
21012110
continue;
21022111
}
21032112

2104-
const fallbackTokenAccountId = extractAccountId(fallbackAuth.access);
2105-
const fallbackAccountId = resolveRequestAccountId(
2106-
fallbackAccount.accountId,
2107-
fallbackAccount.accountIdSource,
2108-
fallbackTokenAccountId,
2109-
);
2113+
const fallbackStoredAccountId = fallbackAccount.accountId;
2114+
const fallbackStoredAccountIdSource = fallbackAccount.accountIdSource;
2115+
const fallbackStoredEmail = fallbackAccount.email;
2116+
const hadFallbackAccountId = !!fallbackStoredAccountId;
2117+
const fallbackRuntimeIdentity = resolveRuntimeRequestIdentity({
2118+
storedAccountId: fallbackStoredAccountId,
2119+
source: fallbackStoredAccountIdSource,
2120+
storedEmail: fallbackStoredEmail,
2121+
accessToken: fallbackAuth.access,
2122+
idToken: fallbackAuth.idToken,
2123+
});
2124+
const fallbackTokenAccountId = fallbackRuntimeIdentity.tokenAccountId;
2125+
const fallbackAccountId = fallbackRuntimeIdentity.accountId;
21102126
if (!fallbackAccountId) {
21112127
continue;
21122128
}
2129+
const fallbackResolvedEmail = fallbackRuntimeIdentity.email;
2130+
const fallbackEntitlementAccountKey = resolveEntitlementAccountKey({
2131+
accountId: fallbackStoredAccountId ?? fallbackAccountId,
2132+
email: fallbackResolvedEmail,
2133+
refreshToken: fallbackAccount.refreshToken,
2134+
index: fallbackAccount.index,
2135+
});
2136+
const fallbackEntitlementBlock = entitlementCache.isBlocked(
2137+
fallbackEntitlementAccountKey,
2138+
model ?? modelFamily,
2139+
);
2140+
if (fallbackEntitlementBlock.blocked) {
2141+
runtimeMetrics.accountRotations++;
2142+
runtimeMetrics.lastError =
2143+
`Entitlement cached block for account ${fallbackAccount.index + 1}`;
2144+
logWarn(
2145+
`Skipping account ${fallbackAccount.index + 1} due to cached entitlement block (${formatWaitTime(fallbackEntitlementBlock.waitMs)} remaining).`,
2146+
);
2147+
continue;
2148+
}
21132149

21142150
if (!accountManager.consumeToken(fallbackAccount, modelFamily, model)) {
21152151
continue;
21162152
}
2153+
fallbackAccount.accountId = fallbackAccountId;
2154+
if (
2155+
!hadFallbackAccountId &&
2156+
fallbackTokenAccountId &&
2157+
fallbackAccountId === fallbackTokenAccountId
2158+
) {
2159+
fallbackAccount.accountIdSource =
2160+
fallbackStoredAccountIdSource ?? "token";
2161+
}
2162+
if (fallbackResolvedEmail) {
2163+
fallbackAccount.email = fallbackResolvedEmail;
2164+
}
21172165

21182166
const fallbackHeaders = createCodexHeaders(
21192167
requestInit,
@@ -2155,7 +2203,7 @@ while (attempted.size < Math.max(1, accountCount)) {
21552203
);
21562204
if (fallbackSnapshot) {
21572205
preemptiveQuotaScheduler.update(
2158-
`${resolveEntitlementAccountKey(fallbackAccount)}:${model ?? modelFamily}`,
2206+
`${fallbackEntitlementAccountKey}:${model ?? modelFamily}`,
21592207
fallbackSnapshot,
21602208
);
21612209
}
@@ -2180,13 +2228,14 @@ while (attempted.size < Math.max(1, accountCount)) {
21802228
accountManager.recordFailure(fallbackAccount, modelFamily, model);
21812229
}
21822230
capabilityPolicyStore.recordFailure(
2183-
resolveEntitlementAccountKey(fallbackAccount),
2231+
fallbackEntitlementAccountKey,
21842232
capabilityModelKey,
21852233
);
21862234
continue;
21872235
}
21882236

21892237
successAccountForResponse = fallbackAccount;
2238+
successEntitlementAccountKey = fallbackEntitlementAccountKey;
21902239
runtimeMetrics.streamFailoverRecoveries += 1;
21912240
if (fallbackAccount.index !== account.index) {
21922241
runtimeMetrics.streamFailoverCrossAccountRecoveries += 1;
@@ -2206,7 +2255,7 @@ while (attempted.size < Math.max(1, accountCount)) {
22062255
accountManager.refundToken(fallbackAccount, modelFamily, model);
22072256
accountManager.recordFailure(fallbackAccount, modelFamily, model);
22082257
capabilityPolicyStore.recordFailure(
2209-
resolveEntitlementAccountKey(fallbackAccount),
2258+
fallbackEntitlementAccountKey,
22102259
capabilityModelKey,
22112260
);
22122261
logWarn(
@@ -2296,10 +2345,7 @@ while (attempted.size < Math.max(1, accountCount)) {
22962345
if (successAccountForResponse.index !== account.index) {
22972346
accountManager.markSwitched(successAccountForResponse, "rotation", modelFamily);
22982347
}
2299-
const successAccountKey =
2300-
successAccountForResponse.index === account.index
2301-
? entitlementAccountKey
2302-
: resolveEntitlementAccountKey(successAccountForResponse);
2348+
const successAccountKey = successEntitlementAccountKey;
23032349
accountManager.recordSuccess(successAccountForResponse, modelFamily, model);
23042350
capabilityPolicyStore.recordSuccess(
23052351
successAccountKey,

‎lib/accounts.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ export {
3030
extractAccountEmail,
3131
getAccountIdCandidates,
3232
selectBestAccountCandidate,
33+
resolveRuntimeRequestIdentity,
3334
shouldUpdateAccountIdFromToken,
3435
resolveRequestAccountId,
3536
sanitizeEmail,

‎lib/auth/token-utils.ts‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,37 @@ export function resolveRequestAccountId(
383383
return tokenAccountId ?? storedAccountId;
384384
}
385385

386+
export interface RuntimeRequestIdentity {
387+
accountId?: string;
388+
email?: string;
389+
tokenAccountId?: string;
390+
}
391+
392+
/**
393+
* Resolve the live request identity for an account using the stored workspace binding
394+
* plus the freshest token-derived hints available for this request.
395+
*/
396+
export function resolveRuntimeRequestIdentity(input: {
397+
storedAccountId?: string;
398+
source?: AccountIdSource;
399+
storedEmail?: string;
400+
accessToken?: string;
401+
idToken?: string;
402+
}): RuntimeRequestIdentity {
403+
const tokenAccountId = extractAccountId(input.accessToken);
404+
return {
405+
accountId: resolveRequestAccountId(
406+
input.storedAccountId,
407+
input.source,
408+
tokenAccountId,
409+
),
410+
email:
411+
sanitizeEmail(extractAccountEmail(input.accessToken, input.idToken)) ??
412+
sanitizeEmail(input.storedEmail),
413+
tokenAccountId,
414+
};
415+
}
416+
386417
/**
387418
* Sanitizes an email address by trimming whitespace and lowercasing.
388419
* @param email - Email string to sanitize

‎lib/codex-manager.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1586,7 +1586,7 @@ async function runOAuthFlow(
15861586
);
15871587
}
15881588

1589-
const waitingForCallback = oauthServer?.ready === true;
1589+
const waitingForCallback = signInMode === "browser" && oauthServer?.ready === true;
15901590
if (waitingForCallback && oauthServer) {
15911591
console.log(stylePromptText(UI_COPY.oauth.waitingCallback, "muted"));
15921592
const callbackResult = await oauthServer.waitForCode(state);

‎test/codex-manager-cli.test.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3301,7 +3301,6 @@ describe("codex manager cli commands", () => {
33013301
expect(exchangeAuthorizationCodeMock).not.toHaveBeenCalled();
33023302
expect(storageState.accounts).toHaveLength(0);
33033303
});
3304-
33053304
it("preserves distinct same-email workspaces when oauth login reuses a refresh token", async () => {
33063305
const now = Date.now();
33073306
let storageState = {

‎test/index-retry.test.ts‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ vi.mock("../lib/request/request-transformer.js", () => ({
3535
applyFastSessionDefaults: <T>(config: T) => config,
3636
}));
3737

38-
vi.mock("../lib/accounts.js", () => {
38+
vi.mock("../lib/accounts.js", async () => {
39+
const tokenUtils = await vi.importActual("../lib/auth/token-utils.js");
3940
class AccountManager {
4041
private calls = 0;
4142

@@ -114,7 +115,38 @@ vi.mock("../lib/accounts.js", () => {
114115
extractAccountEmail: () => "user@example.com",
115116
extractAccountId: () => "account-1",
116117
selectBestAccountCandidate: (candidates: Array<{ accountId: string }>) => candidates[0] ?? null,
117-
resolveRequestAccountId: (_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) => tokenId,
118+
resolveRuntimeRequestIdentity: ({
119+
storedAccountId,
120+
source,
121+
storedEmail,
122+
accessToken,
123+
idToken,
124+
}: {
125+
storedAccountId?: string;
126+
source?: string;
127+
storedEmail?: string;
128+
accessToken?: string;
129+
idToken?: string;
130+
}) => {
131+
const tokenUtilsModule = tokenUtils as typeof import("../lib/auth/token-utils.js");
132+
const tokenAccountId = accessToken ? "account-1" : undefined;
133+
const tokenEmail = tokenUtilsModule.sanitizeEmail(
134+
tokenUtilsModule.extractAccountEmail(accessToken, idToken),
135+
);
136+
const sanitizedStoredEmail = tokenUtilsModule.sanitizeEmail(storedEmail);
137+
return {
138+
accountId: tokenUtilsModule.resolveRequestAccountId(
139+
storedAccountId,
140+
source as never,
141+
tokenAccountId,
142+
),
143+
email: tokenEmail ?? sanitizedStoredEmail,
144+
tokenAccountId,
145+
};
146+
},
147+
resolveRequestAccountId: (
148+
tokenUtils as typeof import("../lib/auth/token-utils.js")
149+
).resolveRequestAccountId,
118150
formatAccountLabel: (_account: any, index: number) => `Account ${index + 1}`,
119151
formatCooldown: (ms: number) => `${ms}ms`,
120152
formatWaitTime: (ms: number) => `${ms}ms`,

0 commit comments

Comments
 (0)