diff --git a/index.ts b/index.ts index 45882d05..6e2f8ef6 100644 --- a/index.ts +++ b/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,18 +528,25 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { }; const invalidateAccountManagerCache = (): void => { - const next = invalidateAccountManagerCacheState(); - cachedAccountManager = next.cachedAccountManager; - accountManagerPromise = next.accountManagerPromise; + invalidateAccountManagerCacheEntry({ + invalidateAccountManagerCacheState, + setCachedAccountManager: (manager) => { + cachedAccountManager = manager; + }, + setAccountManagerPromise: (promise) => { + accountManagerPromise = promise; + }, + }); }; const reloadAccountManagerFromDisk = async ( authFallback?: OAuthAuthDetails, - ): Promise => { - accountReloadInFlight = reloadAccountManagerFromDiskState({ + ): Promise => + reloadAccountManagerFromDiskEntry({ + authFallback, currentReloadInFlight: accountReloadInFlight, + reloadAccountManagerFromDiskState, loadFromDisk: (fallback) => AccountManager.loadFromDisk(fallback), - authFallback, onLoaded: (reloaded) => { cachedAccountManager = reloaded; accountManagerPromise = Promise.resolve(reloaded); @@ -543,9 +554,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { onSettled: () => { accountReloadInFlight = null; }, + setReloadInFlight: (promise) => { + accountReloadInFlight = promise; + }, }); - return accountReloadInFlight; - }; const applyAccountStorageScope = ( pluginConfig: ReturnType, diff --git a/lib/runtime/account-manager-cache-entry.ts b/lib/runtime/account-manager-cache-entry.ts new file mode 100644 index 00000000..ebb4ffa4 --- /dev/null +++ b/lib/runtime/account-manager-cache-entry.ts @@ -0,0 +1,40 @@ +import type { OAuthAuthDetails } from "../types.js"; + +export function invalidateAccountManagerCacheEntry(params: { + invalidateAccountManagerCacheState: () => { + cachedAccountManager: null; + accountManagerPromise: null; + }; + setCachedAccountManager: (manager: TManager | null) => void; + setAccountManagerPromise: (promise: Promise | null) => void; +}): void { + const next = params.invalidateAccountManagerCacheState(); + params.setCachedAccountManager(next.cachedAccountManager); + params.setAccountManagerPromise(next.accountManagerPromise); +} + +export async function reloadAccountManagerFromDiskEntry(params: { + authFallback?: OAuthAuthDetails; + currentReloadInFlight: Promise | null; + reloadAccountManagerFromDiskState: (args: { + currentReloadInFlight: Promise | null; + loadFromDisk: (authFallback?: OAuthAuthDetails) => Promise; + authFallback?: OAuthAuthDetails; + onLoaded: (manager: TManager) => void; + onSettled: () => void; + }) => Promise; + loadFromDisk: (authFallback?: OAuthAuthDetails) => Promise; + onLoaded: (manager: TManager) => void; + onSettled: () => void; + setReloadInFlight: (promise: Promise) => void; +}): Promise { + const inFlight = params.reloadAccountManagerFromDiskState({ + currentReloadInFlight: params.currentReloadInFlight, + loadFromDisk: params.loadFromDisk, + authFallback: params.authFallback, + onLoaded: params.onLoaded, + onSettled: params.onSettled, + }); + params.setReloadInFlight(inFlight); + return inFlight; +} diff --git a/test/account-manager-cache-entry.test.ts b/test/account-manager-cache-entry.test.ts new file mode 100644 index 00000000..0cdf821c --- /dev/null +++ b/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 }); + }); +});