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 <clio-agent@sisyphuslabs.ai>
2026-04-18 01:53:03 +09:00
parent 0d10498a11
commit 81b68d828c
2 changed files with 179 additions and 93 deletions
@@ -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<string, string>
+  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<typeof getNextReachableFallback>
+  clearPendingModelFallback: (sessionID: string) => void
+  hasPendingModelFallback: (sessionID: string) => boolean
+  getFallbackState: (sessionID: string) => ModelFallbackStateLike | undefined
+  reset: () => void
+}
+
+export function createModelFallbackStateController(input: {
+  pendingModelFallbacks: Map<string, ModelFallbackStateLike>
+  lastToastKey: Map<string, string>
+  sessionFallbackChains: Map<string, FallbackEntry[]>
+}): 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<typeof getNextReachableFallback> {
+    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,
+  }
+}
@@ -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<string, ModelFallbackState>()
-const lastToastKey = new Map<string, string>()
-const sessionFallbackChains = new Map<string, FallbackEntry[]>()
+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<string, ModelFallbackState>()
+    const lastToastKey = new Map<string, string>()
+    const sessionFallbackChains = new Map<string, FallbackEntry[]>()
+
+    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()
 }