refactor: extract session affinity wrapper

Author: ndycode

index.ts

@@ -220,6 +220,7 @@ import {
 	ensureRefreshGuardianState,
 	ensureSessionAffinityState,
 } from "./lib/runtime/runtime-services.js";
+import { ensureSessionAffinityEntry } from "./lib/runtime/session-affinity-entry.js";
 import { applyAccountStorageScopeFromConfig } from "./lib/runtime/storage-scope.js";
 import {
 	applyUiRuntimeFromConfig,
@@ -610,14 +611,16 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
 	const ensureSessionAffinity = (
 		pluginConfig: ReturnType<typeof loadPluginConfig>,
 	): void => {
-		const next = ensureSessionAffinityState({
-			enabled: getSessionAffinity(pluginConfig),
-			ttlMs: getSessionAffinityTtlMs(pluginConfig),
-			maxEntries: getSessionAffinityMaxEntries(pluginConfig),
+		const next = ensureSessionAffinityEntry({
+			pluginConfig,
 			currentStore: sessionAffinityStore,
 			currentConfigKey: sessionAffinityConfigKey,
+			getSessionAffinity,
+			getSessionAffinityTtlMs,
+			getSessionAffinityMaxEntries,
 			createStore: ({ ttlMs, maxEntries }) =>
 				new SessionAffinityStore({ ttlMs, maxEntries }),
+			ensureSessionAffinityState,
 		});
 		sessionAffinityStore = next.sessionAffinityStore;
 		sessionAffinityConfigKey = next.sessionAffinityConfigKey;

lib/runtime/session-affinity-entry.ts

@@ -0,0 +1,38 @@
+export function ensureSessionAffinityEntry<TStore>(params: {
+	pluginConfig: ReturnType<typeof import("../config.js").loadPluginConfig>;
+	currentStore: TStore | null;
+	currentConfigKey: string | null;
+	getSessionAffinity: (
+		config: ReturnType<typeof import("../config.js").loadPluginConfig>,
+	) => boolean;
+	getSessionAffinityTtlMs: (
+		config: ReturnType<typeof import("../config.js").loadPluginConfig>,
+	) => number;
+	getSessionAffinityMaxEntries: (
+		config: ReturnType<typeof import("../config.js").loadPluginConfig>,
+	) => number;
+	createStore: (options: { ttlMs: number; maxEntries: number }) => TStore;
+	ensureSessionAffinityState: (args: {
+		enabled: boolean;
+		ttlMs: number;
+		maxEntries: number;
+		currentStore: TStore | null;
+		currentConfigKey: string | null;
+		createStore: (options: { ttlMs: number; maxEntries: number }) => TStore;
+	}) => {
+		sessionAffinityStore: TStore | null;
+		sessionAffinityConfigKey: string | null;
+	};
+}): {
+	sessionAffinityStore: TStore | null;
+	sessionAffinityConfigKey: string | null;
+} {
+	return params.ensureSessionAffinityState({
+		enabled: params.getSessionAffinity(params.pluginConfig),
+		ttlMs: params.getSessionAffinityTtlMs(params.pluginConfig),
+		maxEntries: params.getSessionAffinityMaxEntries(params.pluginConfig),
+		currentStore: params.currentStore,
+		currentConfigKey: params.currentConfigKey,
+		createStore: params.createStore,
+	});
+}

test/session-affinity-entry.test.ts

@@ -0,0 +1,35 @@
+import { describe, expect, it, vi } from "vitest";
+import { ensureSessionAffinityEntry } from "../lib/runtime/session-affinity-entry.js";
+
+describe("session affinity entry", () => {
+	it("delegates config-derived arguments into the session affinity state helper", () => {
+		const ensureSessionAffinityState = vi.fn(() => ({
+			sessionAffinityStore: { id: 1 },
+			sessionAffinityConfigKey: "1000:10",
+		}));
+
+		const result = ensureSessionAffinityEntry({
+			pluginConfig: {} as never,
+			currentStore: null,
+			currentConfigKey: null,
+			getSessionAffinity: () => true,
+			getSessionAffinityTtlMs: () => 1000,
+			getSessionAffinityMaxEntries: () => 10,
+			createStore: vi.fn(() => ({ id: 1 })),
+			ensureSessionAffinityState,
+		});
+
+		expect(ensureSessionAffinityState).toHaveBeenCalledWith({
+			enabled: true,
+			ttlMs: 1000,
+			maxEntries: 10,
+			currentStore: null,
+			currentConfigKey: null,
+			createStore: expect.any(Function),
+		});
+		expect(result).toEqual({
+			sessionAffinityStore: { id: 1 },
+			sessionAffinityConfigKey: "1000:10",
+		});
+	});
+});