diff --git a/index.ts b/index.ts index bcf20b13..9240b6aa 100644 --- a/index.ts +++ b/index.ts @@ -177,6 +177,7 @@ import { import { applyFastSessionDefaults } from "./lib/request/request-transformer.js"; import { applyResponseCompaction } from "./lib/request/response-compaction.js"; import { isEmptyResponse } from "./lib/request/response-handler.js"; +import { CodexAuthError } from "./lib/errors.js"; import { parseRetryAfterHintMs, sanitizeResponseHeadersForLog, @@ -997,9 +998,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { accountAuth, client, )) as OAuthAuthDetails; - accountManager.updateFromAuth(account, accountAuth); - accountManager.clearAuthFailures(account); - accountManager.saveToDiskDebounced(); + await accountManager.commitRefreshedAuth( + account, + accountAuth, + ); } } catch (err) { logDebug( @@ -1010,18 +1012,48 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { runtimeMetrics.accountRotations++; runtimeMetrics.lastError = (err as Error)?.message ?? String(err); - const failures = - accountManager.incrementAuthFailures(account); + const isRetryableRefreshError = + err instanceof CodexAuthError && err.retryable; const accountLabel = formatAccountLabel( account, account.index, ); + sessionAffinityStore?.forgetSession(sessionAffinityKey); + + if (isRetryableRefreshError) { + runtimeMetrics.networkErrors++; + const retryableRefreshPolicy = evaluateFailurePolicy( + { kind: "network", failoverMode }, + { networkCooldownMs: networkErrorCooldownMs }, + ); + if (retryableRefreshPolicy.recordFailure) { + accountManager.recordFailure( + account, + modelFamily, + model, + ); + } + if ( + typeof retryableRefreshPolicy.cooldownMs === "number" && + retryableRefreshPolicy.cooldownReason + ) { + accountManager.markAccountCoolingDown( + account, + retryableRefreshPolicy.cooldownMs, + retryableRefreshPolicy.cooldownReason, + ); + accountManager.saveToDiskDebounced(); + } + continue; + } + + const failures = + accountManager.incrementAuthFailures(account); const authFailurePolicy = evaluateFailurePolicy({ kind: "auth-refresh", consecutiveAuthFailures: failures, }); - sessionAffinityStore?.forgetSession(sessionAffinityKey); if (authFailurePolicy.removeAccount) { const removedIndex = account.index; @@ -1945,14 +1977,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { fallbackAuth, client, )) as OAuthAuthDetails; - accountManager.updateFromAuth( + await accountManager.commitRefreshedAuth( fallbackAccount, fallbackAuth, ); - accountManager.clearAuthFailures( - fallbackAccount, - ); - accountManager.saveToDiskDebounced(); } } catch (refreshError) { logWarn( diff --git a/lib/accounts.ts b/lib/accounts.ts index 4bbb035e..6f63d8bb 100644 --- a/lib/accounts.ts +++ b/lib/accounts.ts @@ -7,6 +7,7 @@ import { type CooldownReason, type RateLimitStateV3, findMatchingAccountIndex, + withAccountStorageTransaction, } from "./storage.js"; import type { AccountIdSource, OAuthAuthDetails } from "./types.js"; import { MODEL_FAMILIES, type ModelFamily } from "./prompts/codex.js"; @@ -18,6 +19,8 @@ import { type HybridSelectionOptions, } from "./rotation.js"; import { nowMs } from "./utils.js"; +import { ERROR_MESSAGES, HTTP_STATUS } from "./constants.js"; +import { CodexAuthError } from "./errors.js"; import { loadCodexCliState, type CodexCliTokenCacheEntry, @@ -81,6 +84,112 @@ function initFamilyState(defaultValue: number): Record { ) as Record; } +type AccountIdentityCandidate = Pick< + ManagedAccount, + "accountId" | "email" | "refreshToken" +> & { + index?: number; +}; + +function getAuthIdentityCandidate( + auth: OAuthAuthDetails | undefined, +): AccountIdentityCandidate { + const accountId = extractAccountId(auth?.access)?.trim() || undefined; + const email = sanitizeEmail(extractAccountEmail(auth?.access)); + return { + accountId, + email, + refreshToken: auth?.refresh, + }; +} + +function buildAccountIdentityCandidates( + source: AccountIdentityCandidate, + auth?: OAuthAuthDetails, +): AccountIdentityCandidate[] { + const derived = getAuthIdentityCandidate(auth); + const candidates: AccountIdentityCandidate[] = []; + const seen = new Set(); + + const pushCandidate = (candidate: AccountIdentityCandidate): void => { + const key = `${candidate.accountId ?? ""}|${candidate.email ?? ""}|${candidate.refreshToken ?? ""}`; + if (seen.has(key)) return; + seen.add(key); + candidates.push(candidate); + }; + + pushCandidate(source); + pushCandidate({ + accountId: source.accountId ?? derived.accountId, + email: source.email ?? derived.email, + refreshToken: source.refreshToken, + index: source.index, + }); + pushCandidate({ + accountId: derived.accountId ?? source.accountId, + email: derived.email ?? source.email, + refreshToken: source.refreshToken, + index: source.index, + }); + pushCandidate({ + accountId: derived.accountId ?? source.accountId, + email: derived.email ?? source.email, + refreshToken: derived.refreshToken ?? source.refreshToken, + index: source.index, + }); + + return candidates; +} + +function findAccountIndexByIdentity< + T extends Pick, +>( + accounts: readonly T[], + source: AccountIdentityCandidate, + auth?: OAuthAuthDetails, +): number | undefined { + for (const candidate of buildAccountIdentityCandidates(source, auth)) { + const matchIndex = findMatchingAccountIndex(accounts, candidate, { + allowUniqueAccountIdFallbackWithoutEmail: true, + }); + if (matchIndex !== undefined) { + return matchIndex; + } + } + return undefined; +} + +const RETRYABLE_AUTH_PERSISTENCE_CODES = new Set(["EAGAIN", "EBUSY", "EPERM"]); + +function isRetryableAuthPersistenceError(error: unknown): boolean { + if (!error || typeof error !== "object") { + return false; + } + + const candidate = error as { + code?: unknown; + status?: unknown; + cause?: unknown; + }; + const code = + typeof candidate.code === "string" + ? candidate.code.toUpperCase() + : undefined; + if (code && RETRYABLE_AUTH_PERSISTENCE_CODES.has(code)) { + return true; + } + + if (candidate.status === HTTP_STATUS.TOO_MANY_REQUESTS) { + return true; + } + + if (candidate.cause && candidate.cause !== error) { + return isRetryableAuthPersistenceError(candidate.cause); + } + + return false; +} + export interface Workspace { id: string; name?: string; @@ -642,6 +751,17 @@ export class AccountManager { account.consecutiveAuthFailures = 0; } + getAccountByIdentity( + candidate: AccountIdentityCandidate, + auth?: OAuthAuthDetails, + ): ManagedAccount | null { + const index = findAccountIndexByIdentity(this.accounts, candidate, auth); + if (index === undefined) { + return null; + } + return this.accounts[index] ?? null; + } + shouldShowAccountToast(accountIndex: number, debounceMs = 30000): boolean { const now = nowMs(); if (accountIndex === this.lastToastAccountIndex && now - this.lastToastTime < debounceMs) { @@ -659,7 +779,7 @@ export class AccountManager { account.refreshToken = auth.refresh; account.access = auth.access; account.expires = auth.expires; - const tokenAccountId = extractAccountId(auth.access); + const tokenAccountId = extractAccountId(auth.access)?.trim() || undefined; if ( tokenAccountId && (shouldUpdateAccountIdFromToken(account.accountIdSource, account.accountId)) @@ -670,6 +790,161 @@ export class AccountManager { account.email = sanitizeEmail(extractAccountEmail(auth.access)) ?? account.email; } + private buildStorageSnapshot(): AccountStorageV3 { + const activeIndexByFamily: Partial> = {}; + for (const family of MODEL_FAMILIES) { + const raw = this.currentAccountIndexByFamily[family]; + activeIndexByFamily[family] = clampNonNegativeInt(raw, 0); + } + + const activeIndex = clampNonNegativeInt(activeIndexByFamily.codex, 0); + + return { + version: 3, + accounts: this.accounts.map((account) => ({ + accountId: account.accountId, + accountIdSource: account.accountIdSource, + accountLabel: account.accountLabel, + email: account.email, + refreshToken: account.refreshToken, + accessToken: account.access, + expiresAt: account.expires, + enabled: account.enabled === false ? false : undefined, + addedAt: account.addedAt, + lastUsed: account.lastUsed, + lastSwitchReason: account.lastSwitchReason, + rateLimitResetTimes: + Object.keys(account.rateLimitResetTimes).length > 0 ? account.rateLimitResetTimes : undefined, + coolingDownUntil: account.coolingDownUntil, + cooldownReason: account.cooldownReason, + workspaces: account.workspaces, + currentWorkspaceIndex: account.currentWorkspaceIndex, + })), + activeIndex, + activeIndexByFamily, + }; + } + + async commitRefreshedAuth( + source: Pick< + ManagedAccount, + "index" | "accountId" | "email" | "refreshToken" + >, + auth: OAuthAuthDetails, + ): Promise { + const nextAccountId = extractAccountId(auth.access)?.trim() || undefined; + const nextEmail = sanitizeEmail(extractAccountEmail(auth.access)); + try { + return await withAccountStorageTransaction(async (_current, persist) => { + // Snapshot the live in-memory pool under the storage lock so refresh + // persistence merges against the latest account state. + const nextStorage = structuredClone( + this.buildStorageSnapshot(), + ) as AccountStorageV3; + const storageIndex = findAccountIndexByIdentity( + nextStorage.accounts, + source, + auth, + ); + if (storageIndex === undefined) { + log.warn("Unable to resolve refreshed account for persistence", { + sourceIndex: source.index, + }); + return null; + } + + const storedAccount = nextStorage.accounts[storageIndex]; + if (!storedAccount) { + return null; + } + + storedAccount.refreshToken = auth.refresh; + storedAccount.accessToken = auth.access; + storedAccount.expiresAt = auth.expires; + if ( + nextAccountId && + shouldUpdateAccountIdFromToken( + storedAccount.accountIdSource, + storedAccount.accountId, + ) + ) { + storedAccount.accountId = nextAccountId; + storedAccount.accountIdSource = "token"; + } + if (nextEmail) { + storedAccount.email = nextEmail; + } + storedAccount.enabled = undefined; + delete storedAccount.coolingDownUntil; + delete storedAccount.cooldownReason; + + const liveAccount = this.getAccountByIdentity(source, auth); + if (liveAccount) { + const previousLiveAccountState = { + access: liveAccount.access, + refreshToken: liveAccount.refreshToken, + expires: liveAccount.expires, + accountId: liveAccount.accountId, + accountIdSource: liveAccount.accountIdSource, + email: liveAccount.email, + enabled: liveAccount.enabled, + coolingDownUntil: liveAccount.coolingDownUntil, + cooldownReason: liveAccount.cooldownReason, + consecutiveAuthFailures: liveAccount.consecutiveAuthFailures, + }; + + this.updateFromAuth(liveAccount, auth); + liveAccount.enabled = true; + this.clearAccountCooldown(liveAccount); + this.clearAuthFailures(liveAccount); + + try { + await persist(nextStorage); + } catch (error) { + liveAccount.access = previousLiveAccountState.access; + liveAccount.refreshToken = previousLiveAccountState.refreshToken; + liveAccount.expires = previousLiveAccountState.expires; + liveAccount.accountId = previousLiveAccountState.accountId; + liveAccount.accountIdSource = + previousLiveAccountState.accountIdSource; + liveAccount.email = previousLiveAccountState.email; + liveAccount.enabled = previousLiveAccountState.enabled; + liveAccount.consecutiveAuthFailures = + previousLiveAccountState.consecutiveAuthFailures; + if ( + previousLiveAccountState.coolingDownUntil === undefined + ) { + delete liveAccount.coolingDownUntil; + } else { + liveAccount.coolingDownUntil = + previousLiveAccountState.coolingDownUntil; + } + if (previousLiveAccountState.cooldownReason === undefined) { + delete liveAccount.cooldownReason; + } else { + liveAccount.cooldownReason = + previousLiveAccountState.cooldownReason; + } + throw error; + } + + return liveAccount; + } + + await persist(nextStorage); + log.warn("Unable to resolve refreshed live account after persistence", { + sourceIndex: source.index, + }); + return null; + }); + } catch (error) { + throw new CodexAuthError(ERROR_MESSAGES.TOKEN_REFRESH_FAILED, { + retryable: isRetryableAuthPersistenceError(error), + cause: error, + }); + } + } + toAuthDetails(account: ManagedAccount): Auth { return { type: "oauth", @@ -780,40 +1055,9 @@ export class AccountManager { } async saveToDisk(): Promise { - const activeIndexByFamily: Partial> = {}; - for (const family of MODEL_FAMILIES) { - const raw = this.currentAccountIndexByFamily[family]; - activeIndexByFamily[family] = clampNonNegativeInt(raw, 0); - } - - const activeIndex = clampNonNegativeInt(activeIndexByFamily.codex, 0); - - const storage: AccountStorageV3 = { - version: 3, - accounts: this.accounts.map((account) => ({ - accountId: account.accountId, - accountIdSource: account.accountIdSource, - accountLabel: account.accountLabel, - email: account.email, - refreshToken: account.refreshToken, - accessToken: account.access, - expiresAt: account.expires, - enabled: account.enabled === false ? false : undefined, - addedAt: account.addedAt, - lastUsed: account.lastUsed, - lastSwitchReason: account.lastSwitchReason, - rateLimitResetTimes: - Object.keys(account.rateLimitResetTimes).length > 0 ? account.rateLimitResetTimes : undefined, - coolingDownUntil: account.coolingDownUntil, - cooldownReason: account.cooldownReason, - workspaces: account.workspaces, - currentWorkspaceIndex: account.currentWorkspaceIndex, - })), - activeIndex, - activeIndexByFamily, - }; - - await saveAccounts(storage); + await withAccountStorageTransaction(async (_current, persist) => { + await persist(this.buildStorageSnapshot()); + }); } saveToDiskDebounced(delayMs = 500): void { diff --git a/lib/proactive-refresh.ts b/lib/proactive-refresh.ts index b80e40be..4e5d4a4b 100644 --- a/lib/proactive-refresh.ts +++ b/lib/proactive-refresh.ts @@ -14,6 +14,12 @@ import { queuedRefresh } from "./refresh-queue.js"; import { createLogger } from "./logger.js"; import type { ManagedAccount } from "./accounts.js"; import type { TokenResult } from "./types.js"; +import { + extractAccountEmail, + extractAccountId, + sanitizeEmail, + shouldUpdateAccountIdFromToken, +} from "./auth/token-utils.js"; const log = createLogger("proactive-refresh"); @@ -43,16 +49,16 @@ export function shouldRefreshProactively( account: ManagedAccount, bufferMs: number = DEFAULT_PROACTIVE_BUFFER_MS, ): boolean { - // No expiry set - can't determine if refresh is needed - if (account.expires === undefined) { - return false; - } - // No access token - definitely needs refresh if (!account.access) { return true; } + // No expiry set - can't determine if refresh is needed + if (account.expires === undefined) { + return false; + } + // Clamp buffer to minimum const safeBufferMs = Math.max(MIN_PROACTIVE_BUFFER_MS, bufferMs); @@ -135,6 +141,10 @@ export async function proactiveRefreshAccount( export async function refreshExpiringAccounts( accounts: ManagedAccount[], bufferMs: number = DEFAULT_PROACTIVE_BUFFER_MS, + onResult?: ( + account: ManagedAccount, + result: ProactiveRefreshResult, + ) => Promise | void, ): Promise> { const results = new Map(); const accountsToRefresh = accounts.filter((a) => @@ -151,6 +161,7 @@ export async function refreshExpiringAccounts( // Refresh in parallel for efficiency const refreshPromises = accountsToRefresh.map(async (account) => { const result = await proactiveRefreshAccount(account, bufferMs); + await onResult?.(account, result); return { index: account.index, result }; }); @@ -194,4 +205,13 @@ export function applyRefreshResult( if (result.refresh !== account.refreshToken) { account.refreshToken = result.refresh; } + const tokenAccountId = extractAccountId(result.access)?.trim() || undefined; + if ( + tokenAccountId && + shouldUpdateAccountIdFromToken(account.accountIdSource, account.accountId) + ) { + account.accountId = tokenAccountId; + account.accountIdSource = "token"; + } + account.email = sanitizeEmail(extractAccountEmail(result.access)) ?? account.email; } diff --git a/lib/refresh-guardian.ts b/lib/refresh-guardian.ts index b40bddd8..9be4414c 100644 --- a/lib/refresh-guardian.ts +++ b/lib/refresh-guardian.ts @@ -1,6 +1,7 @@ import { createLogger } from "./logger.js"; -import { applyRefreshResult, refreshExpiringAccounts } from "./proactive-refresh.js"; +import { refreshExpiringAccounts } from "./proactive-refresh.js"; import type { AccountManager } from "./accounts.js"; +import { CodexAuthError } from "./errors.js"; import type { CooldownReason } from "./storage.js"; import type { TokenResult } from "./types.js"; @@ -95,6 +96,104 @@ export class RefreshGuardian { return "network-error"; } + private async applyRefreshOutcome( + manager: AccountManager, + sourceAccount: ReturnType[number], + result: Awaited> extends Map< + number, + infer TValue + > + ? TValue + : never, + ): Promise { + switch (result.reason) { + case "success": { + if (result.tokenResult?.type !== "success") { + const account = manager.getAccountByIdentity(sourceAccount); + if (!account) return false; + manager.markAccountCoolingDown(account, this.bufferMs, "network-error"); + this.stats.failed += 1; + this.stats.networkFailed += 1; + return true; + } + + const refreshedAuth = { + type: "oauth" as const, + access: result.tokenResult.access, + refresh: result.tokenResult.refresh, + expires: result.tokenResult.expires, + multiAccount: true, + }; + try { + const committedAccount = await manager.commitRefreshedAuth( + sourceAccount, + refreshedAuth, + ); + if (!committedAccount) { + const account = + manager.getAccountByIdentity(sourceAccount, refreshedAuth) ?? + manager.getAccountByIdentity(sourceAccount); + if (account) { + manager.markAccountCoolingDown( + account, + this.bufferMs, + "network-error", + ); + } + this.stats.failed += 1; + this.stats.networkFailed += 1; + return !!account; + } + } catch (error) { + log.warn("Refresh guardian commit failed", { + sourceIndex: sourceAccount.index, + error: error instanceof Error ? error.message : String(error), + }); + const account = + manager.getAccountByIdentity(sourceAccount, refreshedAuth) ?? + manager.getAccountByIdentity(sourceAccount); + const cooldownReason: CooldownReason = + error instanceof CodexAuthError && !error.retryable + ? "auth-failure" + : "network-error"; + if (account) { + manager.markAccountCoolingDown(account, this.bufferMs, cooldownReason); + } + this.stats.failed += 1; + if (cooldownReason === "auth-failure") this.stats.authFailed += 1; + else this.stats.networkFailed += 1; + return !!account; + } + this.stats.refreshed += 1; + return false; + } + case "failed": { + const account = manager.getAccountByIdentity(sourceAccount); + if (!account) return false; + const cooldownReason = this.classifyFailureReason(result.tokenResult); + manager.markAccountCoolingDown(account, this.bufferMs, cooldownReason); + this.stats.failed += 1; + if (cooldownReason === "rate-limit") this.stats.rateLimited += 1; + else if (cooldownReason === "auth-failure") this.stats.authFailed += 1; + else this.stats.networkFailed += 1; + return true; + } + case "not_needed": + this.stats.notNeeded += 1; + return false; + case "no_refresh_token": { + const account = manager.getAccountByIdentity(sourceAccount); + if (!account) return false; + manager.markAccountCoolingDown(account, this.bufferMs, "auth-failure"); + manager.setAccountEnabled(account.index, false); + this.stats.noRefreshToken += 1; + this.stats.failed += 1; + this.stats.authFailed += 1; + return true; + } + } + } + async tick(): Promise { if (this.running) return; const manager = this.getAccountManager(); @@ -106,63 +205,37 @@ export class RefreshGuardian { return; } - const refreshResults = await refreshExpiringAccounts(snapshot, this.bufferMs); - if (refreshResults.size === 0) { + const eligibleSnapshot = snapshot.filter((account) => !manager.isAccountCoolingDown(account)); + if (eligibleSnapshot.length === 0) { this.stats.runs += 1; this.stats.lastRunAt = Date.now(); return; } - const snapshotByIndex = new Map(); - for (const candidate of snapshot) { - snapshotByIndex.set(candidate.index, candidate); - } - - for (const [accountIndex, result] of refreshResults.entries()) { - const sourceAccount = snapshotByIndex.get(accountIndex); - if (!sourceAccount) continue; - const liveAccounts = manager.getAccountsSnapshot(); - const resolvedIndex = liveAccounts.findIndex( - (candidate) => candidate.refreshToken === sourceAccount.refreshToken, - ); - const account = resolvedIndex >= 0 ? manager.getAccountByIndex(resolvedIndex) : null; - if (!account) continue; - - switch (result.reason) { - case "success": - if (result.tokenResult?.type === "success") { - applyRefreshResult(account, result.tokenResult); - manager.clearAuthFailures(account); - this.stats.refreshed += 1; - } else { - manager.markAccountCoolingDown(account, this.bufferMs, "network-error"); - this.stats.failed += 1; - this.stats.networkFailed += 1; - } - break; - case "failed": { - const cooldownReason = this.classifyFailureReason(result.tokenResult); - manager.markAccountCoolingDown(account, this.bufferMs, cooldownReason); - this.stats.failed += 1; - if (cooldownReason === "rate-limit") this.stats.rateLimited += 1; - else if (cooldownReason === "auth-failure") this.stats.authFailed += 1; - else this.stats.networkFailed += 1; - break; + let requiresSave = false; + const refreshResults = await refreshExpiringAccounts( + eligibleSnapshot, + this.bufferMs, + async (sourceAccount, result) => { + const saveNeeded = await this.applyRefreshOutcome( + manager, + sourceAccount, + result, + ); + if (saveNeeded) { + requiresSave = true; } - case "not_needed": - this.stats.notNeeded += 1; - break; - case "no_refresh_token": - manager.markAccountCoolingDown(account, this.bufferMs, "auth-failure"); - manager.setAccountEnabled(account.index, false); - this.stats.noRefreshToken += 1; - this.stats.failed += 1; - this.stats.authFailed += 1; - break; - } + }, + ); + if (refreshResults.size === 0) { + this.stats.runs += 1; + this.stats.lastRunAt = Date.now(); + return; } - manager.saveToDiskDebounced(); + if (requiresSave) { + manager.saveToDiskDebounced(); + } this.stats.runs += 1; this.stats.lastRunAt = Date.now(); } catch (error) { diff --git a/lib/request/fetch-helpers.ts b/lib/request/fetch-helpers.ts index ac0a7efb..e62241d2 100644 --- a/lib/request/fetch-helpers.ts +++ b/lib/request/fetch-helpers.ts @@ -19,7 +19,7 @@ import { convertSseToJson, ensureContentType, } from "./response-handler.js"; -import type { UserConfig, RequestBody } from "../types.js"; +import type { UserConfig, RequestBody, TokenResult } from "../types.js"; import { registerCleanup } from "../shutdown.js"; import { CodexAuthError } from "../errors.js"; import { isRecord } from "../utils.js"; @@ -428,6 +428,56 @@ export function shouldRefreshToken(auth: Auth, skewMs = 0): boolean { return auth.expires <= Date.now() + safeSkewMs; } +function isRetryableRefreshFailure( + result: Extract, +): boolean { + switch (result.reason) { + case "network_error": + case "unknown": + case "invalid_response": + return true; + case "missing_refresh": + return false; + case "http_error": + return !( + result.statusCode === HTTP_STATUS.BAD_REQUEST || + result.statusCode === HTTP_STATUS.UNAUTHORIZED || + result.statusCode === HTTP_STATUS.FORBIDDEN + ); + default: + return false; + } +} + +function isRetryableAuthSetterError(error: unknown): boolean { + if (!error || typeof error !== "object") { + return false; + } + + const candidate = error as { + code?: unknown; + status?: unknown; + cause?: unknown; + }; + const code = + typeof candidate.code === "string" + ? candidate.code.toUpperCase() + : undefined; + if (code === "EAGAIN" || code === "EBUSY" || code === "EPERM") { + return true; + } + + if (candidate.status === HTTP_STATUS.TOO_MANY_REQUESTS) { + return true; + } + + if (candidate.cause && candidate.cause !== error) { + return isRetryableAuthSetterError(candidate.cause); + } + + return false; +} + /** * Refreshes the OAuth token and updates stored credentials * @param currentAuth - Current auth state @@ -440,26 +490,41 @@ export async function refreshAndUpdateToken( ): Promise { const authSetter = (client as Partial).auth; if (!authSetter || typeof authSetter.set !== "function") { - throw new CodexAuthError(ERROR_MESSAGES.TOKEN_REFRESH_FAILED, { retryable: false }); + throw new CodexAuthError(ERROR_MESSAGES.TOKEN_REFRESH_FAILED, { + retryable: false, + }); } const refreshToken = currentAuth.type === "oauth" ? currentAuth.refresh : ""; const refreshResult = await queuedRefresh(refreshToken); if (refreshResult.type === "failed") { - throw new CodexAuthError(ERROR_MESSAGES.TOKEN_REFRESH_FAILED, { retryable: false }); + throw new CodexAuthError(ERROR_MESSAGES.TOKEN_REFRESH_FAILED, { + retryable: isRetryableRefreshFailure(refreshResult), + context: { + refreshFailureReason: refreshResult.reason, + statusCode: refreshResult.statusCode, + }, + }); } - await authSetter.set({ - path: { id: "openai" }, - body: { - type: "oauth", - access: refreshResult.access, - refresh: refreshResult.refresh, - expires: refreshResult.expires, - multiAccount: true, - }, - }); + try { + await authSetter.set({ + path: { id: "openai" }, + body: { + type: "oauth", + access: refreshResult.access, + refresh: refreshResult.refresh, + expires: refreshResult.expires, + multiAccount: true, + }, + }); + } catch (error) { + throw new CodexAuthError(ERROR_MESSAGES.TOKEN_REFRESH_FAILED, { + retryable: isRetryableAuthSetterError(error), + cause: error, + }); + } // Update current auth reference if it's OAuth type if (currentAuth.type === "oauth") { diff --git a/test/accounts-edge.test.ts b/test/accounts-edge.test.ts index 31c34b07..e9621076 100644 --- a/test/accounts-edge.test.ts +++ b/test/accounts-edge.test.ts @@ -3,6 +3,7 @@ import type { OAuthAuthDetails } from "../lib/types.js"; const mockLoadAccounts = vi.fn(); const mockSaveAccounts = vi.fn(); +const mockWithAccountStorageTransaction = vi.fn(); const mockLoadCodexCliState = vi.fn(); const mockSyncAccountStorageFromCodexCli = vi.fn(); const mockSetCodexCliActiveSelection = vi.fn(); @@ -14,6 +15,7 @@ vi.mock("../lib/storage.js", async (importOriginal) => { ...actual, loadAccounts: mockLoadAccounts, saveAccounts: mockSaveAccounts, + withAccountStorageTransaction: mockWithAccountStorageTransaction, }; }); @@ -81,6 +83,11 @@ describe("accounts edge branches", () => { vi.clearAllMocks(); mockLoadAccounts.mockResolvedValue(null); mockSaveAccounts.mockResolvedValue(undefined); + mockWithAccountStorageTransaction.mockImplementation(async (handler) => + handler(null, async (storage) => { + await mockSaveAccounts(storage); + }), + ); mockLoadCodexCliState.mockResolvedValue(null); mockSyncAccountStorageFromCodexCli.mockImplementation(async (storage) => ({ storage, diff --git a/test/accounts-load-from-disk.test.ts b/test/accounts-load-from-disk.test.ts index 61c2b8b0..1976fab8 100644 --- a/test/accounts-load-from-disk.test.ts +++ b/test/accounts-load-from-disk.test.ts @@ -1,12 +1,23 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { AccountManager } from "../lib/accounts.js"; +const { + mockLoadAccounts, + mockSaveAccounts, + mockWithAccountStorageTransaction, +} = vi.hoisted(() => ({ + mockLoadAccounts: vi.fn(), + mockSaveAccounts: vi.fn(), + mockWithAccountStorageTransaction: vi.fn(), +})); + vi.mock("../lib/storage.js", async (importOriginal) => { const actual = await importOriginal(); return { ...actual, - loadAccounts: vi.fn(), - saveAccounts: vi.fn().mockResolvedValue(undefined), + loadAccounts: mockLoadAccounts, + saveAccounts: mockSaveAccounts, + withAccountStorageTransaction: mockWithAccountStorageTransaction, }; }); @@ -22,7 +33,6 @@ vi.mock("../lib/codex-cli/writer.js", () => ({ setCodexCliActiveSelection: vi.fn().mockResolvedValue(undefined), })); -import { loadAccounts, saveAccounts } from "../lib/storage.js"; import { syncAccountStorageFromCodexCli } from "../lib/codex-cli/sync.js"; import { loadCodexCliState } from "../lib/codex-cli/state.js"; import { setCodexCliActiveSelection } from "../lib/codex-cli/writer.js"; @@ -30,8 +40,13 @@ import { setCodexCliActiveSelection } from "../lib/codex-cli/writer.js"; describe("AccountManager loadFromDisk", () => { beforeEach(() => { vi.clearAllMocks(); - vi.mocked(loadAccounts).mockResolvedValue(null); - vi.mocked(saveAccounts).mockResolvedValue(undefined); + mockLoadAccounts.mockResolvedValue(null); + mockSaveAccounts.mockResolvedValue(undefined); + mockWithAccountStorageTransaction.mockImplementation(async (handler) => + handler(null, async (storage) => { + await mockSaveAccounts(storage); + }), + ); vi.mocked(syncAccountStorageFromCodexCli).mockResolvedValue({ changed: false, storage: null, @@ -56,7 +71,7 @@ describe("AccountManager loadFromDisk", () => { ], }; - vi.mocked(loadAccounts).mockResolvedValue(stored); + mockLoadAccounts.mockResolvedValue(stored); vi.mocked(syncAccountStorageFromCodexCli).mockResolvedValue({ changed: true, storage: synced, @@ -64,7 +79,7 @@ describe("AccountManager loadFromDisk", () => { const manager = await AccountManager.loadFromDisk(); - expect(saveAccounts).toHaveBeenCalledWith(synced); + expect(mockSaveAccounts).toHaveBeenCalledWith(synced); expect(manager.getAccountCount()).toBe(2); expect(manager.getCurrentAccount()?.refreshToken).toBe("stored-refresh"); }); @@ -81,7 +96,7 @@ describe("AccountManager loadFromDisk", () => { changed: true, storage: synced, }); - vi.mocked(saveAccounts).mockRejectedValueOnce(new Error("forced persist failure")); + mockSaveAccounts.mockRejectedValueOnce(new Error("forced persist failure")); const manager = await AccountManager.loadFromDisk(); @@ -91,7 +106,7 @@ describe("AccountManager loadFromDisk", () => { it("hydrates missing access/accountId fields from Codex CLI token cache", async () => { const now = Date.now(); - vi.mocked(loadAccounts).mockResolvedValue({ + mockLoadAccounts.mockResolvedValue({ version: 3 as const, activeIndex: 0, accounts: [ @@ -122,12 +137,12 @@ describe("AccountManager loadFromDisk", () => { expect(account?.expires).toBe(now + 120_000); expect(account?.accountId).toBe("acct-123"); expect(account?.accountIdSource).toBe("token"); - expect(saveAccounts).toHaveBeenCalledTimes(1); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); }); it("skips expired Codex CLI cache entries and does not persist", async () => { const now = Date.now(); - vi.mocked(loadAccounts).mockResolvedValue({ + mockLoadAccounts.mockResolvedValue({ version: 3 as const, activeIndex: 0, accounts: [ @@ -156,7 +171,7 @@ describe("AccountManager loadFromDisk", () => { expect(account?.access).toBeUndefined(); expect(account?.accountId).toBeUndefined(); - expect(saveAccounts).not.toHaveBeenCalled(); + expect(mockSaveAccounts).not.toHaveBeenCalled(); }); it("syncCodexCliActiveSelectionForIndex ignores invalid indices and syncs a valid one", async () => { diff --git a/test/accounts.test.ts b/test/accounts.test.ts index 54b44e61..af2f5dc6 100644 --- a/test/accounts.test.ts +++ b/test/accounts.test.ts @@ -11,6 +11,7 @@ import { getAccountIdCandidates, } from "../lib/accounts.js"; import { getHealthTracker, getTokenTracker, resetTrackers } from "../lib/rotation.js"; +import { CodexAuthError } from "../lib/errors.js"; import type { OAuthAuthDetails } from "../lib/types.js"; vi.mock("../lib/storage.js", async (importOriginal) => { @@ -18,9 +19,33 @@ vi.mock("../lib/storage.js", async (importOriginal) => { return { ...actual, saveAccounts: vi.fn().mockResolvedValue(undefined), + withAccountStorageTransaction: vi.fn(), }; }); +beforeEach(async () => { + resetTrackers(); + const { saveAccounts, withAccountStorageTransaction } = await import( + "../lib/storage.js" + ); + const mockSaveAccounts = vi.mocked(saveAccounts); + const mockWithAccountStorageTransaction = vi.mocked( + withAccountStorageTransaction, + ); + + mockSaveAccounts.mockReset(); + mockSaveAccounts.mockResolvedValue(undefined); + mockWithAccountStorageTransaction.mockReset(); + mockWithAccountStorageTransaction.mockImplementation(async (handler) => { + let current = null; + const persist = async (storage: Parameters[0]) => { + current = structuredClone(storage); + await mockSaveAccounts(storage); + }; + return handler(current as never, persist); + }); +}); + describe("parseRateLimitReason", () => { it("returns quota for quota-related codes", () => { expect(parseRateLimitReason("usage_limit_reached")).toBe("quota"); @@ -1049,6 +1074,394 @@ describe("AccountManager", () => { }); }); + describe("commitRefreshedAuth", () => { + it("persists refreshed auth transactionally and updates the live account", async () => { + const { withAccountStorageTransaction } = await import("../lib/storage.js"); + const mockWithAccountStorageTransaction = vi.mocked( + withAccountStorageTransaction, + ); + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "old-refresh", + accessToken: "old-access", + expiresAt: now, + addedAt: now, + lastUsed: now, + email: "old@example.com", + accountId: "old-account-id", + accountIdSource: "token" as const, + enabled: false, + coolingDownUntil: now + 30_000, + cooldownReason: "auth-failure" as const, + }, + ], + }; + const manager = new AccountManager(undefined, stored as any); + const account = manager.getAccountByIndex(0)!; + account.enabled = false; + manager.markAccountCoolingDown(account, 30_000, "auth-failure"); + manager.incrementAuthFailures(account); + + const payload = Buffer.from( + JSON.stringify({ + email: "new@example.com", + "https://api.openai.com/auth": { + chatgpt_account_id: "new-account-id", + }, + }), + ).toString("base64url"); + const refreshedAuth: OAuthAuthDetails = { + type: "oauth", + access: `header.${payload}.signature`, + refresh: "new-refresh", + expires: now + 3_600_000, + }; + + mockWithAccountStorageTransaction.mockImplementationOnce(async (handler) => { + const persist = vi.fn().mockResolvedValue(undefined); + const result = await handler(stored as any, persist); + + expect(persist).toHaveBeenCalledTimes(1); + const persistedStorage = persist.mock.calls[0]?.[0]; + expect(persistedStorage?.accounts[0]?.refreshToken).toBe("new-refresh"); + expect(persistedStorage?.accounts[0]?.accessToken).toBe( + refreshedAuth.access, + ); + expect(persistedStorage?.accounts[0]?.expiresAt).toBe( + refreshedAuth.expires, + ); + expect(persistedStorage?.accounts[0]?.accountId).toBe("new-account-id"); + expect(persistedStorage?.accounts[0]?.accountIdSource).toBe("token"); + expect(persistedStorage?.accounts[0]?.email).toBe("new@example.com"); + expect(persistedStorage?.accounts[0]?.enabled).toBeUndefined(); + expect(persistedStorage?.accounts[0]?.coolingDownUntil).toBeUndefined(); + expect(persistedStorage?.accounts[0]?.cooldownReason).toBeUndefined(); + + return result; + }); + + const updated = await manager.commitRefreshedAuth(account, refreshedAuth); + + expect(updated).toBe(account); + expect(account.refreshToken).toBe("new-refresh"); + expect(account.access).toBe(refreshedAuth.access); + expect(account.expires).toBe(refreshedAuth.expires); + expect(account.accountId).toBe("new-account-id"); + expect(account.accountIdSource).toBe("token"); + expect(account.email).toBe("new@example.com"); + expect(account.enabled).toBe(true); + expect(account.coolingDownUntil).toBeUndefined(); + expect(account.cooldownReason).toBeUndefined(); + expect(account.consecutiveAuthFailures).toBe(0); + }); + + it("preserves unsaved live pool state when persisting refreshed auth", async () => { + const { withAccountStorageTransaction } = await import("../lib/storage.js"); + const mockWithAccountStorageTransaction = vi.mocked( + withAccountStorageTransaction, + ); + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + activeIndexByFamily: { + codex: 0, + }, + accounts: [ + { + refreshToken: "old-refresh-a", + accessToken: "old-access-a", + expiresAt: now, + addedAt: now, + lastUsed: now, + email: "a@example.com", + }, + { + refreshToken: "old-refresh-b", + accessToken: "old-access-b", + expiresAt: now, + addedAt: now, + lastUsed: now, + email: "b@example.com", + }, + ], + }; + const manager = new AccountManager(undefined, stored as any); + const accountA = manager.getAccountByIndex(0)!; + const accountB = manager.getAccountByIndex(1)!; + manager.markSwitched(accountB, "rotation", "codex"); + manager.markRateLimited(accountB, 45_000, "codex"); + manager.markAccountCoolingDown(accountB, 30_000, "network-error"); + + const refreshedAuth: OAuthAuthDetails = { + type: "oauth", + access: "header.payload.signature", + refresh: "new-refresh-a", + expires: now + 3_600_000, + }; + + mockWithAccountStorageTransaction.mockImplementationOnce(async (handler) => { + const persist = vi.fn().mockResolvedValue(undefined); + const result = await handler(stored as any, persist); + + expect(persist).toHaveBeenCalledTimes(1); + const persistedStorage = persist.mock.calls[0]?.[0]; + expect(persistedStorage?.activeIndex).toBe(1); + expect(persistedStorage?.activeIndexByFamily?.codex).toBe(1); + expect(persistedStorage?.accounts[0]?.refreshToken).toBe("new-refresh-a"); + expect(persistedStorage?.accounts[1]?.rateLimitResetTimes?.codex).toBeGreaterThan( + now, + ); + expect(persistedStorage?.accounts[1]?.cooldownReason).toBe( + "network-error", + ); + expect(persistedStorage?.accounts[1]?.coolingDownUntil).toBeGreaterThan( + now, + ); + + return result; + }); + + await manager.commitRefreshedAuth(accountA, refreshedAuth); + }); + + it("keeps trimmed token accountId identical in memory and persisted storage", async () => { + const { withAccountStorageTransaction } = await import("../lib/storage.js"); + const mockWithAccountStorageTransaction = vi.mocked( + withAccountStorageTransaction, + ); + const now = Date.now(); + const payload = Buffer.from( + JSON.stringify({ + "https://api.openai.com/auth": { + chatgpt_account_id: " matching-account-id ", + }, + }), + ).toString("base64url"); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "old-refresh", + accessToken: "old-access", + expiresAt: now, + addedAt: now, + lastUsed: now, + accountId: "matching-account-id", + accountIdSource: "token" as const, + }, + ], + }; + const manager = new AccountManager(undefined, stored as any); + const account = manager.getAccountByIndex(0)!; + const refreshedAuth: OAuthAuthDetails = { + type: "oauth", + access: `header.${payload}.signature`, + refresh: "new-refresh", + expires: now + 3_600_000, + }; + + mockWithAccountStorageTransaction.mockImplementationOnce(async (handler) => { + const persist = vi.fn().mockResolvedValue(undefined); + const result = await handler(stored as any, persist); + + expect(persist).toHaveBeenCalledTimes(1); + const persistedStorage = persist.mock.calls[0]?.[0]; + expect(persistedStorage?.accounts[0]?.accountId).toBe("matching-account-id"); + + return result; + }); + + await manager.commitRefreshedAuth(account, refreshedAuth); + + expect(account.accountId).toBe("matching-account-id"); + }); + + it("propagates storage write failure as retryable CodexAuthError", async () => { + const { withAccountStorageTransaction } = await import("../lib/storage.js"); + const mockWithAccountStorageTransaction = vi.mocked( + withAccountStorageTransaction, + ); + mockWithAccountStorageTransaction.mockRejectedValueOnce( + Object.assign(new Error("EBUSY"), { code: "EBUSY" }), + ); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "old-refresh", + accessToken: "old-access", + expiresAt: now, + addedAt: now, + lastUsed: now, + }, + ], + }; + const manager = new AccountManager(undefined, stored as any); + const account = manager.getAccountByIndex(0)!; + const refreshedAuth: OAuthAuthDetails = { + type: "oauth", + access: "header.payload.signature", + refresh: "new-refresh", + expires: now + 3_600_000, + }; + + const error = await manager.commitRefreshedAuth(account, refreshedAuth).catch( + (err) => err as CodexAuthError, + ); + + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(true); + expect(account.refreshToken).toBe("old-refresh"); + }); + + it("propagates non-transient storage write failure as terminal CodexAuthError", async () => { + const { withAccountStorageTransaction } = await import("../lib/storage.js"); + const mockWithAccountStorageTransaction = vi.mocked( + withAccountStorageTransaction, + ); + mockWithAccountStorageTransaction.mockRejectedValueOnce( + Object.assign(new Error("EACCES"), { code: "EACCES" }), + ); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "old-refresh", + accessToken: "old-access", + expiresAt: now, + addedAt: now, + lastUsed: now, + }, + ], + }; + const manager = new AccountManager(undefined, stored as any); + const account = manager.getAccountByIndex(0)!; + const refreshedAuth: OAuthAuthDetails = { + type: "oauth", + access: "header.payload.signature", + refresh: "new-refresh", + expires: now + 3_600_000, + }; + + const error = await manager.commitRefreshedAuth(account, refreshedAuth).catch( + (err) => err as CodexAuthError, + ); + + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(false); + expect(account.refreshToken).toBe("old-refresh"); + }); + + it("prevents debounced saves from writing stale auth during refresh persistence", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts, withAccountStorageTransaction } = await import( + "../lib/storage.js" + ); + const mockSaveAccounts = vi.mocked(saveAccounts); + const mockWithAccountStorageTransaction = vi.mocked( + withAccountStorageTransaction, + ); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "old-refresh", + accessToken: "old-access", + expiresAt: now, + addedAt: now, + lastUsed: now, + }, + ], + }; + const manager = new AccountManager(undefined, stored as any); + const account = manager.getAccountByIndex(0)!; + const refreshedAuth: OAuthAuthDetails = { + type: "oauth", + access: "new-access", + refresh: "new-refresh", + expires: now + 3_600_000, + }; + + let storageState = structuredClone(stored) as typeof stored; + let lock = Promise.resolve(); + let releasePersist!: () => void; + const persistBlocked = new Promise((resolve) => { + releasePersist = resolve; + }); + let notifyPersistStarted!: () => void; + const persistStarted = new Promise((resolve) => { + notifyPersistStarted = resolve; + }); + + mockWithAccountStorageTransaction.mockImplementation((handler) => { + const run = async () => { + const current = structuredClone(storageState) as typeof storageState; + const persist = async ( + nextStorage: Parameters[0], + ) => { + if ( + nextStorage.accounts[0]?.refreshToken === "new-refresh" && + nextStorage.accounts[0]?.accessToken === "new-access" + ) { + notifyPersistStarted(); + await persistBlocked; + } + storageState = structuredClone(nextStorage) as typeof storageState; + await mockSaveAccounts(nextStorage); + }; + return handler(current as never, persist); + }; + + const result = lock.then(run, run); + lock = result.then( + () => undefined, + () => undefined, + ); + return result; + }); + + const commitPromise = manager.commitRefreshedAuth(account, refreshedAuth); + await persistStarted; + + manager.saveToDiskDebounced(0); + await vi.advanceTimersByTimeAsync(0); + + releasePersist(); + await commitPromise; + await manager.flushPendingSave(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + expect(mockSaveAccounts.mock.calls[0]?.[0]?.accounts[0]?.refreshToken).toBe( + "new-refresh", + ); + expect(mockSaveAccounts.mock.calls[1]?.[0]?.accounts[0]?.refreshToken).toBe( + "new-refresh", + ); + expect(mockSaveAccounts.mock.calls[1]?.[0]?.accounts[0]?.accessToken).toBe( + "new-access", + ); + } finally { + vi.useRealTimers(); + } + }); + }); + describe("toAuthDetails", () => { it("converts account to Auth object", () => { const now = Date.now(); diff --git a/test/fetch-helpers.test.ts b/test/fetch-helpers.test.ts index 1363c8fd..15313e9a 100644 --- a/test/fetch-helpers.test.ts +++ b/test/fetch-helpers.test.ts @@ -23,6 +23,7 @@ import * as loggerModule from '../lib/logger.js'; import type { Auth } from '../lib/types.js'; import type { CreateCodexHeadersParams } from '../lib/request/fetch-helpers.js'; import { URL_PATHS, OPENAI_HEADERS, OPENAI_HEADER_VALUES, CODEX_BASE_URL } from '../lib/constants.js'; +import { CodexAuthError } from '../lib/errors.js'; describe('Fetch Helpers Module', () => { afterEach(async () => { @@ -81,7 +82,11 @@ describe('Fetch Helpers Module', () => { const client = {} as any; const refreshSpy = vi.spyOn(refreshQueueModule, 'queuedRefresh'); - await expect(refreshAndUpdateToken(auth, client)).rejects.toThrow(); + const error = await refreshAndUpdateToken(auth, client).catch( + (err) => err as CodexAuthError, + ); + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(false); expect(refreshSpy).not.toHaveBeenCalled(); }); @@ -90,16 +95,61 @@ describe('Fetch Helpers Module', () => { const client = { auth: { set: 'bad' } } as any; const refreshSpy = vi.spyOn(refreshQueueModule, 'queuedRefresh'); - await expect(refreshAndUpdateToken(auth, client)).rejects.toThrow(); + const error = await refreshAndUpdateToken(auth, client).catch( + (err) => err as CodexAuthError, + ); + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(false); expect(refreshSpy).not.toHaveBeenCalled(); }); - it('throws when refresh fails', async () => { + it('throws retryable auth errors for transient refresh failures', async () => { const auth: Auth = { type: 'oauth', access: 'old', refresh: 'bad', expires: 0 }; const client = { auth: { set: vi.fn() } } as any; - vi.spyOn(refreshQueueModule, 'queuedRefresh').mockResolvedValue({ type: 'failed' } as any); + vi.spyOn(refreshQueueModule, 'queuedRefresh').mockResolvedValue({ + type: 'failed', + reason: 'network_error', + message: 'temporary network issue', + } as any); - await expect(refreshAndUpdateToken(auth, client)).rejects.toThrow(); + const error = await refreshAndUpdateToken(auth, client).catch( + (err) => err as CodexAuthError, + ); + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(true); + }); + + it('throws terminal auth errors for explicit invalid-refresh responses', async () => { + const auth: Auth = { type: 'oauth', access: 'old', refresh: 'bad', expires: 0 }; + const client = { auth: { set: vi.fn() } } as any; + vi.spyOn(refreshQueueModule, 'queuedRefresh').mockResolvedValue({ + type: 'failed', + reason: 'http_error', + statusCode: 401, + message: 'refresh token rejected', + } as any); + + const error = await refreshAndUpdateToken(auth, client).catch( + (err) => err as CodexAuthError, + ); + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(false); + }); + + it('treats missing refresh tokens as terminal auth errors', async () => { + const auth: Auth = { type: 'oauth', access: 'old', refresh: '', expires: 0 }; + const client = { auth: { set: vi.fn() } } as any; + vi.spyOn(refreshQueueModule, 'queuedRefresh').mockResolvedValue({ + type: 'failed', + reason: 'missing_refresh', + message: 'missing refresh token', + } as any); + + const error = await refreshAndUpdateToken(auth, client).catch( + (err) => err as CodexAuthError, + ); + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(false); }); it('updates stored auth on success', async () => { @@ -128,6 +178,52 @@ describe('Fetch Helpers Module', () => { expect(updated.refresh).toBe('newr'); expect(updated.expires).toBe(123); }); + + it('throws retryable auth errors when auth persistence fails', async () => { + const auth: Auth = { type: 'oauth', access: 'old', refresh: 'oldr', expires: 0 }; + const client = { + auth: { + set: vi.fn().mockRejectedValue( + Object.assign(new Error('persist failed'), { code: 'EBUSY' }), + ), + }, + } as any; + vi.spyOn(refreshQueueModule, 'queuedRefresh').mockResolvedValue({ + type: 'success', + access: 'new', + refresh: 'newr', + expires: 123, + } as any); + + const error = await refreshAndUpdateToken(auth, client).catch( + (err) => err as CodexAuthError, + ); + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(true); + }); + + it('throws terminal auth errors when auth persistence fails permanently', async () => { + const auth: Auth = { type: 'oauth', access: 'old', refresh: 'oldr', expires: 0 }; + const client = { + auth: { + set: vi.fn().mockRejectedValue( + Object.assign(new Error('persist failed'), { code: 'EACCES' }), + ), + }, + } as any; + vi.spyOn(refreshQueueModule, 'queuedRefresh').mockResolvedValue({ + type: 'success', + access: 'new', + refresh: 'newr', + expires: 123, + } as any); + + const error = await refreshAndUpdateToken(auth, client).catch( + (err) => err as CodexAuthError, + ); + expect(error).toBeInstanceOf(CodexAuthError); + expect(error.retryable).toBe(false); + }); }); describe('extractRequestUrl', () => { diff --git a/test/proactive-refresh.test.ts b/test/proactive-refresh.test.ts index 9528c102..582d653e 100644 --- a/test/proactive-refresh.test.ts +++ b/test/proactive-refresh.test.ts @@ -39,7 +39,10 @@ describe("proactive-refresh", () => { describe("shouldRefreshProactively", () => { it("returns false when no expiry is set", () => { - const account = createMockAccount({ expires: undefined }); + const account = createMockAccount({ + access: "test-access", + expires: undefined, + }); expect(shouldRefreshProactively(account)).toBe(false); }); @@ -51,6 +54,14 @@ describe("proactive-refresh", () => { expect(shouldRefreshProactively(account)).toBe(true); }); + it("returns true when access token is missing even if expiry is undefined", () => { + const account = createMockAccount({ + access: undefined, + expires: undefined, + }); + expect(shouldRefreshProactively(account)).toBe(true); + }); + it("returns true when token expires within buffer window", () => { const account = createMockAccount({ access: "test-access", @@ -360,6 +371,47 @@ describe("proactive-refresh", () => { expect(results.get(0)?.reason).toBe("success"); expect(results.get(1)?.reason).toBe("failed"); }); + + it("invokes onResult for each refreshed account", async () => { + const accounts = [ + createMockAccount({ + index: 0, + access: "access-0", + expires: Date.now() + 60_000, + refreshToken: "refresh-0", + }), + createMockAccount({ + index: 1, + access: "access-1", + expires: Date.now() + 10 * 60 * 1000, + refreshToken: "refresh-1", + }), + ]; + const onResult = vi.fn(); + + vi.mocked(refreshQueue.queuedRefresh).mockResolvedValue({ + type: "success" as const, + access: "new-access", + refresh: "new-refresh", + expires: Date.now() + 3_600_000, + }); + + const results = await refreshExpiringAccounts( + accounts, + DEFAULT_PROACTIVE_BUFFER_MS, + onResult, + ); + + expect(results.size).toBe(1); + expect(onResult).toHaveBeenCalledTimes(1); + expect(onResult).toHaveBeenCalledWith( + accounts[0], + expect.objectContaining({ + refreshed: true, + reason: "success", + }), + ); + }); }); describe("applyRefreshResult", () => { @@ -399,6 +451,36 @@ describe("proactive-refresh", () => { expect(account.refreshToken).toBe("same-refresh"); }); + + it("updates account identity fields from refreshed access tokens", () => { + const account = createMockAccount({ + access: "old-access", + expires: Date.now(), + refreshToken: "old-refresh", + accountId: "old-account-id", + accountIdSource: "token", + email: "old@example.com", + }); + const payload = Buffer.from( + JSON.stringify({ + email: "new@example.com", + "https://api.openai.com/auth": { + chatgpt_account_id: "new-account-id", + }, + }), + ).toString("base64url"); + + applyRefreshResult(account, { + type: "success", + access: `header.${payload}.signature`, + refresh: "new-refresh", + expires: Date.now() + 3_600_000, + }); + + expect(account.accountId).toBe("new-account-id"); + expect(account.accountIdSource).toBe("token"); + expect(account.email).toBe("new@example.com"); + }); }); describe("constants", () => { diff --git a/test/refresh-guardian.test.ts b/test/refresh-guardian.test.ts index 9584d5fe..c2ad3024 100644 --- a/test/refresh-guardian.test.ts +++ b/test/refresh-guardian.test.ts @@ -1,11 +1,33 @@ import { describe, expect, it, beforeEach, afterEach, vi } from "vitest"; import type { AccountManager, ManagedAccount } from "../lib/accounts.js"; +import { + extractAccountEmail, + extractAccountId, + sanitizeEmail, +} from "../lib/auth/token-utils.js"; +import { CodexAuthError } from "../lib/errors.js"; +import { findMatchingAccountIndex } from "../lib/storage.js"; +import type { OAuthAuthDetails } from "../lib/types.js"; const refreshExpiringAccountsMock = vi.fn(); const applyRefreshResultMock = vi.fn(); vi.mock("../lib/proactive-refresh.js", () => ({ - refreshExpiringAccounts: refreshExpiringAccountsMock, + refreshExpiringAccounts: vi.fn(async (accounts, bufferMs, onResult) => { + const results = await refreshExpiringAccountsMock(accounts, bufferMs, onResult); + if (results instanceof Map && typeof onResult === "function") { + for (const [accountIndex, result] of results.entries()) { + const sourceAccount = accounts.find( + (account) => account.index === accountIndex, + ); + if (!sourceAccount) { + continue; + } + await onResult(sourceAccount, result); + } + } + return results; + }), applyRefreshResult: applyRefreshResultMock, })); @@ -20,6 +42,57 @@ function createManagedAccount(index: number): ManagedAccount { }; } +function findAccountByIdentity( + accounts: ManagedAccount[], + candidate: Partial, + auth?: OAuthAuthDetails, +): ManagedAccount | null { + const derived = { + accountId: extractAccountId(auth?.access)?.trim() || undefined, + email: sanitizeEmail(extractAccountEmail(auth?.access)), + refreshToken: auth?.refresh, + }; + const lookupCandidates = [ + candidate, + { + accountId: candidate.accountId ?? derived.accountId, + email: candidate.email ?? derived.email, + refreshToken: candidate.refreshToken, + }, + { + accountId: derived.accountId ?? candidate.accountId, + email: derived.email ?? candidate.email, + refreshToken: candidate.refreshToken, + }, + { + accountId: derived.accountId ?? candidate.accountId, + email: derived.email ?? candidate.email, + refreshToken: derived.refreshToken ?? candidate.refreshToken, + }, + ]; + + const seen = new Set(); + for (const lookupCandidate of lookupCandidates) { + const key = `${lookupCandidate.accountId ?? ""}|${lookupCandidate.email ?? ""}|${lookupCandidate.refreshToken ?? ""}`; + if (seen.has(key)) { + continue; + } + seen.add(key); + const matchIndex = findMatchingAccountIndex(accounts, { + accountId: lookupCandidate.accountId, + email: lookupCandidate.email, + refreshToken: lookupCandidate.refreshToken, + }, { + allowUniqueAccountIdFallbackWithoutEmail: true, + }); + if (matchIndex !== undefined) { + return accounts[matchIndex] ?? null; + } + } + + return null; +} + function createManagerMock(accounts: ManagedAccount[]): AccountManager { return { getAccountsSnapshot: vi.fn(() => accounts), @@ -27,6 +100,17 @@ function createManagerMock(accounts: ManagedAccount[]): AccountManager { (index: number) => accounts.find((account) => account.index === index) ?? null, ), + isAccountCoolingDown: vi.fn( + (account: ManagedAccount) => + typeof account.coolingDownUntil === "number" && + account.coolingDownUntil > Date.now(), + ), + getAccountByIdentity: vi.fn((candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity(accounts, candidate, auth), + ), + commitRefreshedAuth: vi.fn(async (candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity(accounts, candidate, auth), + ), clearAuthFailures: vi.fn(), markAccountCoolingDown: vi.fn(), setAccountEnabled: vi.fn(), @@ -108,6 +192,33 @@ describe("refresh-guardian", () => { expect(stats.lastRunAt).not.toBeNull(); }); + it("skips accounts that are already cooling down", async () => { + const coolingAccount = { + ...createManagedAccount(0), + coolingDownUntil: Date.now() + 60_000, + cooldownReason: "auth-failure" as const, + }; + const readyAccount = createManagedAccount(1); + const manager = createManagerMock([coolingAccount, readyAccount]); + const { RefreshGuardian } = await import("../lib/refresh-guardian.js"); + const guardian = new RefreshGuardian(() => manager, { + intervalMs: 5_000, + bufferMs: 60_000, + }); + + refreshExpiringAccountsMock.mockResolvedValue(new Map()); + + await guardian.tick(); + + expect(refreshExpiringAccountsMock).toHaveBeenCalledWith( + [readyAccount], + 60_000, + expect.any(Function), + ); + expect(manager.markAccountCoolingDown as ReturnType).not.toHaveBeenCalled(); + expect(guardian.getStats().runs).toBe(1); + }); + it("applies refresh outcomes and updates stats", async () => { const accountA = createManagedAccount(0); const accountB = createManagedAccount(1); @@ -151,14 +262,20 @@ describe("refresh-guardian", () => { await guardian.tick(); expect(refreshExpiringAccountsMock).toHaveBeenCalledTimes(1); - expect(applyRefreshResultMock).toHaveBeenCalledTimes(1); - expect(applyRefreshResultMock).toHaveBeenCalledWith( - accountA, - expect.objectContaining({ type: "success" }), - ); + expect(applyRefreshResultMock).not.toHaveBeenCalled(); expect( manager.clearAuthFailures as ReturnType, - ).toHaveBeenCalledWith(accountA); + ).not.toHaveBeenCalled(); + expect( + manager.commitRefreshedAuth as ReturnType, + ).toHaveBeenCalledWith( + accountA, + expect.objectContaining({ + type: "oauth", + access: "access-0", + refresh: "refresh-0-new", + }), + ); expect( manager.markAccountCoolingDown as ReturnType, ).toHaveBeenCalledWith(accountB, 60_000, "auth-failure"); @@ -232,6 +349,13 @@ describe("refresh-guardian", () => { (index: number) => [liveB, liveA].find((account) => account.index === index) ?? null, ), + isAccountCoolingDown: vi.fn(() => false), + getAccountByIdentity: vi.fn((candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity([liveB, liveA], candidate, auth), + ), + commitRefreshedAuth: vi.fn(async (candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity([liveB, liveA], candidate, auth), + ), clearAuthFailures: vi.fn(), markAccountCoolingDown: vi.fn(), saveToDiskDebounced: vi.fn(), @@ -262,14 +386,17 @@ describe("refresh-guardian", () => { await guardian.tick(); - expect(applyRefreshResultMock).toHaveBeenCalledTimes(1); - expect(applyRefreshResultMock).toHaveBeenCalledWith( - liveB, - expect.objectContaining({ type: "success" }), - ); + expect(applyRefreshResultMock).not.toHaveBeenCalled(); expect( - manager.clearAuthFailures as ReturnType, - ).toHaveBeenCalledWith(liveB); + manager.commitRefreshedAuth as ReturnType, + ).toHaveBeenCalledWith( + originalB, + expect.objectContaining({ + type: "oauth", + access: "access-shifted", + refresh: "refresh-shifted", + }), + ); }); it("classifies failure reasons and handles no-op branches", async () => { @@ -395,6 +522,13 @@ describe("refresh-guardian", () => { (index: number) => liveSnapshot.find((account) => account.index === index) ?? null, ), + isAccountCoolingDown: vi.fn(() => false), + getAccountByIdentity: vi.fn((candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity(liveSnapshot, candidate, auth), + ), + commitRefreshedAuth: vi.fn(async (candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity(liveSnapshot, candidate, auth), + ), clearAuthFailures: vi.fn(), markAccountCoolingDown: vi.fn(), setAccountEnabled: vi.fn(), @@ -560,6 +694,13 @@ describe("refresh-guardian", () => { getAccountByIndex: vi.fn( (index: number) => liveAfterRemoval[index] ?? null, ), + isAccountCoolingDown: vi.fn(() => false), + getAccountByIdentity: vi.fn((candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity(liveAfterRemoval, candidate, auth), + ), + commitRefreshedAuth: vi.fn(async (candidate: Partial, auth?: OAuthAuthDetails) => + findAccountByIdentity(liveAfterRemoval, candidate, auth), + ), clearAuthFailures: vi.fn(), markAccountCoolingDown: vi.fn(), setAccountEnabled: vi.fn(), @@ -604,10 +745,7 @@ describe("refresh-guardian", () => { ); await expect(guardian.tick()).resolves.toBeUndefined(); - expect(applyRefreshResultMock).not.toHaveBeenCalledWith( - expect.objectContaining({ refreshToken: originalA.refreshToken }), - expect.anything(), - ); + expect(applyRefreshResultMock).not.toHaveBeenCalled(); expect( manager.markAccountCoolingDown as ReturnType, ).toHaveBeenCalledWith( @@ -616,4 +754,168 @@ describe("refresh-guardian", () => { "rate-limit", ); }); + + it("treats null commit results as retryable network cooldowns", async () => { + const accountA = createManagedAccount(0); + const manager = createManagerMock([accountA]); + const commitRefreshedAuthMock = manager + .commitRefreshedAuth as ReturnType; + commitRefreshedAuthMock.mockResolvedValueOnce(null); + + const { RefreshGuardian } = await import("../lib/refresh-guardian.js"); + const guardian = new RefreshGuardian(() => manager, { + bufferMs: 60_000, + intervalMs: 5_000, + }); + + refreshExpiringAccountsMock.mockResolvedValue( + new Map([ + [ + 0, + { + refreshed: true, + reason: "success", + tokenResult: { + type: "success", + access: "access-0", + refresh: "refresh-0-new", + expires: Date.now() + 3_600_000, + }, + }, + ], + ]), + ); + + await guardian.tick(); + + expect(applyRefreshResultMock).not.toHaveBeenCalled(); + expect(commitRefreshedAuthMock).toHaveBeenCalledTimes(1); + expect( + manager.markAccountCoolingDown as ReturnType, + ).toHaveBeenCalledWith(accountA, 60_000, "network-error"); + + const stats = guardian.getStats(); + expect(stats.runs).toBe(1); + expect(stats.refreshed).toBe(0); + expect(stats.failed).toBe(1); + expect(stats.networkFailed).toBe(1); + }); + + it("treats commit failures as retryable network cooldowns and continues the batch", async () => { + const accountA = createManagedAccount(0); + const accountB = createManagedAccount(1); + const manager = createManagerMock([accountA, accountB]); + const commitRefreshedAuthMock = manager + .commitRefreshedAuth as ReturnType; + commitRefreshedAuthMock.mockImplementation( + async (candidate: Partial, auth?: OAuthAuthDetails) => { + if (candidate.refreshToken === accountA.refreshToken) { + throw Object.assign(new Error("EBUSY"), { code: "EBUSY" }); + } + return manager.getAccountByIdentity(candidate, auth); + }, + ); + + const { RefreshGuardian } = await import("../lib/refresh-guardian.js"); + const guardian = new RefreshGuardian(() => manager, { + bufferMs: 60_000, + intervalMs: 5_000, + }); + + refreshExpiringAccountsMock.mockResolvedValue( + new Map([ + [ + 0, + { + refreshed: true, + reason: "success", + tokenResult: { + type: "success", + access: "access-0", + refresh: "refresh-0-new", + expires: Date.now() + 3_600_000, + }, + }, + ], + [ + 1, + { + refreshed: true, + reason: "failed", + tokenResult: { + type: "failed", + reason: "http_error", + statusCode: 429, + message: "rate limited", + }, + }, + ], + ]), + ); + + await guardian.tick(); + + expect(applyRefreshResultMock).not.toHaveBeenCalled(); + expect(commitRefreshedAuthMock).toHaveBeenCalledTimes(1); + expect( + manager.markAccountCoolingDown as ReturnType, + ).toHaveBeenNthCalledWith(1, accountA, 60_000, "network-error"); + expect( + manager.markAccountCoolingDown as ReturnType, + ).toHaveBeenNthCalledWith(2, accountB, 60_000, "rate-limit"); + expect( + manager.saveToDiskDebounced as ReturnType, + ).toHaveBeenCalledTimes(1); + + const stats = guardian.getStats(); + expect(stats.runs).toBe(1); + expect(stats.refreshed).toBe(0); + expect(stats.failed).toBe(2); + expect(stats.networkFailed).toBe(1); + expect(stats.rateLimited).toBe(1); + }); + + it("treats non-retryable commit failures as auth cooldowns", async () => { + const accountA = createManagedAccount(0); + const manager = createManagerMock([accountA]); + const commitRefreshedAuthMock = manager + .commitRefreshedAuth as ReturnType; + commitRefreshedAuthMock.mockRejectedValueOnce( + new CodexAuthError("refresh persistence failed", { retryable: false }), + ); + + const { RefreshGuardian } = await import("../lib/refresh-guardian.js"); + const guardian = new RefreshGuardian(() => manager, { + bufferMs: 60_000, + intervalMs: 5_000, + }); + + refreshExpiringAccountsMock.mockResolvedValue( + new Map([ + [ + 0, + { + refreshed: true, + reason: "success", + tokenResult: { + type: "success", + access: "access-0", + refresh: "refresh-0-new", + expires: Date.now() + 3_600_000, + }, + }, + ], + ]), + ); + + await guardian.tick(); + + expect( + manager.markAccountCoolingDown as ReturnType, + ).toHaveBeenCalledWith(accountA, 60_000, "auth-failure"); + const stats = guardian.getStats(); + expect(stats.failed).toBe(1); + expect(stats.authFailed).toBe(1); + expect(stats.networkFailed).toBe(0); + }); });