Skip to content

Commit 8c4646c

Browse files
committed
refactor: extract runtime cache and event helpers
1 parent 16acf57 commit 8c4646c

8 files changed

Lines changed: 410 additions & 91 deletions

index.ts

Lines changed: 55 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,10 @@ import {
193193
clampActiveIndices,
194194
isFlaggableFailure,
195195
} from "./lib/runtime/account-check-helpers.js";
196+
import {
197+
invalidateAccountManagerCacheState,
198+
reloadAccountManagerFromDiskState,
199+
} from "./lib/runtime/account-manager-cache.js";
196200
import {
197201
type TokenSuccessWithAccount as AccountPoolTokenSuccessWithAccount,
198202
persistAccountPoolResults,
@@ -204,7 +208,12 @@ import {
204208
resolveActiveIndex,
205209
} from "./lib/runtime/account-status.js";
206210
import { runBrowserOAuthFlow } from "./lib/runtime/browser-oauth-flow.js";
211+
import { handleRuntimeEvent } from "./lib/runtime/event-handler.js";
207212
import { buildManualOAuthFlow } from "./lib/runtime/manual-oauth-flow.js";
213+
import {
214+
applyPreemptiveQuotaSettingsFromConfig,
215+
resolveUiRuntimeFromConfig,
216+
} from "./lib/runtime/quota-settings.js";
208217
import {
209218
ensureLiveAccountSyncState,
210219
ensureRefreshGuardianState,
@@ -500,31 +509,33 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
500509
};
501510

502511
const resolveUiRuntime = (): UiRuntimeOptions => {
503-
return applyUiRuntimeFromConfig(loadPluginConfig(), setUiRuntimeOptions);
512+
return resolveUiRuntimeFromConfig(loadPluginConfig, (pluginConfig) =>
513+
applyUiRuntimeFromConfig(pluginConfig, setUiRuntimeOptions),
514+
);
504515
};
505516

506517
const invalidateAccountManagerCache = (): void => {
507-
cachedAccountManager = null;
508-
accountManagerPromise = null;
518+
const next = invalidateAccountManagerCacheState();
519+
cachedAccountManager = next.cachedAccountManager;
520+
accountManagerPromise = next.accountManagerPromise;
509521
};
510522

511523
const reloadAccountManagerFromDisk = async (
512524
authFallback?: OAuthAuthDetails,
513525
): Promise<AccountManager> => {
514-
if (accountReloadInFlight) {
515-
return accountReloadInFlight;
516-
}
517-
accountReloadInFlight = (async () => {
518-
const reloaded = await AccountManager.loadFromDisk(authFallback);
519-
cachedAccountManager = reloaded;
520-
accountManagerPromise = Promise.resolve(reloaded);
521-
return reloaded;
522-
})();
523-
try {
524-
return await accountReloadInFlight;
525-
} finally {
526-
accountReloadInFlight = null;
527-
}
526+
accountReloadInFlight = reloadAccountManagerFromDiskState({
527+
currentReloadInFlight: accountReloadInFlight,
528+
loadFromDisk: (fallback) => AccountManager.loadFromDisk(fallback),
529+
authFallback,
530+
onLoaded: (reloaded) => {
531+
cachedAccountManager = reloaded;
532+
accountManagerPromise = Promise.resolve(reloaded);
533+
},
534+
onSettled: () => {
535+
accountReloadInFlight = null;
536+
},
537+
});
538+
return accountReloadInFlight;
528539
};
529540

530541
const applyAccountStorageScope = (
@@ -611,85 +622,38 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
611622

612623
const applyPreemptiveQuotaSettings = (
613624
pluginConfig: ReturnType<typeof loadPluginConfig>,
614-
): void => {
615-
preemptiveQuotaScheduler.configure({
616-
enabled: getPreemptiveQuotaEnabled(pluginConfig),
617-
remainingPercentThresholdPrimary:
618-
getPreemptiveQuotaRemainingPercent5h(pluginConfig),
619-
remainingPercentThresholdSecondary:
620-
getPreemptiveQuotaRemainingPercent7d(pluginConfig),
621-
maxDeferralMs: getPreemptiveQuotaMaxDeferralMs(pluginConfig),
625+
): void =>
626+
applyPreemptiveQuotaSettingsFromConfig(pluginConfig, {
627+
configure: (options) => preemptiveQuotaScheduler.configure(options),
628+
getPreemptiveQuotaEnabled,
629+
getPreemptiveQuotaRemainingPercent5h,
630+
getPreemptiveQuotaRemainingPercent7d,
631+
getPreemptiveQuotaMaxDeferralMs,
622632
});
623-
};
624633

625634
// Event handler for session recovery and account selection
626635
const eventHandler = async (input: {
627636
event: { type: string; properties?: unknown };
628-
}) => {
629-
try {
630-
const { event } = input;
631-
// Handle TUI account selection events
632-
// Accepts generic selection events with an index property
633-
if (
634-
event.type === "account.select" ||
635-
event.type === "openai.account.select"
636-
) {
637-
const props = event.properties as {
638-
index?: number;
639-
accountIndex?: number;
640-
provider?: string;
641-
};
642-
// Filter by provider if specified
643-
if (
644-
props.provider &&
645-
props.provider !== "openai" &&
646-
props.provider !== PROVIDER_ID
647-
) {
648-
return;
649-
}
650-
651-
const index = props.index ?? props.accountIndex;
652-
if (typeof index === "number") {
653-
const storage = await loadAccounts();
654-
if (!storage || index < 0 || index >= storage.accounts.length) {
655-
return;
656-
}
657-
658-
const now = Date.now();
659-
const account = storage.accounts[index];
660-
if (account) {
661-
account.lastUsed = now;
662-
account.lastSwitchReason = "rotation";
663-
}
664-
storage.activeIndex = index;
665-
storage.activeIndexByFamily = storage.activeIndexByFamily ?? {};
666-
for (const family of MODEL_FAMILIES) {
667-
storage.activeIndexByFamily[family] = index;
668-
}
669-
670-
await saveAccounts(storage);
671-
if (cachedAccountManager) {
672-
await cachedAccountManager.syncCodexCliActiveSelectionForIndex(
673-
index,
674-
);
675-
}
676-
lastCodexCliActiveSyncIndex = index;
677-
678-
// Reload manager from disk so we don't overwrite newer rotated
679-
// refresh tokens with stale in-memory state.
680-
if (cachedAccountManager) {
681-
await reloadAccountManagerFromDisk();
682-
}
683-
684-
await showToast(`Switched to account ${index + 1}`, "info");
685-
}
686-
}
687-
} catch (error) {
688-
logDebug(
689-
`[${PLUGIN_NAME}] Event handler error: ${error instanceof Error ? error.message : String(error)}`,
690-
);
691-
}
692-
};
637+
}) =>
638+
handleRuntimeEvent({
639+
input,
640+
providerId: PROVIDER_ID,
641+
modelFamilies: MODEL_FAMILIES,
642+
loadAccounts,
643+
saveAccounts,
644+
hasCachedAccountManager: () => !!cachedAccountManager,
645+
syncCodexCliActiveSelectionForIndex: async (index) => {
646+
if (!cachedAccountManager) return;
647+
await cachedAccountManager.syncCodexCliActiveSelectionForIndex(index);
648+
},
649+
setLastCodexCliActiveSyncIndex: (index) => {
650+
lastCodexCliActiveSyncIndex = index;
651+
},
652+
reloadAccountManagerFromDisk: () => reloadAccountManagerFromDisk(),
653+
showToast: (message, variant) => showToast(message, variant),
654+
logDebug,
655+
pluginName: PLUGIN_NAME,
656+
});
693657

694658
// Initialize runtime UI settings once on plugin load; auth/tools refresh this dynamically.
695659
resolveUiRuntime();
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import type { OAuthAuthDetails } from "../types.js";
2+
3+
export function invalidateAccountManagerCacheState(): {
4+
cachedAccountManager: null;
5+
accountManagerPromise: null;
6+
} {
7+
return {
8+
cachedAccountManager: null,
9+
accountManagerPromise: null,
10+
};
11+
}
12+
13+
export async function reloadAccountManagerFromDiskState<TManager>(params: {
14+
currentReloadInFlight: Promise<TManager> | null;
15+
loadFromDisk: (authFallback?: OAuthAuthDetails) => Promise<TManager>;
16+
authFallback?: OAuthAuthDetails;
17+
onLoaded: (manager: TManager) => void;
18+
onSettled: () => void;
19+
}): Promise<TManager> {
20+
if (params.currentReloadInFlight) {
21+
return params.currentReloadInFlight;
22+
}
23+
24+
const inFlight = (async () => {
25+
const reloaded = await params.loadFromDisk(params.authFallback);
26+
params.onLoaded(reloaded);
27+
return reloaded;
28+
})();
29+
30+
try {
31+
return await inFlight;
32+
} finally {
33+
params.onSettled();
34+
}
35+
}

lib/runtime/event-handler.ts

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import type { ModelFamily } from "../prompts/codex.js";
2+
import type { AccountStorageV3 } from "../storage.js";
3+
4+
export async function handleRuntimeEvent(params: {
5+
input: { event: { type: string; properties?: unknown } };
6+
providerId: string;
7+
modelFamilies: readonly ModelFamily[];
8+
loadAccounts: () => Promise<AccountStorageV3 | null>;
9+
saveAccounts: (storage: AccountStorageV3) => Promise<void>;
10+
hasCachedAccountManager: () => boolean;
11+
syncCodexCliActiveSelectionForIndex: (index: number) => Promise<void>;
12+
setLastCodexCliActiveSyncIndex: (index: number) => void;
13+
reloadAccountManagerFromDisk: () => Promise<unknown>;
14+
showToast: (message: string, variant: "info") => Promise<void>;
15+
logDebug: (message: string) => void;
16+
pluginName: string;
17+
}): Promise<void> {
18+
try {
19+
const { event } = params.input;
20+
if (
21+
event.type !== "account.select" &&
22+
event.type !== "openai.account.select"
23+
) {
24+
return;
25+
}
26+
27+
const props = event.properties as {
28+
index?: number;
29+
accountIndex?: number;
30+
provider?: string;
31+
};
32+
if (
33+
props.provider &&
34+
props.provider !== "openai" &&
35+
props.provider !== params.providerId
36+
) {
37+
return;
38+
}
39+
40+
const index = props.index ?? props.accountIndex;
41+
if (typeof index !== "number") return;
42+
43+
const storage = await params.loadAccounts();
44+
if (!storage || index < 0 || index >= storage.accounts.length) {
45+
return;
46+
}
47+
48+
const now = Date.now();
49+
const account = storage.accounts[index];
50+
if (account) {
51+
account.lastUsed = now;
52+
account.lastSwitchReason = "rotation";
53+
}
54+
storage.activeIndex = index;
55+
storage.activeIndexByFamily = storage.activeIndexByFamily ?? {};
56+
for (const family of params.modelFamilies) {
57+
storage.activeIndexByFamily[family] = index;
58+
}
59+
60+
await params.saveAccounts(storage);
61+
if (params.hasCachedAccountManager()) {
62+
await params.syncCodexCliActiveSelectionForIndex(index);
63+
}
64+
params.setLastCodexCliActiveSyncIndex(index);
65+
66+
if (params.hasCachedAccountManager()) {
67+
await params.reloadAccountManagerFromDisk();
68+
}
69+
70+
await params.showToast(`Switched to account ${index + 1}`, "info");
71+
} catch (error) {
72+
params.logDebug(
73+
`[${params.pluginName}] Event handler error: ${error instanceof Error ? error.message : String(error)}`,
74+
);
75+
}
76+
}

lib/runtime/quota-settings.ts

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
export function applyPreemptiveQuotaSettingsFromConfig(
2+
pluginConfig: ReturnType<typeof import("../config.js").loadPluginConfig>,
3+
deps: {
4+
configure: (options: {
5+
enabled: boolean;
6+
remainingPercentThresholdPrimary: number;
7+
remainingPercentThresholdSecondary: number;
8+
maxDeferralMs: number;
9+
}) => void;
10+
getPreemptiveQuotaEnabled: (
11+
config: ReturnType<typeof import("../config.js").loadPluginConfig>,
12+
) => boolean;
13+
getPreemptiveQuotaRemainingPercent5h: (
14+
config: ReturnType<typeof import("../config.js").loadPluginConfig>,
15+
) => number;
16+
getPreemptiveQuotaRemainingPercent7d: (
17+
config: ReturnType<typeof import("../config.js").loadPluginConfig>,
18+
) => number;
19+
getPreemptiveQuotaMaxDeferralMs: (
20+
config: ReturnType<typeof import("../config.js").loadPluginConfig>,
21+
) => number;
22+
},
23+
): void {
24+
deps.configure({
25+
enabled: deps.getPreemptiveQuotaEnabled(pluginConfig),
26+
remainingPercentThresholdPrimary:
27+
deps.getPreemptiveQuotaRemainingPercent5h(pluginConfig),
28+
remainingPercentThresholdSecondary:
29+
deps.getPreemptiveQuotaRemainingPercent7d(pluginConfig),
30+
maxDeferralMs: deps.getPreemptiveQuotaMaxDeferralMs(pluginConfig),
31+
});
32+
}
33+
34+
export function resolveUiRuntimeFromConfig(
35+
loadPluginConfig: () => ReturnType<
36+
typeof import("../config.js").loadPluginConfig
37+
>,
38+
applyUiRuntimeFromConfig: (
39+
pluginConfig: ReturnType<typeof import("../config.js").loadPluginConfig>,
40+
) => ReturnType<typeof import("../ui/runtime.js").getUiRuntimeOptions>,
41+
): ReturnType<typeof import("../ui/runtime.js").getUiRuntimeOptions> {
42+
return applyUiRuntimeFromConfig(loadPluginConfig());
43+
}

test/account-manager-cache.test.ts

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
import { describe, expect, it, vi } from "vitest";
2+
import {
3+
invalidateAccountManagerCacheState,
4+
reloadAccountManagerFromDiskState,
5+
} from "../lib/runtime/account-manager-cache.js";
6+
7+
describe("account manager cache helpers", () => {
8+
it("invalidates cache state", () => {
9+
expect(invalidateAccountManagerCacheState()).toEqual({
10+
cachedAccountManager: null,
11+
accountManagerPromise: null,
12+
});
13+
});
14+
15+
it("reuses in-flight reload and updates callbacks on success", async () => {
16+
const onLoaded = vi.fn();
17+
const onSettled = vi.fn();
18+
const inFlight = Promise.resolve({ id: 1 });
19+
await expect(
20+
reloadAccountManagerFromDiskState({
21+
currentReloadInFlight: inFlight,
22+
loadFromDisk: vi.fn(),
23+
onLoaded,
24+
onSettled,
25+
}),
26+
).resolves.toEqual({ id: 1 });
27+
28+
expect(onLoaded).not.toHaveBeenCalled();
29+
expect(onSettled).not.toHaveBeenCalled();
30+
31+
const loadFromDisk = vi.fn(async () => ({ id: 2 }));
32+
await expect(
33+
reloadAccountManagerFromDiskState({
34+
currentReloadInFlight: null,
35+
loadFromDisk,
36+
onLoaded,
37+
onSettled,
38+
}),
39+
).resolves.toEqual({ id: 2 });
40+
expect(onLoaded).toHaveBeenCalledWith({ id: 2 });
41+
expect(onSettled).toHaveBeenCalled();
42+
});
43+
});

0 commit comments

Comments
 (0)