Skip to content

Commit 22db10b

Browse files
ndycodendycode
committed
chore: finalize 1.2.2 release integration
1 parent a01e56d commit 22db10b

23 files changed

Lines changed: 1461 additions & 217 deletions

‎lib/accounts.ts‎

Lines changed: 124 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,11 @@ import {
2727
} from "./codex-cli/state.js";
2828
import { syncAccountStorageFromCodexCli } from "./codex-cli/sync.js";
2929
import { setCodexCliActiveSelection } from "./codex-cli/writer.js";
30+
import {
31+
getAccountIdentityKey,
32+
getRuntimeAccountIdentityKey,
33+
} from "./storage/identity.js";
34+
import { getCircuitBreaker } from "./circuit-breaker.js";
3035

3136
export {
3237
extractAccountId,
@@ -77,7 +82,27 @@ import {
7782
} from "./accounts/rate-limits.js";
7883

7984
const log = createLogger("accounts");
80-
const ACCOUNT_SELECTION_FRESH_WINDOW_MS = 5 * 60_000;
85+
let nextRuntimeCircuitKeyId = 0;
86+
87+
function getAccountCircuitKey(account: ManagedAccount): string {
88+
if (!account.circuitKeyId) {
89+
account.circuitKeyId =
90+
getAccountIdentityKey(account) ?? `circuit:${nextRuntimeCircuitKeyId++}`;
91+
}
92+
return account.circuitKeyId;
93+
}
94+
95+
export function getRuntimeTrackerKey(
96+
account: ManagedAccount,
97+
): string | number {
98+
if (account._runtimeTrackerKey !== undefined) {
99+
return account._runtimeTrackerKey;
100+
}
101+
102+
const trackerKey = getRuntimeAccountIdentityKey(account) ?? account.index;
103+
account._runtimeTrackerKey = trackerKey;
104+
return trackerKey;
105+
}
81106

82107
function initFamilyState(defaultValue: number): Record<ModelFamily, number> {
83108
return Object.fromEntries(
@@ -201,6 +226,8 @@ export interface Workspace {
201226

202227
export interface ManagedAccount {
203228
index: number;
229+
_runtimeTrackerKey?: string | number;
230+
circuitKeyId?: string;
204231
accountId?: string;
205232
accountIdSource?: AccountIdSource;
206233
accountLabel?: string;
@@ -285,24 +312,20 @@ export class AccountManager {
285312
const cached = cache.get(email);
286313
if (!cached) continue;
287314

288-
const cachedAccessUsable =
289-
typeof cached.expiresAt !== "number" || cached.expiresAt > now;
315+
if (typeof cached.expiresAt === "number" && cached.expiresAt <= now) {
316+
continue;
317+
}
290318

291319
const missingOrExpired =
292320
!account.access || account.expires === undefined || account.expires <= now;
293-
if (missingOrExpired && cachedAccessUsable) {
321+
if (missingOrExpired) {
294322
account.access = cached.accessToken;
295323
if (typeof cached.expiresAt === "number") {
296324
account.expires = cached.expiresAt;
297325
}
298326
changed = true;
299327
}
300328

301-
if (cached.refreshToken && !account.refreshToken) {
302-
account.refreshToken = cached.refreshToken;
303-
changed = true;
304-
}
305-
306329
if (
307330
!account.accountId &&
308331
cached.accountId &&
@@ -482,10 +505,16 @@ export class AccountManager {
482505
}
483506

484507
getAccountsSnapshot(): ManagedAccount[] {
485-
return this.accounts.map((account) => ({
486-
...account,
487-
rateLimitResetTimes: { ...account.rateLimitResetTimes },
488-
}));
508+
return this.accounts.map((account) => {
509+
const trackerKey = getRuntimeTrackerKey(account);
510+
const circuitKeyId = getAccountCircuitKey(account);
511+
return {
512+
...account,
513+
_runtimeTrackerKey: trackerKey,
514+
circuitKeyId,
515+
rateLimitResetTimes: { ...account.rateLimitResetTimes },
516+
};
517+
});
489518
}
490519

491520
getAccountByIndex(index: number): ManagedAccount | null {
@@ -495,45 +524,16 @@ export class AccountManager {
495524
return account ?? null;
496525
}
497526

498-
private hasFreshAccessToken(account: ManagedAccount): boolean {
499-
if (!account.access) return false;
500-
if (typeof account.expires !== "number" || !Number.isFinite(account.expires)) {
501-
return false;
502-
}
503-
return account.expires > nowMs() + ACCOUNT_SELECTION_FRESH_WINDOW_MS;
504-
}
505-
506-
private isAccountSelectableForFamily(
507-
account: ManagedAccount,
508-
family: ModelFamily,
509-
model?: string | null,
510-
requireFresh = false,
511-
): boolean {
512-
if (account.enabled === false) return false;
513-
clearExpiredRateLimits(account);
514-
if (isRateLimitedForFamily(account, family, model) || this.isAccountCoolingDown(account)) {
515-
return false;
516-
}
517-
return !requireFresh || this.hasFreshAccessToken(account);
518-
}
519-
520-
private hasFreshAvailableAccountForFamily(
521-
family: ModelFamily,
522-
model?: string | null,
523-
): boolean {
524-
return this.accounts.some(
525-
(account) =>
526-
Boolean(account) &&
527-
this.isAccountSelectableForFamily(account, family, model) &&
528-
this.hasFreshAccessToken(account),
529-
);
530-
}
531-
532527
isAccountAvailableForFamily(index: number, family: ModelFamily, model?: string | null): boolean {
533528
const account = this.getAccountByIndex(index);
534529
if (!account) return false;
535-
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
536-
return this.isAccountSelectableForFamily(account, family, model, requireFresh);
530+
if (account.enabled === false) return false;
531+
clearExpiredRateLimits(account);
532+
return (
533+
!isRateLimitedForFamily(account, family, model) &&
534+
!this.isAccountCoolingDown(account) &&
535+
this.isCircuitAvailable(account)
536+
);
537537
}
538538

539539
setActiveIndex(index: number): ManagedAccount | null {
@@ -593,13 +593,19 @@ export class AccountManager {
593593
if (count === 0) return null;
594594

595595
const cursor = this.cursorByFamily[family];
596-
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
597596

598597
for (let i = 0; i < count; i++) {
599598
const idx = (cursor + i) % count;
600599
const account = this.accounts[idx];
601600
if (!account) continue;
602-
if (!this.isAccountSelectableForFamily(account, family, model, requireFresh)) {
601+
if (account.enabled === false) continue;
602+
603+
clearExpiredRateLimits(account);
604+
if (
605+
isRateLimitedForFamily(account, family, model) ||
606+
this.isAccountCoolingDown(account) ||
607+
!this.isCircuitAvailable(account)
608+
) {
603609
continue;
604610
}
605611

@@ -617,13 +623,19 @@ export class AccountManager {
617623
if (count === 0) return null;
618624

619625
const cursor = this.cursorByFamily[family];
620-
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
621626

622627
for (let i = 0; i < count; i++) {
623628
const idx = (cursor + i) % count;
624629
const account = this.accounts[idx];
625630
if (!account) continue;
626-
if (!this.isAccountSelectableForFamily(account, family, model, requireFresh)) {
631+
if (account.enabled === false) continue;
632+
633+
clearExpiredRateLimits(account);
634+
if (
635+
isRateLimitedForFamily(account, family, model) ||
636+
this.isAccountCoolingDown(account) ||
637+
!this.isCircuitAvailable(account)
638+
) {
627639
continue;
628640
}
629641

@@ -638,29 +650,6 @@ export class AccountManager {
638650
getCurrentOrNextForFamilyHybrid(family: ModelFamily, model?: string | null, options?: HybridSelectionOptions): ManagedAccount | null {
639651
const count = this.accounts.length;
640652
if (count === 0) return null;
641-
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
642-
643-
const currentIndex = this.currentAccountIndexByFamily[family];
644-
if (currentIndex >= 0 && currentIndex < count) {
645-
const currentAccount = this.accounts[currentIndex];
646-
if (currentAccount) {
647-
if (currentAccount.enabled === false) {
648-
// Fall through to hybrid selection.
649-
} else {
650-
if (
651-
this.isAccountSelectableForFamily(
652-
currentAccount,
653-
family,
654-
model,
655-
requireFresh,
656-
)
657-
) {
658-
currentAccount.lastUsed = nowMs();
659-
return currentAccount;
660-
}
661-
}
662-
}
663-
}
664653

665654
const quotaKey = model ? `${family}:${model}` : family;
666655
const healthTracker = getHealthTracker();
@@ -669,14 +658,16 @@ export class AccountManager {
669658
const accountsWithMetrics: AccountWithMetrics[] = this.accounts
670659
.map((account): AccountWithMetrics | null => {
671660
if (!account) return null;
661+
if (account.enabled === false) return null;
662+
clearExpiredRateLimits(account);
663+
const isAvailable =
664+
!isRateLimitedForFamily(account, family, model) &&
665+
!this.isAccountCoolingDown(account) &&
666+
this.isCircuitAvailable(account);
672667
return {
673668
index: account.index,
674-
isAvailable: this.isAccountSelectableForFamily(
675-
account,
676-
family,
677-
model,
678-
requireFresh,
679-
),
669+
trackerKey: getRuntimeTrackerKey(account),
670+
isAvailable,
680671
lastUsed: account.lastUsed,
681672
};
682673
})
@@ -697,7 +688,7 @@ export class AccountManager {
697688
recordSuccess(account: ManagedAccount, family: ModelFamily, model?: string | null): void {
698689
const quotaKey = model ? `${family}:${model}` : family;
699690
const healthTracker = getHealthTracker();
700-
healthTracker.recordSuccess(account.index, quotaKey);
691+
healthTracker.recordSuccess(getRuntimeTrackerKey(account), quotaKey);
701692
const hadCooldownMetadata =
702693
account.coolingDownUntil !== undefined || account.cooldownReason !== undefined;
703694
const hadAuthFailures = (account.consecutiveAuthFailures ?? 0) > 0;
@@ -717,26 +708,40 @@ export class AccountManager {
717708
if (healed) {
718709
this.saveToDiskDebounced();
719710
}
711+
getCircuitBreaker(getAccountCircuitKey(account)).recordSuccess();
720712
}
721713

722714
recordRateLimit(account: ManagedAccount, family: ModelFamily, model?: string | null): void {
723715
const quotaKey = model ? `${family}:${model}` : family;
724716
const healthTracker = getHealthTracker();
725717
const tokenTracker = getTokenTracker();
726-
healthTracker.recordRateLimit(account.index, quotaKey);
727-
tokenTracker.drain(account.index, quotaKey);
718+
const trackerKey = getRuntimeTrackerKey(account);
719+
healthTracker.recordRateLimit(trackerKey, quotaKey);
720+
tokenTracker.drain(trackerKey, quotaKey);
728721
}
729722

730723
recordFailure(account: ManagedAccount, family: ModelFamily, model?: string | null): void {
731724
const quotaKey = model ? `${family}:${model}` : family;
732725
const healthTracker = getHealthTracker();
733-
healthTracker.recordFailure(account.index, quotaKey);
726+
healthTracker.recordFailure(getRuntimeTrackerKey(account), quotaKey);
727+
getCircuitBreaker(getAccountCircuitKey(account)).recordFailure();
734728
}
735729

736730
consumeToken(account: ManagedAccount, family: ModelFamily, model?: string | null): boolean {
737731
const quotaKey = model ? `${family}:${model}` : family;
738732
const tokenTracker = getTokenTracker();
739-
return tokenTracker.tryConsume(account.index, quotaKey);
733+
const trackerKey = getRuntimeTrackerKey(account);
734+
if (!tokenTracker.tryConsume(trackerKey, quotaKey)) {
735+
return false;
736+
}
737+
738+
try {
739+
getCircuitBreaker(getAccountCircuitKey(account)).canExecute();
740+
return true;
741+
} catch {
742+
tokenTracker.refundToken(trackerKey, quotaKey);
743+
return false;
744+
}
740745
}
741746

742747
/**
@@ -747,7 +752,7 @@ export class AccountManager {
747752
refundToken(account: ManagedAccount, family: ModelFamily, model?: string | null): boolean {
748753
const quotaKey = model ? `${family}:${model}` : family;
749754
const tokenTracker = getTokenTracker();
750-
return tokenTracker.refundToken(account.index, quotaKey);
755+
return tokenTracker.refundToken(getRuntimeTrackerKey(account), quotaKey);
751756
}
752757

753758
markSwitched(account: ManagedAccount, reason: "rate-limit" | "initial" | "rotation", family: ModelFamily): void {
@@ -770,9 +775,14 @@ export class AccountManager {
770775
const resetAt = nowMs() + retryMs;
771776

772777
const baseKey = getQuotaKey(family);
773-
account.rateLimitResetTimes[baseKey] = resetAt;
778+
if (!model || reason === "quota" || reason === "unknown") {
779+
account.rateLimitResetTimes[baseKey] = resetAt;
780+
}
774781

775-
if (model) {
782+
if (
783+
model &&
784+
(reason === "tokens" || reason === "concurrent" || reason === "unknown")
785+
) {
776786
const modelKey = getQuotaKey(family, model);
777787
account.rateLimitResetTimes[modelKey] = resetAt;
778788
}
@@ -800,6 +810,10 @@ export class AccountManager {
800810
delete account.cooldownReason;
801811
}
802812

813+
private isCircuitAvailable(account: ManagedAccount): boolean {
814+
return getCircuitBreaker(getAccountCircuitKey(account)).isAvailable();
815+
}
816+
803817
incrementAuthFailures(account: ManagedAccount): number {
804818
account.consecutiveAuthFailures = (account.consecutiveAuthFailures ?? 0) + 1;
805819
return account.consecutiveAuthFailures;
@@ -1021,7 +1035,11 @@ export class AccountManager {
10211035
const enabledAccounts = this.accounts.filter((account) => account.enabled !== false);
10221036
const available = enabledAccounts.filter((account) => {
10231037
clearExpiredRateLimits(account);
1024-
return !isRateLimitedForFamily(account, family, model) && !this.isAccountCoolingDown(account);
1038+
return (
1039+
!isRateLimitedForFamily(account, family, model) &&
1040+
!this.isAccountCoolingDown(account) &&
1041+
this.isCircuitAvailable(account)
1042+
);
10251043
});
10261044
if (available.length > 0) return 0;
10271045
if (enabledAccounts.length === 0) return 0;
@@ -1031,20 +1049,32 @@ export class AccountManager {
10311049
const modelKey = model ? getQuotaKey(family, model) : null;
10321050

10331051
for (const account of enabledAccounts) {
1052+
const perAccountWaitTimes: number[] = [];
10341053
const baseResetAt = account.rateLimitResetTimes[baseKey];
10351054
if (typeof baseResetAt === "number") {
1036-
waitTimes.push(Math.max(0, baseResetAt - now));
1055+
perAccountWaitTimes.push(Math.max(0, baseResetAt - now));
10371056
}
10381057

10391058
if (modelKey) {
10401059
const modelResetAt = account.rateLimitResetTimes[modelKey];
10411060
if (typeof modelResetAt === "number") {
1042-
waitTimes.push(Math.max(0, modelResetAt - now));
1061+
perAccountWaitTimes.push(Math.max(0, modelResetAt - now));
10431062
}
10441063
}
10451064

10461065
if (typeof account.coolingDownUntil === "number") {
1047-
waitTimes.push(Math.max(0, account.coolingDownUntil - now));
1066+
perAccountWaitTimes.push(Math.max(0, account.coolingDownUntil - now));
1067+
}
1068+
1069+
const breakerWait = getCircuitBreaker(
1070+
getAccountCircuitKey(account),
1071+
).getTimeUntilAvailable();
1072+
if (breakerWait > 0) {
1073+
perAccountWaitTimes.push(breakerWait);
1074+
}
1075+
1076+
if (perAccountWaitTimes.length > 0) {
1077+
waitTimes.push(Math.max(...perAccountWaitTimes));
10481078
}
10491079
}
10501080

0 commit comments

Comments
 (0)