From 81b68d828cfa50910750efdc6180b2c75f314c1f Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sat, 18 Apr 2026 01:53:03 +0900 Subject: [PATCH] refactor(model-fallback): move fallback state into factory closure and split hook.ts Move the pending fallback, toast, and session-chain maps behind a shared controller initialized from the hook factory. This preserves the existing singleton semantics because exported helpers and hook instances still resolve the same lazily initialized controller while hook.ts stays under the 200-line limit. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../fallback-state-controller.ts | 135 +++++++++++++++++ src/hooks/model-fallback/hook.ts | 137 ++++++------------ 2 files changed, 179 insertions(+), 93 deletions(-) create mode 100644 src/hooks/model-fallback/fallback-state-controller.ts diff --git a/src/hooks/model-fallback/fallback-state-controller.ts b/src/hooks/model-fallback/fallback-state-controller.ts new file mode 100644 index 000000000..b2e6831a0 --- /dev/null +++ b/src/hooks/model-fallback/fallback-state-controller.ts @@ -0,0 +1,135 @@ +import type { FallbackEntry } from "../../shared/model-requirements" +import { getAgentConfigKey } from "../../shared/agent-display-names" +import { AGENT_MODEL_REQUIREMENTS } from "../../shared/model-requirements" +import { log } from "../../shared/logger" +import { getNextReachableFallback } from "./next-fallback" + +type ModelFallbackStateLike = { + providerID: string + modelID: string + fallbackChain: FallbackEntry[] + attemptCount: number + pending: boolean +} + +export type ModelFallbackStateController = { + lastToastKey: Map + setSessionFallbackChain: (sessionID: string, fallbackChain: FallbackEntry[] | undefined) => void + clearSessionFallbackChain: (sessionID: string) => void + setPendingModelFallback: ( + sessionID: string, + agentName: string, + currentProviderID: string, + currentModelID: string, + ) => boolean + getNextFallback: (sessionID: string) => ReturnType + clearPendingModelFallback: (sessionID: string) => void + hasPendingModelFallback: (sessionID: string) => boolean + getFallbackState: (sessionID: string) => ModelFallbackStateLike | undefined + reset: () => void +} + +export function createModelFallbackStateController(input: { + pendingModelFallbacks: Map + lastToastKey: Map + sessionFallbackChains: Map +}): ModelFallbackStateController { + const { pendingModelFallbacks, lastToastKey, sessionFallbackChains } = input + + function setSessionFallbackChain(sessionID: string, fallbackChain: FallbackEntry[] | undefined): void { + if (!sessionID) return + sessionFallbackChains.set(sessionID, fallbackChain?.length ? fallbackChain : []) + } + + function clearSessionFallbackChain(sessionID: string): void { + sessionFallbackChains.delete(sessionID) + } + + function setPendingModelFallback( + sessionID: string, + agentName: string, + currentProviderID: string, + currentModelID: string, + ): boolean { + const agentKey = getAgentConfigKey(agentName) + const requirements = AGENT_MODEL_REQUIREMENTS[agentKey] + const fallbackChain = sessionFallbackChains.has(sessionID) + ? sessionFallbackChains.get(sessionID) + : requirements?.fallbackChain + + if (!fallbackChain?.length) { + log("[model-fallback] No fallback chain for agent: " + agentName + " (key: " + agentKey + ")") + return false + } + + const existing = pendingModelFallbacks.get(sessionID) + if (existing) { + if (existing.pending) { + log("[model-fallback] Pending fallback already armed for session: " + sessionID) + return false + } + existing.providerID = currentProviderID + existing.modelID = currentModelID + existing.pending = true + if (existing.attemptCount >= existing.fallbackChain.length) { + log("[model-fallback] Fallback chain exhausted for session: " + sessionID) + return false + } + log("[model-fallback] Re-armed pending fallback for session: " + sessionID) + return true + } + + pendingModelFallbacks.set(sessionID, { + providerID: currentProviderID, + modelID: currentModelID, + fallbackChain, + attemptCount: 0, + pending: true, + }) + log("[model-fallback] Set pending fallback for session: " + sessionID + ", agent: " + agentName) + return true + } + + function getNextFallback(sessionID: string): ReturnType { + const state = pendingModelFallbacks.get(sessionID) + if (!state?.pending) return null + + const fallback = getNextReachableFallback(sessionID, state) + if (fallback) return fallback + + log("[model-fallback] No more fallbacks for session: " + sessionID) + pendingModelFallbacks.delete(sessionID) + return null + } + + function clearPendingModelFallback(sessionID: string): void { + pendingModelFallbacks.delete(sessionID) + lastToastKey.delete(sessionID) + } + + function hasPendingModelFallback(sessionID: string): boolean { + return pendingModelFallbacks.get(sessionID)?.pending === true + } + + function getFallbackState(sessionID: string): ModelFallbackStateLike | undefined { + return pendingModelFallbacks.get(sessionID) + } + + function reset(): void { + pendingModelFallbacks.clear() + lastToastKey.clear() + sessionFallbackChains.clear() + } + + return { + lastToastKey, + setSessionFallbackChain, + clearSessionFallbackChain, + setPendingModelFallback, + getNextFallback, + clearPendingModelFallback, + hasPendingModelFallback, + getFallbackState, + reset, + } +} diff --git a/src/hooks/model-fallback/hook.ts b/src/hooks/model-fallback/hook.ts index b188bd48d..191a58e3a 100644 --- a/src/hooks/model-fallback/hook.ts +++ b/src/hooks/model-fallback/hook.ts @@ -1,13 +1,10 @@ import type { FallbackEntry } from "../../shared/model-requirements" -import { getAgentConfigKey } from "../../shared/agent-display-names" -import { AGENT_MODEL_REQUIREMENTS } from "../../shared/model-requirements" -import { readConnectedProvidersCache, readProviderModelsCache } from "../../shared/connected-providers-cache" -import { selectFallbackProvider } from "../../shared/model-error-classifier" -import { transformModelForProvider } from "../../shared/provider-model-id-transform" -import { log } from "../../shared/logger" import type { ChatMessageInput, ChatMessageHandlerOutput } from "../../plugin/chat-message" import { applyFallbackToChatMessage } from "./chat-message-fallback-handler" -import { getNextReachableFallback } from "./next-fallback" +import { + createModelFallbackStateController, + type ModelFallbackStateController, +} from "./fallback-state-controller" type FallbackToast = (input: { title: string @@ -31,30 +28,26 @@ export type ModelFallbackState = { pending: boolean } -/** - * Map of sessionID -> pending model fallback state - * When a model error occurs, we store the fallback info here. - * The next chat.message call will use this to switch to the fallback model. - */ -const pendingModelFallbacks = new Map() -const lastToastKey = new Map() -const sessionFallbackChains = new Map() +const modelFallbackControllerRef: { current?: ModelFallbackStateController } = {} + +function getOrCreateModelFallbackController(): ModelFallbackStateController { + if (!modelFallbackControllerRef.current) { + createModelFallbackHook() + } + + const controller = modelFallbackControllerRef.current + if (!controller) { + throw new Error("Model fallback controller should be initialized") + } + return controller +} export function setSessionFallbackChain(sessionID: string, fallbackChain: FallbackEntry[] | undefined): void { - if (!sessionID) return - if (!fallbackChain) { - sessionFallbackChains.set(sessionID, []) - return - } - if (fallbackChain.length === 0) { - sessionFallbackChains.set(sessionID, []) - return - } - sessionFallbackChains.set(sessionID, fallbackChain) + getOrCreateModelFallbackController().setSessionFallbackChain(sessionID, fallbackChain) } export function clearSessionFallbackChain(sessionID: string): void { - sessionFallbackChains.delete(sessionID) + getOrCreateModelFallbackController().clearSessionFallbackChain(sessionID) } /** @@ -67,51 +60,12 @@ export function setPendingModelFallback( currentProviderID: string, currentModelID: string, ): boolean { - const agentKey = getAgentConfigKey(agentName) - const requirements = AGENT_MODEL_REQUIREMENTS[agentKey] - const hasSessionFallback = sessionFallbackChains.has(sessionID) - const sessionFallback = sessionFallbackChains.get(sessionID) - const fallbackChain = hasSessionFallback - ? sessionFallback - : requirements?.fallbackChain - - if (!fallbackChain || fallbackChain.length === 0) { - log("[model-fallback] No fallback chain for agent: " + agentName + " (key: " + agentKey + ")") - return false - } - - const existing = pendingModelFallbacks.get(sessionID) - - if (existing) { - if (existing.pending) { - log("[model-fallback] Pending fallback already armed for session: " + sessionID) - return false - } - - // Preserve progression across repeated session.error retries in same session. - // We only mark the next turn as pending fallback application. - existing.providerID = currentProviderID - existing.modelID = currentModelID - existing.pending = true - if (existing.attemptCount >= existing.fallbackChain.length) { - log("[model-fallback] Fallback chain exhausted for session: " + sessionID) - return false - } - log("[model-fallback] Re-armed pending fallback for session: " + sessionID) - return true - } - - const state: ModelFallbackState = { - providerID: currentProviderID, - modelID: currentModelID, - fallbackChain, - attemptCount: 0, - pending: true, - } - - pendingModelFallbacks.set(sessionID, state) - log("[model-fallback] Set pending fallback for session: " + sessionID + ", agent: " + agentName) - return true + return getOrCreateModelFallbackController().setPendingModelFallback( + sessionID, + agentName, + currentProviderID, + currentModelID, + ) } /** @@ -121,19 +75,7 @@ export function setPendingModelFallback( export function getNextFallback( sessionID: string, ): { providerID: string; modelID: string; variant?: string } | null { - const state = pendingModelFallbacks.get(sessionID) - if (!state) return null - - if (!state.pending) return null - - const fallback = getNextReachableFallback(sessionID, state) - if (fallback) { - return fallback - } - - log("[model-fallback] No more fallbacks for session: " + sessionID) - pendingModelFallbacks.delete(sessionID) - return null + return getOrCreateModelFallbackController().getNextFallback(sessionID) } /** @@ -141,29 +83,40 @@ export function getNextFallback( * Called after fallback is successfully applied. */ export function clearPendingModelFallback(sessionID: string): void { - pendingModelFallbacks.delete(sessionID) - lastToastKey.delete(sessionID) + getOrCreateModelFallbackController().clearPendingModelFallback(sessionID) } /** * Checks if there's a pending fallback for a session. */ export function hasPendingModelFallback(sessionID: string): boolean { - const state = pendingModelFallbacks.get(sessionID) - return state?.pending === true + return getOrCreateModelFallbackController().hasPendingModelFallback(sessionID) } /** * Gets the current fallback state for a session (for debugging). */ export function getFallbackState(sessionID: string): ModelFallbackState | undefined { - return pendingModelFallbacks.get(sessionID) + return getOrCreateModelFallbackController().getFallbackState(sessionID) } /** * Creates a chat.message hook that applies model fallbacks when pending. */ export function createModelFallbackHook(args?: { toast?: FallbackToast; onApplied?: FallbackCallback }) { + if (!modelFallbackControllerRef.current) { + const pendingModelFallbacks = new Map() + const lastToastKey = new Map() + const sessionFallbackChains = new Map() + + modelFallbackControllerRef.current = createModelFallbackStateController({ + pendingModelFallbacks, + lastToastKey, + sessionFallbackChains, + }) + } + + const controller = getOrCreateModelFallbackController() const toast = args?.toast const onApplied = args?.onApplied @@ -184,7 +137,7 @@ export function createModelFallbackHook(args?: { toast?: FallbackToast; onApplie fallback, toast, onApplied, - lastToastKey, + lastToastKey: controller.lastToastKey, }) }, } @@ -195,7 +148,5 @@ export function createModelFallbackHook(args?: { toast?: FallbackToast; onApplie * Clears pending fallbacks, toast keys, and session chains. */ export function _resetForTesting(): void { - pendingModelFallbacks.clear() - lastToastKey.clear() - sessionFallbackChains.clear() + getOrCreateModelFallbackController().reset() }