From af53e38110c3eca6c5812ea1791267540b9fbc88 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 21 May 2026 02:30:50 +0000 Subject: [PATCH] refactor(web): use Zustand store for current platform assistant (LUM-1718) Replace the hand-rolled `useSyncExternalStore` + module-level listener Set in `use-current-platform-assistant` with a Zustand `persist`-backed store. A custom `StateStorage` adapter keeps the existing per-org `vellum_current_assistant_id__{orgId}` localStorage key format so the on-disk shape stays backward compatible. React Query still owns the platform assistants list; the hook is now a thin composition of the store + the query. --- .../current-platform-assistant-store.ts | 168 ++++++++++++++++++ .../hooks/use-current-platform-assistant.ts | 88 ++------- 2 files changed, 187 insertions(+), 69 deletions(-) create mode 100644 apps/web/src/domains/settings/current-platform-assistant-store.ts diff --git a/apps/web/src/domains/settings/current-platform-assistant-store.ts b/apps/web/src/domains/settings/current-platform-assistant-store.ts new file mode 100644 index 00000000000..86b2fc432f7 --- /dev/null +++ b/apps/web/src/domains/settings/current-platform-assistant-store.ts @@ -0,0 +1,168 @@ +/** + * Zustand store for the per-organization "current platform assistant" + * selection. Persists the selected assistant ID for each org to + * localStorage so a reload or new tab restores the prior selection. + * + * **Storage model — one localStorage key per org:** + * + * Each org's selection is stored under + * `vellum_current_assistant_id__{orgId}`. A custom `StateStorage` + * adapter wired into the `persist` middleware reads/writes those keys + * directly, so the on-disk format stays compatible with the prior + * hand-rolled implementation. + * + * **Cross-tab sync:** + * + * The persist middleware doesn't subscribe to `storage` events on its + * own. We listen for any `storage` event whose key carries the + * per-org prefix and trigger `persist.rehydrate()` so the store + * picks up writes from other tabs. + * + * References: + * - {@link https://zustand.docs.pmnd.rs/} + * - {@link https://zustand.docs.pmnd.rs/integrations/persisting-store-data} + */ + +import { create } from "zustand"; +import { + createJSONStorage, + persist, + type StateStorage, +} from "zustand/middleware"; + +import { createSelectors } from "@/utils/create-selectors.js"; + +export const PLATFORM_ASSISTANT_STORAGE_PREFIX = + "vellum_current_assistant_id__"; + +export interface CurrentPlatformAssistantState { + /** orgId → selected assistant ID. Absent entries mean "no selection yet". */ + byOrg: Record; +} + +export interface CurrentPlatformAssistantActions { + setAssistantId: (orgId: string, id: string | null) => void; + getAssistantId: (orgId: string) => string | null; +} + +export type CurrentPlatformAssistantStore = CurrentPlatformAssistantState & + CurrentPlatformAssistantActions; + +function readByOrgFromLocalStorage(): Record { + if (typeof window === "undefined") return {}; + const byOrg: Record = {}; + try { + for (let i = 0; i < window.localStorage.length; i++) { + const key = window.localStorage.key(i); + if (key && key.startsWith(PLATFORM_ASSISTANT_STORAGE_PREFIX)) { + const orgId = key.slice(PLATFORM_ASSISTANT_STORAGE_PREFIX.length); + const value = window.localStorage.getItem(key); + if (value != null) byOrg[orgId] = value; + } + } + } catch { + // ignore storage failures + } + return byOrg; +} + +function writeByOrgToLocalStorage(byOrg: Record): void { + if (typeof window === "undefined") return; + try { + const existing: string[] = []; + for (let i = 0; i < window.localStorage.length; i++) { + const key = window.localStorage.key(i); + if (key && key.startsWith(PLATFORM_ASSISTANT_STORAGE_PREFIX)) { + existing.push(key); + } + } + for (const key of existing) { + const orgId = key.slice(PLATFORM_ASSISTANT_STORAGE_PREFIX.length); + if (byOrg[orgId] == null) { + window.localStorage.removeItem(key); + } + } + for (const [orgId, id] of Object.entries(byOrg)) { + window.localStorage.setItem( + `${PLATFORM_ASSISTANT_STORAGE_PREFIX}${orgId}`, + id, + ); + } + } catch { + // ignore storage failures + } +} + +/** + * Translates the single-name `name` view that `persist` expects into + * per-org reads and writes against the existing + * `vellum_current_assistant_id__{orgId}` localStorage keys. + */ +const perOrgStorage: StateStorage = { + getItem: () => { + const byOrg = readByOrgFromLocalStorage(); + return JSON.stringify({ state: { byOrg }, version: 0 }); + }, + setItem: (_name, value) => { + let parsed: { state?: { byOrg?: Record } }; + try { + parsed = JSON.parse(value) as typeof parsed; + } catch { + return; + } + const byOrg = parsed.state?.byOrg ?? {}; + writeByOrgToLocalStorage(byOrg); + }, + removeItem: () => { + writeByOrgToLocalStorage({}); + }, +}; + +const CURRENT_PLATFORM_ASSISTANT_STORE_NAME = + "vellum:current-platform-assistant"; + +const useCurrentPlatformAssistantStoreBase = create()( + persist( + (set, get) => ({ + byOrg: readByOrgFromLocalStorage(), + + setAssistantId: (orgId, id) => + set((state) => { + const next = { ...state.byOrg }; + if (id == null) { + delete next[orgId]; + } else { + next[orgId] = id; + } + return { byOrg: next }; + }), + + getAssistantId: (orgId) => get().byOrg[orgId] ?? null, + }), + { + name: CURRENT_PLATFORM_ASSISTANT_STORE_NAME, + storage: createJSONStorage(() => perOrgStorage), + partialize: (state) => ({ byOrg: state.byOrg }), + }, + ), +); + +export const useCurrentPlatformAssistantStore = createSelectors( + useCurrentPlatformAssistantStoreBase, +); + +// --------------------------------------------------------------------------- +// Cross-tab sync +// --------------------------------------------------------------------------- + +if (typeof window !== "undefined") { + window.addEventListener("storage", (event) => { + if (event.key === null) { + void useCurrentPlatformAssistantStoreBase.persist.rehydrate(); + return; + } + if (event.key.startsWith(PLATFORM_ASSISTANT_STORAGE_PREFIX)) { + void useCurrentPlatformAssistantStoreBase.persist.rehydrate(); + } + }); +} diff --git a/apps/web/src/domains/settings/hooks/use-current-platform-assistant.ts b/apps/web/src/domains/settings/hooks/use-current-platform-assistant.ts index 0806c3d4ba1..069747540e7 100644 --- a/apps/web/src/domains/settings/hooks/use-current-platform-assistant.ts +++ b/apps/web/src/domains/settings/hooks/use-current-platform-assistant.ts @@ -1,72 +1,15 @@ import { useQuery } from "@tanstack/react-query"; -import { useCallback, useEffect, useSyncExternalStore } from "react"; +import { useCallback, useEffect } from "react"; import { assistantsListOptions } from "@/generated/api/@tanstack/react-query.gen.js"; import type { Assistant } from "@/generated/api/types.gen.js"; +import { useCurrentPlatformAssistantStore } from "@/domains/settings/current-platform-assistant-store.js"; import { useOrganizationStore } from "@/stores/organization-store.js"; -const PLATFORM_ASSISTANT_STORAGE_PREFIX = "vellum_current_assistant_id__"; - const PLATFORM_LIST_OPTIONS = assistantsListOptions({ query: { hosting: "platform" }, }); -function storageKeyForOrg(orgId: string | null): string | null { - if (!orgId) return null; - return `${PLATFORM_ASSISTANT_STORAGE_PREFIX}${orgId}`; -} - -type Listener = () => void; -const listeners = new Set(); - -function notify(): void { - for (const listener of listeners) listener(); -} - -if (typeof window !== "undefined") { - window.addEventListener("storage", (event: StorageEvent) => { - if (event.key === null) { - notify(); - return; - } - if (event.key.startsWith(PLATFORM_ASSISTANT_STORAGE_PREFIX)) { - notify(); - } - }); -} - -function subscribe(listener: Listener): () => void { - listeners.add(listener); - return () => { - listeners.delete(listener); - }; -} - -function getSnapshot(orgId: string | null): string | null { - const key = storageKeyForOrg(orgId); - if (!key) return null; - try { - return window.localStorage.getItem(key); - } catch { - return null; - } -} - -function writeStoredId(orgId: string | null, id: string | null): void { - const key = storageKeyForOrg(orgId); - if (!key) return; - try { - if (id == null) { - window.localStorage.removeItem(key); - } else { - window.localStorage.setItem(key, id); - } - } catch { - // ignore storage failures - } - notify(); -} - export interface UseCurrentPlatformAssistantResult { assistantId: string | null; assistant: Assistant | null; @@ -78,12 +21,11 @@ export interface UseCurrentPlatformAssistantResult { export function useCurrentPlatformAssistant(): UseCurrentPlatformAssistantResult { const orgId = useOrganizationStore.use.currentOrganizationId(); + const byOrg = useCurrentPlatformAssistantStore.use.byOrg(); + const setAssistantIdAction = + useCurrentPlatformAssistantStore.use.setAssistantId(); - const storedId = useSyncExternalStore( - subscribe, - useCallback(() => getSnapshot(orgId), [orgId]), - () => null, - ); + const storedId = orgId ? (byOrg[orgId] ?? null) : null; const listQuery = useQuery(PLATFORM_LIST_OPTIONS); @@ -109,16 +51,24 @@ export function useCurrentPlatformAssistant(): UseCurrentPlatformAssistantResult if (!isListLoaded) return; if (platformAssistants.length === 0) return; if (resolvedId === storedId) return; - if (resolvedId != null) { - writeStoredId(orgId, resolvedId); + if (resolvedId != null && orgId) { + setAssistantIdAction(orgId, resolvedId); } - }, [isListLoaded, platformAssistants.length, resolvedId, storedId, orgId]); + }, [ + isListLoaded, + platformAssistants.length, + resolvedId, + storedId, + orgId, + setAssistantIdAction, + ]); const setAssistantId = useCallback( (id: string | null) => { - writeStoredId(orgId, id); + if (!orgId) return; + setAssistantIdAction(orgId, id); }, - [orgId], + [orgId, setAssistantIdAction], ); return {