refactor: extract runtime account pool persistence

Author: ndycode

index.ts

@@ -122,7 +122,6 @@ import {
 	getCodexInstructions,
 	getModelFamily,
 	MODEL_FAMILIES,
-	type ModelFamily,
 	prewarmCodexInstructions,
 } from "./lib/prompts/codex.js";
 import { prewarmHostCodexPrompt } from "./lib/prompts/host-codex-prompt.js";
@@ -161,6 +160,7 @@ import { applyFastSessionDefaults } from "./lib/request/request-transformer.js";
 import { isEmptyResponse } from "./lib/request/response-handler.js";
 import { withStreamingFailover } from "./lib/request/stream-failover.js";
 import { addJitter } from "./lib/rotation.js";
+import { persistAccountPool } from "./lib/runtime/account-pool.js";
 import {
 	resolveAccountSelection,
 	type TokenSuccessWithAccount,
@@ -292,182 +292,18 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 			logWarn: (message) => logWarn(`[${PLUGIN_NAME}] ${message}`),
 		});
 
-	const persistAccountPool = async (
+	const persistAccounts = async (
 		results: TokenSuccessWithAccount[],
 		replaceAll: boolean = false,
-	): Promise<void> => {
-		if (results.length === 0) return;
-		await withAccountStorageTransaction(async (loadedStorage, persist) => {
-			const now = Date.now();
-			const stored = replaceAll ? null : loadedStorage;
-			const accounts = stored?.accounts ? [...stored.accounts] : [];
-
-			for (const result of results) {
-				const accountId =
-					result.accountIdOverride ?? extractAccountId(result.access);
-				const accountIdSource = accountId
-					? (result.accountIdSource ??
-						(result.accountIdOverride ? "manual" : "token"))
-					: undefined;
-				const accountLabel = result.accountLabel;
-				const accountEmail = sanitizeEmail(
-					extractAccountEmail(result.access, result.idToken),
-				);
-				const existingIndex = findMatchingAccountIndex(
-					accounts,
-					{
-						accountId,
-						email: accountEmail,
-						refreshToken: result.refresh,
-					},
-					{
-						allowUniqueAccountIdFallbackWithoutEmail: true,
-					},
-				);
-
-				if (existingIndex === undefined) {
-					const initialWorkspaceIndex =
-						result.workspaces && result.workspaces.length > 0
-							? (() => {
-									if (accountId) {
-										const matchingWorkspaceIndex = result.workspaces.findIndex(
-											(workspace) => workspace.id === accountId,
-										);
-										if (matchingWorkspaceIndex >= 0) {
-											return matchingWorkspaceIndex;
-										}
-									}
-									const firstEnabledWorkspaceIndex =
-										result.workspaces.findIndex(
-											(workspace) => workspace.enabled !== false,
-										);
-									return firstEnabledWorkspaceIndex >= 0
-										? firstEnabledWorkspaceIndex
-										: 0;
-								})()
-							: undefined;
-					accounts.push({
-						accountId,
-						accountIdSource,
-						accountLabel,
-						email: accountEmail,
-						refreshToken: result.refresh,
-						accessToken: result.access,
-						expiresAt: result.expires,
-						addedAt: now,
-						lastUsed: now,
-						workspaces: result.workspaces,
-						currentWorkspaceIndex: initialWorkspaceIndex,
-					});
-					continue;
-				}
-
-				const existing = accounts[existingIndex];
-				if (!existing) continue;
-
-				const nextEmail = accountEmail ?? sanitizeEmail(existing.email);
-				const nextAccountId = accountId ?? existing.accountId;
-				const nextAccountIdSource = accountId
-					? (accountIdSource ?? existing.accountIdSource)
-					: existing.accountIdSource;
-				const nextAccountLabel = accountLabel ?? existing.accountLabel;
-				// Preserve tracked workspace state when auth refreshes do not return workspace metadata.
-				const mergedWorkspaces = result.workspaces
-					? result.workspaces.map((newWs) => {
-							const existingWs = existing.workspaces?.find(
-								(w) => w.id === newWs.id,
-							);
-							return existingWs
-								? {
-										...newWs,
-										enabled: existingWs.enabled,
-										disabledAt: existingWs.disabledAt,
-									}
-								: newWs;
-						})
-					: existing.workspaces;
-				const currentWorkspaceId =
-					existing.workspaces?.[
-						typeof existing.currentWorkspaceIndex === "number"
-							? existing.currentWorkspaceIndex
-							: 0
-					]?.id;
-				const nextCurrentWorkspaceIndex =
-					mergedWorkspaces && mergedWorkspaces.length > 0
-						? (() => {
-								if (currentWorkspaceId) {
-									const matchingWorkspaceIndex = mergedWorkspaces.findIndex(
-										(workspace) => workspace.id === currentWorkspaceId,
-									);
-									if (matchingWorkspaceIndex >= 0) {
-										return matchingWorkspaceIndex;
-									}
-								}
-								const defaultWorkspaceIndex = mergedWorkspaces.findIndex(
-									(workspace) => workspace.isDefault === true,
-								);
-								if (defaultWorkspaceIndex >= 0) {
-									return defaultWorkspaceIndex;
-								}
-								const firstEnabledWorkspaceIndex = mergedWorkspaces.findIndex(
-									(workspace) => workspace.enabled !== false,
-								);
-								return firstEnabledWorkspaceIndex >= 0
-									? firstEnabledWorkspaceIndex
-									: 0;
-							})()
-						: existing.currentWorkspaceIndex;
-				accounts[existingIndex] = {
-					...existing,
-					accountId: nextAccountId,
-					accountIdSource: nextAccountIdSource,
-					accountLabel: nextAccountLabel,
-					email: nextEmail,
-					refreshToken: result.refresh,
-					accessToken: result.access,
-					expiresAt: result.expires,
-					lastUsed: now,
-					workspaces: mergedWorkspaces,
-					currentWorkspaceIndex: nextCurrentWorkspaceIndex,
-				};
-			}
-
-			if (accounts.length === 0) return;
-
-			const activeIndex = replaceAll
-				? 0
-				: typeof stored?.activeIndex === "number" &&
-						Number.isFinite(stored.activeIndex)
-					? stored.activeIndex
-					: 0;
-
-			const clampedActiveIndex = Math.max(
-				0,
-				Math.min(activeIndex, accounts.length - 1),
-			);
-			const activeIndexByFamily: Partial<Record<ModelFamily, number>> = {};
-			for (const family of MODEL_FAMILIES) {
-				const storedFamilyIndex = stored?.activeIndexByFamily?.[family];
-				const rawFamilyIndex = replaceAll
-					? 0
-					: typeof storedFamilyIndex === "number" &&
-							Number.isFinite(storedFamilyIndex)
-						? storedFamilyIndex
-						: clampedActiveIndex;
-				activeIndexByFamily[family] = Math.max(
-					0,
-					Math.min(Math.floor(rawFamilyIndex), accounts.length - 1),
-				);
-			}
-
-			await persist({
-				version: 3,
-				accounts,
-				activeIndex: clampedActiveIndex,
-				activeIndexByFamily,
-			});
+	): Promise<void> =>
+		persistAccountPool(results, replaceAll, {
+			withAccountStorageTransaction,
+			extractAccountId,
+			extractAccountEmail,
+			sanitizeEmail,
+			findMatchingAccountIndex,
+			MODEL_FAMILIES,
 		});
-	};
 
 	const showToast = async (
 		message: string,
@@ -3485,7 +3321,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 							}
 
 							if (restored.length > 0) {
-								await persistAccountPool(restored, false);
+								await persistAccounts(restored, false);
 								invalidateAccountManagerCache();
 							}
 
@@ -3710,7 +3546,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 								logInfo,
 								onSuccess: async (tokens: TokenSuccessWithAccount) => {
 									try {
-										await persistAccountPool([tokens], startFresh);
+										await persistAccounts([tokens], startFresh);
 										invalidateAccountManagerCache();
 									} catch (err) {
 										const storagePath = getStoragePath();
@@ -3814,10 +3650,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 
 							try {
 								const isFirstAccount = accounts.length === 1;
-								await persistAccountPool(
-									[resolved],
-									isFirstAccount && startFresh,
-								);
+								await persistAccounts([resolved], isFirstAccount && startFresh);
 								invalidateAccountManagerCache();
 							} catch (err) {
 								const storagePath = getStoragePath();
@@ -3905,7 +3738,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 							logInfo,
 							onSuccess: async (tokens: TokenSuccessWithAccount) => {
 								try {
-									await persistAccountPool([tokens], false);
+									await persistAccounts([tokens], false);
 								} catch (err) {
 									const storagePath = getStoragePath();
 									const errorCode =