Skip to content

Commit 66d2d63

Browse files
ndycodendycode
committed
fix: prefer fresh access tokens during account selection
1 parent 53488ef commit 66d2d63

3 files changed

Lines changed: 198 additions & 19 deletions

File tree

‎lib/accounts.ts‎

Lines changed: 54 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ import {
7474
} from "./accounts/rate-limits.js";
7575

7676
const log = createLogger("accounts");
77+
const ACCOUNT_SELECTION_FRESH_WINDOW_MS = 5 * 60_000;
7778

7879
function initFamilyState(defaultValue: number): Record<ModelFamily, number> {
7980
return Object.fromEntries(
@@ -381,12 +382,45 @@ export class AccountManager {
381382
return account ?? null;
382383
}
383384

385+
private hasFreshAccessToken(account: ManagedAccount): boolean {
386+
if (!account.access) return false;
387+
if (typeof account.expires !== "number" || !Number.isFinite(account.expires)) {
388+
return false;
389+
}
390+
return account.expires > nowMs() + ACCOUNT_SELECTION_FRESH_WINDOW_MS;
391+
}
392+
393+
private isAccountSelectableForFamily(
394+
account: ManagedAccount,
395+
family: ModelFamily,
396+
model?: string | null,
397+
requireFresh = false,
398+
): boolean {
399+
if (account.enabled === false) return false;
400+
clearExpiredRateLimits(account);
401+
if (isRateLimitedForFamily(account, family, model) || this.isAccountCoolingDown(account)) {
402+
return false;
403+
}
404+
return !requireFresh || this.hasFreshAccessToken(account);
405+
}
406+
407+
private hasFreshAvailableAccountForFamily(
408+
family: ModelFamily,
409+
model?: string | null,
410+
): boolean {
411+
return this.accounts.some(
412+
(account) =>
413+
Boolean(account) &&
414+
this.isAccountSelectableForFamily(account, family, model) &&
415+
this.hasFreshAccessToken(account),
416+
);
417+
}
418+
384419
isAccountAvailableForFamily(index: number, family: ModelFamily, model?: string | null): boolean {
385420
const account = this.getAccountByIndex(index);
386421
if (!account) return false;
387-
if (account.enabled === false) return false;
388-
clearExpiredRateLimits(account);
389-
return !isRateLimitedForFamily(account, family, model) && !this.isAccountCoolingDown(account);
422+
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
423+
return this.isAccountSelectableForFamily(account, family, model, requireFresh);
390424
}
391425

392426
setActiveIndex(index: number): ManagedAccount | null {
@@ -446,15 +480,13 @@ export class AccountManager {
446480
if (count === 0) return null;
447481

448482
const cursor = this.cursorByFamily[family];
483+
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
449484

450485
for (let i = 0; i < count; i++) {
451486
const idx = (cursor + i) % count;
452487
const account = this.accounts[idx];
453488
if (!account) continue;
454-
if (account.enabled === false) continue;
455-
456-
clearExpiredRateLimits(account);
457-
if (isRateLimitedForFamily(account, family, model) || this.isAccountCoolingDown(account)) {
489+
if (!this.isAccountSelectableForFamily(account, family, model, requireFresh)) {
458490
continue;
459491
}
460492

@@ -472,15 +504,13 @@ export class AccountManager {
472504
if (count === 0) return null;
473505

474506
const cursor = this.cursorByFamily[family];
507+
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
475508

476509
for (let i = 0; i < count; i++) {
477510
const idx = (cursor + i) % count;
478511
const account = this.accounts[idx];
479512
if (!account) continue;
480-
if (account.enabled === false) continue;
481-
482-
clearExpiredRateLimits(account);
483-
if (isRateLimitedForFamily(account, family, model) || this.isAccountCoolingDown(account)) {
513+
if (!this.isAccountSelectableForFamily(account, family, model, requireFresh)) {
484514
continue;
485515
}
486516

@@ -495,6 +525,7 @@ export class AccountManager {
495525
getCurrentOrNextForFamilyHybrid(family: ModelFamily, model?: string | null, options?: HybridSelectionOptions): ManagedAccount | null {
496526
const count = this.accounts.length;
497527
if (count === 0) return null;
528+
const requireFresh = this.hasFreshAvailableAccountForFamily(family, model);
498529

499530
const currentIndex = this.currentAccountIndexByFamily[family];
500531
if (currentIndex >= 0 && currentIndex < count) {
@@ -503,10 +534,13 @@ export class AccountManager {
503534
if (currentAccount.enabled === false) {
504535
// Fall through to hybrid selection.
505536
} else {
506-
clearExpiredRateLimits(currentAccount);
507537
if (
508-
!isRateLimitedForFamily(currentAccount, family, model) &&
509-
!this.isAccountCoolingDown(currentAccount)
538+
this.isAccountSelectableForFamily(
539+
currentAccount,
540+
family,
541+
model,
542+
requireFresh,
543+
)
510544
) {
511545
currentAccount.lastUsed = nowMs();
512546
return currentAccount;
@@ -522,13 +556,14 @@ export class AccountManager {
522556
const accountsWithMetrics: AccountWithMetrics[] = this.accounts
523557
.map((account): AccountWithMetrics | null => {
524558
if (!account) return null;
525-
if (account.enabled === false) return null;
526-
clearExpiredRateLimits(account);
527-
const isAvailable =
528-
!isRateLimitedForFamily(account, family, model) && !this.isAccountCoolingDown(account);
529559
return {
530560
index: account.index,
531-
isAvailable,
561+
isAvailable: this.isAccountSelectableForFamily(
562+
account,
563+
family,
564+
model,
565+
requireFresh,
566+
),
532567
lastUsed: account.lastUsed,
533568
};
534569
})

‎test/accounts-load-from-disk.test.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,4 +253,31 @@ describe("AccountManager loadFromDisk", () => {
253253

254254
expect(manager.getNextForFamily("codex")).toBeNull();
255255
});
256+
257+
it("getNextForFamily falls back to stale accounts when no fresh option exists", () => {
258+
const now = Date.now();
259+
const manager = new AccountManager(undefined, {
260+
version: 3 as const,
261+
activeIndex: 0,
262+
accounts: [
263+
{
264+
refreshToken: "stale-1",
265+
accessToken: "access-1",
266+
expiresAt: now + 60_000,
267+
addedAt: now,
268+
lastUsed: now,
269+
},
270+
{
271+
refreshToken: "stale-2",
272+
accessToken: "access-2",
273+
expiresAt: now + 120_000,
274+
addedAt: now,
275+
lastUsed: now - 5_000,
276+
},
277+
],
278+
} as never);
279+
280+
const selected = manager.getNextForFamily("codex");
281+
expect(selected?.refreshToken).toBe("stale-1");
282+
});
256283
});

‎test/accounts.test.ts‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,62 @@ describe("AccountManager", () => {
406406
expect(account?.rateLimitResetTimes?.codex).toBeUndefined();
407407
});
408408

409+
it("prefers accounts with fresh access tokens in availability checks", () => {
410+
const now = Date.now();
411+
const stored = {
412+
version: 3 as const,
413+
activeIndex: 0,
414+
accounts: [
415+
{
416+
refreshToken: "token-stale",
417+
accessToken: "access-stale",
418+
expiresAt: now + 60_000,
419+
addedAt: now,
420+
lastUsed: now,
421+
},
422+
{
423+
refreshToken: "token-fresh",
424+
accessToken: "access-fresh",
425+
expiresAt: now + 10 * 60_000,
426+
addedAt: now,
427+
lastUsed: now,
428+
},
429+
],
430+
};
431+
const manager = new AccountManager(undefined, stored);
432+
433+
expect(manager.isAccountAvailableForFamily(0, "codex")).toBe(false);
434+
expect(manager.isAccountAvailableForFamily(1, "codex")).toBe(true);
435+
});
436+
437+
it("keeps stale accounts available when the whole pool is stale", () => {
438+
const now = Date.now();
439+
const stored = {
440+
version: 3 as const,
441+
activeIndex: 0,
442+
accounts: [
443+
{
444+
refreshToken: "token-stale-1",
445+
accessToken: "access-stale-1",
446+
expiresAt: now + 60_000,
447+
addedAt: now,
448+
lastUsed: now,
449+
},
450+
{
451+
refreshToken: "token-stale-2",
452+
accessToken: "access-stale-2",
453+
expiresAt: now + 120_000,
454+
addedAt: now,
455+
lastUsed: now,
456+
},
457+
],
458+
};
459+
const manager = new AccountManager(undefined, stored);
460+
461+
expect(manager.isAccountAvailableForFamily(0, "codex")).toBe(true);
462+
expect(manager.isAccountAvailableForFamily(1, "codex")).toBe(true);
463+
});
464+
409465
it("rotates when the active account is rate-limited", () => {
410466
const now = Date.now();
411467
const stored = {
@@ -614,6 +670,36 @@ describe("AccountManager", () => {
614670
expect(gpt51Second?.refreshToken).toBe("token-2");
615671
});
616672

673+
it("skips a stale active account when a fresher account is available", () => {
674+
const now = Date.now();
675+
const stored = {
676+
version: 3 as const,
677+
activeIndex: 0,
678+
activeIndexByFamily: { codex: 0 },
679+
accounts: [
680+
{
681+
refreshToken: "token-stale",
682+
accessToken: "access-stale",
683+
expiresAt: now + 60_000,
684+
addedAt: now,
685+
lastUsed: now,
686+
},
687+
{
688+
refreshToken: "token-fresh",
689+
accessToken: "access-fresh",
690+
expiresAt: now + 10 * 60_000,
691+
addedAt: now,
692+
lastUsed: now - 5_000,
693+
},
694+
],
695+
};
696+
697+
const manager = new AccountManager(undefined, stored as never);
698+
const selected = manager.getCurrentOrNextForFamily("codex");
699+
700+
expect(selected?.refreshToken).toBe("token-fresh");
701+
});
702+
617703
it("hybrid selection prefers active index when available", () => {
618704
const now = Date.now();
619705
const stored = {
@@ -2085,6 +2171,37 @@ describe("AccountManager", () => {
20852171
expect(secondCall?.index).toBe(selected?.index);
20862172
});
20872173

2174+
it("prefers a fresh alternate account over a stale current account", () => {
2175+
const now = Date.now();
2176+
const stored = {
2177+
version: 3 as const,
2178+
activeIndex: 0,
2179+
activeIndexByFamily: { codex: 0 },
2180+
accounts: [
2181+
{
2182+
refreshToken: "token-stale",
2183+
accessToken: "access-stale",
2184+
expiresAt: now + 60_000,
2185+
addedAt: now,
2186+
lastUsed: now,
2187+
},
2188+
{
2189+
refreshToken: "token-fresh",
2190+
accessToken: "access-fresh",
2191+
expiresAt: now + 10 * 60_000,
2192+
addedAt: now,
2193+
lastUsed: now - 5_000,
2194+
},
2195+
],
2196+
};
2197+
2198+
const manager = new AccountManager(undefined, stored as never);
2199+
const selected = manager.getCurrentOrNextForFamilyHybrid("codex");
2200+
2201+
expect(selected?.refreshToken).toBe("token-fresh");
2202+
expect(selected?.index).toBe(1);
2203+
});
2204+
20882205
it("falls back to least-recently-used when all accounts are rate-limited", () => {
20892206
const now = Date.now();
20902207
const stored = {

0 commit comments

Comments
 (0)