Skip to content

Commit efd08ad

Browse files
ndycodendycode
committed
fix: harden refresh persistence follow-ups
1 parent c809a99 commit efd08ad

6 files changed

Lines changed: 345 additions & 79 deletions

File tree

‎lib/accounts.ts‎

Lines changed: 62 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ import {
1919
type HybridSelectionOptions,
2020
} from "./rotation.js";
2121
import { nowMs } from "./utils.js";
22+
import { ERROR_MESSAGES } from "./constants.js";
23+
import { CodexAuthError } from "./errors.js";
2224
import {
2325
loadCodexCliState,
2426
type CodexCliTokenCacheEntry,
@@ -801,66 +803,70 @@ export class AccountManager {
801803
): Promise<ManagedAccount | null> {
802804
const nextAccountId = extractAccountId(auth.access)?.trim() || undefined;
803805
const nextEmail = sanitizeEmail(extractAccountEmail(auth.access));
806+
try {
807+
return await withAccountStorageTransaction(async (current, persist) => {
808+
const nextStorage = structuredClone(
809+
current ?? this.buildStorageSnapshot(),
810+
) as AccountStorageV3;
811+
const storageIndex = findAccountIndexByIdentity(
812+
nextStorage.accounts,
813+
source,
814+
auth,
815+
);
816+
if (storageIndex === undefined) {
817+
log.warn("Unable to resolve refreshed account for persistence", {
818+
sourceIndex: source.index,
819+
});
820+
return null;
821+
}
804822

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-
}
823+
const storedAccount = nextStorage.accounts[storageIndex];
824+
if (!storedAccount) {
825+
return null;
826+
}
821827

822-
const storedAccount = nextStorage.accounts[storageIndex];
823-
if (!storedAccount) {
824-
return;
825-
}
828+
storedAccount.refreshToken = auth.refresh;
829+
storedAccount.accessToken = auth.access;
830+
storedAccount.expiresAt = auth.expires;
831+
if (
832+
nextAccountId &&
833+
shouldUpdateAccountIdFromToken(
834+
storedAccount.accountIdSource,
835+
storedAccount.accountId,
836+
)
837+
) {
838+
storedAccount.accountId = nextAccountId;
839+
storedAccount.accountIdSource = "token";
840+
}
841+
if (nextEmail) {
842+
storedAccount.email = nextEmail;
843+
}
844+
storedAccount.enabled = undefined;
845+
delete storedAccount.coolingDownUntil;
846+
delete storedAccount.cooldownReason;
826847

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;
848+
await persist(nextStorage);
846849

847-
await persist(nextStorage);
848-
});
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+
});
855+
return null;
856+
}
849857

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,
858+
this.updateFromAuth(liveAccount, auth);
859+
liveAccount.enabled = true;
860+
this.clearAccountCooldown(liveAccount);
861+
this.clearAuthFailures(liveAccount);
862+
return liveAccount;
863+
});
864+
} catch (error) {
865+
throw new CodexAuthError(ERROR_MESSAGES.TOKEN_REFRESH_FAILED, {
866+
retryable: true,
867+
cause: error,
855868
});
856-
return null;
857869
}
858-
859-
this.updateFromAuth(liveAccount, auth);
860-
liveAccount.enabled = true;
861-
this.clearAccountCooldown(liveAccount);
862-
this.clearAuthFailures(liveAccount);
863-
return liveAccount;
864870
}
865871

866872
toAuthDetails(account: ManagedAccount): Auth {
@@ -973,7 +979,9 @@ export class AccountManager {
973979
}
974980

975981
async saveToDisk(): Promise<void> {
976-
await saveAccounts(this.buildStorageSnapshot());
982+
await withAccountStorageTransaction(async (_current, persist) => {
983+
await persist(this.buildStorageSnapshot());
984+
});
977985
}
978986

979987
saveToDiskDebounced(delayMs = 500): void {

‎lib/refresh-guardian.ts‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { createLogger } from "./logger.js";
2-
import { applyRefreshResult, refreshExpiringAccounts } from "./proactive-refresh.js";
2+
import { refreshExpiringAccounts } from "./proactive-refresh.js";
33
import type { AccountManager } from "./accounts.js";
44
import type { CooldownReason } from "./storage.js";
55
import type { TokenResult } from "./types.js";
@@ -123,12 +123,21 @@ export class RefreshGuardian {
123123
expires: result.tokenResult.expires,
124124
multiAccount: true,
125125
};
126-
const account = manager.getAccountByIdentity(sourceAccount, refreshedAuth);
127-
if (account) {
128-
applyRefreshResult(account, result.tokenResult);
129-
manager.clearAuthFailures(account);
126+
try {
127+
await manager.commitRefreshedAuth(sourceAccount, refreshedAuth);
128+
} catch (error) {
129+
log.warn("Refresh guardian commit failed", {
130+
sourceIndex: sourceAccount.index,
131+
error: error instanceof Error ? error.message : String(error),
132+
});
133+
const account = manager.getAccountByIdentity(sourceAccount, refreshedAuth);
134+
if (account) {
135+
manager.markAccountCoolingDown(account, this.bufferMs, "network-error");
136+
}
137+
this.stats.failed += 1;
138+
this.stats.networkFailed += 1;
139+
return !!account;
130140
}
131-
await manager.commitRefreshedAuth(sourceAccount, refreshedAuth);
132141
this.stats.refreshed += 1;
133142
return false;
134143
}

‎lib/request/fetch-helpers.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -435,8 +435,9 @@ function isRetryableRefreshFailure(
435435
case "network_error":
436436
case "unknown":
437437
case "invalid_response":
438-
case "missing_refresh":
439438
return true;
439+
case "missing_refresh":
440+
return false;
440441
case "http_error":
441442
return !(
442443
result.statusCode === HTTP_STATUS.BAD_REQUEST ||

‎test/accounts.test.ts‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
getAccountIdCandidates,
1212
} from "../lib/accounts.js";
1313
import { getHealthTracker, getTokenTracker, resetTrackers } from "../lib/rotation.js";
14+
import { CodexAuthError } from "../lib/errors.js";
1415
import type { OAuthAuthDetails } from "../lib/types.js";
1516

1617
vi.mock("../lib/storage.js", async (importOriginal) => {
@@ -22,6 +23,29 @@ vi.mock("../lib/storage.js", async (importOriginal) => {
2223
};
2324
});
2425

26+
beforeEach(async () => {
27+
resetTrackers();
28+
const { saveAccounts, withAccountStorageTransaction } = await import(
29+
"../lib/storage.js"
30+
);
31+
const mockSaveAccounts = vi.mocked(saveAccounts);
32+
const mockWithAccountStorageTransaction = vi.mocked(
33+
withAccountStorageTransaction,
34+
);
35+
36+
mockSaveAccounts.mockReset();
37+
mockSaveAccounts.mockResolvedValue(undefined);
38+
mockWithAccountStorageTransaction.mockReset();
39+
mockWithAccountStorageTransaction.mockImplementation(async (handler) => {
40+
let current = null;
41+
const persist = async (storage: Parameters<typeof saveAccounts>[0]) => {
42+
current = structuredClone(storage);
43+
await mockSaveAccounts(storage);
44+
};
45+
return handler(current as never, persist);
46+
});
47+
});
48+
2549
describe("parseRateLimitReason", () => {
2650
it("returns quota for quota-related codes", () => {
2751
expect(parseRateLimitReason("usage_limit_reached")).toBe("quota");
@@ -1134,6 +1158,144 @@ describe("AccountManager", () => {
11341158
expect(account.cooldownReason).toBeUndefined();
11351159
expect(account.consecutiveAuthFailures).toBe(0);
11361160
});
1161+
1162+
it("propagates storage write failure as retryable CodexAuthError", async () => {
1163+
const { withAccountStorageTransaction } = await import("../lib/storage.js");
1164+
const mockWithAccountStorageTransaction = vi.mocked(
1165+
withAccountStorageTransaction,
1166+
);
1167+
mockWithAccountStorageTransaction.mockRejectedValueOnce(
1168+
Object.assign(new Error("EBUSY"), { code: "EBUSY" }),
1169+
);
1170+
1171+
const now = Date.now();
1172+
const stored = {
1173+
version: 3 as const,
1174+
activeIndex: 0,
1175+
accounts: [
1176+
{
1177+
refreshToken: "old-refresh",
1178+
accessToken: "old-access",
1179+
expiresAt: now,
1180+
addedAt: now,
1181+
lastUsed: now,
1182+
},
1183+
],
1184+
};
1185+
const manager = new AccountManager(undefined, stored as any);
1186+
const account = manager.getAccountByIndex(0)!;
1187+
const refreshedAuth: OAuthAuthDetails = {
1188+
type: "oauth",
1189+
access: "header.payload.signature",
1190+
refresh: "new-refresh",
1191+
expires: now + 3_600_000,
1192+
};
1193+
1194+
const error = await manager.commitRefreshedAuth(account, refreshedAuth).catch(
1195+
(err) => err as CodexAuthError,
1196+
);
1197+
1198+
expect(error).toBeInstanceOf(CodexAuthError);
1199+
expect(error.retryable).toBe(true);
1200+
expect(account.refreshToken).toBe("old-refresh");
1201+
});
1202+
1203+
it("prevents debounced saves from writing stale auth during refresh persistence", async () => {
1204+
vi.useFakeTimers();
1205+
try {
1206+
const { saveAccounts, withAccountStorageTransaction } = await import(
1207+
"../lib/storage.js"
1208+
);
1209+
const mockSaveAccounts = vi.mocked(saveAccounts);
1210+
const mockWithAccountStorageTransaction = vi.mocked(
1211+
withAccountStorageTransaction,
1212+
);
1213+
1214+
const now = Date.now();
1215+
const stored = {
1216+
version: 3 as const,
1217+
activeIndex: 0,
1218+
accounts: [
1219+
{
1220+
refreshToken: "old-refresh",
1221+
accessToken: "old-access",
1222+
expiresAt: now,
1223+
addedAt: now,
1224+
lastUsed: now,
1225+
},
1226+
],
1227+
};
1228+
const manager = new AccountManager(undefined, stored as any);
1229+
const account = manager.getAccountByIndex(0)!;
1230+
const refreshedAuth: OAuthAuthDetails = {
1231+
type: "oauth",
1232+
access: "new-access",
1233+
refresh: "new-refresh",
1234+
expires: now + 3_600_000,
1235+
};
1236+
1237+
let storageState = structuredClone(stored) as typeof stored;
1238+
let lock = Promise.resolve();
1239+
let releasePersist!: () => void;
1240+
const persistBlocked = new Promise<void>((resolve) => {
1241+
releasePersist = resolve;
1242+
});
1243+
let notifyPersistStarted!: () => void;
1244+
const persistStarted = new Promise<void>((resolve) => {
1245+
notifyPersistStarted = resolve;
1246+
});
1247+
1248+
mockWithAccountStorageTransaction.mockImplementation((handler) => {
1249+
const run = async () => {
1250+
const current = structuredClone(storageState) as typeof storageState;
1251+
const persist = async (
1252+
nextStorage: Parameters<typeof saveAccounts>[0],
1253+
) => {
1254+
if (
1255+
nextStorage.accounts[0]?.refreshToken === "new-refresh" &&
1256+
nextStorage.accounts[0]?.accessToken === "new-access"
1257+
) {
1258+
notifyPersistStarted();
1259+
await persistBlocked;
1260+
}
1261+
storageState = structuredClone(nextStorage) as typeof storageState;
1262+
await mockSaveAccounts(nextStorage);
1263+
};
1264+
return handler(current as never, persist);
1265+
};
1266+
1267+
const result = lock.then(run, run);
1268+
lock = result.then(
1269+
() => undefined,
1270+
() => undefined,
1271+
);
1272+
return result;
1273+
});
1274+
1275+
const commitPromise = manager.commitRefreshedAuth(account, refreshedAuth);
1276+
await persistStarted;
1277+
1278+
manager.saveToDiskDebounced(0);
1279+
await vi.advanceTimersByTimeAsync(0);
1280+
1281+
releasePersist();
1282+
await commitPromise;
1283+
await manager.flushPendingSave();
1284+
1285+
expect(mockSaveAccounts).toHaveBeenCalledTimes(2);
1286+
expect(mockSaveAccounts.mock.calls[0]?.[0]?.accounts[0]?.refreshToken).toBe(
1287+
"new-refresh",
1288+
);
1289+
expect(mockSaveAccounts.mock.calls[1]?.[0]?.accounts[0]?.refreshToken).toBe(
1290+
"new-refresh",
1291+
);
1292+
expect(mockSaveAccounts.mock.calls[1]?.[0]?.accounts[0]?.accessToken).toBe(
1293+
"new-access",
1294+
);
1295+
} finally {
1296+
vi.useRealTimers();
1297+
}
1298+
});
11371299
});
11381300

11391301
describe("toAuthDetails", () => {

‎test/fetch-helpers.test.ts‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,22 @@ describe('Fetch Helpers Module', () => {
136136
expect(error.retryable).toBe(false);
137137
});
138138

139+
it('treats missing refresh tokens as terminal auth errors', async () => {
140+
const auth: Auth = { type: 'oauth', access: 'old', refresh: '', expires: 0 };
141+
const client = { auth: { set: vi.fn() } } as any;
142+
vi.spyOn(refreshQueueModule, 'queuedRefresh').mockResolvedValue({
143+
type: 'failed',
144+
reason: 'missing_refresh',
145+
message: 'missing refresh token',
146+
} as any);
147+
148+
const error = await refreshAndUpdateToken(auth, client).catch(
149+
(err) => err as CodexAuthError,
150+
);
151+
expect(error).toBeInstanceOf(CodexAuthError);
152+
expect(error.retryable).toBe(false);
153+
});
154+
139155
it('updates stored auth on success', async () => {
140156
const auth: Auth = { type: 'oauth', access: 'old', refresh: 'oldr', expires: 0 };
141157
const client = { auth: { set: vi.fn() } } as any;

0 commit comments

Comments
 (0)