refactor: extract account manager cache entry wrapper

index.ts

@@ -197,6 +197,10 @@ import {
 	invalidateAccountManagerCacheState,
 	reloadAccountManagerFromDiskState,
 } from "./lib/runtime/account-manager-cache.js";
+import {
+	invalidateAccountManagerCacheEntry,
+	reloadAccountManagerFromDiskEntry,
+} from "./lib/runtime/account-manager-cache-entry.js";
 import {
 	type TokenSuccessWithAccount as AccountPoolTokenSuccessWithAccount,
 	persistAccountPoolResults,
@@ -524,28 +528,36 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 	};
 
 	const invalidateAccountManagerCache = (): void => {
-		const next = invalidateAccountManagerCacheState();
-		cachedAccountManager = next.cachedAccountManager;
-		accountManagerPromise = next.accountManagerPromise;
+		invalidateAccountManagerCacheEntry<AccountManager>({
+			invalidateAccountManagerCacheState,
+			setCachedAccountManager: (manager) => {
+				cachedAccountManager = manager;
+			},
+			setAccountManagerPromise: (promise) => {
+				accountManagerPromise = promise;
+			},
+		});
 	};
 
 	const reloadAccountManagerFromDisk = async (
 		authFallback?: OAuthAuthDetails,
-	): Promise<AccountManager> => {
-		accountReloadInFlight = reloadAccountManagerFromDiskState({
+	): Promise<AccountManager> =>
+		reloadAccountManagerFromDiskEntry<AccountManager>({
+			authFallback,
 			currentReloadInFlight: accountReloadInFlight,
+			reloadAccountManagerFromDiskState,
 			loadFromDisk: (fallback) => AccountManager.loadFromDisk(fallback),
-			authFallback,
 			onLoaded: (reloaded) => {
 				cachedAccountManager = reloaded;
 				accountManagerPromise = Promise.resolve(reloaded);
 			},
 			onSettled: () => {
 				accountReloadInFlight = null;
 			},
+			setReloadInFlight: (promise) => {
+				accountReloadInFlight = promise;
+			},
 		});
-		return accountReloadInFlight;
-	};
 
 	const applyAccountStorageScope = (
 		pluginConfig: ReturnType<typeof loadPluginConfig>,

lib/runtime/account-manager-cache-entry.ts

@@ -0,0 +1,40 @@
+import type { OAuthAuthDetails } from "../types.js";
+
+export function invalidateAccountManagerCacheEntry<TManager>(params: {
+	invalidateAccountManagerCacheState: () => {
+		cachedAccountManager: null;
+		accountManagerPromise: null;
+	};
+	setCachedAccountManager: (manager: TManager | null) => void;
+	setAccountManagerPromise: (promise: Promise<TManager> | null) => void;
+}): void {
+	const next = params.invalidateAccountManagerCacheState();
+	params.setCachedAccountManager(next.cachedAccountManager);
+	params.setAccountManagerPromise(next.accountManagerPromise);
+}
+
+export async function reloadAccountManagerFromDiskEntry<TManager>(params: {
+	authFallback?: OAuthAuthDetails;
+	currentReloadInFlight: Promise<TManager> | null;
+	reloadAccountManagerFromDiskState: (args: {
+		currentReloadInFlight: Promise<TManager> | null;
+		loadFromDisk: (authFallback?: OAuthAuthDetails) => Promise<TManager>;
+		authFallback?: OAuthAuthDetails;
+		onLoaded: (manager: TManager) => void;
+		onSettled: () => void;
+	}) => Promise<TManager>;
+	loadFromDisk: (authFallback?: OAuthAuthDetails) => Promise<TManager>;
+	onLoaded: (manager: TManager) => void;
+	onSettled: () => void;
+	setReloadInFlight: (promise: Promise<TManager>) => void;
+}): Promise<TManager> {
+	const inFlight = params.reloadAccountManagerFromDiskState({
+		currentReloadInFlight: params.currentReloadInFlight,
+		loadFromDisk: params.loadFromDisk,
+		authFallback: params.authFallback,
+		onLoaded: params.onLoaded,
+		onSettled: params.onSettled,
+	});
+	params.setReloadInFlight(inFlight);
+	return inFlight;
+}

test/account-manager-cache-entry.test.ts

@@ -0,0 +1,42 @@
+import { describe, expect, it, vi } from "vitest";
+import {
+	invalidateAccountManagerCacheEntry,
+	reloadAccountManagerFromDiskEntry,
+} from "../lib/runtime/account-manager-cache-entry.js";
+
+describe("account manager cache entry", () => {
+	it("delegates cache invalidation state into the setter callbacks", () => {
+		const setCachedAccountManager = vi.fn();
+		const setAccountManagerPromise = vi.fn();
+
+		invalidateAccountManagerCacheEntry({
+			invalidateAccountManagerCacheState: () => ({
+				cachedAccountManager: null,
+				accountManagerPromise: null,
+			}),
+			setCachedAccountManager,
+			setAccountManagerPromise,
+		});
+
+		expect(setCachedAccountManager).toHaveBeenCalledWith(null);
+		expect(setAccountManagerPromise).toHaveBeenCalledWith(null);
+	});
+
+	it("delegates reload state into the injected runtime callbacks", async () => {
+		const reloadState = vi.fn(async () => ({ id: 1 }));
+		const setReloadInFlight = vi.fn();
+
+		const result = await reloadAccountManagerFromDiskEntry({
+			currentReloadInFlight: null,
+			reloadAccountManagerFromDiskState: reloadState,
+			loadFromDisk: vi.fn(async () => ({ id: 1 })),
+			onLoaded: vi.fn(),
+			onSettled: vi.fn(),
+			setReloadInFlight,
+		});
+
+		expect(reloadState).toHaveBeenCalled();
+		expect(setReloadInFlight).toHaveBeenCalledWith(expect.any(Promise));
+		expect(result).toEqual({ id: 1 });
+	});
+});