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,10 +33,19 @@ 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
+const MAX_PERSISTED_SESSIONS = 150
+
+export const useSelectionStore = create<SelectionState>()(
+  persist(
+    (set, get) => ({
       sessionModelSelections: new Map(),
       sessionAgentSelections: new Map(),
       sessionAgentModelSelections: new Map(),
@@ -33,6 +54,7 @@ export const useSelectionStore = create<SelectionState>()((set, get) => ({
       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 } }
         }),
@@ -43,6 +65,7 @@ export const useSelectionStore = create<SelectionState>()((set, get) => ({
         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 }
         }),
@@ -55,8 +78,11 @@ export const useSelectionStore = create<SelectionState>()((set, get) => ({
           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 }
         }),
 
@@ -96,4 +122,47 @@ export const useSelectionStore = create<SelectionState>()((set, get) => ({
         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())])
+
+        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))
+          })
+        }
+
+        return {
+          ...currentState,
+          lastUsedProvider: persisted?.lastUsedProvider ?? currentState.lastUsedProvider,
+          sessionModelSelections: new Map(persisted?.sessionModelSelections ?? []),
+          sessionAgentSelections: new Map(persisted?.sessionAgentSelections ?? []),
+          sessionAgentModelSelections: agentModelSelections,
+        }
+      },
+      migrate: (persistedState: unknown) => {
+        // Scaffold for future schema migrations
+        return persistedState
+      }
+    }
+  )
+)