fix(ui): persist selection-store state across reloads (#1229)

* fix(ui): persist selection-store state across reloads

* fix(ui): resolve selection store persistence build

---------

Co-authored-by: Bohdan Triapitsyn <artmore@protonmail.com>

packages/ui/src/sync/selection-store.ts

@@ -4,12 +4,24 @@
  */
 
 import { create } from "zustand"
+import { persist, createJSONStorage } from "zustand/middleware"
+import { getSafeStorage } from "@/stores/utils/safeStorage"
+
+type ModelSelection = { providerId: string; modelId: string }
+type LastUsedProvider = { providerID: string; modelID: string }
+type AgentModelSelectionEntries = [string, [string, ModelSelection][]][]
+type PersistedSelectionState = {
+  sessionModelSelections?: [string, ModelSelection][]
+  sessionAgentSelections?: [string, string][]
+  sessionAgentModelSelections?: AgentModelSelectionEntries
+  lastUsedProvider?: LastUsedProvider | null
+}
 
 export type SelectionState = {
-  sessionModelSelections: Map<string, { providerId: string; modelId: string }>
+  sessionModelSelections: Map<string, ModelSelection>
   sessionAgentSelections: Map<string, string>
-  sessionAgentModelSelections: Map<string, Map<string, { providerId: string; modelId: string }>>
+  sessionAgentModelSelections: Map<string, Map<string, ModelSelection>>
-  lastUsedProvider: { providerID: string; modelID: string } | null
+  lastUsedProvider: LastUsedProvider | null
 
   saveSessionModelSelection: (sessionId: string, providerId: string, modelId: string) => void
   getSessionModelSelection: (sessionId: string) => { providerId: string; modelId: string } | null
@@ -21,79 +33,136 @@ export type SelectionState = {
   getAgentModelVariantForSession: (sessionId: string, agentName: string, providerId: string, modelId: string) => string | undefined
 }
 
+const isPersistedSelectionState = (state: unknown): state is PersistedSelectionState => (
+  typeof state === "object" && state !== null
+)
+
 // In-memory variant storage (not persisted)
 const agentModelVariantSelections = new Map<string, Map<string, Map<string, string>>>()
 
-export const useSelectionStore = create<SelectionState>()((set, get) => ({
+// Maximum number of sessions to persist to local storage to prevent unbounded growth
-  sessionModelSelections: new Map(),
+const MAX_PERSISTED_SESSIONS = 150
-  sessionAgentSelections: new Map(),
-  sessionAgentModelSelections: new Map(),
-  lastUsedProvider: null,
 
-  saveSessionModelSelection: (sessionId, providerId, modelId) =>
+export const useSelectionStore = create<SelectionState>()(
-    set((s) => {
+  persist(
-      const map = new Map(s.sessionModelSelections)
+    (set, get) => ({
-      map.set(sessionId, { providerId, modelId })
+      sessionModelSelections: new Map(),
-      return { sessionModelSelections: map, lastUsedProvider: { providerID: providerId, modelID: modelId } }
+      sessionAgentSelections: new Map(),
+      sessionAgentModelSelections: new Map(),
+      lastUsedProvider: null,
+
+      saveSessionModelSelection: (sessionId, providerId, modelId) =>
+        set((s) => {
+          const map = new Map(s.sessionModelSelections)
+          map.delete(sessionId) // Delete first to ensure it moves to the end of insertion order (MRU)
+          map.set(sessionId, { providerId, modelId })
+          return { sessionModelSelections: map, lastUsedProvider: { providerID: providerId, modelID: modelId } }
+        }),
+
+      getSessionModelSelection: (sessionId) => get().sessionModelSelections.get(sessionId) ?? null,
+
+      saveSessionAgentSelection: (sessionId, agentName) =>
+        set((s) => {
+          if (s.sessionAgentSelections.get(sessionId) === agentName) return s
+          const map = new Map(s.sessionAgentSelections)
+          map.delete(sessionId) // Delete first to ensure it moves to the end of insertion order (MRU)
+          map.set(sessionId, agentName)
+          return { sessionAgentSelections: map }
+        }),
+
+      getSessionAgentSelection: (sessionId) => get().sessionAgentSelections.get(sessionId) ?? null,
+
+      saveAgentModelForSession: (sessionId, agentName, providerId, modelId) =>
+        set((s) => {
+          const existing = s.sessionAgentModelSelections.get(sessionId)?.get(agentName)
+          if (existing?.providerId === providerId && existing?.modelId === modelId) return s
+          const outer = new Map(s.sessionAgentModelSelections)
+          const inner = new Map(outer.get(sessionId) ?? new Map())
+
+          outer.delete(sessionId) // Delete first to ensure it moves to the end of insertion order (MRU)
+          inner.set(agentName, { providerId, modelId })
+          outer.set(sessionId, inner)
+
+          return { sessionAgentModelSelections: outer }
+        }),
+
+      getAgentModelForSession: (sessionId, agentName) =>
+        get().sessionAgentModelSelections.get(sessionId)?.get(agentName) ?? null,
+
+      saveAgentModelVariantForSession: (sessionId, agentName, providerId, modelId, variant) => {
+        const key = `${providerId}/${modelId}`
+        let agentMap = agentModelVariantSelections.get(sessionId)
+        if (!agentMap && variant) {
+          agentMap = new Map()
+          agentModelVariantSelections.set(sessionId, agentMap)
+        }
+        if (!agentMap) return
+        let modelMap = agentMap.get(agentName)
+        if (!modelMap && variant) {
+          modelMap = new Map()
+          agentMap.set(agentName, modelMap)
+        }
+        if (!modelMap) return
+
+        if (!variant) {
+          modelMap.delete(key)
+          if (modelMap.size === 0) {
+            agentMap.delete(agentName)
+          }
+          if (agentMap.size === 0) {
+            agentModelVariantSelections.delete(sessionId)
+          }
+          return
+        }
+
+        modelMap.set(key, variant)
+      },
+
+      getAgentModelVariantForSession: (sessionId, agentName, providerId, modelId) => {
+        const key = `${providerId}/${modelId}`
+        return agentModelVariantSelections.get(sessionId)?.get(agentName)?.get(key)
+      },
     }),
+    {
+      name: "selection-store",
+      version: 1,
+      storage: createJSONStorage(() => getSafeStorage()),
+      partialize: (state) => {
+        // Convert Maps to arrays and slice to keep only the most recent MAX_PERSISTED_SESSIONS
+        const models = Array.from(state.sessionModelSelections.entries()).slice(-MAX_PERSISTED_SESSIONS)
+        const agents = Array.from(state.sessionAgentSelections.entries()).slice(-MAX_PERSISTED_SESSIONS)
+        const agentModels = Array.from(state.sessionAgentModelSelections.entries())
+          .slice(-MAX_PERSISTED_SESSIONS)
+          .map(([sessionId, agentMap]) => [sessionId, Array.from(agentMap.entries())])
 
-  getSessionModelSelection: (sessionId) => get().sessionModelSelections.get(sessionId) ?? null,
+        return {
+          sessionModelSelections: models,
+          sessionAgentSelections: agents,
+          sessionAgentModelSelections: agentModels,
+          lastUsedProvider: state.lastUsedProvider,
+        }
+      },
+      merge: (persistedState: unknown, currentState) => {
+        const persisted = isPersistedSelectionState(persistedState) ? persistedState : undefined
+        const agentModelSelections = new Map<string, Map<string, ModelSelection>>()
+        if (Array.isArray(persisted?.sessionAgentModelSelections)) {
+          persisted.sessionAgentModelSelections.forEach(([sessionId, agentArray]) => {
+            agentModelSelections.set(sessionId, new Map(agentArray))
+          })
+        }
 
-  saveSessionAgentSelection: (sessionId, agentName) =>
+        return {
-    set((s) => {
+          ...currentState,
-      if (s.sessionAgentSelections.get(sessionId) === agentName) return s
+          lastUsedProvider: persisted?.lastUsedProvider ?? currentState.lastUsedProvider,
-      const map = new Map(s.sessionAgentSelections)
+          sessionModelSelections: new Map(persisted?.sessionModelSelections ?? []),
-      map.set(sessionId, agentName)
+          sessionAgentSelections: new Map(persisted?.sessionAgentSelections ?? []),
-      return { sessionAgentSelections: map }
+          sessionAgentModelSelections: agentModelSelections,
-    }),
+        }
-
+      },
-  getSessionAgentSelection: (sessionId) => get().sessionAgentSelections.get(sessionId) ?? null,
+      migrate: (persistedState: unknown) => {
-
+        // Scaffold for future schema migrations
-  saveAgentModelForSession: (sessionId, agentName, providerId, modelId) =>
+        return persistedState
-    set((s) => {
-      const existing = s.sessionAgentModelSelections.get(sessionId)?.get(agentName)
-      if (existing?.providerId === providerId && existing?.modelId === modelId) return s
-      const outer = new Map(s.sessionAgentModelSelections)
-      const inner = new Map(outer.get(sessionId) ?? new Map())
-      inner.set(agentName, { providerId, modelId })
-      outer.set(sessionId, inner)
-      return { sessionAgentModelSelections: outer }
-    }),
-
-  getAgentModelForSession: (sessionId, agentName) =>
-    get().sessionAgentModelSelections.get(sessionId)?.get(agentName) ?? null,
-
-  saveAgentModelVariantForSession: (sessionId, agentName, providerId, modelId, variant) => {
-    const key = `${providerId}/${modelId}`
-    let agentMap = agentModelVariantSelections.get(sessionId)
-    if (!agentMap && variant) {
-      agentMap = new Map()
-      agentModelVariantSelections.set(sessionId, agentMap)
-    }
-    if (!agentMap) return
-    let modelMap = agentMap.get(agentName)
-    if (!modelMap && variant) {
-      modelMap = new Map()
-      agentMap.set(agentName, modelMap)
-    }
-    if (!modelMap) return
-
-    if (!variant) {
-      modelMap.delete(key)
-      if (modelMap.size === 0) {
-        agentMap.delete(agentName)
       }
-      if (agentMap.size === 0) {
-        agentModelVariantSelections.delete(sessionId)
-      }
-      return
     }
-
+  )
-    modelMap.set(key, variant)
+)
-  },
-
-  getAgentModelVariantForSession: (sessionId, agentName, providerId, modelId) => {
-    const key = `${providerId}/${modelId}`
-    return agentModelVariantSelections.get(sessionId)?.get(agentName)?.get(key)
-  },
-}))