From 9347f97a984bd6be47efef2976b00f682dc94755 Mon Sep 17 00:00:00 2001 From: chutastic Date: Wed, 13 May 2026 08:06:20 -0400 Subject: [PATCH] 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 --- packages/ui/src/sync/selection-store.ts | 209 ++++++++++++++++-------- 1 file changed, 139 insertions(+), 70 deletions(-) diff --git a/packages/ui/src/sync/selection-store.ts b/packages/ui/src/sync/selection-store.ts index 2e3e6383..7b3058e3 100644 --- a/packages/ui/src/sync/selection-store.ts +++ b/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 + sessionModelSelections: Map sessionAgentSelections: Map - sessionAgentModelSelections: Map> - lastUsedProvider: { providerID: string; modelID: string } | null + sessionAgentModelSelections: Map> + 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>>() -export const useSelectionStore = create()((set, get) => ({ - sessionModelSelections: new Map(), - sessionAgentSelections: new Map(), - sessionAgentModelSelections: new Map(), - lastUsedProvider: null, +// Maximum number of sessions to persist to local storage to prevent unbounded growth +const MAX_PERSISTED_SESSIONS = 150 - saveSessionModelSelection: (sessionId, providerId, modelId) => - set((s) => { - const map = new Map(s.sessionModelSelections) - map.set(sessionId, { providerId, modelId }) - return { sessionModelSelections: map, lastUsedProvider: { providerID: providerId, modelID: modelId } } +export const useSelectionStore = create()( + persist( + (set, get) => ({ + sessionModelSelections: new Map(), + 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>() + if (Array.isArray(persisted?.sessionAgentModelSelections)) { + persisted.sessionAgentModelSelections.forEach(([sessionId, agentArray]) => { + agentModelSelections.set(sessionId, new Map(agentArray)) + }) + } - saveSessionAgentSelection: (sessionId, agentName) => - set((s) => { - if (s.sessionAgentSelections.get(sessionId) === agentName) return s - const map = new Map(s.sessionAgentSelections) - 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()) - 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) + 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 } - 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) - }, -})) + ) +)