codex-multi-auth 1.2.1 → 1.2.2

package/dist/lib/accounts.js

@@ -1,20 +1,112 @@
 import { createLogger } from "./logger.js";
-import { loadAccounts, saveAccounts, findMatchingAccountIndex, } from "./storage.js";
+import { loadAccounts, saveAccounts, findMatchingAccountIndex, withAccountStorageTransaction, } from "./storage.js";
 import { MODEL_FAMILIES } from "./prompts/codex.js";
 import { getHealthTracker, getTokenTracker, selectHybridAccount, } from "./rotation.js";
 import { nowMs } from "./utils.js";
+import { ERROR_MESSAGES, HTTP_STATUS } from "./constants.js";
+import { CodexAuthError } from "./errors.js";
 import { loadCodexCliState, } from "./codex-cli/state.js";
 import { syncAccountStorageFromCodexCli } from "./codex-cli/sync.js";
 import { setCodexCliActiveSelection } from "./codex-cli/writer.js";
+import { getAccountIdentityKey, getRuntimeAccountIdentityKey, } from "./storage/identity.js";
+import { getCircuitBreaker } from "./circuit-breaker.js";
 export { extractAccountId, extractAccountEmail, getAccountIdCandidates, selectBestAccountCandidate, resolveRuntimeRequestIdentity, shouldUpdateAccountIdFromToken, resolveRequestAccountId, sanitizeEmail, } from "./auth/token-utils.js";
 export { parseRateLimitReason, getQuotaKey, clampNonNegativeInt, clearExpiredRateLimits, isRateLimitedForQuotaKey, isRateLimitedForFamily, formatWaitTime, } from "./accounts/rate-limits.js";
 export { lookupCodexCliTokensByEmail, isCodexCliSyncEnabled, } from "./codex-cli/state.js";
 import { extractAccountId, extractAccountEmail, shouldUpdateAccountIdFromToken, sanitizeEmail, } from "./auth/token-utils.js";
 import { clampNonNegativeInt, getQuotaKey, clearExpiredRateLimits, isRateLimitedForFamily, formatWaitTime, } from "./accounts/rate-limits.js";
 const log = createLogger("accounts");
+let nextRuntimeCircuitKeyId = 0;
+function getAccountCircuitKey(account) {
+    if (!account.circuitKeyId) {
+        account.circuitKeyId =
+            getAccountIdentityKey(account) ?? `circuit:${nextRuntimeCircuitKeyId++}`;
+    }
+    return account.circuitKeyId;
+}
+export function getRuntimeTrackerKey(account) {
+    if (account._runtimeTrackerKey !== undefined) {
+        return account._runtimeTrackerKey;
+    }
+    const trackerKey = getRuntimeAccountIdentityKey(account) ?? account.index;
+    account._runtimeTrackerKey = trackerKey;
+    return trackerKey;
+}
 function initFamilyState(defaultValue) {
     return Object.fromEntries(MODEL_FAMILIES.map((family) => [family, defaultValue]));
 }
+function getAuthIdentityCandidate(auth) {
+    const accountId = extractAccountId(auth?.access)?.trim() || undefined;
+    const email = sanitizeEmail(extractAccountEmail(auth?.access));
+    return {
+        accountId,
+        email,
+        refreshToken: auth?.refresh,
+    };
+}
+function buildAccountIdentityCandidates(source, auth) {
+    const derived = getAuthIdentityCandidate(auth);
+    const candidates = [];
+    const seen = new Set();
+    const pushCandidate = (candidate) => {
+        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(accounts, source, auth) {
+    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) {
+    if (!error || typeof error !== "object") {
+        return false;
+    }
+    const candidate = error;
+    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 class AccountManager {
     accounts = [];
     cursorByFamily = initFamilyState(0);
@@ -227,10 +319,16 @@ export class AccountManager {
         return index;
     }
     getAccountsSnapshot() {
-        return this.accounts.map((account) => ({
-            ...account,
-            rateLimitResetTimes: { ...account.rateLimitResetTimes },
-        }));
+        return this.accounts.map((account) => {
+            const trackerKey = getRuntimeTrackerKey(account);
+            const circuitKeyId = getAccountCircuitKey(account);
+            return {
+                ...account,
+                _runtimeTrackerKey: trackerKey,
+                circuitKeyId,
+                rateLimitResetTimes: { ...account.rateLimitResetTimes },
+            };
+        });
     }
     getAccountByIndex(index) {
         if (!Number.isFinite(index))
@@ -247,7 +345,9 @@ export class AccountManager {
         if (account.enabled === false)
             return false;
         clearExpiredRateLimits(account);
-        return !isRateLimitedForFamily(account, family, model) && !this.isAccountCoolingDown(account);
+        return (!isRateLimitedForFamily(account, family, model) &&
+            !this.isAccountCoolingDown(account) &&
+            this.isCircuitAvailable(account));
     }
     setActiveIndex(index) {
         if (!Number.isFinite(index))
@@ -314,7 +414,9 @@ export class AccountManager {
             if (account.enabled === false)
                 continue;
             clearExpiredRateLimits(account);
-            if (isRateLimitedForFamily(account, family, model) || this.isAccountCoolingDown(account)) {
+            if (isRateLimitedForFamily(account, family, model) ||
+                this.isAccountCoolingDown(account) ||
+                !this.isCircuitAvailable(account)) {
                 continue;
             }
             this.cursorByFamily[family] = (idx + 1) % count;
@@ -337,7 +439,9 @@ export class AccountManager {
             if (account.enabled === false)
                 continue;
             clearExpiredRateLimits(account);
-            if (isRateLimitedForFamily(account, family, model) || this.isAccountCoolingDown(account)) {
+            if (isRateLimitedForFamily(account, family, model) ||
+                this.isAccountCoolingDown(account) ||
+                !this.isCircuitAvailable(account)) {
                 continue;
             }
             this.cursorByFamily[family] = (idx + 1) % count;
@@ -350,23 +454,6 @@ export class AccountManager {
         const count = this.accounts.length;
         if (count === 0)
             return null;
-        const currentIndex = this.currentAccountIndexByFamily[family];
-        if (currentIndex >= 0 && currentIndex < count) {
-            const currentAccount = this.accounts[currentIndex];
-            if (currentAccount) {
-                if (currentAccount.enabled === false) {
-                    // Fall through to hybrid selection.
-                }
-                else {
-                    clearExpiredRateLimits(currentAccount);
-                    if (!isRateLimitedForFamily(currentAccount, family, model) &&
-                        !this.isAccountCoolingDown(currentAccount)) {
-                        currentAccount.lastUsed = nowMs();
-                        return currentAccount;
-                    }
-                }
-            }
-        }
         const quotaKey = model ? `${family}:${model}` : family;
         const healthTracker = getHealthTracker();
         const tokenTracker = getTokenTracker();
@@ -377,9 +464,12 @@ export class AccountManager {
             if (account.enabled === false)
                 return null;
             clearExpiredRateLimits(account);
-            const isAvailable = !isRateLimitedForFamily(account, family, model) && !this.isAccountCoolingDown(account);
+            const isAvailable = !isRateLimitedForFamily(account, family, model) &&
+                !this.isAccountCoolingDown(account) &&
+                this.isCircuitAvailable(account);
             return {
                 index: account.index,
+                trackerKey: getRuntimeTrackerKey(account),
                 isAvailable,
                 lastUsed: account.lastUsed,
             };
@@ -399,24 +489,53 @@ export class AccountManager {
     recordSuccess(account, family, model) {
         const quotaKey = model ? `${family}:${model}` : family;
         const healthTracker = getHealthTracker();
-        healthTracker.recordSuccess(account.index, quotaKey);
+        healthTracker.recordSuccess(getRuntimeTrackerKey(account), quotaKey);
+        const hadCooldownMetadata = account.coolingDownUntil !== undefined || account.cooldownReason !== undefined;
+        const hadAuthFailures = (account.consecutiveAuthFailures ?? 0) > 0;
+        const isCoolingDown = this.isAccountCoolingDown(account);
+        let healed = false;
+        if (!isCoolingDown && hadCooldownMetadata) {
+            this.clearAccountCooldown(account);
+            healed = true;
+        }
+        if (!isCoolingDown && hadAuthFailures) {
+            this.clearAuthFailures(account);
+            healed = true;
+        }
+        if (healed) {
+            this.saveToDiskDebounced();
+        }
+        getCircuitBreaker(getAccountCircuitKey(account)).recordSuccess();
     }
     recordRateLimit(account, family, model) {
         const quotaKey = model ? `${family}:${model}` : family;
         const healthTracker = getHealthTracker();
         const tokenTracker = getTokenTracker();
-        healthTracker.recordRateLimit(account.index, quotaKey);
-        tokenTracker.drain(account.index, quotaKey);
+        const trackerKey = getRuntimeTrackerKey(account);
+        healthTracker.recordRateLimit(trackerKey, quotaKey);
+        tokenTracker.drain(trackerKey, quotaKey);
     }
     recordFailure(account, family, model) {
         const quotaKey = model ? `${family}:${model}` : family;
         const healthTracker = getHealthTracker();
-        healthTracker.recordFailure(account.index, quotaKey);
+        healthTracker.recordFailure(getRuntimeTrackerKey(account), quotaKey);
+        getCircuitBreaker(getAccountCircuitKey(account)).recordFailure();
     }
     consumeToken(account, family, model) {
         const quotaKey = model ? `${family}:${model}` : family;
         const tokenTracker = getTokenTracker();
-        return tokenTracker.tryConsume(account.index, quotaKey);
+        const trackerKey = getRuntimeTrackerKey(account);
+        if (!tokenTracker.tryConsume(trackerKey, quotaKey)) {
+            return false;
+        }
+        try {
+            getCircuitBreaker(getAccountCircuitKey(account)).canExecute();
+            return true;
+        }
+        catch {
+            tokenTracker.refundToken(trackerKey, quotaKey);
+            return false;
+        }
     }
     /**
      * Refund a token consumed within the refund window (30 seconds).
@@ -426,7 +545,7 @@ export class AccountManager {
     refundToken(account, family, model) {
         const quotaKey = model ? `${family}:${model}` : family;
         const tokenTracker = getTokenTracker();
-        return tokenTracker.refundToken(account.index, quotaKey);
+        return tokenTracker.refundToken(getRuntimeTrackerKey(account), quotaKey);
     }
     markSwitched(account, reason, family) {
         account.lastSwitchReason = reason;
@@ -439,8 +558,11 @@ export class AccountManager {
         const retryMs = Math.max(0, Math.floor(retryAfterMs));
         const resetAt = nowMs() + retryMs;
         const baseKey = getQuotaKey(family);
-        account.rateLimitResetTimes[baseKey] = resetAt;
-        if (model) {
+        if (!model || reason === "quota" || reason === "unknown") {
+            account.rateLimitResetTimes[baseKey] = resetAt;
+        }
+        if (model &&
+            (reason === "tokens" || reason === "concurrent" || reason === "unknown")) {
             const modelKey = getQuotaKey(family, model);
             account.rateLimitResetTimes[modelKey] = resetAt;
         }
@@ -464,6 +586,9 @@ export class AccountManager {
         delete account.coolingDownUntil;
         delete account.cooldownReason;
     }
+    isCircuitAvailable(account) {
+        return getCircuitBreaker(getAccountCircuitKey(account)).isAvailable();
+    }
     incrementAuthFailures(account) {
         account.consecutiveAuthFailures = (account.consecutiveAuthFailures ?? 0) + 1;
         return account.consecutiveAuthFailures;
@@ -471,6 +596,13 @@ export class AccountManager {
     clearAuthFailures(account) {
         account.consecutiveAuthFailures = 0;
     }
+    getAccountByIdentity(candidate, auth) {
+        const index = findAccountIndexByIdentity(this.accounts, candidate, auth);
+        if (index === undefined) {
+            return null;
+        }
+        return this.accounts[index] ?? null;
+    }
     shouldShowAccountToast(accountIndex, debounceMs = 30000) {
         const now = nowMs();
         if (accountIndex === this.lastToastAccountIndex && now - this.lastToastTime < debounceMs) {
@@ -486,7 +618,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))) {
             account.accountId = tokenAccountId;
@@ -494,6 +626,134 @@ export class AccountManager {
         }
         account.email = sanitizeEmail(extractAccountEmail(auth.access)) ?? account.email;
     }
+    buildStorageSnapshot() {
+        const activeIndexByFamily = {};
+        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, auth) {
+        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());
+                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) {
         return {
             type: "oauth",
@@ -510,7 +770,9 @@ export class AccountManager {
         const enabledAccounts = this.accounts.filter((account) => account.enabled !== false);
         const available = enabledAccounts.filter((account) => {
             clearExpiredRateLimits(account);
-            return !isRateLimitedForFamily(account, family, model) && !this.isAccountCoolingDown(account);
+            return (!isRateLimitedForFamily(account, family, model) &&
+                !this.isAccountCoolingDown(account) &&
+                this.isCircuitAvailable(account));
         });
         if (available.length > 0)
             return 0;
@@ -520,18 +782,26 @@ export class AccountManager {
         const baseKey = getQuotaKey(family);
         const modelKey = model ? getQuotaKey(family, model) : null;
         for (const account of enabledAccounts) {
+            const perAccountWaitTimes = [];
             const baseResetAt = account.rateLimitResetTimes[baseKey];
             if (typeof baseResetAt === "number") {
-                waitTimes.push(Math.max(0, baseResetAt - now));
+                perAccountWaitTimes.push(Math.max(0, baseResetAt - now));
             }
             if (modelKey) {
                 const modelResetAt = account.rateLimitResetTimes[modelKey];
                 if (typeof modelResetAt === "number") {
-                    waitTimes.push(Math.max(0, modelResetAt - now));
+                    perAccountWaitTimes.push(Math.max(0, modelResetAt - now));
                 }
             }
             if (typeof account.coolingDownUntil === "number") {
-                waitTimes.push(Math.max(0, account.coolingDownUntil - now));
+                perAccountWaitTimes.push(Math.max(0, account.coolingDownUntil - now));
+            }
+            const breakerWait = getCircuitBreaker(getAccountCircuitKey(account)).getTimeUntilAvailable();
+            if (breakerWait > 0) {
+                perAccountWaitTimes.push(breakerWait);
+            }
+            if (perAccountWaitTimes.length > 0) {
+                waitTimes.push(Math.max(...perAccountWaitTimes));
             }
         }
         return waitTimes.length > 0 ? Math.min(...waitTimes) : 0;
@@ -596,36 +866,9 @@ export class AccountManager {
         return account;
     }
     async saveToDisk() {
-        const activeIndexByFamily = {};
-        for (const family of MODEL_FAMILIES) {
-            const raw = this.currentAccountIndexByFamily[family];
-            activeIndexByFamily[family] = clampNonNegativeInt(raw, 0);
-        }
-        const activeIndex = clampNonNegativeInt(activeIndexByFamily.codex, 0);
-        const storage = {
-            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) {
         if (this.saveDebounceTimer) {