Skip to content

Commit 029f539

Browse files
ndycodendycode
committed
Fix shared-account runtime cache identity
1 parent 63a7a2f commit 029f539

5 files changed

Lines changed: 386 additions & 34 deletions

File tree

‎index.ts‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1361,6 +1361,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
13611361
const accountCount = accountManager.getAccountCount();
13621362
const attempted = new Set<number>();
13631363
let restartAccountTraversalWithFallback = false;
1364+
let retryNextAccountBeforeFallback = false;
13641365
let usedPreferredSessionAccount = false;
13651366
const capabilityBoostByAccount: Record<number, number> = {};
13661367
type AccountSnapshotCandidate = {
@@ -1501,11 +1502,6 @@ while (attempted.size < Math.max(1, accountCount)) {
15011502
account.accountIdSource,
15021503
tokenAccountId,
15031504
);
1504-
const entitlementAccountKey = resolveEntitlementAccountKey({
1505-
accountId: hadAccountId ? account.accountId : undefined,
1506-
email: account.email,
1507-
index: account.index,
1508-
});
15091505
if (!accountId) {
15101506
accountManager.markAccountCoolingDown(
15111507
account,
@@ -1515,12 +1511,19 @@ while (attempted.size < Math.max(1, accountCount)) {
15151511
accountManager.saveToDiskDebounced();
15161512
continue;
15171513
}
1514+
const resolvedEmail =
1515+
extractAccountEmail(accountAuth.access) ?? account.email;
1516+
const entitlementAccountKey = resolveEntitlementAccountKey({
1517+
accountId: account.accountId ?? accountId,
1518+
email: account.email ?? resolvedEmail,
1519+
refreshToken: account.refreshToken,
1520+
index: account.index,
1521+
});
15181522
account.accountId = accountId;
15191523
if (!hadAccountId && tokenAccountId && accountId === tokenAccountId) {
15201524
account.accountIdSource = account.accountIdSource ?? "token";
15211525
}
1522-
account.email =
1523-
extractAccountEmail(accountAuth.access) ?? account.email;
1526+
account.email = resolvedEmail;
15241527
const entitlementBlock = entitlementCache.isBlocked(
15251528
entitlementAccountKey,
15261529
model ?? modelFamily,
@@ -1771,6 +1774,7 @@ while (attempted.size < Math.max(1, accountCount)) {
17711774
fallbackReason: "unsupported-model-entitlement",
17721775
},
17731776
);
1777+
retryNextAccountBeforeFallback = true;
17741778
break;
17751779
}
17761780

@@ -2292,7 +2296,10 @@ while (attempted.size < Math.max(1, accountCount)) {
22922296
if (successAccountForResponse.index !== account.index) {
22932297
accountManager.markSwitched(successAccountForResponse, "rotation", modelFamily);
22942298
}
2295-
const successAccountKey = resolveEntitlementAccountKey(successAccountForResponse);
2299+
const successAccountKey =
2300+
successAccountForResponse.index === account.index
2301+
? entitlementAccountKey
2302+
: resolveEntitlementAccountKey(successAccountForResponse);
22962303
accountManager.recordSuccess(successAccountForResponse, modelFamily, model);
22972304
capabilityPolicyStore.recordSuccess(
22982305
successAccountKey,
@@ -2311,6 +2318,11 @@ while (attempted.size < Math.max(1, accountCount)) {
23112318
}
23122319
return successResponse;
23132320
}
2321+
if (retryNextAccountBeforeFallback) {
2322+
retryNextAccountBeforeFallback = false;
2323+
continue;
2324+
}
2325+
23142326
if (restartAccountTraversalWithFallback) {
23152327
break;
23162328
}

‎lib/codex-manager.ts‎

Lines changed: 83 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -416,6 +416,27 @@ function normalizeQuotaEmail(email: string | undefined): string | null {
416416
return normalized && normalized.length > 0 ? normalized : null;
417417
}
418418

419+
function normalizeQuotaAccountId(accountId: string | undefined): string | null {
420+
if (typeof accountId !== "string") return null;
421+
const trimmed = accountId.trim();
422+
return trimmed.length > 0 ? trimmed : null;
423+
}
424+
425+
function hasUniqueQuotaAccountId(
426+
accounts: readonly Pick<AccountMetadataV3, "accountId">[],
427+
account: Pick<AccountMetadataV3, "accountId">,
428+
): boolean {
429+
const accountId = normalizeQuotaAccountId(account.accountId);
430+
if (!accountId) return false;
431+
let matchCount = 0;
432+
for (const candidate of accounts) {
433+
if (normalizeQuotaAccountId(candidate.accountId) !== accountId) continue;
434+
matchCount += 1;
435+
if (matchCount > 1) return false;
436+
}
437+
return matchCount === 1;
438+
}
439+
419440
function quotaCacheEntryToSnapshot(entry: QuotaCacheEntry): CodexQuotaSnapshot {
420441
return {
421442
status: entry.status,
@@ -490,21 +511,28 @@ function formatAccountQuotaSummary(entry: QuotaCacheEntry): string {
490511
function getQuotaCacheEntryForAccount(
491512
cache: QuotaCacheData,
492513
account: Pick<AccountMetadataV3, "accountId" | "email">,
514+
accounts: readonly Pick<AccountMetadataV3, "accountId">[],
493515
): QuotaCacheEntry | null {
494-
if (account.accountId && cache.byAccountId[account.accountId]) {
495-
return cache.byAccountId[account.accountId] ?? null;
496-
}
497516
const email = normalizeQuotaEmail(account.email);
498517
if (email && cache.byEmail[email]) {
499518
return cache.byEmail[email] ?? null;
500519
}
520+
const accountId = normalizeQuotaAccountId(account.accountId);
521+
if (
522+
accountId &&
523+
hasUniqueQuotaAccountId(accounts, account) &&
524+
cache.byAccountId[accountId]
525+
) {
526+
return cache.byAccountId[accountId] ?? null;
527+
}
501528
return null;
502529
}
503530

504531
function updateQuotaCacheForAccount(
505532
cache: QuotaCacheData,
506533
account: Pick<AccountMetadataV3, "accountId" | "email">,
507534
snapshot: CodexQuotaSnapshot,
535+
accounts: readonly Pick<AccountMetadataV3, "accountId">[],
508536
): boolean {
509537
const nextEntry: QuotaCacheEntry = {
510538
updatedAt: Date.now(),
@@ -524,14 +552,16 @@ function updateQuotaCacheForAccount(
524552
};
525553

526554
let changed = false;
527-
if (account.accountId) {
528-
cache.byAccountId[account.accountId] = nextEntry;
529-
changed = true;
530-
}
531555
const email = normalizeQuotaEmail(account.email);
532556
if (email) {
533557
cache.byEmail[email] = nextEntry;
534558
changed = true;
559+
return changed;
560+
}
561+
const accountId = normalizeQuotaAccountId(account.accountId);
562+
if (accountId && hasUniqueQuotaAccountId(accounts, account)) {
563+
cache.byAccountId[accountId] = nextEntry;
564+
changed = true;
535565
}
536566
return changed;
537567
}
@@ -550,11 +580,12 @@ function resolveMenuQuotaProbeInput(
550580
cache: QuotaCacheData,
551581
maxAgeMs: number,
552582
now: number,
583+
accounts: readonly Pick<AccountMetadataV3, "accountId">[],
553584
): { accountId: string; accessToken: string } | null {
554585
if (account.enabled === false) return null;
555586
if (!hasUsableAccessToken(account, now)) return null;
556587

557-
const existing = getQuotaCacheEntryForAccount(cache, account);
588+
const existing = getQuotaCacheEntryForAccount(cache, account, accounts);
558589
if (
559590
existing &&
560591
typeof existing.updatedAt === "number" &&
@@ -580,7 +611,13 @@ function collectMenuQuotaRefreshTargets(
580611
): MenuQuotaProbeTarget[] {
581612
const targets: MenuQuotaProbeTarget[] = [];
582613
for (const account of storage.accounts) {
583-
const probeInput = resolveMenuQuotaProbeInput(account, cache, maxAgeMs, now);
614+
const probeInput = resolveMenuQuotaProbeInput(
615+
account,
616+
cache,
617+
maxAgeMs,
618+
now,
619+
storage.accounts,
620+
);
584621
if (!probeInput) continue;
585622
targets.push({
586623
account,
@@ -599,7 +636,7 @@ function countMenuQuotaRefreshTargets(
599636
): number {
600637
let count = 0;
601638
for (const account of storage.accounts) {
602-
if (resolveMenuQuotaProbeInput(account, cache, maxAgeMs, now)) {
639+
if (resolveMenuQuotaProbeInput(account, cache, maxAgeMs, now, storage.accounts)) {
603640
count += 1;
604641
}
605642
}
@@ -632,7 +669,9 @@ async function refreshQuotaCacheForMenu(
632669
accessToken: target.accessToken,
633670
model: MENU_QUOTA_REFRESH_MODEL,
634671
});
635-
changed = updateQuotaCacheForAccount(cache, target.account, snapshot) || changed;
672+
changed =
673+
updateQuotaCacheForAccount(cache, target.account, snapshot, storage.accounts) ||
674+
changed;
636675
} catch {
637676
// Keep existing cached values if probing fails.
638677
}
@@ -783,7 +822,9 @@ function toExistingAccountInfo(
783822
const activeIndex = resolveActiveIndex(storage, "codex");
784823
const layoutMode = resolveMenuLayoutMode(displaySettings);
785824
const baseAccounts = storage.accounts.map((account, index) => {
786-
const entry = quotaCache ? getQuotaCacheEntryForAccount(quotaCache, account) : null;
825+
const entry = quotaCache
826+
? getQuotaCacheEntryForAccount(quotaCache, account, storage.accounts)
827+
: null;
787828
return {
788829
index,
789830
sourceIndex: index,
@@ -1477,7 +1518,12 @@ async function runHealthCheck(options: HealthCheckOptions = {}): Promise<void> {
14771518
});
14781519
if (quotaCache) {
14791520
quotaCacheChanged =
1480-
updateQuotaCacheForAccount(quotaCache, account, snapshot) || quotaCacheChanged;
1521+
updateQuotaCacheForAccount(
1522+
quotaCache,
1523+
account,
1524+
snapshot,
1525+
storage.accounts,
1526+
) || quotaCacheChanged;
14811527
}
14821528
healthDetail = formatQuotaSnapshotForDashboard(snapshot, display);
14831529
} catch (error) {
@@ -1550,7 +1596,12 @@ async function runHealthCheck(options: HealthCheckOptions = {}): Promise<void> {
15501596
});
15511597
if (quotaCache) {
15521598
quotaCacheChanged =
1553-
updateQuotaCacheForAccount(quotaCache, account, snapshot) || quotaCacheChanged;
1599+
updateQuotaCacheForAccount(
1600+
quotaCache,
1601+
account,
1602+
snapshot,
1603+
storage.accounts,
1604+
) || quotaCacheChanged;
15541605
}
15551606
healthyMessage = formatQuotaSnapshotForDashboard(snapshot, display);
15561607
} catch (error) {
@@ -2048,7 +2099,12 @@ async function runForecast(args: string[]): Promise<number> {
20482099
const account = storage.accounts[i];
20492100
if (account) {
20502101
quotaCacheChanged =
2051-
updateQuotaCacheForAccount(quotaCache, account, liveQuota) || quotaCacheChanged;
2102+
updateQuotaCacheForAccount(
2103+
quotaCache,
2104+
account,
2105+
liveQuota,
2106+
storage.accounts,
2107+
) || quotaCacheChanged;
20522108
}
20532109
}
20542110
} catch (error) {
@@ -2803,7 +2859,12 @@ async function runFix(args: string[]): Promise<number> {
28032859
});
28042860
if (quotaCache) {
28052861
quotaCacheChanged =
2806-
updateQuotaCacheForAccount(quotaCache, account, snapshot) || quotaCacheChanged;
2862+
updateQuotaCacheForAccount(
2863+
quotaCache,
2864+
account,
2865+
snapshot,
2866+
storage.accounts,
2867+
) || quotaCacheChanged;
28072868
}
28082869
reports.push({
28092870
index: i,
@@ -2881,7 +2942,12 @@ async function runFix(args: string[]): Promise<number> {
28812942
});
28822943
if (quotaCache) {
28832944
quotaCacheChanged =
2884-
updateQuotaCacheForAccount(quotaCache, account, snapshot) || quotaCacheChanged;
2945+
updateQuotaCacheForAccount(
2946+
quotaCache,
2947+
account,
2948+
snapshot,
2949+
storage.accounts,
2950+
) || quotaCacheChanged;
28852951
}
28862952
reports.push({
28872953
index: i,

‎lib/entitlement-cache.ts‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ const MAX_ACCOUNT_BUCKETS = 512;
1919
export interface EntitlementAccountRef {
2020
accountId?: string;
2121
email?: string;
22+
refreshToken?: string;
2223
index?: number;
2324
}
2425

@@ -39,25 +40,41 @@ function normalizeModel(model: string | undefined): string | null {
3940
return stripped.replace(/-(none|minimal|low|medium|high|xhigh)$/i, "");
4041
}
4142

43+
function normalizeEntitlementEmail(email: string | undefined): string | undefined {
44+
if (!email) return undefined;
45+
const trimmed = email.trim();
46+
if (!trimmed) return undefined;
47+
return trimmed.toLowerCase();
48+
}
49+
4250
/**
4351
* Derives a stable cache key for an entitlement account reference.
4452
*
45-
* Produces one of three deterministic keys:
46-
* - `id:<trimmed accountId>` when `accountId` is present,
47-
* - `email:<lowercased trimmed email>` when `email` is present and no `accountId`,
53+
* Produces one of six deterministic keys:
54+
* - `account:<trimmed accountId>::email:<lowercased trimmed email>` when both are present,
55+
* - `email:<lowercased trimmed email>` when only `email` is present,
56+
* - `account:<trimmed accountId>::idx:<non-negative integer>` when `accountId` is present without email,
57+
* - `account:<trimmed accountId>` when only `accountId` is present and no index is available,
58+
* - `refresh:<trimmed refreshToken>` when no accountId/email exists but a refresh token is available,
4859
* - `idx:<non-negative integer>` otherwise (index defaults to 0).
4960
*
5061
* This function is pure and concurrency-safe; it performs no I/O and is not affected by Windows filesystem semantics. It does not redact secrets or tokens — values are only trimmed and, for emails, lowercased.
5162
*
5263
* @param ref - Reference identifying an account (may include `accountId`, `email`, or `index`)
53-
* @returns A deterministic string key prefixed with `id:`, `email:`, or `idx:` as described above
64+
* @returns A deterministic string key prefixed with `account:`, `email:`, `refresh:`, or `idx:` as described above
5465
*/
5566
export function resolveEntitlementAccountKey(ref: EntitlementAccountRef): string {
5667
const accountId = typeof ref.accountId === "string" ? ref.accountId.trim() : "";
57-
if (accountId) return `id:${accountId}`;
58-
const email = typeof ref.email === "string" ? ref.email.trim().toLowerCase() : "";
68+
const hasIndex = Number.isFinite(ref.index);
69+
const index = hasIndex ? Math.max(0, Math.floor(ref.index ?? 0)) : 0;
70+
const email = normalizeEntitlementEmail(ref.email);
71+
const refreshToken =
72+
typeof ref.refreshToken === "string" ? ref.refreshToken.trim() : "";
73+
if (accountId && email) return `account:${accountId}::email:${email}`;
5974
if (email) return `email:${email}`;
60-
const index = Number.isFinite(ref.index) ? Math.max(0, Math.floor(ref.index ?? 0)) : 0;
75+
if (accountId && hasIndex) return `account:${accountId}::idx:${index}`;
76+
if (accountId) return `account:${accountId}`;
77+
if (refreshToken) return `refresh:${refreshToken}`;
6178
return `idx:${index}`;
6279
}
6380

0 commit comments

Comments
 (0)