diff --git a/index.ts b/index.ts index e29c9179..900621ec 100644 --- a/index.ts +++ b/index.ts @@ -88,6 +88,7 @@ import { } from "./lib/logger.js"; import { checkAndNotify } from "./lib/auto-update-checker.js"; import { handleContextOverflow } from "./lib/context-overflow.js"; +import { registerCleanup, unregisterCleanup } from "./lib/shutdown.js"; import { AccountManager, type AccountSelectionExplainability, @@ -118,6 +119,7 @@ import { loadFlaggedAccounts, saveFlaggedAccounts, clearFlaggedAccounts, + type AccountDisabledReason, StorageError, formatStorageErrorHint, type AccountStorageV3, @@ -199,6 +201,14 @@ import { * } * ``` */ +// Shared across plugin instances so shutdown cleanup can flush debounced saves +// even when multiple plugin objects coexist during reloads or tests. +let accountManagerCleanupHook: (() => Promise) | null = null; +// `active...` retains only the shared currently cached manager. Older managers move to +// `tracked...` only while they still have pending disk work and self-prune once settled. +const activeAccountManagersForCleanup = new Set(); +const trackedAccountManagersForCleanup = new Set(); + // eslint-disable-next-line @typescript-eslint/require-await export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { initLogger(client); @@ -471,15 +481,27 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { ); }; + type PersistAccountPoolOptions = { + reviveMatchingDisabledAccounts?: boolean; + }; + const persistAccountPool = async ( results: TokenSuccessWithAccount[], replaceAll: boolean = false, + options: PersistAccountPoolOptions = {}, ): Promise => { if (results.length === 0) return; await withAccountStorageTransaction(async (loadedStorage, persist) => { const now = Date.now(); const stored = replaceAll ? null : loadedStorage; let accounts = stored?.accounts ? [...stored.accounts] : []; + const refreshTokensToRevive = options.reviveMatchingDisabledAccounts + ? new Set( + results + .map((result) => result.refresh.trim()) + .filter((refreshToken) => refreshToken.length > 0), + ) + : null; const pushIndex = ( map: Map, @@ -548,6 +570,21 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { target.enabled === false || source.enabled === false ? false : target.enabled ?? source.enabled; + const mergedDisabledReason = (() => { + if (mergedEnabled !== false) { + return undefined; + } + const reasons = [target.disabledReason, source.disabledReason]; + if (reasons.includes("user")) { + return "user" satisfies AccountDisabledReason; + } + if (reasons.includes("auth-failure")) { + return "auth-failure" satisfies AccountDisabledReason; + } + // Unknown disabled reasons are treated as legacy/manual disables. + // Explicit login revival is reserved for accounts we know were disabled by auth failure. + return "user" satisfies AccountDisabledReason; + })(); const targetCoolingDownUntil = typeof target.coolingDownUntil === "number" && Number.isFinite(target.coolingDownUntil) ? target.coolingDownUntil @@ -585,6 +622,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { accessToken: newer.accessToken || older.accessToken, expiresAt: newer.expiresAt ?? older.expiresAt, enabled: mergedEnabled, + disabledReason: mergedDisabledReason, addedAt: Math.max(target.addedAt ?? 0, source.addedAt ?? 0), lastUsed: Math.max(target.lastUsed ?? 0, source.lastUsed ?? 0), lastSwitchReason: target.lastSwitchReason ?? source.lastSwitchReason, @@ -1005,6 +1043,27 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { pruneRefreshTokenCollisions(); + if (refreshTokensToRevive && refreshTokensToRevive.size > 0) { + accounts = accounts.map((account) => { + const refreshToken = account.refreshToken?.trim(); + if ( + !refreshToken || + !refreshTokensToRevive.has(refreshToken) || + account.disabledReason !== "auth-failure" + ) { + return account; + } + + return { + ...account, + enabled: undefined, + disabledReason: undefined, + coolingDownUntil: undefined, + cooldownReason: undefined, + }; + }); + } + if (accounts.length === 0) return; const resolveIndexByIdentityKeys = (identityKeys: string[] | undefined): number | undefined => { @@ -1069,7 +1128,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { message: string, variant: "info" | "success" | "warning" | "error" = "success", options?: { title?: string; duration?: number }, - ): Promise => { + ): Promise => { try { await client.tui.showToast({ body: { @@ -1079,8 +1138,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { ...(options?.duration && { duration: options.duration }), }, }); + return true; } catch { // Ignore when TUI is not available. + return false; } }; @@ -1627,11 +1688,94 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { } }; + const managerHasPendingSave = (candidate: AccountManager): boolean => + typeof candidate.hasPendingSave === "function" ? candidate.hasPendingSave() : true; + const trackedManagerSettledWaits = new WeakMap>(); + + const scheduleTrackedManagerPrune = (manager: AccountManager): void => { + if ( + !managerHasPendingSave(manager) || + trackedManagerSettledWaits.has(manager) || + typeof manager.waitForPendingSaveToSettle !== "function" + ) { + return; + } + + // Re-registering this waiter after each settle is intentional: a manager can + // move from "debounce only" to "in-flight save" repeatedly, and we do not + // want a transient idle check to drop shutdown tracking before the next save. + const waitForSettle = manager + .waitForPendingSaveToSettle() + .finally(() => { + trackedManagerSettledWaits.delete(manager); + if (activeAccountManagersForCleanup.has(manager)) { + return; + } + if (managerHasPendingSave(manager)) { + trackedAccountManagersForCleanup.add(manager); + scheduleTrackedManagerPrune(manager); + return; + } + trackedAccountManagersForCleanup.delete(manager); + }); + + trackedManagerSettledWaits.set(manager, waitForSettle); + }; + + const setCachedAccountManager = (manager: AccountManager): AccountManager => { + if (cachedAccountManager && cachedAccountManager !== manager) { + activeAccountManagersForCleanup.delete(cachedAccountManager); + // Detached managers can still enqueue their first debounced save after the + // cache switches. Keep them tracked until shutdown so that late work is flushed. + trackedAccountManagersForCleanup.add(cachedAccountManager); + if (managerHasPendingSave(cachedAccountManager)) { + scheduleTrackedManagerPrune(cachedAccountManager); + } + } + activeAccountManagersForCleanup.add(manager); + trackedAccountManagersForCleanup.delete(manager); + cachedAccountManager = manager; + accountManagerPromise = Promise.resolve(manager); + return manager; + }; + const invalidateAccountManagerCache = (): void => { + if (cachedAccountManager) { + activeAccountManagersForCleanup.delete(cachedAccountManager); + // Once detached, this manager may still persist state from in-flight request + // handlers. Keep it in the shutdown tracking set until cleanup runs. + trackedAccountManagersForCleanup.add(cachedAccountManager); + if (managerHasPendingSave(cachedAccountManager)) { + scheduleTrackedManagerPrune(cachedAccountManager); + } + } cachedAccountManager = null; accountManagerPromise = null; }; + accountManagerCleanupHook ??= async (context?: { deadlineMs?: number }) => { + const managersToFlush = new Set([ + ...activeAccountManagersForCleanup, + ...trackedAccountManagersForCleanup, + ]); + for (const manager of managersToFlush) { + try { + await manager.flushPendingSave( + context?.deadlineMs !== undefined ? { deadlineMs: context.deadlineMs } : undefined, + ); + } catch (error) { + logWarn("[shutdown] flushPendingSave failed; disabled state may not be persisted", { + error: error instanceof Error ? error.message : String(error), + }); + } finally { + activeAccountManagersForCleanup.delete(manager); + trackedAccountManagersForCleanup.delete(manager); + } + } + }; + unregisterCleanup(accountManagerCleanupHook); + registerCleanup(accountManagerCleanupHook); + // Event handler for session recovery and account selection const eventHandler = async (input: { event: { type: string; properties?: unknown } }) => { try { @@ -1673,8 +1817,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { // refresh tokens with stale in-memory state. if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } await showToast(`Switched to account ${index + 1}`, "info"); @@ -1743,8 +1886,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { if (!accountManagerPromise) { accountManagerPromise = AccountManager.loadFromDisk(authFallback); } - let accountManager = await accountManagerPromise; - cachedAccountManager = accountManager; + let accountManager = setCachedAccountManager(await accountManagerPromise); const refreshToken = authFallback?.refresh ?? ""; const needsPersist = refreshToken && @@ -2083,7 +2225,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { } while (true) { - let accountCount = accountManager.getAccountCount(); + let accountCount = accountManager.getEnabledAccountCount(); const attempted = new Set(); let restartAccountTraversalWithFallback = false; @@ -2116,7 +2258,7 @@ while (attempted.size < Math.max(1, accountCount)) { } // Log account selection for debugging rotation logDebug( - `Using account ${account.index + 1}/${accountCount}: ${account.email ?? "unknown"} for ${modelFamily}`, + `Using account ${account.index + 1} (enabled ${attempted.size}/${accountCount}): ${account.email ?? "unknown"} for ${modelFamily}`, ); let accountAuth = accountManager.toAuthDetails(account) as OAuthAuthDetails; @@ -2161,11 +2303,11 @@ while (attempted.size < Math.max(1, accountCount)) { const failures = accountManager.incrementAuthFailures(account); const accountLabel = formatAccountLabel(account, account.index); - if (failures >= ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_REMOVAL) { - const removedCount = accountManager.removeAccountsWithSameRefreshToken(account); - if (removedCount <= 0) { + if (failures >= ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_DISABLE) { + const disabledCount = accountManager.disableAccountsWithSameRefreshToken(account); + if (disabledCount <= 0) { logWarn( - `[${PLUGIN_NAME}] Expected grouped account removal after auth failures, but removed ${removedCount}.`, + `[${PLUGIN_NAME}] Expected grouped account disable after auth failures, but disabled ${disabledCount}.`, ); const cooledCount = accountManager.markAccountsWithRefreshTokenCoolingDown( account.refreshToken, @@ -2181,17 +2323,17 @@ while (attempted.size < Math.max(1, accountCount)) { continue; } accountManager.saveToDiskDebounced(); - const removalMessage = removedCount > 1 - ? `Removed ${removedCount} accounts (same refresh token) after ${failures} consecutive auth failures. Run 'opencode auth login' to re-add.` - : `Removed ${accountLabel} after ${failures} consecutive auth failures. Run 'opencode auth login' to re-add.`; + const disableMessage = disabledCount > 1 + ? `Disabled ${disabledCount} accounts (same refresh token) after ${failures} consecutive auth failures. Run 'opencode auth login' to re-enable.` + : `Disabled ${accountLabel} after ${failures} consecutive auth failures. Run 'opencode auth login' to re-enable.`; await showToast( - removalMessage, + disableMessage, "error", { duration: toastDurationMs * 2 }, ); // Restart traversal: clear attempted and refresh accountCount to avoid skipping healthy accounts attempted.clear(); - accountCount = accountManager.getAccountCount(); + accountCount = accountManager.getEnabledAccountCount(); continue; } @@ -2236,7 +2378,7 @@ while (attempted.size < Math.max(1, accountCount)) { ) { const accountLabel = formatAccountLabel(account, account.index); await showToast( - `Using ${accountLabel} (${account.index + 1}/${accountCount})`, + `Using ${accountLabel} (enabled ${attempted.size}/${accountCount})`, "info", ); accountManager.markToastShown(account.index); @@ -2549,7 +2691,7 @@ while (attempted.size < Math.max(1, accountCount)) { ); if ( - accountManager.getAccountCount() > 1 && + accountManager.getEnabledAccountCount() > 1 && accountManager.shouldShowAccountToast( account.index, rateLimitToastDebounceMs, @@ -2632,11 +2774,12 @@ while (attempted.size < Math.max(1, accountCount)) { } const waitMs = accountManager.getMinWaitTimeForFamily(modelFamily, model); - const count = accountManager.getAccountCount(); + const enabledCount = accountManager.getEnabledAccountCount(); + const totalCount = accountManager.getAccountCount(); if ( retryAllAccountsRateLimited && - count > 0 && + enabledCount > 0 && waitMs > 0 && (retryAllAccountsMaxWaitMs === 0 || waitMs <= retryAllAccountsMaxWaitMs) && @@ -2646,19 +2789,40 @@ while (attempted.size < Math.max(1, accountCount)) { `All accounts rate-limited wait ${waitMs}ms`, ) ) { - const countdownMessage = `All ${count} account(s) rate-limited. Waiting`; + const countdownMessage = `All ${enabledCount} enabled account(s) rate-limited. Waiting`; await sleepWithCountdown(addJitter(waitMs, 0.2), countdownMessage); allRateLimitedRetries++; continue; } const waitLabel = waitMs > 0 ? formatWaitTime(waitMs) : "a bit"; + const disabledSnapshot = + enabledCount === 0 ? accountManager.getAccountsSnapshot() : []; + const hasUserDisabled = disabledSnapshot.some( + (account) => + account.enabled === false && + account.disabledReason === "user", + ); + const hasAuthFailureDisabled = disabledSnapshot.some( + (account) => + account.enabled === false && + account.disabledReason === "auth-failure", + ); + const disabledMessage = hasUserDisabled && hasAuthFailureDisabled + ? "All stored Codex accounts are disabled. Re-enable user-disabled accounts from account management, or run `opencode auth login` to restore auth-failure disables." + : hasAuthFailureDisabled + ? "All stored Codex accounts are disabled after repeated auth failures. Run `opencode auth login` to restore access." + : hasUserDisabled + ? "All stored Codex accounts are user-disabled. Re-enable them from account management." + : "All stored Codex accounts are disabled. Re-enable any manually disabled accounts from account management, or run `opencode auth login` to restore access."; const message = - count === 0 + totalCount === 0 ? "No Codex accounts configured. Run `opencode auth login`." + : enabledCount === 0 + ? disabledMessage : waitMs > 0 - ? `All ${count} account(s) are rate-limited. Try again in ${waitLabel} or add another account with \`opencode auth login\`.` - : `All ${count} account(s) failed (server errors or auth issues). Check account health with \`codex-health\`.`; + ? `All ${enabledCount} enabled account(s) are rate-limited. Try again in ${waitLabel} or add another account with \`opencode auth login\`.` + : `All ${enabledCount} enabled account(s) failed (server errors or auth issues). Check account health with \`codex-health\`.`; runtimeMetrics.failedRequests++; runtimeMetrics.lastError = message; runtimeMetrics.lastErrorCategory = waitMs > 0 ? "rate-limit" : "account-failure"; @@ -3323,7 +3487,9 @@ while (attempted.size < Math.max(1, accountCount)) { } if (restored.length > 0) { - await persistAccountPool(restored, false); + await persistAccountPool(restored, false, { + reviveMatchingDisabledAccounts: true, + }); invalidateAccountManagerCache(); } @@ -3437,7 +3603,34 @@ while (attempted.size < Math.max(1, accountCount)) { if (typeof menuResult.toggleAccountIndex === "number") { const target = workingStorage.accounts[menuResult.toggleAccountIndex]; if (target) { - target.enabled = target.enabled === false ? true : false; + const shouldEnable = target.enabled === false; + if (shouldEnable && target.disabledReason === "auth-failure") { + const message = + "This account was disabled after repeated auth failures. Run 'opencode auth login' to re-enable with fresh credentials."; + const logContext = { + accountIndex: menuResult.toggleAccountIndex + 1, + }; + logWarn("[account-menu] blocked re-enable for auth-failure disabled account", { + ...logContext, + }); + logInfo("[account-menu] prompted re-auth for auth-failure disabled account", { + ...logContext, + }); + const toastShown = await showToast(message, "warning"); + if (!toastShown) { + process.stdout.write("\nRun 'opencode auth login' to re-enable this account.\n\n"); + } + continue; + } + if (shouldEnable) { + delete target.enabled; + delete target.disabledReason; + delete target.coolingDownUntil; + delete target.cooldownReason; + } else { + target.enabled = false; + target.disabledReason = "user"; + } await saveAccounts(workingStorage); invalidateAccountManagerCache(); console.log( @@ -3507,7 +3700,9 @@ while (attempted.size < Math.max(1, accountCount)) { const { pkce, state, url } = await createAuthorizationFlow(); return buildManualOAuthFlow(pkce, url, state, async (selection) => { try { - await persistAccountPool(selection.variantsForPersistence, startFresh); + await persistAccountPool(selection.variantsForPersistence, startFresh, { + reviveMatchingDisabledAccounts: true, + }); invalidateAccountManagerCache(); } catch (err) { const storagePath = getStoragePath(); @@ -3582,7 +3777,9 @@ while (attempted.size < Math.max(1, accountCount)) { const isFirstAccount = accounts.length === 1; const entriesToPersist = variantsForPersistence.length > 0 ? variantsForPersistence : [resolved]; - await persistAccountPool(entriesToPersist, isFirstAccount && startFresh); + await persistAccountPool(entriesToPersist, isFirstAccount && startFresh, { + reviveMatchingDisabledAccounts: true, + }); invalidateAccountManagerCache(); } catch (err) { const storagePath = getStoragePath(); @@ -3667,7 +3864,10 @@ while (attempted.size < Math.max(1, accountCount)) { const { pkce, state, url } = await createAuthorizationFlow(); return buildManualOAuthFlow(pkce, url, state, async (selection) => { try { - await persistAccountPool(selection.variantsForPersistence, false); + await persistAccountPool(selection.variantsForPersistence, false, { + reviveMatchingDisabledAccounts: true, + }); + invalidateAccountManagerCache(); } catch (err) { const storagePath = getStoragePath(); const errorCode = (err as NodeJS.ErrnoException)?.code || "UNKNOWN"; @@ -3956,10 +4156,9 @@ while (attempted.size < Math.max(1, accountCount)) { return `Switched to ${label} but failed to persist. Changes may be lost on restart.`; } - if (cachedAccountManager) { + if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } const label = formatCommandAccountLabel(account, targetIndex); @@ -5016,8 +5215,7 @@ while (attempted.size < Math.max(1, accountCount)) { if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } } @@ -5261,8 +5459,7 @@ while (attempted.size < Math.max(1, accountCount)) { if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } const accountLabel = formatCommandAccountLabel(account, targetIndex); @@ -5362,8 +5559,7 @@ while (attempted.size < Math.max(1, accountCount)) { if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } const accountLabel = formatCommandAccountLabel(account, targetIndex); @@ -5440,8 +5636,7 @@ while (attempted.size < Math.max(1, accountCount)) { if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } const accountLabel = formatCommandAccountLabel(account, targetIndex); @@ -5757,8 +5952,7 @@ while (attempted.size < Math.max(1, accountCount)) { if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } const remaining = storage.accounts.length; @@ -5852,8 +6046,7 @@ while (attempted.size < Math.max(1, accountCount)) { await saveAccounts(storage); if (cachedAccountManager) { const reloadedManager = await AccountManager.loadFromDisk(); - cachedAccountManager = reloadedManager; - accountManagerPromise = Promise.resolve(reloadedManager); + setCachedAccountManager(reloadedManager); } results.push(""); results.push(`Summary: ${refreshedCount} refreshed, ${failedCount} failed`); diff --git a/lib/accounts.ts b/lib/accounts.ts index 286bc6f2..eb7c67a6 100644 --- a/lib/accounts.ts +++ b/lib/accounts.ts @@ -6,6 +6,7 @@ import { createLogger } from "./logger.js"; import { loadAccounts, saveAccounts, + type AccountDisabledReason, type AccountStorageV3, type CooldownReason, type RateLimitStateV3, @@ -72,6 +73,10 @@ export type CodexCliTokenCacheEntry = { accountId?: string; }; +type FlushPendingSaveOptions = { + deadlineMs?: number; +}; + const CODEX_CLI_ACCOUNTS_PATH = join(homedir(), ".codex", "accounts.json"); const CODEX_CLI_CACHE_TTL_MS = 5_000; let codexCliTokenCache: Map | null = null; @@ -181,6 +186,7 @@ export interface ManagedAccount { email?: string; refreshToken: string; enabled?: boolean; + disabledReason?: AccountDisabledReason; access?: string; expires?: number; addedAt: number; @@ -206,6 +212,12 @@ export interface AccountSelectionExplainability { lastUsed: number; } +export type SetAccountEnabledFailureReason = "auth-failure-blocked" | "invalid-index"; + +export type SetAccountEnabledResult = + | { ok: true; account: ManagedAccount } + | { ok: false; reason: SetAccountEnabledFailureReason }; + export class AccountManager { private accounts: ManagedAccount[] = []; private cursorByFamily: Record = initFamilyState(0); @@ -214,7 +226,11 @@ export class AccountManager { private lastToastTime = 0; private saveDebounceTimer: ReturnType | null = null; private pendingSave: Promise | null = null; + private saveRetryNeeded = false; private authFailuresByRefreshToken: Map = new Map(); + private pendingSaveSettledWaiters = new Set<() => void>(); + private saveFinalizationTick = 0; + private saveFinalizationWaiters = new Set<{ afterTick: number; resolve: () => void }>(); static async loadFromDisk(authFallback?: OAuthAuthDetails): Promise { const stored = await loadAccounts(); @@ -307,8 +323,11 @@ export class AccountManager { email: matchesFallback ? fallbackAccountEmail ?? sanitizeEmail(account.email) : sanitizeEmail(account.email), + // Storage only persists `enabled: false`; in memory we normalize to a concrete boolean. refreshToken, enabled: account.enabled !== false, + disabledReason: + account.enabled === false ? account.disabledReason ?? "user" : undefined, access: matchesFallback && authFallback ? authFallback.access : account.accessToken, expires: matchesFallback && authFallback ? authFallback.expires : account.expiresAt, addedAt: clampNonNegativeInt(account.addedAt, baseNow), @@ -392,6 +411,16 @@ export class AccountManager { return this.accounts.length; } + getEnabledAccountCount(): number { + let enabledCount = 0; + for (const account of this.accounts) { + if (account.enabled !== false) { + enabledCount += 1; + } + } + return enabledCount; + } + getActiveIndex(): number { return this.getActiveIndexForFamily("codex"); } @@ -872,6 +901,31 @@ export class AccountManager { return this.removeAccount(account); } + /** + * Disable all accounts that share the same refreshToken as the given account. + * This keeps org/workspace variants visible in the pool while preventing reuse. + * @returns Number of accounts newly disabled + */ + disableAccountsWithSameRefreshToken(account: ManagedAccount): number { + const refreshToken = account.refreshToken; + let disabledCount = 0; + + for (const accountToDisable of this.accounts) { + if (accountToDisable.refreshToken !== refreshToken) continue; + if (accountToDisable.enabled === false) continue; + this.clearAccountCooldown(accountToDisable); + accountToDisable.enabled = false; + accountToDisable.disabledReason = "auth-failure"; + disabledCount++; + } + + // Reset any accumulated auth failures for this token, even if matching accounts + // were already manually disabled, so a later manual re-enable starts clean. + this.authFailuresByRefreshToken.delete(refreshToken); + + return disabledCount; + } + /** * Remove all accounts that share the same refreshToken as the given account. * This is used when auth refresh fails to remove all org variants together. @@ -895,16 +949,53 @@ export class AccountManager { return removedCount; } - setAccountEnabled(index: number, enabled: boolean): ManagedAccount | null { - if (!Number.isFinite(index)) return null; - if (index < 0 || index >= this.accounts.length) return null; + trySetAccountEnabled( + index: number, + enabled: boolean, + reason?: AccountDisabledReason, + ): SetAccountEnabledResult { + if (!Number.isFinite(index)) return { ok: false, reason: "invalid-index" }; + if (index < 0 || index >= this.accounts.length) return { ok: false, reason: "invalid-index" }; const account = this.accounts[index]; - if (!account) return null; - account.enabled = enabled; - return account; + if (!account) return { ok: false, reason: "invalid-index" }; + if (enabled && account.disabledReason === "auth-failure") { + return { ok: false, reason: "auth-failure-blocked" }; + } + // Once an account is auth-failure disabled, callers must refresh credentials + // instead of downgrading the reason to a manually re-enableable state. + if (!enabled && account.disabledReason === "auth-failure") { + account.enabled = false; + return { ok: true, account }; + } + if (enabled) { + const wasDisabled = account.enabled === false; + account.enabled = true; + if (wasDisabled) { + delete account.disabledReason; + this.clearAccountCooldown(account); + } + } else if (reason) { + account.enabled = false; + account.disabledReason = reason; + } else if (account.disabledReason !== "auth-failure") { + account.enabled = false; + account.disabledReason = "user"; + } else { + account.enabled = false; + } + return { ok: true, account }; } - async saveToDisk(): Promise { + setAccountEnabled( + index: number, + enabled: boolean, + reason?: AccountDisabledReason, + ): ManagedAccount | null { + const result = this.trySetAccountEnabled(index, enabled, reason); + return result.ok ? result.account : null; + } + + private async persistToDisk(): Promise { const activeIndexByFamily: Partial> = {}; for (const family of MODEL_FAMILIES) { const raw = this.currentAccountIndexByFamily[family]; @@ -926,7 +1017,10 @@ export class AccountManager { refreshToken: account.refreshToken, accessToken: account.access, expiresAt: account.expires, + // Persist enabled accounts by omitting the flag to preserve the storage convention. enabled: account.enabled === false ? false : undefined, + disabledReason: + account.enabled === false ? account.disabledReason ?? "user" : undefined, addedAt: account.addedAt, lastUsed: account.lastUsed, lastSwitchReason: account.lastSwitchReason, @@ -942,7 +1036,93 @@ export class AccountManager { await saveAccounts(storage); } + async saveToDisk(): Promise { + return this.enqueueSave(() => this.persistToDisk()); + } + + private enqueueSave(saveOperation: () => Promise): Promise { + const previousSave = this.pendingSave; + const nextSave = (async () => { + if (previousSave) { + try { + await previousSave; + } catch (error) { + log.warn("Continuing queued save after previous save failure", { + error: error instanceof Error ? error.message : String(error), + }); + } + } + try { + await saveOperation(); + this.saveRetryNeeded = false; + } catch (error) { + this.saveRetryNeeded = true; + throw error; + } + })().finally(() => { + if (this.pendingSave === nextSave) { + this.pendingSave = null; + } + this.resolvePendingSaveSettledWaiters(); + this.notifySaveFinalization(); + }); + this.pendingSave = nextSave; + return nextSave; + } + + hasPendingSave(): boolean { + return this.saveDebounceTimer !== null || this.pendingSave !== null; + } + + waitForPendingSaveToSettle(): Promise { + if (!this.hasPendingSave()) { + return Promise.resolve(); + } + + // This intentionally treats debounce-only timers as pending work too. Callers + // such as scheduleTrackedManagerPrune wait across those gaps so a later re-arm + // cannot drop shutdown tracking before the deferred save actually runs. If a + // manager is abandoned before the debounce is replayed through flushPendingSave + // or the timer itself fires, cleanup-time pruning is the final backstop. + return new Promise((resolve) => { + this.pendingSaveSettledWaiters.add(resolve); + }); + } + + private waitForSaveFinalization(afterTick: number): Promise { + if (this.saveFinalizationTick > afterTick) { + return Promise.resolve(); + } + + return new Promise((resolve) => { + this.saveFinalizationWaiters.add({ afterTick, resolve }); + }); + } + + private notifySaveFinalization(): void { + this.saveFinalizationTick += 1; + for (const waiter of [...this.saveFinalizationWaiters]) { + if (this.saveFinalizationTick <= waiter.afterTick) { + continue; + } + this.saveFinalizationWaiters.delete(waiter); + waiter.resolve(); + } + } + + private resolvePendingSaveSettledWaiters(): void { + if (this.hasPendingSave() || this.pendingSaveSettledWaiters.size === 0) { + return; + } + + for (const resolve of [...this.pendingSaveSettledWaiters]) { + this.pendingSaveSettledWaiters.delete(resolve); + resolve(); + } + } + saveToDiskDebounced(delayMs = 500): void { + this.saveRetryNeeded = true; if (this.saveDebounceTimer) { clearTimeout(this.saveDebounceTimer); } @@ -950,13 +1130,7 @@ export class AccountManager { this.saveDebounceTimer = null; const doSave = async () => { try { - if (this.pendingSave) { - await this.pendingSave; - } - this.pendingSave = this.saveToDisk().finally(() => { - this.pendingSave = null; - }); - await this.pendingSave; + await this.saveToDisk(); } catch (error) { log.warn("Debounced save failed", { error: error instanceof Error ? error.message : String(error) }); } @@ -965,14 +1139,95 @@ export class AccountManager { }, delayMs); } - async flushPendingSave(): Promise { - if (this.saveDebounceTimer) { - clearTimeout(this.saveDebounceTimer); - this.saveDebounceTimer = null; - await this.saveToDisk(); - } - if (this.pendingSave) { - await this.pendingSave; + async flushPendingSave(options?: FlushPendingSaveOptions): Promise { + const MAX_FLUSH_ITERATIONS = 20; + let flushIterations = 0; + const deadlineMs = + typeof options?.deadlineMs === "number" && Number.isFinite(options.deadlineMs) + ? options.deadlineMs + : undefined; + const throwIfDeadlineExceeded = (): void => { + if (deadlineMs !== undefined && Date.now() >= deadlineMs) { + throw new Error("flushPendingSave: shutdown deadline exceeded before save state settled"); + } + }; + + while (true) { + throwIfDeadlineExceeded(); + flushIterations += 1; + if (flushIterations > MAX_FLUSH_ITERATIONS) { + // This is intentionally far above realistic debounce re-arm chains; if we + // still haven't converged, shutdown callers log the failure and continue exit. + log.warn("flushPendingSave exceeded max iterations; possible save loop", { + iterations: flushIterations - 1, + }); + throw new Error("flushPendingSave: exceeded max flush iterations; save state may be incomplete"); + } + const hadDebouncedSave = !!this.saveDebounceTimer; + if (this.saveDebounceTimer) { + clearTimeout(this.saveDebounceTimer); + this.saveDebounceTimer = null; + } + if (this.pendingSave) { + const pendingSaveTick = this.saveFinalizationTick; + let pendingSaveError: unknown; + try { + await this.pendingSave; + } catch (error) { + pendingSaveError = error; + } + await this.waitForSaveFinalization(pendingSaveTick); + if (this.saveDebounceTimer !== null || this.pendingSave !== null) { + throwIfDeadlineExceeded(); + continue; + } + const shouldRetryAfterPendingSave = hadDebouncedSave || this.saveRetryNeeded; + if (pendingSaveError) { + if (!shouldRetryAfterPendingSave) { + throw pendingSaveError; + } + log.warn("flushPendingSave: retrying after save failure to flush latest state", { + error: + pendingSaveError instanceof Error + ? pendingSaveError.message + : String(pendingSaveError), + }); + } + } + if (!hadDebouncedSave && !this.saveRetryNeeded) { + return; + } + if (this.saveDebounceTimer !== null || this.pendingSave !== null) { + throwIfDeadlineExceeded(); + continue; + } + const flushSaveTick = this.saveFinalizationTick; + const flushSave = this.saveToDisk(); + let flushSaveError: unknown; + try { + await flushSave; + } catch (error) { + flushSaveError = error; + } + await this.waitForSaveFinalization(flushSaveTick); + if (this.saveDebounceTimer !== null || this.pendingSave !== null) { + if (flushSaveError) { + log.warn("flushPendingSave: retrying after flush save failure while newer save is queued", { + error: + flushSaveError instanceof Error + ? flushSaveError.message + : String(flushSaveError), + }); + } + throwIfDeadlineExceeded(); + continue; + } + if (flushSaveError) { + throw flushSaveError; + } + if (this.saveDebounceTimer === null && this.pendingSave === null) { + return; + } } } } diff --git a/lib/constants.ts b/lib/constants.ts index 69e7f5bf..a601b7a4 100644 --- a/lib/constants.ts +++ b/lib/constants.ts @@ -89,6 +89,6 @@ export const ACCOUNT_LIMITS = { MAX_ACCOUNTS: 20, /** Cooldown period (ms) after auth failure before retrying account */ AUTH_FAILURE_COOLDOWN_MS: 30_000, - /** Number of consecutive auth failures before auto-removing account */ - MAX_AUTH_FAILURES_BEFORE_REMOVAL: 3, + /** Number of consecutive auth failures before auto-disabling account */ + MAX_AUTH_FAILURES_BEFORE_DISABLE: 3, } as const; diff --git a/lib/schemas.ts b/lib/schemas.ts index 6028246d..28928bc7 100644 --- a/lib/schemas.ts +++ b/lib/schemas.ts @@ -5,6 +5,7 @@ */ import { z } from "zod"; import { MODEL_FAMILIES, type ModelFamily } from "./prompts/codex.js"; +import { normalizeAccountDisabledReason } from "./storage/migrations.js"; // ============================================================================ // Plugin Configuration Schema @@ -74,6 +75,21 @@ export const CooldownReasonSchema = z.enum(["auth-failure", "network-error"]); export type CooldownReasonFromSchema = z.infer; +const AccountDisabledReasonValueSchema = z.enum(["user", "auth-failure"]); + +// Storage normalization strips unknown disabled reasons later; keep schema parsing +// lenient so legacy/downgraded files don't warn or fail before that step runs. +export const AccountDisabledReasonSchema = z.preprocess( + normalizeAccountDisabledReason, + AccountDisabledReasonValueSchema.optional(), +); + +// Preserve the older export name for callers while using a more specific schema name internally. +export const DisabledReasonSchema = AccountDisabledReasonSchema; + +export type AccountDisabledReasonFromSchema = z.infer; +export type DisabledReasonFromSchema = AccountDisabledReasonFromSchema; + /** * Last switch reason for account rotation tracking. */ @@ -117,6 +133,7 @@ export const AccountMetadataV3Schema = z.object({ accessToken: z.string().optional(), expiresAt: z.number().optional(), enabled: z.boolean().optional(), + disabledReason: AccountDisabledReasonSchema.optional(), addedAt: z.number(), lastUsed: z.number(), lastSwitchReason: SwitchReasonSchema.optional(), @@ -164,6 +181,7 @@ export const AccountMetadataV1Schema = z.object({ accessToken: z.string().optional(), expiresAt: z.number().optional(), enabled: z.boolean().optional(), + disabledReason: AccountDisabledReasonSchema.optional(), addedAt: z.number(), lastUsed: z.number(), lastSwitchReason: SwitchReasonSchema.optional(), diff --git a/lib/shutdown.ts b/lib/shutdown.ts index 66965627..2b0ad19b 100644 --- a/lib/shutdown.ts +++ b/lib/shutdown.ts @@ -1,7 +1,20 @@ -type CleanupFn = () => void | Promise; +export type CleanupContext = { + deadlineMs?: number; +}; + +type CleanupFn = (context?: CleanupContext) => void | Promise; + +// Allow enough time for serialized manager flushes to finish on Windows hosts +// where AV/file-indexer contention can hold the accounts file for multiple seconds. +const SIGNAL_CLEANUP_TIMEOUT_MS = 10_000; const cleanupFunctions: CleanupFn[] = []; let shutdownRegistered = false; +let sigintHandler: (() => void) | null = null; +let sigtermHandler: (() => void) | null = null; +let beforeExitHandler: (() => void) | null = null; +let cleanupInFlight: Promise | null = null; +let signalExitPending = false; export function registerCleanup(fn: CleanupFn): void { cleanupFunctions.push(fn); @@ -13,18 +26,58 @@ export function unregisterCleanup(fn: CleanupFn): void { if (index !== -1) { cleanupFunctions.splice(index, 1); } + if ( + cleanupFunctions.length === 0 && + !cleanupInFlight && + !signalExitPending && + shutdownRegistered + ) { + removeShutdownHandlers(); + shutdownRegistered = false; + } } -export async function runCleanup(): Promise { - const fns = [...cleanupFunctions]; - cleanupFunctions.length = 0; +export function runCleanup(context?: CleanupContext): Promise { + if (cleanupInFlight) { + return cleanupInFlight; + } + + cleanupInFlight = (async () => { + while (cleanupFunctions.length > 0) { + const fns = [...cleanupFunctions]; + cleanupFunctions.length = 0; - for (const fn of fns) { - try { - await fn(); - } catch { - // Ignore cleanup errors during shutdown + for (const fn of fns) { + try { + await fn(context); + } catch { + // Ignore cleanup errors during shutdown + } + } + } + })().finally(() => { + cleanupInFlight = null; + if (!signalExitPending) { + removeShutdownHandlers(); + shutdownRegistered = false; } + if (cleanupFunctions.length > 0 && !shutdownRegistered) { + ensureShutdownHandler(); + } + }); + + return cleanupInFlight; +} + +function resetSignalStateAfterInterceptedExit(): void { + if (!signalExitPending) { + return; + } + + signalExitPending = false; + if (cleanupFunctions.length === 0 && !cleanupInFlight && shutdownRegistered) { + removeShutdownHandlers(); + shutdownRegistered = false; } } @@ -33,16 +86,62 @@ function ensureShutdownHandler(): void { shutdownRegistered = true; const handleSignal = () => { - void runCleanup().finally(() => { + if (signalExitPending) { + return; + } + signalExitPending = true; + const deadlineMs = Date.now() + SIGNAL_CLEANUP_TIMEOUT_MS; + let exitRequested = false; + const requestExit = () => { + if (exitRequested) { + return; + } + exitRequested = true; + if (timeoutHandle) { + clearTimeout(timeoutHandle); + timeoutHandle = null; + } process.exit(0); + }; + let timeoutHandle: ReturnType | null = setTimeout(() => { + requestExit(); + }, Math.max(0, deadlineMs - Date.now())); + void runCleanup({ deadlineMs }).finally(() => { + try { + requestExit(); + } finally { + // `process.exit()` normally terminates immediately. This only runs when exit + // is intercepted (for example in tests), so we need to unlatch signal state. + resetSignalStateAfterInterceptedExit(); + } }); }; + sigintHandler = handleSignal; + sigtermHandler = handleSignal; + beforeExitHandler = () => { + if (!cleanupInFlight) { + void runCleanup(); + } + }; - process.once("SIGINT", handleSignal); - process.once("SIGTERM", handleSignal); - process.once("beforeExit", () => { - void runCleanup(); - }); + process.on("SIGINT", sigintHandler); + process.on("SIGTERM", sigtermHandler); + process.on("beforeExit", beforeExitHandler); +} + +function removeShutdownHandlers(): void { + if (sigintHandler) { + process.off("SIGINT", sigintHandler); + sigintHandler = null; + } + if (sigtermHandler) { + process.off("SIGTERM", sigtermHandler); + sigtermHandler = null; + } + if (beforeExitHandler) { + process.off("beforeExit", beforeExitHandler); + beforeExitHandler = null; + } } export function getCleanupCount(): number { diff --git a/lib/storage.ts b/lib/storage.ts index 151e2213..e3ac36ee 100644 --- a/lib/storage.ts +++ b/lib/storage.ts @@ -7,16 +7,27 @@ import { MODEL_FAMILIES, type ModelFamily } from "./prompts/codex.js"; import { AnyAccountStorageSchema, getValidationErrors } from "./schemas.js"; import { getConfigDir, getProjectConfigDir, getProjectGlobalConfigDir, findProjectRoot, resolvePath } from "./storage/paths.js"; import { - migrateV1ToV3, - type CooldownReason, - type RateLimitStateV3, - type AccountMetadataV1, - type AccountStorageV1, - type AccountMetadataV3, - type AccountStorageV3, + migrateV1ToV3, + normalizeStoredAccountDisabledReason, + normalizeStoredEnabled, + type AccountDisabledReason, + type CooldownReason, + type RateLimitStateV3, + type AccountMetadataV1, + type AccountStorageV1, + type AccountMetadataV3, + type AccountStorageV3, } from "./storage/migrations.js"; -export type { CooldownReason, RateLimitStateV3, AccountMetadataV1, AccountStorageV1, AccountMetadataV3, AccountStorageV3 }; +export type { + AccountDisabledReason, + CooldownReason, + RateLimitStateV3, + AccountMetadataV1, + AccountStorageV1, + AccountMetadataV3, + AccountStorageV3, +}; const log = createLogger("storage"); const ACCOUNTS_FILE_NAME = "openai-codex-accounts.json"; @@ -612,10 +623,27 @@ export function normalizeAccountStorage(data: unknown): AccountStorageV3 | null ? migrateV1ToV3(data as unknown as AccountStorageV1) : (data as unknown as AccountStorageV3); - const validAccounts = rawAccounts.filter( - (account): account is AccountMetadataV3 => - isRecord(account) && typeof account.refreshToken === "string" && !!account.refreshToken.trim(), - ); + const validAccounts: AccountMetadataV3[] = baseStorage.accounts + .filter( + (account): account is AccountMetadataV3 => + isRecord(account) && typeof account.refreshToken === "string" && !!account.refreshToken.trim(), + ) + .map((account) => { + const normalizedAccount: AccountMetadataV3 = { ...account }; + const enabled = normalizeStoredEnabled(account.enabled); + const disabledReason = normalizeStoredAccountDisabledReason(account.enabled, account.disabledReason); + if (enabled === undefined) { + delete normalizedAccount.enabled; + } else { + normalizedAccount.enabled = enabled; + } + if (disabledReason === undefined) { + delete normalizedAccount.disabledReason; + } else { + normalizedAccount.disabledReason = disabledReason; + } + return normalizedAccount; + }); const deduplicatedAccounts = deduplicateAccountsForStorage(validAccounts); @@ -972,6 +1000,10 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 { const cooldownReason = isCooldownReason(rawAccount.cooldownReason) ? rawAccount.cooldownReason : undefined; + const disabledReason = normalizeStoredAccountDisabledReason( + rawAccount.enabled, + rawAccount.disabledReason, + ); const accountTags = normalizeTags(rawAccount.accountTags); const accountNote = typeof rawAccount.accountNote === "string" && rawAccount.accountNote.trim() @@ -1000,6 +1032,9 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 { flaggedReason: typeof rawAccount.flaggedReason === "string" ? rawAccount.flaggedReason : undefined, lastError: typeof rawAccount.lastError === "string" ? rawAccount.lastError : undefined, }; + if (disabledReason !== undefined) { + normalized.disabledReason = disabledReason; + } byRefreshToken.set(refreshToken, normalized); } diff --git a/lib/storage/migrations.ts b/lib/storage/migrations.ts index 092782cd..498fbbbe 100644 --- a/lib/storage/migrations.ts +++ b/lib/storage/migrations.ts @@ -7,6 +7,32 @@ import { MODEL_FAMILIES, type ModelFamily } from "../prompts/codex.js"; import type { AccountIdSource } from "../types.js"; export type CooldownReason = "auth-failure" | "network-error"; +export type AccountDisabledReason = "user" | "auth-failure"; + +export function isAccountDisabledReason(value: unknown): value is AccountDisabledReason { + return value === "user" || value === "auth-failure"; +} + +export function normalizeAccountDisabledReason( + value: unknown, +): AccountDisabledReason | undefined { + return isAccountDisabledReason(value) ? value : undefined; +} + +export function normalizeStoredAccountDisabledReason( + enabled: unknown, + disabledReason: unknown, +): AccountDisabledReason | undefined { + // Unlike V1 migration, V3 normalization intentionally leaves missing/invalid + // reasons unset so legacy-disabled accounts are not auto-revived as auth failures. + return enabled === false + ? normalizeAccountDisabledReason(disabledReason) + : undefined; +} + +export function normalizeStoredEnabled(enabled: unknown): false | undefined { + return enabled === false ? false : undefined; +} export interface RateLimitStateV3 { [key: string]: number | undefined; @@ -26,6 +52,8 @@ export interface AccountMetadataV1 { /** Optional access token expiry timestamp (ms since epoch). */ expiresAt?: number; enabled?: boolean; + /** Never set by V1 storage; migration defaults disabled V1 accounts to "user". */ + disabledReason?: AccountDisabledReason; addedAt: number; lastUsed: number; lastSwitchReason?: "rate-limit" | "initial" | "rotation"; @@ -54,6 +82,7 @@ export interface AccountMetadataV3 { /** Optional access token expiry timestamp (ms since epoch). */ expiresAt?: number; enabled?: boolean; + disabledReason?: AccountDisabledReason; addedAt: number; lastUsed: number; lastSwitchReason?: "rate-limit" | "initial" | "rotation"; @@ -95,7 +124,11 @@ export function migrateV1ToV3(v1: AccountStorageV1): AccountStorageV3 { refreshToken: account.refreshToken, accessToken: account.accessToken, expiresAt: account.expiresAt, - enabled: account.enabled, + enabled: normalizeStoredEnabled(account.enabled), + disabledReason: + account.enabled === false + ? normalizeAccountDisabledReason(account.disabledReason) ?? "user" + : undefined, addedAt: account.addedAt, lastUsed: account.lastUsed, lastSwitchReason: account.lastSwitchReason, diff --git a/test/accounts.test.ts b/test/accounts.test.ts index 68e6da9c..57135c9c 100644 --- a/test/accounts.test.ts +++ b/test/accounts.test.ts @@ -13,6 +13,10 @@ import { import { getHealthTracker, getTokenTracker, resetTrackers } from "../lib/rotation.js"; import type { OAuthAuthDetails } from "../lib/types.js"; +const { accountLoggerWarn } = vi.hoisted(() => ({ + accountLoggerWarn: vi.fn(), +})); + vi.mock("../lib/storage.js", async (importOriginal) => { const actual = await importOriginal(); return { @@ -21,6 +25,21 @@ vi.mock("../lib/storage.js", async (importOriginal) => { }; }); +vi.mock("../lib/logger.js", () => ({ + createLogger: vi.fn(() => ({ + debug: vi.fn(), + info: vi.fn(), + warn: accountLoggerWarn, + error: vi.fn(), + time: vi.fn(() => vi.fn(() => 0)), + timeEnd: vi.fn(), + })), +})); + +beforeEach(() => { + accountLoggerWarn.mockClear(); +}); + describe("parseRateLimitReason", () => { it("returns quota for quota-related codes", () => { expect(parseRateLimitReason("usage_limit_reached")).toBe("quota"); @@ -856,6 +875,120 @@ describe("AccountManager", () => { expect(manager.getAccountsSnapshot()[0].refreshToken).toBe("token-2"); }); + it("disables all accounts with the same refreshToken", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + { refreshToken: "token-1", organizationId: "org-1", addedAt: now, lastUsed: now }, + { refreshToken: "token-1", organizationId: "org-2", addedAt: now, lastUsed: now }, + { refreshToken: "token-2", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const accounts = manager.getAccountsSnapshot(); + expect(accounts).toHaveLength(4); + + const disabledCount = manager.disableAccountsWithSameRefreshToken(accounts[0]); + expect(disabledCount).toBe(3); + expect(manager.getAccountCount()).toBe(4); + + const updated = manager.getAccountsSnapshot(); + expect(updated.slice(0, 3).every((account) => account.enabled === false)).toBe(true); + expect(updated.slice(0, 3).every((account) => account.disabledReason === "auth-failure")).toBe(true); + expect(updated[3]?.enabled).not.toBe(false); + }); + + it("preserves a manual disable reason when auth-failure disabling sibling variants", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", accountId: "workspace-a", addedAt: now, lastUsed: now }, + { + refreshToken: "token-1", + accountId: "workspace-user-disabled", + enabled: false, + disabledReason: "user" as const, + coolingDownUntil: now + 60_000, + cooldownReason: "network-error" as const, + addedAt: now, + lastUsed: now, + }, + { refreshToken: "token-1", accountId: "workspace-b", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const accounts = manager.getAccountsSnapshot(); + const disabledCount = manager.disableAccountsWithSameRefreshToken(accounts[0]); + + expect(disabledCount).toBe(2); + const updated = manager.getAccountsSnapshot(); + expect(updated[0]).toMatchObject({ + accountId: "workspace-a", + enabled: false, + disabledReason: "auth-failure", + }); + expect(updated[1]).toMatchObject({ + accountId: "workspace-user-disabled", + enabled: false, + disabledReason: "user", + coolingDownUntil: now + 60_000, + cooldownReason: "network-error", + }); + expect(updated[2]).toMatchObject({ + accountId: "workspace-b", + enabled: false, + disabledReason: "auth-failure", + }); + }); + + it("defaults legacy disabled accounts to a user disable reason", async () => { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + mockSaveAccounts.mockResolvedValueOnce(); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + enabled: false, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + expect(manager.getAccountsSnapshot()[0]).toMatchObject({ + enabled: false, + disabledReason: "user", + }); + + await manager.saveToDisk(); + + expect(mockSaveAccounts).toHaveBeenCalledWith( + expect.objectContaining({ + accounts: [ + expect.objectContaining({ + enabled: false, + disabledReason: "user", + }), + ], + }), + ); + }); + it("clears auth failure counter when removing accounts with same refreshToken", () => { const now = Date.now(); const stored = { @@ -884,13 +1017,322 @@ describe("AccountManager", () => { const removedCount = manager.removeAccountsWithSameRefreshToken(account1); expect(removedCount).toBe(2); expect(manager.getAccountCount()).toBe(0); - const failuresByRefreshToken = Reflect.get( - manager, - "authFailuresByRefreshToken", - ) as Map; - expect(failuresByRefreshToken.has("token-1")).toBe(false); expect(manager.incrementAuthFailures(account1)).toBe(1); }); + + it("clears auth failure counter when disabling accounts with same refreshToken", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + coolingDownUntil: now + 60_000, + cooldownReason: "auth-failure" as const, + addedAt: now, + lastUsed: now, + }, + { + refreshToken: "token-1", + organizationId: "org-1", + coolingDownUntil: now + 60_000, + cooldownReason: "auth-failure" as const, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const accounts = manager.getAccountsSnapshot(); + expect(accounts).toHaveLength(2); + + expect(manager.incrementAuthFailures(accounts[0])).toBe(1); + expect(manager.incrementAuthFailures(accounts[1])).toBe(2); + expect(manager.incrementAuthFailures(accounts[0])).toBe(3); + + const disabledCount = manager.disableAccountsWithSameRefreshToken(accounts[0]); + expect(disabledCount).toBe(2); + expect(manager.getAccountsSnapshot().every((account) => account.enabled === false)).toBe(true); + expect( + manager.getAccountsSnapshot().every((account) => account.disabledReason === "auth-failure"), + ).toBe(true); + expect( + manager.getAccountsSnapshot().every((account) => account.coolingDownUntil === undefined), + ).toBe(true); + expect( + manager.getAccountsSnapshot().every((account) => account.cooldownReason === undefined), + ).toBe(true); + expect(manager.incrementAuthFailures(accounts[0])).toBe(1); + }); + + it("clears auth failure counter even when matching accounts were already user-disabled", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + accountId: "workspace-a", + enabled: false, + disabledReason: "user" as const, + coolingDownUntil: now + 60_000, + cooldownReason: "network-error" as const, + addedAt: now, + lastUsed: now, + }, + { + refreshToken: "token-1", + accountId: "workspace-b", + enabled: false, + disabledReason: "user" as const, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const accounts = manager.getAccountsSnapshot(); + + expect(manager.incrementAuthFailures(accounts[0])).toBe(1); + expect(manager.incrementAuthFailures(accounts[1])).toBe(2); + + const disabledCount = manager.disableAccountsWithSameRefreshToken(accounts[0]); + expect(disabledCount).toBe(0); + expect(manager.getAccountsSnapshot()).toMatchObject([ + { + accountId: "workspace-a", + enabled: false, + disabledReason: "user", + coolingDownUntil: now + 60_000, + cooldownReason: "network-error", + }, + { + accountId: "workspace-b", + enabled: false, + disabledReason: "user", + }, + ]); + expect(manager.incrementAuthFailures(accounts[0])).toBe(1); + }); + + it("records disable reason when setAccountEnabled disables an account", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + coolingDownUntil: now + 60_000, + cooldownReason: "network-error" as const, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + const disabled = manager.trySetAccountEnabled(0, false, "user"); + expect(disabled).toMatchObject({ + ok: true, + account: { + enabled: false, + disabledReason: "user", + }, + }); + + const reenabled = manager.trySetAccountEnabled(0, true); + expect(reenabled).toMatchObject({ + ok: true, + account: { + enabled: true, + }, + }); + if (!reenabled.ok) { + throw new Error("expected setAccountEnabled to re-enable account"); + } + expect(reenabled.account.disabledReason).toBeUndefined(); + expect(reenabled.account.coolingDownUntil).toBeUndefined(); + expect(reenabled.account.cooldownReason).toBeUndefined(); + }); + + it("defaults disable reason to user when disabling without an explicit reason", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + const disabled = manager.trySetAccountEnabled(0, false); + expect(disabled).toMatchObject({ + ok: true, + account: { + enabled: false, + disabledReason: "user", + }, + }); + }); + + it("preserves auth-failure disabledReason when disabling without an explicit reason", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + enabled: false, + disabledReason: "auth-failure" as const, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + const disabled = manager.trySetAccountEnabled(0, false); + expect(disabled).toMatchObject({ + ok: true, + account: { + enabled: false, + disabledReason: "auth-failure", + }, + }); + }); + + it("blocks re-enabling auth-failure disabled accounts through setAccountEnabled", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + enabled: false, + disabledReason: "auth-failure" as const, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + const reenabled = manager.trySetAccountEnabled(0, true); + expect(reenabled).toEqual({ + ok: false, + reason: "auth-failure-blocked", + }); + expect(manager.getAccountsSnapshot()[0]).toMatchObject({ + enabled: false, + disabledReason: "auth-failure", + }); + }); + + it("preserves cooldown metadata when enabling an already enabled account", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + enabled: true, + coolingDownUntil: now + 60_000, + cooldownReason: "network-error" as const, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + const result = manager.trySetAccountEnabled(0, true); + expect(result).toMatchObject({ + ok: true, + account: { + enabled: true, + coolingDownUntil: now + 60_000, + cooldownReason: "network-error", + }, + }); + }); + + it("ignores explicit disable-reason overrides for auth-failure disabled accounts", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + enabled: false, + disabledReason: "auth-failure" as const, + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + const disabled = manager.trySetAccountEnabled(0, false, "user"); + expect(disabled).toMatchObject({ + ok: true, + account: { + enabled: false, + disabledReason: "auth-failure", + }, + }); + + expect(manager.trySetAccountEnabled(0, true)).toEqual({ + ok: false, + reason: "auth-failure-blocked", + }); + expect(manager.setAccountEnabled(0, true)).toBeNull(); + }); + + it("preserves the legacy setAccountEnabled account-or-null contract", () => { + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "token-1", + addedAt: now, + lastUsed: now, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + expect(manager.setAccountEnabled(99, false)).toBeNull(); + + const userDisabled = manager.setAccountEnabled(0, false, "user"); + expect(userDisabled).toMatchObject({ + enabled: false, + disabledReason: "user", + }); + + const reenabled = manager.setAccountEnabled(0, true); + expect(reenabled).toMatchObject({ + enabled: true, + }); + }); }); describe("getMinWaitTimeForFamily", () => { @@ -1277,6 +1719,7 @@ describe("AccountManager", () => { it("saves accounts with all fields", async () => { const { saveAccounts } = await import("../lib/storage.js"); const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); mockSaveAccounts.mockResolvedValueOnce(); const now = Date.now(); @@ -1397,6 +1840,181 @@ describe("AccountManager", () => { expect(mockSaveAccounts).toHaveBeenCalledTimes(2); }); + + it("continues a queued save after the previous save fails", async () => { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + const savedSnapshots: Array< + Array<{ + refreshToken?: string; + enabled?: false; + disabledReason?: string; + }> + > = []; + + let rejectFirstSave!: (error: Error) => void; + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + await new Promise((_, reject) => { + rejectFirstSave = reject; + }); + }); + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + }); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + manager.saveToDiskDebounced(50); + await vi.advanceTimersByTimeAsync(60); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + const account = manager.getAccountsSnapshot()[0]!; + manager.disableAccountsWithSameRefreshToken(account); + manager.saveToDiskDebounced(50); + await vi.advanceTimersByTimeAsync(60); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + rejectFirstSave(Object.assign(new Error("EBUSY: file in use"), { code: "EBUSY" })); + await vi.advanceTimersByTimeAsync(0); + for (let turn = 0; turn < 6; turn += 1) { + await Promise.resolve(); + } + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + expect(savedSnapshots[1]).toEqual([ + expect.objectContaining({ + refreshToken: "token-1", + enabled: false, + disabledReason: "auth-failure", + }), + ]); + }); + + it("persists newer disabled state after two queued save failures once a later flush requeues it", async () => { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + const savedSnapshots: Array< + Array<{ + refreshToken?: string; + enabled?: false; + disabledReason?: string; + }> + > = []; + + let rejectFirstSave!: (error: Error) => void; + let rejectSecondSave!: (error: Error) => void; + const firstSave = new Promise((_resolve, reject) => { + rejectFirstSave = reject; + }); + const secondSave = new Promise((_resolve, reject) => { + rejectSecondSave = reject; + }); + + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + await firstSave; + }); + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + await secondSave; + }); + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + }); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + manager.saveToDiskDebounced(0); + await vi.advanceTimersByTimeAsync(0); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + const account = manager.getAccountsSnapshot()[0]!; + manager.disableAccountsWithSameRefreshToken(account); + manager.saveToDiskDebounced(50); + await vi.advanceTimersByTimeAsync(60); + + rejectFirstSave(Object.assign(new Error("EBUSY: file in use"), { code: "EBUSY" })); + await vi.advanceTimersByTimeAsync(0); + for (let turn = 0; turn < 6; turn += 1) { + await Promise.resolve(); + } + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + + rejectSecondSave(Object.assign(new Error("EBUSY: file in use"), { code: "EBUSY" })); + await vi.advanceTimersByTimeAsync(0); + for (let turn = 0; turn < 6; turn += 1) { + await Promise.resolve(); + } + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + + manager.saveToDiskDebounced(0); + await manager.flushPendingSave(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(3); + expect(savedSnapshots[2]).toEqual([ + expect.objectContaining({ + refreshToken: "token-1", + enabled: false, + disabledReason: "auth-failure", + }), + ]); + }); }); describe("constructor edge cases", () => { @@ -1689,6 +2307,12 @@ describe("AccountManager", () => { }); describe("flushPendingSave", () => { + const drainMicrotasks = async (turns = 10): Promise => { + for (let turn = 0; turn < turns; turn++) { + await Promise.resolve(); + } + }; + it("flushes pending debounced save", async () => { const { saveAccounts } = await import("../lib/storage.js"); const mockSaveAccounts = vi.mocked(saveAccounts); @@ -1785,10 +2409,626 @@ describe("AccountManager", () => { const savePromise = manager.saveToDisk(); queueMicrotask(() => resolveInFlight!()); - + await manager.flushPendingSave(); await savePromise; }); + + it("tracks direct saveToDisk work so flushPendingSave waits for it", async () => { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + let resolveInFlight: (() => void) | undefined; + const inFlightSave = new Promise((resolve) => { + resolveInFlight = resolve; + }); + mockSaveAccounts.mockImplementation(() => inFlightSave); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const savePromise = manager.saveToDisk(); + const flushPromise = manager.flushPendingSave(); + + expect(manager.hasPendingSave()).toBe(true); + + let flushSettled = false; + void flushPromise.then(() => { + flushSettled = true; + }); + await Promise.resolve(); + expect(flushSettled).toBe(false); + + resolveInFlight?.(); + await savePromise; + await flushPromise; + + expect(manager.hasPendingSave()).toBe(false); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + }); + + it("waits for in-flight save before flushing a pending debounce", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + let resolveFirst!: () => void; + const firstSave = new Promise((resolve) => { + resolveFirst = resolve; + }); + mockSaveAccounts.mockImplementationOnce(() => firstSave); + mockSaveAccounts.mockResolvedValue(); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + manager.saveToDiskDebounced(50); + await vi.advanceTimersByTimeAsync(60); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + manager.saveToDiskDebounced(50); + const flushPromise = manager.flushPendingSave(); + await Promise.resolve(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + resolveFirst(); + await flushPromise; + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); + + it("resolves settle waiters when flushPendingSave cancels a pending debounce", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + mockSaveAccounts.mockResolvedValue(); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + manager.saveToDiskDebounced(50); + const settlePromise = manager.waitForPendingSaveToSettle(); + const flushPromise = manager.flushPendingSave(); + + await flushPromise; + await expect(settlePromise).resolves.toBeUndefined(); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it("retries a failed debounced save during flush even after pending state clears", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + const savedSnapshots: Array< + Array<{ + refreshToken?: string; + enabled?: false; + disabledReason?: string; + }> + > = []; + + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + throw Object.assign(new Error("EBUSY: file in use"), { code: "EBUSY" }); + }); + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + }); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const account = manager.getAccountsSnapshot()[0]!; + + manager.disableAccountsWithSameRefreshToken(account); + manager.saveToDiskDebounced(0); + await vi.advanceTimersByTimeAsync(0); + await drainMicrotasks(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + expect(manager.hasPendingSave()).toBe(false); + + await manager.flushPendingSave(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + expect(savedSnapshots[1]).toEqual([ + expect.objectContaining({ + refreshToken: "token-1", + enabled: false, + disabledReason: "auth-failure", + }), + ]); + } finally { + vi.useRealTimers(); + } + }); + + it("drains saves queued during flush without overlapping writes", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + let activeWrites = 0; + let maxConcurrentWrites = 0; + const settleWrites: Array<() => void> = []; + const writePromises: Promise[] = []; + + mockSaveAccounts.mockImplementation(() => { + activeWrites++; + maxConcurrentWrites = Math.max(maxConcurrentWrites, activeWrites); + const writePromise = new Promise((resolve) => { + settleWrites.push(() => { + activeWrites--; + resolve(); + }); + }); + writePromises.push(writePromise); + return writePromise; + }); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + manager.saveToDiskDebounced(50); + await vi.advanceTimersByTimeAsync(60); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + const armGapSave = writePromises[0]!.then(() => { + manager.saveToDiskDebounced(0); + }); + + manager.saveToDiskDebounced(50); + const flushPromise = manager.flushPendingSave(); + + settleWrites[0]!(); + await armGapSave; + await drainMicrotasks(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + + manager.saveToDiskDebounced(0); + await vi.advanceTimersByTimeAsync(0); + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + + settleWrites[1]!(); + await drainMicrotasks(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(3); + + settleWrites[2]!(); + await flushPromise; + + expect(maxConcurrentWrites).toBe(1); + } finally { + vi.useRealTimers(); + } + }); + + it("persists disabled auth-failure state after flushing over an in-flight save", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + let activeWrites = 0; + let maxConcurrentWrites = 0; + const settleWrites: Array<() => void> = []; + const savedSnapshots: Array< + Array<{ + refreshToken?: string; + enabled?: false; + disabledReason?: string; + coolingDownUntil?: number; + cooldownReason?: string; + }> + > = []; + + mockSaveAccounts.mockImplementation(async (storage) => { + activeWrites++; + maxConcurrentWrites = Math.max(maxConcurrentWrites, activeWrites); + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + coolingDownUntil: account.coolingDownUntil, + cooldownReason: account.cooldownReason, + })), + ); + await new Promise((resolve) => { + settleWrites.push(() => { + activeWrites--; + resolve(); + }); + }); + }); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { + refreshToken: "shared-refresh", + addedAt: now, + lastUsed: now, + coolingDownUntil: now + 10_000, + cooldownReason: "auth-failure" as const, + }, + { + refreshToken: "shared-refresh", + addedAt: now + 1, + lastUsed: now + 1, + coolingDownUntil: now + 20_000, + cooldownReason: "auth-failure" as const, + }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + manager.saveToDiskDebounced(50); + await vi.advanceTimersByTimeAsync(60); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + const account = manager.getAccountsSnapshot()[0]!; + manager.disableAccountsWithSameRefreshToken(account); + manager.saveToDiskDebounced(50); + + const flushPromise = manager.flushPendingSave(); + + settleWrites[0]!(); + await drainMicrotasks(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + + const flushedSharedAccounts = savedSnapshots[1]!.filter( + (savedAccount) => savedAccount.refreshToken === "shared-refresh", + ); + expect(flushedSharedAccounts).toEqual([ + expect.objectContaining({ + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: undefined, + cooldownReason: undefined, + }), + expect.objectContaining({ + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: undefined, + cooldownReason: undefined, + }), + ]); + + settleWrites[1]!(); + await flushPromise; + + expect(maxConcurrentWrites).toBe(1); + } finally { + vi.useRealTimers(); + } + }); + + it("flushes newer debounced state after an older pending save fails", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + const savedSnapshots: Array< + Array<{ + refreshToken?: string; + enabled?: false; + disabledReason?: string; + }> + > = []; + + let rejectFirstSave: (error: Error) => void; + const firstSave = new Promise((_resolve, reject) => { + rejectFirstSave = reject; + }); + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + await firstSave; + }); + mockSaveAccounts.mockImplementationOnce(async (storage) => { + savedSnapshots.push( + storage.accounts.map((account) => ({ + refreshToken: account.refreshToken, + enabled: account.enabled, + disabledReason: account.disabledReason, + })), + ); + }); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "shared-refresh", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + manager.saveToDiskDebounced(0); + await vi.advanceTimersByTimeAsync(0); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + + const account = manager.getAccountsSnapshot()[0]!; + manager.disableAccountsWithSameRefreshToken(account); + manager.saveToDiskDebounced(50); + + const flushPromise = manager.flushPendingSave(); + await Promise.resolve(); + + rejectFirstSave!(Object.assign(new Error("EBUSY: file in use"), { code: "EBUSY" })); + await flushPromise; + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + expect(savedSnapshots[1]).toEqual([ + expect.objectContaining({ + refreshToken: "shared-refresh", + enabled: false, + disabledReason: "auth-failure", + }), + ]); + } finally { + vi.useRealTimers(); + } + }); + + it("keeps retrying re-armed flush saves after EBUSY failures until a later write succeeds", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + const savedSnapshots: Array< + Array<{ + refreshToken?: string; + enabled?: false; + disabledReason?: string; + }> + > = []; + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "shared-refresh", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const account = manager.getAccountsSnapshot()[0]!; + manager.disableAccountsWithSameRefreshToken(account); + + let attempts = 0; + mockSaveAccounts.mockImplementation(async (storage) => { + savedSnapshots.push( + storage.accounts.map((savedAccount) => ({ + refreshToken: savedAccount.refreshToken, + enabled: savedAccount.enabled, + disabledReason: savedAccount.disabledReason, + })), + ); + attempts += 1; + if (attempts < 4) { + manager.saveToDiskDebounced(0); + throw Object.assign(new Error("EBUSY: file in use"), { code: "EBUSY" }); + } + }); + + manager.saveToDiskDebounced(0); + + await expect(manager.flushPendingSave()).resolves.toBeUndefined(); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(4); + expect(savedSnapshots.at(-1)).toEqual([ + expect.objectContaining({ + refreshToken: "shared-refresh", + enabled: false, + disabledReason: "auth-failure", + }), + ]); + expect(accountLoggerWarn).toHaveBeenCalledWith( + "flushPendingSave: retrying after flush save failure while newer save is queued", + expect.objectContaining({ error: "EBUSY: file in use" }), + ); + } finally { + vi.useRealTimers(); + } + }); + + it("warns and rejects if flushPendingSave keeps discovering new saves", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + mockSaveAccounts.mockResolvedValue(); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + const originalSaveToDisk = manager.saveToDisk.bind(manager); + let rearmedSaves = 0; + + vi.spyOn(manager, "saveToDisk").mockImplementation(async () => { + if (rearmedSaves < 25) { + rearmedSaves++; + manager.saveToDiskDebounced(0); + } + await originalSaveToDisk(); + }); + + manager.saveToDiskDebounced(0); + await expect(manager.flushPendingSave()).rejects.toThrow( + "flushPendingSave: exceeded max flush iterations; save state may be incomplete", + ); + + expect(accountLoggerWarn).toHaveBeenCalledWith( + "flushPendingSave exceeded max iterations; possible save loop", + expect.objectContaining({ iterations: 20 }), + ); + } finally { + vi.useRealTimers(); + } + }); + + it("stops retrying flushes once the shutdown deadline is exceeded", async () => { + vi.useFakeTimers(); + try { + vi.setSystemTime(new Date(1_000)); + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + + mockSaveAccounts.mockImplementation(async () => { + vi.setSystemTime(new Date(2_000)); + manager.saveToDiskDebounced(0); + }); + + manager.saveToDiskDebounced(0); + + await expect( + manager.flushPendingSave({ deadlineMs: 1_500 }), + ).rejects.toThrow("flushPendingSave: shutdown deadline exceeded before save state settled"); + expect(mockSaveAccounts).toHaveBeenCalledTimes(1); + } finally { + vi.clearAllTimers(); + vi.useRealTimers(); + } + }); + + it("rejects when saveToDisk throws during flush", async () => { + vi.useFakeTimers(); + try { + const { saveAccounts } = await import("../lib/storage.js"); + const mockSaveAccounts = vi.mocked(saveAccounts); + mockSaveAccounts.mockClear(); + mockSaveAccounts.mockRejectedValueOnce(new Error("EBUSY: file in use")); + mockSaveAccounts.mockResolvedValueOnce(); + + const now = Date.now(); + const stored = { + version: 3 as const, + activeIndex: 0, + accounts: [ + { refreshToken: "token-1", addedAt: now, lastUsed: now }, + ], + }; + + const manager = new AccountManager(undefined, stored); + manager.saveToDiskDebounced(0); + + await expect(manager.flushPendingSave()).rejects.toThrow("EBUSY"); + + manager.saveToDiskDebounced(0); + await vi.advanceTimersByTimeAsync(0); + + expect(mockSaveAccounts).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); }); describe("health and token tracking methods", () => { diff --git a/test/index-retry.test.ts b/test/index-retry.test.ts index e4268e6c..aca18064 100644 --- a/test/index-retry.test.ts +++ b/test/index-retry.test.ts @@ -1,4 +1,13 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { runCleanup } from "../lib/shutdown.js"; + +beforeEach(async () => { + await runCleanup(); +}); + +afterEach(async () => { + await runCleanup(); +}); vi.mock("@opencode-ai/plugin/tool", () => { const makeSchema = () => ({ @@ -45,6 +54,10 @@ vi.mock("../lib/accounts.js", () => { return 1; } + getEnabledAccountCount() { + return 1; + } + getCurrentOrNextForFamily() { this.calls += 1; if (this.calls === 1) return null; diff --git a/test/index.test.ts b/test/index.test.ts index daf55c6c..e3308cdc 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -1,4 +1,13 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { getCleanupCount, runCleanup } from "../lib/shutdown.js"; + +beforeEach(async () => { + await runCleanup(); +}); + +afterEach(async () => { + await runCleanup(); +}); vi.mock("@opencode-ai/plugin/tool", () => { const makeSchema = () => ({ @@ -204,6 +213,7 @@ const mockStorage = { accessToken?: string; expiresAt?: number; enabled?: boolean; + disabledReason?: "user" | "auth-failure"; addedAt?: number; lastUsed?: number; coolingDownUntil?: number; @@ -293,6 +303,10 @@ vi.mock("../lib/accounts.js", () => { return this.accounts.length; } + getEnabledAccountCount() { + return this.accounts.filter((account) => account.enabled !== false).length; + } + getCurrentOrNextForFamily() { return this.accounts[0] ?? null; } @@ -344,6 +358,20 @@ vi.mock("../lib/accounts.js", () => { refundToken() {} markSwitched() {} removeAccount() {} + disableAccountsWithSameRefreshToken(account: { refreshToken?: string }) { + const refreshToken = account.refreshToken; + let disabledCount = 0; + for (const storedAccount of this.accounts) { + if (storedAccount.refreshToken !== refreshToken) continue; + if (storedAccount.enabled === false) continue; + delete storedAccount.coolingDownUntil; + delete storedAccount.cooldownReason; + storedAccount.enabled = false; + storedAccount.disabledReason = "auth-failure"; + disabledCount++; + } + return disabledCount; + } removeAccountsWithSameRefreshToken() { return 1; } getMinWaitTimeForFamily() { @@ -504,6 +532,53 @@ describe("OpenAIOAuthPlugin", () => { expect(result.reason).toBe("invalid_response"); expect(vi.mocked(authModule.exchangeAuthorizationCode)).not.toHaveBeenCalled(); }); + + it("invalidates cached account manager after manual OAuth persistence", async () => { + const { AccountManager } = await import("../lib/accounts.js"); + const authModule = await import("../lib/auth/auth.js"); + const accountsModule = await import("../lib/accounts.js"); + const loadFromDiskSpy = vi.spyOn(AccountManager, "loadFromDisk"); + + mockStorage.accounts = [{ refreshToken: "existing-refresh", email: "existing@example.com" }]; + + const getAuth = async () => ({ + type: "oauth" as const, + access: "existing-access", + refresh: "existing-refresh", + expires: Date.now() + 60_000, + multiAccount: true, + }); + + await plugin.auth.loader(getAuth, { options: {}, models: {} }); + loadFromDiskSpy.mockClear(); + + vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({ + type: "success", + access: "manual-access", + refresh: "manual-refresh", + expires: Date.now() + 60_000, + idToken: "manual-id-token", + }); + vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValueOnce([]); + + const manualMethod = plugin.auth.methods[1] as unknown as { + authorize: () => Promise<{ + callback: (input: string) => Promise<{ type: string }>; + }>; + }; + + const flow = await manualMethod.authorize(); + const result = await flow.callback( + "http://127.0.0.1:1455/auth/callback?code=manual-code&state=test-state", + ); + + expect(result.type).toBe("success"); + expect(mockStorage.accounts).toHaveLength(2); + + await plugin.auth.loader(getAuth, { options: {}, models: {} }); + + expect(loadFromDiskSpy).toHaveBeenCalledTimes(1); + }); }); describe("event handler", () => { @@ -2038,7 +2113,122 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expect(await response.text()).toContain("server errors or auth issues"); }); - it("cools down the account when grouped auth removal removes zero entries", async () => { + it("reports disabled pools separately from missing pools", async () => { + const { AccountManager } = await import("../lib/accounts.js"); + + const customManager = { + getAccountCount: () => 2, + getEnabledAccountCount: () => 0, + getCurrentOrNextForFamilyHybrid: () => null, + getSelectionExplainability: () => [], + toAuthDetails: () => ({ + type: "oauth" as const, + access: "access-disabled", + refresh: "refresh-disabled", + expires: Date.now() + 60_000, + }), + hasRefreshToken: () => true, + saveToDiskDebounced: () => {}, + updateFromAuth: () => {}, + clearAuthFailures: () => {}, + incrementAuthFailures: () => 1, + markAccountCoolingDown: () => {}, + markRateLimitedWithReason: () => {}, + recordRateLimit: () => {}, + consumeToken: () => true, + refundToken: () => {}, + markSwitched: () => {}, + removeAccount: () => {}, + recordFailure: () => {}, + recordSuccess: () => {}, + getMinWaitTimeForFamily: () => 0, + shouldShowAccountToast: () => false, + markToastShown: () => {}, + setActiveIndex: () => null, + hasPendingSave: () => false, + waitForPendingSaveToSettle: async () => {}, + flushPendingSave: async () => {}, + getAccountsSnapshot: () => [ + { + index: 0, + email: "user1@example.com", + refreshToken: "refresh-disabled-a", + enabled: false, + disabledReason: "user" as const, + }, + { + index: 1, + email: "user2@example.com", + refreshToken: "refresh-disabled-b", + enabled: false, + disabledReason: "user" as const, + }, + ], + }; + vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(customManager as never); + globalThis.fetch = vi.fn(); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(globalThis.fetch).not.toHaveBeenCalled(); + expect(response.status).toBe(503); + const responseText = await response.text(); + expect(responseText).toContain("All stored Codex accounts are user-disabled"); + expect(responseText).toContain("Re-enable them from account management"); + }); + + it("disables grouped accounts when auth failures hit the threshold", async () => { + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const { AccountManager } = await import("../lib/accounts.js"); + const { ACCOUNT_LIMITS } = await import("../lib/constants.js"); + + vi.spyOn(fetchHelpers, "shouldRefreshToken").mockReturnValue(true); + vi.mocked(fetchHelpers.refreshAndUpdateToken).mockRejectedValue( + new Error("Token expired"), + ); + const incrementAuthFailuresSpy = vi + .spyOn(AccountManager.prototype, "incrementAuthFailures") + .mockReturnValue(ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_DISABLE); + const disableGroupedAccountsSpy = vi + .spyOn(AccountManager.prototype, "disableAccountsWithSameRefreshToken") + .mockReturnValue(2); + vi.spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementationOnce(() => ({ + index: 0, + email: "user1@example.com", + refreshToken: "refresh-1", + }) as never) + .mockImplementation(() => null); + + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "should-not-fetch" }), { status: 200 }), + ); + + const { sdk, mockClient } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(503); + expect(globalThis.fetch).not.toHaveBeenCalled(); + expect(incrementAuthFailuresSpy).toHaveBeenCalledTimes(1); + expect(disableGroupedAccountsSpy).toHaveBeenCalledTimes(1); + expect(mockClient.tui.showToast).toHaveBeenCalledWith( + expect.objectContaining({ + body: expect.objectContaining({ + message: expect.stringContaining("Disabled 2 accounts"), + variant: "error", + }), + }), + ); + }); + + it("cools down the account when grouped auth disable updates zero entries", async () => { const fetchHelpers = await import("../lib/request/fetch-helpers.js"); const { AccountManager } = await import("../lib/accounts.js"); const { ACCOUNT_LIMITS } = await import("../lib/constants.js"); @@ -2049,9 +2239,9 @@ describe("OpenAIOAuthPlugin fetch handler", () => { ); const incrementAuthFailuresSpy = vi .spyOn(AccountManager.prototype, "incrementAuthFailures") - .mockReturnValue(ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_REMOVAL); - const removeGroupedAccountsSpy = vi - .spyOn(AccountManager.prototype, "removeAccountsWithSameRefreshToken") + .mockReturnValue(ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_DISABLE); + const disableGroupedAccountsSpy = vi + .spyOn(AccountManager.prototype, "disableAccountsWithSameRefreshToken") .mockReturnValue(0); const markAccountsWithRefreshTokenCoolingDownSpy = vi.spyOn( AccountManager.prototype, @@ -2071,12 +2261,266 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expect(response.status).toBe(503); expect(globalThis.fetch).not.toHaveBeenCalled(); expect(incrementAuthFailuresSpy).toHaveBeenCalledTimes(1); - expect(removeGroupedAccountsSpy).toHaveBeenCalledTimes(1); + expect(disableGroupedAccountsSpy).toHaveBeenCalledTimes(1); + expect(markAccountsWithRefreshTokenCoolingDownSpy).toHaveBeenCalledWith( + "refresh-1", + ACCOUNT_LIMITS.AUTH_FAILURE_COOLDOWN_MS, + "auth-failure", + ); + }); + + it("surfaces a disabled-pool response when grouped disable and cooldown both no-op", async () => { + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const { AccountManager } = await import("../lib/accounts.js"); + const { ACCOUNT_LIMITS } = await import("../lib/constants.js"); + + vi.spyOn(fetchHelpers, "shouldRefreshToken").mockReturnValue(true); + vi.mocked(fetchHelpers.refreshAndUpdateToken).mockRejectedValue( + new Error("Token expired"), + ); + + let selected = false; + const disableGroupedAccountsSpy = vi.fn(() => 0); + const markAccountsWithRefreshTokenCoolingDownSpy = vi.fn(() => 0); + const saveToDiskDebouncedSpy = vi.fn(); + const customManager = { + getAccountCount: () => 1, + getEnabledAccountCount: () => 0, + getCurrentOrNextForFamilyHybrid: () => { + if (selected) { + return null; + } + selected = true; + return { + index: 0, + email: "user1@example.com", + refreshToken: "refresh-1", + }; + }, + getSelectionExplainability: () => [], + toAuthDetails: () => ({ + type: "oauth" as const, + access: "access-disabled", + refresh: "refresh-disabled", + expires: Date.now() + 60_000, + }), + hasRefreshToken: () => true, + saveToDiskDebounced: saveToDiskDebouncedSpy, + updateFromAuth: () => {}, + clearAuthFailures: () => {}, + incrementAuthFailures: () => ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_DISABLE, + disableAccountsWithSameRefreshToken: disableGroupedAccountsSpy, + markAccountsWithRefreshTokenCoolingDown: markAccountsWithRefreshTokenCoolingDownSpy, + markAccountCoolingDown: () => {}, + markRateLimitedWithReason: () => {}, + recordRateLimit: () => {}, + consumeToken: () => true, + refundToken: () => {}, + markSwitched: () => {}, + removeAccount: () => {}, + recordFailure: () => {}, + recordSuccess: () => {}, + getMinWaitTimeForFamily: () => 0, + shouldShowAccountToast: () => false, + setActiveIndex: () => null, + hasPendingSave: () => false, + waitForPendingSaveToSettle: async () => {}, + flushPendingSave: async () => {}, + getAccountsSnapshot: () => [], + }; + vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(customManager as never); + + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "should-not-fetch" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(503); + expect(globalThis.fetch).not.toHaveBeenCalled(); + expect(disableGroupedAccountsSpy).toHaveBeenCalledTimes(1); expect(markAccountsWithRefreshTokenCoolingDownSpy).toHaveBeenCalledWith( "refresh-1", ACCOUNT_LIMITS.AUTH_FAILURE_COOLDOWN_MS, "auth-failure", ); + expect(saveToDiskDebouncedSpy).toHaveBeenCalledTimes(1); + const responseText = await response.text(); + expect(responseText).toContain( + "All stored Codex accounts are disabled. Re-enable any manually disabled accounts from account management, or run `opencode auth login` to restore access.", + ); + }); + + it("surfaces the disabled-pool response when same-token variants were already user-disabled", async () => { + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const { AccountManager } = await import("../lib/accounts.js"); + const { ACCOUNT_LIMITS } = await import("../lib/constants.js"); + + vi.spyOn(fetchHelpers, "shouldRefreshToken").mockReturnValue(true); + vi.mocked(fetchHelpers.refreshAndUpdateToken).mockRejectedValue( + new Error("Token expired"), + ); + + let selected = false; + const disableGroupedAccountsSpy = vi.fn(() => 0); + const markAccountsWithRefreshTokenCoolingDownSpy = vi.fn(() => 0); + const saveToDiskDebouncedSpy = vi.fn(); + const customManager = { + getAccountCount: () => 2, + getEnabledAccountCount: () => 0, + getCurrentOrNextForFamilyHybrid: () => { + if (selected) { + return null; + } + selected = true; + return { + index: 0, + email: "user1@example.com", + refreshToken: "refresh-shared", + }; + }, + getSelectionExplainability: () => [], + toAuthDetails: () => ({ + type: "oauth" as const, + access: "access-disabled", + refresh: "refresh-shared", + expires: Date.now() + 60_000, + }), + hasRefreshToken: () => true, + saveToDiskDebounced: saveToDiskDebouncedSpy, + updateFromAuth: () => {}, + clearAuthFailures: () => {}, + incrementAuthFailures: () => ACCOUNT_LIMITS.MAX_AUTH_FAILURES_BEFORE_DISABLE, + disableAccountsWithSameRefreshToken: disableGroupedAccountsSpy, + markAccountsWithRefreshTokenCoolingDown: markAccountsWithRefreshTokenCoolingDownSpy, + markAccountCoolingDown: () => {}, + markRateLimitedWithReason: () => {}, + recordRateLimit: () => {}, + consumeToken: () => true, + refundToken: () => {}, + markSwitched: () => {}, + removeAccount: () => {}, + recordFailure: () => {}, + recordSuccess: () => {}, + getMinWaitTimeForFamily: () => 0, + shouldShowAccountToast: () => false, + setActiveIndex: () => null, + hasPendingSave: () => false, + waitForPendingSaveToSettle: async () => {}, + flushPendingSave: async () => {}, + getAccountsSnapshot: () => [ + { + index: 0, + email: "user1@example.com", + refreshToken: "refresh-shared", + enabled: false, + disabledReason: "user" as const, + }, + { + index: 1, + email: "user1@example.com", + refreshToken: "refresh-shared", + enabled: false, + disabledReason: "user" as const, + }, + ], + }; + vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(customManager as never); + + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "should-not-fetch" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(503); + expect(globalThis.fetch).not.toHaveBeenCalled(); + expect(disableGroupedAccountsSpy).toHaveBeenCalledTimes(1); + expect(markAccountsWithRefreshTokenCoolingDownSpy).toHaveBeenCalledWith( + "refresh-shared", + ACCOUNT_LIMITS.AUTH_FAILURE_COOLDOWN_MS, + "auth-failure", + ); + expect(saveToDiskDebouncedSpy).toHaveBeenCalledTimes(1); + const responseText = await response.text(); + expect(responseText).toContain("All stored Codex accounts are user-disabled"); + expect(responseText).toContain("Re-enable them from account management"); + }); + + it("surfaces an auth-failure-specific disabled-pool response when every account is auth-failure disabled", async () => { + const { AccountManager } = await import("../lib/accounts.js"); + + const customManager = { + getAccountCount: () => 2, + getEnabledAccountCount: () => 0, + getCurrentOrNextForFamilyHybrid: () => null, + getSelectionExplainability: () => [], + toAuthDetails: () => ({ + type: "oauth" as const, + access: "access-disabled", + refresh: "refresh-disabled", + expires: Date.now() + 60_000, + }), + hasRefreshToken: () => true, + saveToDiskDebounced: () => {}, + updateFromAuth: () => {}, + clearAuthFailures: () => {}, + incrementAuthFailures: () => 1, + markAccountCoolingDown: () => {}, + markRateLimitedWithReason: () => {}, + recordRateLimit: () => {}, + consumeToken: () => true, + refundToken: () => {}, + markSwitched: () => {}, + removeAccount: () => {}, + recordFailure: () => {}, + recordSuccess: () => {}, + getMinWaitTimeForFamily: () => 0, + shouldShowAccountToast: () => false, + markToastShown: () => {}, + setActiveIndex: () => null, + hasPendingSave: () => false, + waitForPendingSaveToSettle: async () => {}, + flushPendingSave: async () => {}, + getAccountsSnapshot: () => [ + { + index: 0, + email: "user1@example.com", + refreshToken: "refresh-disabled-a", + enabled: false, + disabledReason: "auth-failure" as const, + }, + { + index: 1, + email: "user2@example.com", + refreshToken: "refresh-disabled-b", + enabled: false, + disabledReason: "auth-failure" as const, + }, + ], + }; + vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(customManager as never); + globalThis.fetch = vi.fn(); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(globalThis.fetch).not.toHaveBeenCalled(); + expect(response.status).toBe(503); + const responseText = await response.text(); + expect(responseText).toContain("disabled after repeated auth failures"); + expect(responseText).toContain("Run `opencode auth login` to restore access"); }); it("skips fetch when local token bucket is depleted", async () => { @@ -2290,6 +2734,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { let fallbackSelection = 0; const customManager = { getAccountCount: () => 2, + getEnabledAccountCount: () => 2, getCurrentOrNextForFamilyHybrid: (_family: string, currentModel?: string) => { if (currentModel === "gpt-5-codex") { if (fallbackSelection === 0) { @@ -2459,6 +2904,80 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expect(response.status).toBe(200); }); + it("uses enabled-account positions in selection toasts", async () => { + const { AccountManager } = await import("../lib/accounts.js"); + + const account = { + index: 2, + accountId: "acc-3", + email: "user3@example.com", + refreshToken: "refresh-3", + }; + const markToastShown = vi.fn(); + const customManager = { + getAccountCount: () => 3, + getEnabledAccountCount: () => 2, + getCurrentOrNextForFamilyHybrid: (() => { + let selected = false; + return () => { + if (selected) { + return null; + } + selected = true; + return account; + }; + })(), + getSelectionExplainability: () => [], + toAuthDetails: (selectedAccount: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${selectedAccount.accountId ?? "unknown"}`, + refresh: "refresh-token", + expires: Date.now() + 60_000, + }), + hasRefreshToken: () => true, + saveToDiskDebounced: () => {}, + updateFromAuth: () => {}, + clearAuthFailures: () => {}, + incrementAuthFailures: () => 1, + markAccountCoolingDown: () => {}, + markRateLimitedWithReason: () => {}, + recordRateLimit: () => {}, + consumeToken: () => true, + refundToken: () => {}, + markSwitched: () => {}, + removeAccount: () => {}, + recordFailure: () => {}, + recordSuccess: () => {}, + getMinWaitTimeForFamily: () => 0, + shouldShowAccountToast: () => true, + markToastShown, + setActiveIndex: () => account, + getAccountsSnapshot: () => [account], + }; + vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(customManager as never); + + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "test" }), { status: 200 }), + ); + + const { sdk, mockClient } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(200); + expect(mockClient.tui.showToast).toHaveBeenCalledWith( + expect.objectContaining({ + body: expect.objectContaining({ + message: expect.stringContaining("enabled 1/2"), + variant: "info", + }), + }), + ); + expect(markToastShown).toHaveBeenCalledWith(2); + }); + it("handles empty body in request", async () => { globalThis.fetch = vi.fn().mockResolvedValue( new Response(JSON.stringify({ content: "test" }), { status: 200 }), @@ -3075,7 +3594,6 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => { organizationId: "org-keep", email: "org@example.com", refreshToken: "shared-refresh", - enabled: true, addedAt: 10, lastUsed: 10, }, @@ -3122,28 +3640,35 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => { const mergedOrg = mergedOrgEntries[0]; expect(mergedOrg?.accountId).toBe("org-shared"); expect(mergedOrg?.enabled).toBe(false); + expect(mergedOrg?.disabledReason).toBe("user"); expect(mergedOrg?.coolingDownUntil).toBe(12_000); expect(mergedOrg?.cooldownReason).toBe("auth-failure"); }); - it("preserves same-organization entries when accountId differs", async () => { + it("preserves manual disable reason over auth-failure when collapsing same-organization duplicates", async () => { const accountsModule = await import("../lib/accounts.js"); const authModule = await import("../lib/auth/auth.js"); mockStorage.accounts = [ { - accountId: "org-shared-a", + accountId: "org-shared", organizationId: "org-keep", - email: "org-a@example.com", + email: "manual@example.com", refreshToken: "shared-refresh", + enabled: false, + disabledReason: "user", addedAt: 10, lastUsed: 10, }, { - accountId: "org-shared-b", + accountId: "org-shared", organizationId: "org-keep", - email: "org-b@example.com", + email: "auth-failure@example.com", refreshToken: "shared-refresh", + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: 12_000, + cooldownReason: "auth-failure", addedAt: 20, lastUsed: 20, }, @@ -3151,14 +3676,14 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => { vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({ type: "success", - access: "access-unrelated-preserve", - refresh: "refresh-unrelated-preserve", + access: "access-unrelated-user-disabled", + refresh: "refresh-unrelated-user-disabled", expires: Date.now() + 300_000, - idToken: "id-unrelated-preserve", + idToken: "id-unrelated-user-disabled", }); vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValueOnce([]); vi.mocked(accountsModule.extractAccountId).mockImplementation((accessToken) => - accessToken === "access-unrelated-preserve" ? "unrelated-preserve" : "account-1", + accessToken === "access-unrelated-user-disabled" ? "unrelated-user-disabled" : "account-1", ); const mockClient = createMockClient(); @@ -3172,14 +3697,191 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => { await autoMethod.authorize({ loginMode: "add", accountCount: "1" }); - const orgEntries = mockStorage.accounts.filter( + const mergedOrgEntries = mockStorage.accounts.filter( (account) => account.organizationId === "org-keep", ); - expect(orgEntries).toHaveLength(2); - const accountIds = orgEntries.map((account) => account.accountId).sort(); - expect(accountIds).toEqual(["org-shared-a", "org-shared-b"]); - }); - + expect(mergedOrgEntries).toHaveLength(1); + const mergedOrg = mergedOrgEntries[0]; + expect(mergedOrg).toMatchObject({ + accountId: "org-shared", + enabled: false, + disabledReason: "user", + coolingDownUntil: 12_000, + cooldownReason: "auth-failure", + }); + }); + + it("preserves user-disabled and legacy-disabled variants while reviving auth-failure disables on explicit login", async () => { + const accountsModule = await import("../lib/accounts.js"); + const authModule = await import("../lib/auth/auth.js"); + + mockStorage.accounts = [ + { + accountId: "unrelated-disabled", + organizationId: "org-unrelated", + email: "unrelated@example.com", + refreshToken: "different-refresh", + enabled: false, + disabledReason: "user", + addedAt: 9, + lastUsed: 9, + }, + { + accountId: "org-shared", + organizationId: "org-keep", + email: "org@example.com", + refreshToken: "shared-refresh", + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: 12_000, + cooldownReason: "auth-failure", + addedAt: 10, + lastUsed: 10, + }, + { + accountId: "org-sibling", + organizationId: "org-sibling", + email: "sibling@example.com", + refreshToken: "shared-refresh", + enabled: false, + disabledReason: "user", + addedAt: 11, + lastUsed: 11, + }, + { + accountId: "org-legacy", + organizationId: "org-legacy", + email: "legacy@example.com", + refreshToken: "shared-refresh", + enabled: false, + coolingDownUntil: 18_000, + cooldownReason: "auth-failure", + addedAt: 12, + lastUsed: 12, + }, + ]; + + vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({ + type: "success", + access: "access-shared-refresh", + refresh: "shared-refresh", + expires: Date.now() + 300_000, + idToken: "id-shared-refresh", + }); + vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValueOnce([]); + vi.mocked(accountsModule.extractAccountId).mockImplementation((accessToken) => + accessToken === "access-shared-refresh" ? "org-shared" : "account-1", + ); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + await autoMethod.authorize({ loginMode: "add", accountCount: "1" }); + + const revivedEntries = mockStorage.accounts.filter( + (account) => account.refreshToken === "shared-refresh", + ); + expect(revivedEntries).toHaveLength(3); + expect( + revivedEntries.find((account) => account.accountId === "org-shared"), + ).toMatchObject({ + accountId: "org-shared", + enabled: undefined, + disabledReason: undefined, + coolingDownUntil: undefined, + cooldownReason: undefined, + }); + expect( + revivedEntries.find((account) => account.accountId === "org-sibling"), + ).toMatchObject({ + accountId: "org-sibling", + enabled: false, + disabledReason: "user", + }); + expect( + revivedEntries.find((account) => account.accountId === "org-legacy"), + ).toMatchObject({ + accountId: "org-legacy", + enabled: false, + coolingDownUntil: 18_000, + cooldownReason: "auth-failure", + }); + expect( + revivedEntries.find((account) => account.accountId === "org-legacy")?.disabledReason, + ).toBeUndefined(); + expect( + revivedEntries.find((account) => account.accountId === "org-shared")?.accessToken, + ).toBe("access-shared-refresh"); + expect( + mockStorage.accounts.find((account) => account.accountId === "unrelated-disabled"), + ).toMatchObject({ + accountId: "unrelated-disabled", + refreshToken: "different-refresh", + enabled: false, + disabledReason: "user", + }); + }); + + it("preserves same-organization entries when accountId differs", async () => { + const accountsModule = await import("../lib/accounts.js"); + const authModule = await import("../lib/auth/auth.js"); + + mockStorage.accounts = [ + { + accountId: "org-shared-a", + organizationId: "org-keep", + email: "org-a@example.com", + refreshToken: "shared-refresh", + addedAt: 10, + lastUsed: 10, + }, + { + accountId: "org-shared-b", + organizationId: "org-keep", + email: "org-b@example.com", + refreshToken: "shared-refresh", + addedAt: 20, + lastUsed: 20, + }, + ]; + + vi.mocked(authModule.exchangeAuthorizationCode).mockResolvedValueOnce({ + type: "success", + access: "access-unrelated-preserve", + refresh: "refresh-unrelated-preserve", + expires: Date.now() + 300_000, + idToken: "id-unrelated-preserve", + }); + vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValueOnce([]); + vi.mocked(accountsModule.extractAccountId).mockImplementation((accessToken) => + accessToken === "access-unrelated-preserve" ? "unrelated-preserve" : "account-1", + ); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + await autoMethod.authorize({ loginMode: "add", accountCount: "1" }); + + const orgEntries = mockStorage.accounts.filter( + (account) => account.organizationId === "org-keep", + ); + expect(orgEntries).toHaveLength(2); + const accountIds = orgEntries.map((account) => account.accountId).sort(); + expect(accountIds).toEqual(["org-shared-a", "org-shared-b"]); + }); + it("persists non-team login and updates same record via accountId/refresh fallback", async () => { const accountsModule = await import("../lib/accounts.js"); const authModule = await import("../lib/auth/auth.js"); @@ -3227,6 +3929,23 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => { const accountsModule = await import("../lib/accounts.js"); const refreshQueueModule = await import("../lib/refresh-queue.js"); + mockStorage.accounts = [ + { + refreshToken: "refreshed-refresh", + organizationId: "org-refresh", + accountId: "flagged-live", + accountIdSource: "manual", + accountLabel: "Refresh Workspace", + email: "refresh@example.com", + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: 60_000, + cooldownReason: "auth-failure", + addedAt: Date.now() - 1500, + lastUsed: Date.now() - 1500, + }, + ]; + const flaggedAccounts = [ { refreshToken: "flagged-refresh-cache", @@ -3280,16 +3999,7 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => { } return null; }); - vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValue([ - { - accountId: "token-shared", - source: "token", - label: "Token Shared [id:shared]", - }, - ]); - vi.mocked(accountsModule.selectBestAccountCandidate).mockImplementation( - (candidates) => candidates[0] ?? null, - ); + vi.mocked(accountsModule.getAccountIdCandidates).mockReturnValueOnce([]); const mockClient = createMockClient(); const { OpenAIOAuthPlugin } = await import("../index.js"); @@ -3308,11 +4018,765 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => { expect(new Set(mockStorage.accounts.map((account) => account.organizationId))).toEqual( new Set(["org-cache", "org-refresh"]), ); + expect( + mockStorage.accounts.find((account) => account.organizationId === "org-refresh"), + ).toMatchObject({ + accountId: "flagged-live", + refreshToken: "refreshed-refresh", + enabled: undefined, + disabledReason: undefined, + coolingDownUntil: undefined, + cooldownReason: undefined, + accessToken: "refreshed-access", + }); expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledWith({ version: 1, accounts: [], }); }); + + it("flushes pending account saves during shutdown cleanup", async () => { + const accountsModule = await import("../lib/accounts.js"); + const flushPendingSave = vi.fn(async () => {}); + const manager = await accountsModule.AccountManager.loadFromDisk(); + manager.flushPendingSave = flushPendingSave; + vi.spyOn(accountsModule.AccountManager, "loadFromDisk").mockResolvedValue(manager); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + await runCleanup(); + + expect(flushPendingSave).toHaveBeenCalledTimes(1); + }); + + it("forwards the shutdown deadline to pending account save flushers", async () => { + const accountsModule = await import("../lib/accounts.js"); + const flushPendingSave = vi.fn(async () => {}); + const manager = await accountsModule.AccountManager.loadFromDisk(); + manager.flushPendingSave = flushPendingSave; + vi.spyOn(accountsModule.AccountManager, "loadFromDisk").mockResolvedValue(manager); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + await runCleanup({ deadlineMs: 12_345 }); + + expect(flushPendingSave).toHaveBeenCalledTimes(1); + expect(flushPendingSave).toHaveBeenCalledWith({ deadlineMs: 12_345 }); + }); + + it("logs when shutdown cleanup flush fails", async () => { + const accountsModule = await import("../lib/accounts.js"); + const loggerModule = await import("../lib/logger.js"); + const flushPendingSave = vi.fn(async () => { + throw new Error("EBUSY"); + }); + const manager = await accountsModule.AccountManager.loadFromDisk(); + manager.flushPendingSave = flushPendingSave; + vi.spyOn(accountsModule.AccountManager, "loadFromDisk").mockResolvedValue(manager); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + vi.mocked(loggerModule.logWarn).mockClear(); + await runCleanup(); + + expect(flushPendingSave).toHaveBeenCalledTimes(1); + expect(vi.mocked(loggerModule.logWarn)).toHaveBeenCalledWith( + "[shutdown] flushPendingSave failed; disabled state may not be persisted", + expect.objectContaining({ error: "EBUSY" }), + ); + }); + + it("drops failed tracked managers after shutdown cleanup logs the failure", async () => { + const accountsModule = await import("../lib/accounts.js"); + const loggerModule = await import("../lib/logger.js"); + const failingManager = await accountsModule.AccountManager.loadFromDisk(); + const succeedingManager = await accountsModule.AccountManager.loadFromDisk(); + const failingFlush = vi.fn(async () => { + throw new Error("EBUSY"); + }); + const succeedingFlush = vi.fn(async () => {}); + failingManager.flushPendingSave = failingFlush; + succeedingManager.flushPendingSave = succeedingFlush; + vi.spyOn(accountsModule.AccountManager, "loadFromDisk") + .mockResolvedValueOnce(failingManager) + .mockResolvedValueOnce(succeedingManager); + + const { OpenAIOAuthPlugin } = await import("../index.js"); + const firstPlugin = (await OpenAIOAuthPlugin({ + client: createMockClient(), + } as never)) as unknown as PluginType; + await firstPlugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-1", + refresh: "refresh-token-1", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + vi.mocked(loggerModule.logWarn).mockClear(); + await runCleanup(); + + const secondPlugin = (await OpenAIOAuthPlugin({ + client: createMockClient(), + } as never)) as unknown as PluginType; + await secondPlugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-2", + refresh: "refresh-token-2", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + vi.mocked(loggerModule.logWarn).mockClear(); + await runCleanup(); + + expect(failingFlush).toHaveBeenCalledTimes(1); + expect(succeedingFlush).toHaveBeenCalledTimes(1); + expect(vi.mocked(loggerModule.logWarn)).not.toHaveBeenCalledWith( + "[shutdown] flushPendingSave failed; disabled state may not be persisted", + expect.anything(), + ); + }); + + it("flushes stale invalidated managers alongside the current manager during shutdown cleanup", async () => { + const accountsModule = await import("../lib/accounts.js"); + const cliModule = await import("../lib/cli.js"); + const managerOne = await accountsModule.AccountManager.loadFromDisk(); + const managerTwo = await accountsModule.AccountManager.loadFromDisk(); + const flushManagerOne = vi.fn(async () => {}); + const flushManagerTwo = vi.fn(async () => {}); + managerOne.flushPendingSave = flushManagerOne; + managerTwo.flushPendingSave = flushManagerTwo; + managerOne.hasPendingSave = vi.fn(() => true); + managerTwo.hasPendingSave = vi.fn(() => false); + vi.spyOn(accountsModule.AccountManager, "loadFromDisk") + .mockResolvedValueOnce(managerOne) + .mockResolvedValue(managerTwo); + + mockStorage.accounts = [ + { + accountId: "workspace-managed", + email: "managed@example.com", + refreshToken: "refresh-managed", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-1", + refresh: "refresh-token-1", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + flushManagerOne.mockClear(); + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-2", + refresh: "refresh-token-2", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + await runCleanup(); + + expect(flushManagerOne).toHaveBeenCalledTimes(1); + expect(flushManagerTwo).toHaveBeenCalledTimes(1); + }); + + it("keeps idle detached managers tracked until shutdown in case they enqueue a later save", async () => { + const accountsModule = await import("../lib/accounts.js"); + const cliModule = await import("../lib/cli.js"); + const managerOne = await accountsModule.AccountManager.loadFromDisk(); + const managerTwo = await accountsModule.AccountManager.loadFromDisk(); + const flushManagerOne = vi.fn(async () => {}); + const flushManagerTwo = vi.fn(async () => {}); + let managerOneHasPendingSave = false; + managerOne.flushPendingSave = flushManagerOne; + managerTwo.flushPendingSave = flushManagerTwo; + managerOne.hasPendingSave = vi.fn(() => managerOneHasPendingSave); + managerTwo.hasPendingSave = vi.fn(() => false); + vi.spyOn(accountsModule.AccountManager, "loadFromDisk") + .mockResolvedValueOnce(managerOne) + .mockResolvedValue(managerTwo); + + mockStorage.accounts = [ + { + accountId: "workspace-managed", + email: "managed@example.com", + refreshToken: "refresh-managed", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-1", + refresh: "refresh-token-1", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + flushManagerOne.mockClear(); + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-2", + refresh: "refresh-token-2", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + managerOneHasPendingSave = true; + await runCleanup(); + + expect(flushManagerOne).toHaveBeenCalledTimes(1); + expect(flushManagerTwo).toHaveBeenCalledTimes(1); + }); + + it("prunes tracked invalidated managers once their pending saves settle", async () => { + const accountsModule = await import("../lib/accounts.js"); + const cliModule = await import("../lib/cli.js"); + const managerOne = await accountsModule.AccountManager.loadFromDisk(); + const flushManagerOne = vi.fn(async () => {}); + let managerOneHasPendingSave = true; + let resolveManagerOneSettled: (() => void) | null = null; + managerOne.flushPendingSave = flushManagerOne; + managerOne.hasPendingSave = vi.fn(() => managerOneHasPendingSave); + managerOne.waitForPendingSaveToSettle = vi.fn( + () => + new Promise((resolve) => { + resolveManagerOneSettled = resolve; + }), + ); + vi.spyOn(accountsModule.AccountManager, "loadFromDisk") + .mockResolvedValue(managerOne); + + mockStorage.accounts = [ + { + accountId: "workspace-managed", + email: "managed@example.com", + refreshToken: "refresh-managed", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-1", + refresh: "refresh-token-1", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + flushManagerOne.mockClear(); + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + + managerOneHasPendingSave = false; + resolveManagerOneSettled?.(); + await Promise.resolve(); + await Promise.resolve(); + + await runCleanup(); + + expect(flushManagerOne).not.toHaveBeenCalled(); + }); + + it("flushes a tracked debounce-only manager during shutdown before its settle waiter resolves naturally", async () => { + const accountsModule = await import("../lib/accounts.js"); + const cliModule = await import("../lib/cli.js"); + const managerOne = await accountsModule.AccountManager.loadFromDisk(); + const managerTwo = await accountsModule.AccountManager.loadFromDisk(); + const flushManagerOne = vi.fn(async () => { + managerOneHasPendingSave = false; + resolveManagerOneSettled?.(); + }); + const flushManagerTwo = vi.fn(async () => {}); + let managerOneHasPendingSave = true; + let resolveManagerOneSettled: (() => void) | null = null; + managerOne.flushPendingSave = flushManagerOne; + managerOne.hasPendingSave = vi.fn(() => managerOneHasPendingSave); + managerOne.waitForPendingSaveToSettle = vi.fn( + () => + new Promise((resolve) => { + resolveManagerOneSettled = resolve; + }), + ); + managerTwo.flushPendingSave = flushManagerTwo; + managerTwo.hasPendingSave = vi.fn(() => false); + managerTwo.waitForPendingSaveToSettle = vi.fn(async () => {}); + vi.spyOn(accountsModule.AccountManager, "loadFromDisk") + .mockResolvedValueOnce(managerOne) + .mockResolvedValue(managerTwo); + + mockStorage.accounts = [ + { + accountId: "workspace-managed", + email: "managed@example.com", + refreshToken: "refresh-managed", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-1", + refresh: "refresh-token-1", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + flushManagerOne.mockClear(); + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-2", + refresh: "refresh-token-2", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + expect(managerOne.waitForPendingSaveToSettle).toHaveBeenCalledTimes(1); + + await runCleanup(); + await Promise.resolve(); + await Promise.resolve(); + + expect(flushManagerOne).toHaveBeenCalledTimes(1); + expect(flushManagerTwo).toHaveBeenCalledTimes(1); + expect(managerOneHasPendingSave).toBe(false); + }); + + it("re-arms tracked manager prune waiters when pending saves continue after a settle notification", async () => { + const accountsModule = await import("../lib/accounts.js"); + const cliModule = await import("../lib/cli.js"); + const managerOne = await accountsModule.AccountManager.loadFromDisk(); + const flushManagerOne = vi.fn(async () => { + managerOneHasPendingSave = false; + }); + let managerOneHasPendingSave = true; + const settleResolvers: Array<() => void> = []; + managerOne.flushPendingSave = flushManagerOne; + managerOne.hasPendingSave = vi.fn(() => managerOneHasPendingSave); + managerOne.waitForPendingSaveToSettle = vi.fn( + () => + new Promise((resolve) => { + settleResolvers.push(resolve); + }), + ); + vi.spyOn(accountsModule.AccountManager, "loadFromDisk").mockResolvedValue(managerOne); + + mockStorage.accounts = [ + { + accountId: "workspace-managed", + email: "managed@example.com", + refreshToken: "refresh-managed", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + await plugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-1", + refresh: "refresh-token-1", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + flushManagerOne.mockClear(); + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + expect(managerOne.waitForPendingSaveToSettle).toHaveBeenCalledTimes(1); + + settleResolvers[0]?.(); + await Promise.resolve(); + await Promise.resolve(); + + expect(managerOne.waitForPendingSaveToSettle).toHaveBeenCalledTimes(2); + + await runCleanup(); + + expect(flushManagerOne).toHaveBeenCalledTimes(1); + }); + + it("flushes tracked account managers from multiple plugin instances during shutdown cleanup", async () => { + const accountsModule = await import("../lib/accounts.js"); + const managerOne = await accountsModule.AccountManager.loadFromDisk(); + const managerTwo = await accountsModule.AccountManager.loadFromDisk(); + const flushManagerOne = vi.fn(async () => {}); + const flushManagerTwo = vi.fn(async () => {}); + managerOne.flushPendingSave = flushManagerOne; + managerTwo.flushPendingSave = flushManagerTwo; + vi.spyOn(accountsModule.AccountManager, "loadFromDisk") + .mockResolvedValueOnce(managerOne) + .mockResolvedValueOnce(managerTwo); + + const { OpenAIOAuthPlugin } = await import("../index.js"); + const firstPlugin = (await OpenAIOAuthPlugin({ + client: createMockClient(), + } as never)) as unknown as PluginType; + const secondPlugin = (await OpenAIOAuthPlugin({ + client: createMockClient(), + } as never)) as unknown as PluginType; + + await firstPlugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-1", + refresh: "refresh-token-1", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + await secondPlugin.auth.loader( + async () => ({ + type: "oauth", + access: "access-token-2", + refresh: "refresh-token-2", + expires: Date.now() + 60_000, + }) as never, + {}, + ); + + expect(getCleanupCount()).toBe(1); + + await runCleanup(); + + expect(flushManagerOne).toHaveBeenCalledTimes(1); + expect(flushManagerTwo).toHaveBeenCalledTimes(1); + }); + + it("replaces the account-save shutdown cleanup on plugin re-initialization", async () => { + const { OpenAIOAuthPlugin } = await import("../index.js"); + + expect(getCleanupCount()).toBe(0); + + await OpenAIOAuthPlugin({ client: createMockClient() } as never); + expect(getCleanupCount()).toBe(1); + + await OpenAIOAuthPlugin({ client: createMockClient() } as never); + expect(getCleanupCount()).toBe(1); + }); + + it("writes user disable metadata from the auth manage menu and clears it on manual re-enable", async () => { + const cliModule = await import("../lib/cli.js"); + const storageModule = await import("../lib/storage.js"); + + mockStorage.accounts = [ + { + accountId: "workspace-managed", + email: "managed@example.com", + refreshToken: "refresh-managed", + coolingDownUntil: 30_000, + cooldownReason: "network-error", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + + const saveAccountsMock = vi.mocked(storageModule.saveAccounts); + expect(saveAccountsMock).toHaveBeenCalledTimes(2); + expect(saveAccountsMock.mock.calls[0]?.[0].accounts[0]).toMatchObject({ + accountId: "workspace-managed", + enabled: false, + disabledReason: "user", + coolingDownUntil: 30_000, + cooldownReason: "network-error", + }); + expect(saveAccountsMock.mock.calls[1]?.[0].accounts[0]).toMatchObject({ + accountId: "workspace-managed", + }); + expect(saveAccountsMock.mock.calls[1]?.[0].accounts[0]).not.toHaveProperty("enabled"); + expect(saveAccountsMock.mock.calls[1]?.[0].accounts[0]).not.toHaveProperty("disabledReason"); + expect(saveAccountsMock.mock.calls[1]?.[0].accounts[0]).not.toHaveProperty("coolingDownUntil"); + expect(saveAccountsMock.mock.calls[1]?.[0].accounts[0]).not.toHaveProperty("cooldownReason"); + expect(mockStorage.accounts[0]).toMatchObject({ + accountId: "workspace-managed", + }); + expect(mockStorage.accounts[0]).not.toHaveProperty("enabled"); + expect(mockStorage.accounts[0]).not.toHaveProperty("disabledReason"); + expect(mockStorage.accounts[0]).not.toHaveProperty("coolingDownUntil"); + expect(mockStorage.accounts[0]).not.toHaveProperty("cooldownReason"); + }); + + it("keeps auth-failure disables blocked in the auth manage menu until a fresh login", async () => { + const cliModule = await import("../lib/cli.js"); + const loggerModule = await import("../lib/logger.js"); + const storageModule = await import("../lib/storage.js"); + + mockStorage.accounts = [ + { + accountId: "workspace-auth-failure", + email: "blocked@example.com", + refreshToken: "refresh-blocked", + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: 30_000, + cooldownReason: "auth-failure", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + vi.mocked(loggerModule.logWarn).mockClear(); + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + + expect(vi.mocked(storageModule.saveAccounts)).not.toHaveBeenCalled(); + expect(mockStorage.accounts[0]).toMatchObject({ + accountId: "workspace-auth-failure", + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: 30_000, + cooldownReason: "auth-failure", + }); + expect(mockClient.tui.showToast).toHaveBeenCalledWith({ + body: expect.objectContaining({ + message: + "This account was disabled after repeated auth failures. Run 'opencode auth login' to re-enable with fresh credentials.", + variant: "warning", + }), + }); + const infoCall = vi + .mocked(loggerModule.logInfo) + .mock.calls.find(([message]) => message === "[account-menu] prompted re-auth for auth-failure disabled account"); + const warnCall = vi + .mocked(loggerModule.logWarn) + .mock.calls.find(([message]) => message === "[account-menu] blocked re-enable for auth-failure disabled account"); + expect(infoCall?.[1]).toEqual(expect.objectContaining({ accountIndex: 1 })); + expect(warnCall?.[1]).toEqual(expect.objectContaining({ accountIndex: 1 })); + expect(infoCall?.[1]).not.toHaveProperty("accountId"); + expect(warnCall?.[1]).not.toHaveProperty("accountId"); + }); + + it("falls back to a generic console hint when TUI toast is unavailable for auth-failure disables", async () => { + const cliModule = await import("../lib/cli.js"); + const loggerModule = await import("../lib/logger.js"); + const storageModule = await import("../lib/storage.js"); + const stdoutWrite = vi.spyOn(process.stdout, "write").mockReturnValue(true); + + try { + mockStorage.accounts = [ + { + accountId: "workspace-auth-failure", + email: "blocked@example.com", + refreshToken: "refresh-blocked", + enabled: false, + disabledReason: "auth-failure", + coolingDownUntil: 30_000, + cooldownReason: "auth-failure", + addedAt: 10, + lastUsed: 10, + }, + ]; + + vi.mocked(cliModule.promptLoginMode) + .mockResolvedValueOnce({ mode: "manage", toggleAccountIndex: 0 }) + .mockResolvedValueOnce({ mode: "cancel" }); + + const mockClient = createMockClient(); + mockClient.tui.showToast.mockRejectedValueOnce(new Error("TUI unavailable")); + const { OpenAIOAuthPlugin } = await import("../index.js"); + const plugin = (await OpenAIOAuthPlugin({ + client: mockClient, + } as never)) as unknown as PluginType; + const autoMethod = plugin.auth.methods[0] as unknown as { + authorize: (inputs?: Record) => Promise<{ instructions: string }>; + }; + + vi.mocked(loggerModule.logWarn).mockClear(); + const authResult = await autoMethod.authorize(); + expect(authResult.instructions).toBe("Authentication cancelled"); + + expect(vi.mocked(storageModule.saveAccounts)).not.toHaveBeenCalled(); + expect(stdoutWrite).toHaveBeenCalledWith( + "\nRun 'opencode auth login' to re-enable this account.\n\n", + ); + const infoCall = vi + .mocked(loggerModule.logInfo) + .mock.calls.find(([message]) => message === "[account-menu] prompted re-auth for auth-failure disabled account"); + const warnCall = vi + .mocked(loggerModule.logWarn) + .mock.calls.find(([message]) => message === "[account-menu] blocked re-enable for auth-failure disabled account"); + expect(infoCall?.[1]).toEqual(expect.objectContaining({ accountIndex: 1 })); + expect(warnCall?.[1]).toEqual(expect.objectContaining({ accountIndex: 1 })); + expect(infoCall?.[1]).not.toHaveProperty("accountId"); + expect(warnCall?.[1]).not.toHaveProperty("accountId"); + } finally { + stdoutWrite.mockRestore(); + } + }); }); describe("OpenAIOAuthPlugin showToast error handling", () => { diff --git a/test/schemas.test.ts b/test/schemas.test.ts index 81283aff..fb2e390d 100644 --- a/test/schemas.test.ts +++ b/test/schemas.test.ts @@ -155,6 +155,18 @@ describe("AccountMetadataV3Schema", () => { expect(result.success).toBe(false); }); + it("normalizes invalid disabledReason values to undefined", () => { + const result = AccountMetadataV3Schema.safeParse({ + ...validAccount, + enabled: false, + disabledReason: "future-reason", + }); + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.disabledReason).toBeUndefined(); + } + }); + it("normalizes account tags by trimming, lowercasing, and filtering blanks", () => { const result = AccountMetadataV3Schema.safeParse({ ...validAccount, @@ -541,4 +553,21 @@ describe("getValidationErrors", () => { expect(errors.length).toBeGreaterThan(0); expect(errors[0]).not.toMatch(/^[a-zA-Z0-9_.]+: /); }); + + it("does not warn for unknown disabledReason values in account storage", () => { + const errors = getValidationErrors(AnyAccountStorageSchema, { + version: 3, + accounts: [ + { + refreshToken: "rt", + enabled: false, + disabledReason: "future-reason", + addedAt: 1, + lastUsed: 1, + }, + ], + activeIndex: 0, + }); + expect(errors).toEqual([]); + }); }); diff --git a/test/shutdown.test.ts b/test/shutdown.test.ts index c64ecf41..7e6f4060 100644 --- a/test/shutdown.test.ts +++ b/test/shutdown.test.ts @@ -65,6 +65,21 @@ describe("Graceful shutdown", () => { expect(getCleanupCount()).toBe(0); }); + it("drains cleanup functions registered while cleanup is already running", async () => { + const order: number[] = []; + registerCleanup(async () => { + order.push(1); + registerCleanup(() => { + order.push(2); + }); + }); + + await runCleanup(); + + expect(order).toEqual([1, 2]); + expect(getCleanupCount()).toBe(0); + }); + it("unregister is no-op for non-registered function", () => { const fn = vi.fn(); unregisterCleanup(fn); @@ -75,7 +90,7 @@ describe("Graceful shutdown", () => { it("SIGINT handler runs cleanup and exits with code 0", async () => { const capturedHandlers = new Map void>(); - const processOnceSpy = vi.spyOn(process, "once").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { capturedHandlers.set(String(event), handler); return process; }); @@ -88,23 +103,25 @@ describe("Graceful shutdown", () => { const cleanupFn = vi.fn(); freshRegister(cleanupFn); - const sigintHandler = capturedHandlers.get("SIGINT"); - expect(sigintHandler).toBeDefined(); - - sigintHandler!(); - await new Promise((r) => setTimeout(r, 10)); - - expect(cleanupFn).toHaveBeenCalled(); - expect(processExitSpy).toHaveBeenCalledWith(0); - - processOnceSpy.mockRestore(); - processExitSpy.mockRestore(); + try { + const sigintHandler = capturedHandlers.get("SIGINT"); + expect(sigintHandler).toBeDefined(); + + sigintHandler!(); + await vi.waitFor(() => { + expect(cleanupFn).toHaveBeenCalled(); + expect(processExitSpy).toHaveBeenCalledWith(0); + }); + } finally { + processOnSpy.mockRestore(); + processExitSpy.mockRestore(); + } }); it("SIGTERM handler runs cleanup and exits with code 0", async () => { const capturedHandlers = new Map void>(); - const processOnceSpy = vi.spyOn(process, "once").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { capturedHandlers.set(String(event), handler); return process; }); @@ -117,23 +134,230 @@ describe("Graceful shutdown", () => { const cleanupFn = vi.fn(); freshRegister(cleanupFn); - const sigtermHandler = capturedHandlers.get("SIGTERM"); - expect(sigtermHandler).toBeDefined(); + try { + const sigtermHandler = capturedHandlers.get("SIGTERM"); + expect(sigtermHandler).toBeDefined(); + + sigtermHandler!(); + await vi.waitFor(() => { + expect(cleanupFn).toHaveBeenCalled(); + expect(processExitSpy).toHaveBeenCalledWith(0); + }); + } finally { + processOnSpy.mockRestore(); + processExitSpy.mockRestore(); + } + }); + + it("passes the shutdown deadline to cleanup functions on signal exit", async () => { + const capturedHandlers = new Map void>(); + + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + capturedHandlers.set(String(event), handler); + return process; + }); + const processExitSpy = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + vi.resetModules(); + const { registerCleanup: freshRegister, runCleanup: freshRunCleanup } = await import("../lib/shutdown.js"); + await freshRunCleanup(); + + const cleanupFn = vi.fn(async () => {}); + freshRegister(cleanupFn); + + try { + const sigtermHandler = capturedHandlers.get("SIGTERM"); + expect(sigtermHandler).toBeDefined(); + + sigtermHandler!(); + await vi.waitFor(() => { + expect(cleanupFn).toHaveBeenCalledWith( + expect.objectContaining({ + deadlineMs: expect.any(Number), + }), + ); + expect(processExitSpy).toHaveBeenCalledWith(0); + }); + } finally { + processOnSpy.mockRestore(); + processExitSpy.mockRestore(); + } + }); + + it("keeps shutdown handlers installed until async cleanup completes", async () => { + const capturedHandlers = new Map void>(); + + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + capturedHandlers.set(String(event), handler); + return process; + }); + const processOffSpy = vi.spyOn(process, "off").mockImplementation(() => process); + const processExitSpy = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + vi.resetModules(); + const { registerCleanup: freshRegister, runCleanup: freshRunCleanup } = await import("../lib/shutdown.js"); + await freshRunCleanup(); - sigtermHandler!(); - await new Promise((r) => setTimeout(r, 10)); + let resolveCleanup!: () => void; + const cleanupPromise = new Promise((resolve) => { + resolveCleanup = resolve; + }); + freshRegister(async () => { + await cleanupPromise; + }); - expect(cleanupFn).toHaveBeenCalled(); - expect(processExitSpy).toHaveBeenCalledWith(0); + try { + const sigtermHandler = capturedHandlers.get("SIGTERM"); + expect(sigtermHandler).toBeDefined(); + + sigtermHandler!(); + await Promise.resolve(); + expect(processOffSpy).not.toHaveBeenCalled(); + + resolveCleanup(); + await vi.waitFor(() => { + expect(processExitSpy).toHaveBeenCalledWith(0); + }); + expect(processOffSpy).toHaveBeenCalled(); + } finally { + processOnSpy.mockRestore(); + processOffSpy.mockRestore(); + processExitSpy.mockRestore(); + } + }); - processOnceSpy.mockRestore(); - processExitSpy.mockRestore(); + it("ignores repeated signals while async cleanup is already running", async () => { + const capturedHandlers = new Map void>(); + + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + capturedHandlers.set(String(event), handler); + return process; + }); + const processExitSpy = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + vi.resetModules(); + const { registerCleanup: freshRegister, runCleanup: freshRunCleanup } = await import("../lib/shutdown.js"); + await freshRunCleanup(); + + let resolveCleanup!: () => void; + const cleanupPromise = new Promise((resolve) => { + resolveCleanup = resolve; + }); + const cleanupFn = vi.fn(async () => { + await cleanupPromise; + }); + freshRegister(cleanupFn); + + try { + const sigtermHandler = capturedHandlers.get("SIGTERM"); + expect(sigtermHandler).toBeDefined(); + + sigtermHandler!(); + sigtermHandler!(); + await Promise.resolve(); + + expect(cleanupFn).toHaveBeenCalledTimes(1); + expect(processExitSpy).not.toHaveBeenCalled(); + + resolveCleanup(); + await vi.waitFor(() => { + expect(processExitSpy).toHaveBeenCalledTimes(1); + expect(processExitSpy).toHaveBeenCalledWith(0); + }); + } finally { + processOnSpy.mockRestore(); + processExitSpy.mockRestore(); + } + }); + + it("forces exit if signal cleanup stalls past the timeout", async () => { + vi.useFakeTimers(); + const capturedHandlers = new Map void>(); + + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + capturedHandlers.set(String(event), handler); + return process; + }); + const processExitSpy = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + vi.resetModules(); + const { registerCleanup: freshRegister, runCleanup: freshRunCleanup } = await import("../lib/shutdown.js"); + await freshRunCleanup(); + + const cleanupFn = vi.fn(async () => { + await new Promise(() => { + // Intentionally never resolves so the timeout path can force exit. + }); + }); + freshRegister(cleanupFn); + + try { + const sigtermHandler = capturedHandlers.get("SIGTERM"); + expect(sigtermHandler).toBeDefined(); + + sigtermHandler!(); + await Promise.resolve(); + + expect(cleanupFn).toHaveBeenCalledTimes(1); + expect(processExitSpy).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(9_999); + + expect(processExitSpy).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(1); + + expect(processExitSpy).toHaveBeenCalledTimes(1); + expect(processExitSpy).toHaveBeenCalledWith(0); + } finally { + vi.useRealTimers(); + processOnSpy.mockRestore(); + processExitSpy.mockRestore(); + } + }); + + it("allows later signals again if process exit is intercepted", async () => { + const capturedHandlers = new Map void>(); + + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + capturedHandlers.set(String(event), handler); + return process; + }); + const processExitSpy = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + vi.resetModules(); + const { registerCleanup: freshRegister, runCleanup: freshRunCleanup } = await import("../lib/shutdown.js"); + await freshRunCleanup(); + + const cleanupFn = vi.fn(); + freshRegister(cleanupFn); + + try { + const sigtermHandler = capturedHandlers.get("SIGTERM"); + expect(sigtermHandler).toBeDefined(); + + sigtermHandler!(); + await vi.waitFor(() => { + expect(cleanupFn).toHaveBeenCalledTimes(1); + expect(processExitSpy).toHaveBeenCalledTimes(1); + }); + + sigtermHandler!(); + await vi.waitFor(() => { + expect(processExitSpy).toHaveBeenCalledTimes(2); + }); + + expect(cleanupFn).toHaveBeenCalledTimes(1); + } finally { + processOnSpy.mockRestore(); + processExitSpy.mockRestore(); + } }); it("beforeExit handler runs cleanup without calling exit", async () => { const capturedHandlers = new Map void>(); - const processOnceSpy = vi.spyOn(process, "once").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { capturedHandlers.set(String(event), handler); return process; }); @@ -149,29 +373,138 @@ describe("Graceful shutdown", () => { const beforeExitHandler = capturedHandlers.get("beforeExit"); expect(beforeExitHandler).toBeDefined(); - beforeExitHandler!(); - await new Promise((r) => setTimeout(r, 10)); + try { + beforeExitHandler!(); + await vi.waitFor(() => { + expect(cleanupFn).toHaveBeenCalled(); + }); + expect(processExitSpy).not.toHaveBeenCalled(); + } finally { + processOnSpy.mockRestore(); + processExitSpy.mockRestore(); + } + }); - expect(cleanupFn).toHaveBeenCalled(); - expect(processExitSpy).not.toHaveBeenCalled(); + it("keeps handlers installed when a signal arrives during beforeExit cleanup", async () => { + const capturedHandlers = new Map void>(); - processOnceSpy.mockRestore(); - processExitSpy.mockRestore(); + const processOnSpy = vi.spyOn(process, "on").mockImplementation((event: string | symbol, handler: (...args: unknown[]) => void) => { + capturedHandlers.set(String(event), handler); + return process; + }); + const processOffSpy = vi.spyOn(process, "off").mockImplementation(() => process); + const processExitSpy = vi.spyOn(process, "exit").mockImplementation(() => undefined as never); + + vi.resetModules(); + const { registerCleanup: freshRegister, runCleanup: freshRunCleanup } = await import("../lib/shutdown.js"); + await freshRunCleanup(); + + let resolveCleanup!: () => void; + const cleanupPromise = new Promise((resolve) => { + resolveCleanup = resolve; + }); + const cleanupFn = vi.fn(async () => { + await cleanupPromise; + }); + freshRegister(cleanupFn); + + try { + const beforeExitHandler = capturedHandlers.get("beforeExit"); + const sigtermHandler = capturedHandlers.get("SIGTERM"); + expect(beforeExitHandler).toBeDefined(); + expect(sigtermHandler).toBeDefined(); + + beforeExitHandler!(); + await Promise.resolve(); + sigtermHandler!(); + await Promise.resolve(); + + expect(cleanupFn).toHaveBeenCalledTimes(1); + expect(processExitSpy).not.toHaveBeenCalled(); + expect(processOffSpy).not.toHaveBeenCalled(); + + resolveCleanup(); + await vi.waitFor(() => { + expect(processExitSpy).toHaveBeenCalledWith(0); + }); + expect(processOffSpy).toHaveBeenCalled(); + } finally { + processOnSpy.mockRestore(); + processOffSpy.mockRestore(); + processExitSpy.mockRestore(); + } }); it("signal handlers are only registered once", async () => { - const processOnceSpy = vi.spyOn(process, "once").mockImplementation(() => process); + const processOnSpy = vi.spyOn(process, "on").mockImplementation(() => process); vi.resetModules(); const { registerCleanup: freshRegister } = await import("../lib/shutdown.js"); freshRegister(() => {}); - const firstCallCount = processOnceSpy.mock.calls.length; + const firstCallCount = processOnSpy.mock.calls.length; + + freshRegister(() => {}); + expect(processOnSpy.mock.calls.length).toBe(firstCallCount); + + processOnSpy.mockRestore(); + }); + + it("re-registers signal handlers after cleanup resets shutdown state", async () => { + const processOnSpy = vi.spyOn(process, "on").mockImplementation(() => process); + const processOffSpy = vi.spyOn(process, "off").mockImplementation(() => process); + + vi.resetModules(); + const { + registerCleanup: freshRegister, + runCleanup: freshRunCleanup, + } = await import("../lib/shutdown.js"); + + freshRegister(() => {}); + expect(processOnSpy).toHaveBeenCalledTimes(3); + await freshRunCleanup(); + expect(processOffSpy).toHaveBeenCalledTimes(3); freshRegister(() => {}); - expect(processOnceSpy.mock.calls.length).toBe(firstCallCount); + expect(processOnSpy).toHaveBeenCalledTimes(6); + + processOnSpy.mockRestore(); + processOffSpy.mockRestore(); + }); + + it("reinstalls signal handlers when cleanup is registered during teardown", async () => { + const processOnSpy = vi.spyOn(process, "on").mockImplementation(() => process); + + vi.resetModules(); + const { + registerCleanup: freshRegister, + runCleanup: freshRunCleanup, + getCleanupCount: freshGetCleanupCount, + } = await import("../lib/shutdown.js"); + await freshRunCleanup(); + + let registeredDuringTeardown = false; + const processOffSpy = vi.spyOn(process, "off").mockImplementation((event: string | symbol) => { + if (!registeredDuringTeardown && String(event) === "beforeExit") { + registeredDuringTeardown = true; + freshRegister(() => {}); + } + return process; + }); + + try { + freshRegister(() => {}); + expect(processOnSpy).toHaveBeenCalledTimes(3); + + await freshRunCleanup(); - processOnceSpy.mockRestore(); + expect(freshGetCleanupCount()).toBe(1); + expect(processOffSpy).toHaveBeenCalledTimes(3); + expect(processOnSpy).toHaveBeenCalledTimes(6); + } finally { + processOnSpy.mockRestore(); + processOffSpy.mockRestore(); + } }); }); }); diff --git a/test/storage.test.ts b/test/storage.test.ts index 94d20141..5450269d 100644 --- a/test/storage.test.ts +++ b/test/storage.test.ts @@ -1050,6 +1050,120 @@ describe("storage", () => { expect(result?.accounts).toHaveLength(1); }); + it("defaults migrated legacy disabled accounts to user disabledReason", () => { + const data = { + version: 1, + activeIndex: 0, + accounts: [ + { + refreshToken: "t1", + accountId: "A", + enabled: false, + addedAt: 1, + lastUsed: 1, + }, + ], + }; + + const result = normalizeAccountStorage(data); + expect(result?.version).toBe(3); + expect(result?.accounts[0]).toMatchObject({ + refreshToken: "t1", + enabled: false, + disabledReason: "user", + }); + }); + + it("defaults invalid migrated legacy disabled reasons to user", () => { + const data = { + version: 1, + activeIndex: 0, + accounts: [ + { + refreshToken: "t1", + accountId: "A", + enabled: false, + disabledReason: "future-reason", + addedAt: 1, + lastUsed: 1, + }, + ], + }; + + const result = normalizeAccountStorage(data); + expect(result?.version).toBe(3); + expect(result?.accounts[0]).toMatchObject({ + refreshToken: "t1", + enabled: false, + disabledReason: "user", + }); + }); + + it("omits enabled for enabled accounts during v1 migration", () => { + const data = { + version: 1, + activeIndex: 0, + accounts: [ + { + refreshToken: "t1", + accountId: "A", + enabled: true, + addedAt: 1, + lastUsed: 1, + }, + ], + }; + + const result = normalizeAccountStorage(data); + expect(result?.accounts[0]?.enabled).toBeUndefined(); + expect(result?.accounts[0]).not.toHaveProperty("enabled"); + }); + + it("strips invalid disabledReason values from disabled v3 accounts", () => { + const data = { + version: 3, + activeIndex: 0, + accounts: [ + { + refreshToken: "t1", + accountId: "A", + enabled: false, + disabledReason: "future-reason", + addedAt: 1, + lastUsed: 1, + }, + ], + }; + + const result = normalizeAccountStorage(data); + expect(result?.accounts[0]).toMatchObject({ + refreshToken: "t1", + enabled: false, + }); + expect(result?.accounts[0]?.disabledReason).toBeUndefined(); + expect(result?.accounts[0]).not.toHaveProperty("disabledReason"); + }); + + it("strips enabled: true from v3 accounts during normalization", () => { + const data = { + version: 3, + activeIndex: 0, + accounts: [ + { + refreshToken: "t1", + accountId: "A", + enabled: true, + addedAt: 1, + lastUsed: 1, + }, + ], + }; + + const result = normalizeAccountStorage(data); + expect(result?.accounts[0]?.enabled).toBeUndefined(); + expect(result?.accounts[0]).not.toHaveProperty("enabled"); + }); + it("preserves activeIndexByFamily when valid", () => { const data = { version: 3, @@ -1470,6 +1584,32 @@ describe("storage", () => { expect(loaded.accounts[0]?.accountIdSource).toBe("id_token"); }); + it("strips stale disabledReason from enabled flagged records during load", async () => { + const flaggedPath = join(dirname(testStoragePath), "openai-codex-flagged-accounts.json"); + await fs.writeFile( + flaggedPath, + JSON.stringify({ + version: 1, + accounts: [ + { + refreshToken: "flagged-enabled", + flaggedAt: 123, + addedAt: 123, + lastUsed: 123, + enabled: true, + disabledReason: "auth-failure", + }, + ], + }), + ); + + const loaded = await loadFlaggedAccounts(); + expect(loaded.accounts).toHaveLength(1); + expect(loaded.accounts[0]?.enabled).toBe(true); + expect(loaded.accounts[0]?.disabledReason).toBeUndefined(); + expect(loaded.accounts[0]).not.toHaveProperty("disabledReason"); + }); + it("retries flagged storage rename on EBUSY and succeeds", async () => { const originalRename = fs.rename.bind(fs); let attemptCount = 0;