Skip to content

Commit 9cbc72a

Browse files
ndycodendycode
committed
fix: harden fast-expiry token refresh persistence
1 parent 53488ef commit 9cbc72a

9 files changed

Lines changed: 733 additions & 109 deletions

‎index.ts‎

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,7 @@ import {
177177
import { applyFastSessionDefaults } from "./lib/request/request-transformer.js";
178178
import { applyResponseCompaction } from "./lib/request/response-compaction.js";
179179
import { isEmptyResponse } from "./lib/request/response-handler.js";
180+
import { CodexAuthError } from "./lib/errors.js";
180181
import {
181182
parseRetryAfterHintMs,
182183
sanitizeResponseHeadersForLog,
@@ -997,9 +998,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
997998
accountAuth,
998999
client,
9991000
)) as OAuthAuthDetails;
1000-
accountManager.updateFromAuth(account, accountAuth);
1001-
accountManager.clearAuthFailures(account);
1002-
accountManager.saveToDiskDebounced();
1001+
await accountManager.commitRefreshedAuth(
1002+
account,
1003+
accountAuth,
1004+
);
10031005
}
10041006
} catch (err) {
10051007
logDebug(
@@ -1010,18 +1012,48 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
10101012
runtimeMetrics.accountRotations++;
10111013
runtimeMetrics.lastError =
10121014
(err as Error)?.message ?? String(err);
1013-
const failures =
1014-
accountManager.incrementAuthFailures(account);
1015+
const isRetryableRefreshError =
1016+
err instanceof CodexAuthError && err.retryable;
10151017
const accountLabel = formatAccountLabel(
10161018
account,
10171019
account.index,
10181020
);
10191021

1022+
sessionAffinityStore?.forgetSession(sessionAffinityKey);
1023+
1024+
if (isRetryableRefreshError) {
1025+
runtimeMetrics.networkErrors++;
1026+
const retryableRefreshPolicy = evaluateFailurePolicy(
1027+
{ kind: "network", failoverMode },
1028+
{ networkCooldownMs: networkErrorCooldownMs },
1029+
);
1030+
if (retryableRefreshPolicy.recordFailure) {
1031+
accountManager.recordFailure(
1032+
account,
1033+
modelFamily,
1034+
model,
1035+
);
1036+
}
1037+
if (
1038+
typeof retryableRefreshPolicy.cooldownMs === "number" &&
1039+
retryableRefreshPolicy.cooldownReason
1040+
) {
1041+
accountManager.markAccountCoolingDown(
1042+
account,
1043+
retryableRefreshPolicy.cooldownMs,
1044+
retryableRefreshPolicy.cooldownReason,
1045+
);
1046+
accountManager.saveToDiskDebounced();
1047+
}
1048+
continue;
1049+
}
1050+
1051+
const failures =
1052+
accountManager.incrementAuthFailures(account);
10201053
const authFailurePolicy = evaluateFailurePolicy({
10211054
kind: "auth-refresh",
10221055
consecutiveAuthFailures: failures,
10231056
});
1024-
sessionAffinityStore?.forgetSession(sessionAffinityKey);
10251057

10261058
if (authFailurePolicy.removeAccount) {
10271059
const removedIndex = account.index;
@@ -1945,14 +1977,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
19451977
fallbackAuth,
19461978
client,
19471979
)) as OAuthAuthDetails;
1948-
accountManager.updateFromAuth(
1980+
await accountManager.commitRefreshedAuth(
19491981
fallbackAccount,
19501982
fallbackAuth,
19511983
);
1952-
accountManager.clearAuthFailures(
1953-
fallbackAccount,
1954-
);
1955-
accountManager.saveToDiskDebounced();
19561984
}
19571985
} catch (refreshError) {
19581986
logWarn(

‎lib/accounts.ts‎

Lines changed: 194 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
type CooldownReason,
88
type RateLimitStateV3,
99
findMatchingAccountIndex,
10+
withAccountStorageTransaction,
1011
} from "./storage.js";
1112
import type { AccountIdSource, OAuthAuthDetails } from "./types.js";
1213
import { MODEL_FAMILIES, type ModelFamily } from "./prompts/codex.js";
@@ -81,6 +82,81 @@ function initFamilyState(defaultValue: number): Record<ModelFamily, number> {
8182
) as Record<ModelFamily, number>;
8283
}
8384

85+
type AccountIdentityCandidate = Pick<
86+
ManagedAccount,
87+
"accountId" | "email" | "refreshToken"
88+
> & {
89+
index?: number;
90+
};
91+
92+
function getAuthIdentityCandidate(
93+
auth: OAuthAuthDetails | undefined,
94+
): AccountIdentityCandidate {
95+
const accountId = extractAccountId(auth?.access)?.trim() || undefined;
96+
const email = sanitizeEmail(extractAccountEmail(auth?.access));
97+
return {
98+
accountId,
99+
email,
100+
refreshToken: auth?.refresh,
101+
};
102+
}
103+
104+
function buildAccountIdentityCandidates(
105+
source: AccountIdentityCandidate,
106+
auth?: OAuthAuthDetails,
107+
): AccountIdentityCandidate[] {
108+
const derived = getAuthIdentityCandidate(auth);
109+
const candidates: AccountIdentityCandidate[] = [];
110+
const seen = new Set<string>();
111+
112+
const pushCandidate = (candidate: AccountIdentityCandidate): void => {
113+
const key = `${candidate.accountId ?? ""}|${candidate.email ?? ""}|${candidate.refreshToken ?? ""}`;
114+
if (seen.has(key)) return;
115+
seen.add(key);
116+
candidates.push(candidate);
117+
};
118+
119+
pushCandidate(source);
120+
pushCandidate({
121+
accountId: source.accountId ?? derived.accountId,
122+
email: source.email ?? derived.email,
123+
refreshToken: source.refreshToken,
124+
index: source.index,
125+
});
126+
pushCandidate({
127+
accountId: derived.accountId ?? source.accountId,
128+
email: derived.email ?? source.email,
129+
refreshToken: source.refreshToken,
130+
index: source.index,
131+
});
132+
pushCandidate({
133+
accountId: derived.accountId ?? source.accountId,
134+
email: derived.email ?? source.email,
135+
refreshToken: derived.refreshToken ?? source.refreshToken,
136+
index: source.index,
137+
});
138+
139+
return candidates;
140+
}
141+
142+
function findAccountIndexByIdentity<
143+
T extends Pick<AccountIdentityCandidate, "accountId" | "email" | "refreshToken">,
144+
>(
145+
accounts: readonly T[],
146+
source: AccountIdentityCandidate,
147+
auth?: OAuthAuthDetails,
148+
): number | undefined {
149+
for (const candidate of buildAccountIdentityCandidates(source, auth)) {
150+
const matchIndex = findMatchingAccountIndex(accounts, candidate, {
151+
allowUniqueAccountIdFallbackWithoutEmail: true,
152+
});
153+
if (matchIndex !== undefined) {
154+
return matchIndex;
155+
}
156+
}
157+
return undefined;
158+
}
159+
84160
export interface Workspace {
85161
id: string;
86162
name?: string;
@@ -642,6 +718,17 @@ export class AccountManager {
642718
account.consecutiveAuthFailures = 0;
643719
}
644720

721+
getAccountByIdentity(
722+
candidate: AccountIdentityCandidate,
723+
auth?: OAuthAuthDetails,
724+
): ManagedAccount | null {
725+
const index = findAccountIndexByIdentity(this.accounts, candidate, auth);
726+
if (index === undefined) {
727+
return null;
728+
}
729+
return this.accounts[index] ?? null;
730+
}
731+
645732
shouldShowAccountToast(accountIndex: number, debounceMs = 30000): boolean {
646733
const now = nowMs();
647734
if (accountIndex === this.lastToastAccountIndex && now - this.lastToastTime < debounceMs) {
@@ -670,6 +757,112 @@ export class AccountManager {
670757
account.email = sanitizeEmail(extractAccountEmail(auth.access)) ?? account.email;
671758
}
672759

760+
private buildStorageSnapshot(): AccountStorageV3 {
761+
const activeIndexByFamily: Partial<Record<ModelFamily, number>> = {};
762+
for (const family of MODEL_FAMILIES) {
763+
const raw = this.currentAccountIndexByFamily[family];
764+
activeIndexByFamily[family] = clampNonNegativeInt(raw, 0);
765+
}
766+
767+
const activeIndex = clampNonNegativeInt(activeIndexByFamily.codex, 0);
768+
769+
return {
770+
version: 3,
771+
accounts: this.accounts.map((account) => ({
772+
accountId: account.accountId,
773+
accountIdSource: account.accountIdSource,
774+
accountLabel: account.accountLabel,
775+
email: account.email,
776+
refreshToken: account.refreshToken,
777+
accessToken: account.access,
778+
expiresAt: account.expires,
779+
enabled: account.enabled === false ? false : undefined,
780+
addedAt: account.addedAt,
781+
lastUsed: account.lastUsed,
782+
lastSwitchReason: account.lastSwitchReason,
783+
rateLimitResetTimes:
784+
Object.keys(account.rateLimitResetTimes).length > 0 ? account.rateLimitResetTimes : undefined,
785+
coolingDownUntil: account.coolingDownUntil,
786+
cooldownReason: account.cooldownReason,
787+
workspaces: account.workspaces,
788+
currentWorkspaceIndex: account.currentWorkspaceIndex,
789+
})),
790+
activeIndex,
791+
activeIndexByFamily,
792+
};
793+
}
794+
795+
async commitRefreshedAuth(
796+
source: Pick<
797+
ManagedAccount,
798+
"index" | "accountId" | "email" | "refreshToken"
799+
>,
800+
auth: OAuthAuthDetails,
801+
): Promise<ManagedAccount | null> {
802+
const nextAccountId = extractAccountId(auth.access)?.trim() || undefined;
803+
const nextEmail = sanitizeEmail(extractAccountEmail(auth.access));
804+
805+
await withAccountStorageTransaction(async (current, persist) => {
806+
const nextStorage = structuredClone(
807+
current ?? this.buildStorageSnapshot(),
808+
) as AccountStorageV3;
809+
const storageIndex = findAccountIndexByIdentity(
810+
nextStorage.accounts,
811+
source,
812+
auth,
813+
);
814+
if (storageIndex === undefined) {
815+
log.warn("Unable to resolve refreshed account for persistence", {
816+
sourceIndex: source.index,
817+
email: source.email,
818+
});
819+
return;
820+
}
821+
822+
const storedAccount = nextStorage.accounts[storageIndex];
823+
if (!storedAccount) {
824+
return;
825+
}
826+
827+
storedAccount.refreshToken = auth.refresh;
828+
storedAccount.accessToken = auth.access;
829+
storedAccount.expiresAt = auth.expires;
830+
if (
831+
nextAccountId &&
832+
shouldUpdateAccountIdFromToken(
833+
storedAccount.accountIdSource,
834+
storedAccount.accountId,
835+
)
836+
) {
837+
storedAccount.accountId = nextAccountId;
838+
storedAccount.accountIdSource = "token";
839+
}
840+
if (nextEmail) {
841+
storedAccount.email = nextEmail;
842+
}
843+
storedAccount.enabled = undefined;
844+
delete storedAccount.coolingDownUntil;
845+
delete storedAccount.cooldownReason;
846+
847+
await persist(nextStorage);
848+
});
849+
850+
const liveAccount = this.getAccountByIdentity(source, auth);
851+
if (!liveAccount) {
852+
log.warn("Unable to resolve refreshed live account after persistence", {
853+
sourceIndex: source.index,
854+
email: source.email,
855+
});
856+
return null;
857+
}
858+
859+
this.updateFromAuth(liveAccount, auth);
860+
liveAccount.enabled = true;
861+
this.clearAccountCooldown(liveAccount);
862+
this.clearAuthFailures(liveAccount);
863+
return liveAccount;
864+
}
865+
673866
toAuthDetails(account: ManagedAccount): Auth {
674867
return {
675868
type: "oauth",
@@ -780,40 +973,7 @@ export class AccountManager {
780973
}
781974

782975
async saveToDisk(): Promise<void> {
783-
const activeIndexByFamily: Partial<Record<ModelFamily, number>> = {};
784-
for (const family of MODEL_FAMILIES) {
785-
const raw = this.currentAccountIndexByFamily[family];
786-
activeIndexByFamily[family] = clampNonNegativeInt(raw, 0);
787-
}
788-
789-
const activeIndex = clampNonNegativeInt(activeIndexByFamily.codex, 0);
790-
791-
const storage: AccountStorageV3 = {
792-
version: 3,
793-
accounts: this.accounts.map((account) => ({
794-
accountId: account.accountId,
795-
accountIdSource: account.accountIdSource,
796-
accountLabel: account.accountLabel,
797-
email: account.email,
798-
refreshToken: account.refreshToken,
799-
accessToken: account.access,
800-
expiresAt: account.expires,
801-
enabled: account.enabled === false ? false : undefined,
802-
addedAt: account.addedAt,
803-
lastUsed: account.lastUsed,
804-
lastSwitchReason: account.lastSwitchReason,
805-
rateLimitResetTimes:
806-
Object.keys(account.rateLimitResetTimes).length > 0 ? account.rateLimitResetTimes : undefined,
807-
coolingDownUntil: account.coolingDownUntil,
808-
cooldownReason: account.cooldownReason,
809-
workspaces: account.workspaces,
810-
currentWorkspaceIndex: account.currentWorkspaceIndex,
811-
})),
812-
activeIndex,
813-
activeIndexByFamily,
814-
};
815-
816-
await saveAccounts(storage);
976+
await saveAccounts(this.buildStorageSnapshot());
817977
}
818978

819979
saveToDiskDebounced(delayMs = 500): void {

0 commit comments

Comments
 (0)