diff --git a/packages/core/chat/store.ts b/packages/core/chat/store.ts index 3e46ca5e8..272418367 100644 --- a/packages/core/chat/store.ts +++ b/packages/core/chat/store.ts @@ -1,5 +1,6 @@ import { create } from "zustand"; import type { StorageAdapter } from "../types"; +import { getCurrentWorkspaceId, registerForWorkspaceRehydration } from "../platform/workspace-storage"; const AGENT_STORAGE_KEY = "multica:chat:selectedAgentId"; const SESSION_STORAGE_KEY = "multica:chat:activeSessionId"; @@ -39,12 +40,17 @@ export interface ChatStoreOptions { export function createChatStore(options: ChatStoreOptions) { const { storage } = options; - return create((set) => ({ + const wsKey = (base: string) => { + const wsId = getCurrentWorkspaceId(); + return wsId ? `${base}:${wsId}` : base; + }; + + const store = create((set) => ({ isOpen: false, isFullscreen: false, - activeSessionId: storage.getItem(SESSION_STORAGE_KEY), + activeSessionId: storage.getItem(wsKey(SESSION_STORAGE_KEY)), pendingTaskId: null, - selectedAgentId: storage.getItem(AGENT_STORAGE_KEY), + selectedAgentId: storage.getItem(wsKey(AGENT_STORAGE_KEY)), showHistory: false, timelineItems: [], setOpen: (open) => @@ -57,15 +63,15 @@ export function createChatStore(options: ChatStoreOptions) { toggleFullscreen: () => set((s) => ({ isFullscreen: !s.isFullscreen })), setActiveSession: (id) => { if (id) { - storage.setItem(SESSION_STORAGE_KEY, id); + storage.setItem(wsKey(SESSION_STORAGE_KEY), id); } else { - storage.removeItem(SESSION_STORAGE_KEY); + storage.removeItem(wsKey(SESSION_STORAGE_KEY)); } set({ activeSessionId: id }); }, setPendingTask: (taskId) => set({ pendingTaskId: taskId, timelineItems: [] }), setSelectedAgentId: (id) => { - storage.setItem(AGENT_STORAGE_KEY, id); + storage.setItem(wsKey(AGENT_STORAGE_KEY), id); set({ selectedAgentId: id }); }, setShowHistory: (show) => set({ showHistory: show }), @@ -80,4 +86,14 @@ export function createChatStore(options: ChatStoreOptions) { }), clearTimeline: () => set({ timelineItems: [] }), })); + + registerForWorkspaceRehydration(() => { + store.setState({ + activeSessionId: storage.getItem(wsKey(SESSION_STORAGE_KEY)), + selectedAgentId: storage.getItem(wsKey(AGENT_STORAGE_KEY)), + timelineItems: [], + }); + }); + + return store; } diff --git a/packages/core/issues/stores/draft-store.ts b/packages/core/issues/stores/draft-store.ts index 55296c8cb..7ca50530a 100644 --- a/packages/core/issues/stores/draft-store.ts +++ b/packages/core/issues/stores/draft-store.ts @@ -1,7 +1,7 @@ import { create } from "zustand"; import { createJSONStorage, persist } from "zustand/middleware"; import type { IssueStatus, IssuePriority, IssueAssigneeType } from "../../types"; -import { createWorkspaceAwareStorage } from "../../platform/workspace-storage"; +import { createWorkspaceAwareStorage, registerForWorkspaceRehydration } from "../../platform/workspace-storage"; import { defaultStorage } from "../../platform/storage"; interface IssueDraft { @@ -49,3 +49,5 @@ export const useIssueDraftStore = create()( }, ), ); + +registerForWorkspaceRehydration(() => useIssueDraftStore.persist.rehydrate()); diff --git a/packages/core/issues/stores/issues-scope-store.ts b/packages/core/issues/stores/issues-scope-store.ts index 76a80ee25..f4cfacad3 100644 --- a/packages/core/issues/stores/issues-scope-store.ts +++ b/packages/core/issues/stores/issues-scope-store.ts @@ -2,7 +2,7 @@ import { create } from "zustand"; import { createJSONStorage, persist } from "zustand/middleware"; -import { createWorkspaceAwareStorage } from "../../platform/workspace-storage"; +import { createWorkspaceAwareStorage, registerForWorkspaceRehydration } from "../../platform/workspace-storage"; import { defaultStorage } from "../../platform/storage"; export type IssuesScope = "all" | "members" | "agents"; @@ -24,3 +24,5 @@ export const useIssuesScopeStore = create()( }, ), ); + +registerForWorkspaceRehydration(() => useIssuesScopeStore.persist.rehydrate()); diff --git a/packages/core/issues/stores/my-issues-view-store.ts b/packages/core/issues/stores/my-issues-view-store.ts index 08dff0560..9ffb373c1 100644 --- a/packages/core/issues/stores/my-issues-view-store.ts +++ b/packages/core/issues/stores/my-issues-view-store.ts @@ -7,6 +7,7 @@ import { viewStoreSlice, viewStorePersistOptions, } from "./view-store"; +import { registerForWorkspaceRehydration } from "../../platform/workspace-storage"; export type MyIssuesScope = "assigned" | "created" | "agents"; @@ -17,7 +18,7 @@ export interface MyIssuesViewState extends IssueViewState { const basePersist = viewStorePersistOptions("multica_my_issues_view"); -export const myIssuesViewStore: StoreApi = createStore()( +const _myIssuesViewStore = createStore()( persist( (set) => ({ ...viewStoreSlice(set as unknown as StoreApi["setState"]), @@ -34,3 +35,7 @@ export const myIssuesViewStore: StoreApi = createStore = _myIssuesViewStore; + +registerForWorkspaceRehydration(() => _myIssuesViewStore.persist.rehydrate()); diff --git a/packages/core/issues/stores/view-store.ts b/packages/core/issues/stores/view-store.ts index 7b24a5b6d..d8001ff1c 100644 --- a/packages/core/issues/stores/view-store.ts +++ b/packages/core/issues/stores/view-store.ts @@ -5,7 +5,7 @@ import { createStore, type StoreApi } from "zustand/vanilla"; import { createJSONStorage, persist } from "zustand/middleware"; import type { IssueStatus, IssuePriority } from "../../types"; import { ALL_STATUSES } from "../config"; -import { createWorkspaceAwareStorage } from "../../platform/workspace-storage"; +import { createWorkspaceAwareStorage, registerForWorkspaceRehydration } from "../../platform/workspace-storage"; import { defaultStorage } from "../../platform/storage"; export type ViewMode = "board" | "list"; @@ -183,9 +183,11 @@ export const viewStorePersistOptions = (name: string) => ({ /** Factory: creates a vanilla StoreApi for use with React Context. */ export function createIssueViewStore(persistKey: string): StoreApi { - return createStore()( + const store = createStore()( persist(viewStoreSlice, viewStorePersistOptions(persistKey)) ); + registerForWorkspaceRehydration(() => store.persist.rehydrate()); + return store; } /** Global singleton for the /issues page. */ @@ -193,6 +195,8 @@ export const useIssueViewStore = create()( persist(viewStoreSlice, viewStorePersistOptions("multica_issues_view")) ); +registerForWorkspaceRehydration(() => useIssueViewStore.persist.rehydrate()); + // Clear filters on all registered view stores when workspace switches. const _syncedStores = new Set>(); let _workspaceSyncInitialized = false; diff --git a/packages/core/navigation/store.ts b/packages/core/navigation/store.ts index 2dad77db7..12bcd7964 100644 --- a/packages/core/navigation/store.ts +++ b/packages/core/navigation/store.ts @@ -1,7 +1,9 @@ "use client"; import { create } from "zustand"; -import { persist } from "zustand/middleware"; +import { createJSONStorage, persist } from "zustand/middleware"; +import { createPersistStorage } from "../platform/persist-storage"; +import { defaultStorage } from "../platform/storage"; const EXCLUDED_PREFIXES = ["/login", "/pair/"]; @@ -23,6 +25,7 @@ export const useNavigationStore = create()( }), { name: "multica_navigation", + storage: createJSONStorage(() => createPersistStorage(defaultStorage)), partialize: (state) => ({ lastPath: state.lastPath }), } ) diff --git a/packages/core/platform/index.ts b/packages/core/platform/index.ts index 3f21f71a2..1e59186ae 100644 --- a/packages/core/platform/index.ts +++ b/packages/core/platform/index.ts @@ -3,4 +3,5 @@ export type { CoreProviderProps } from "./types"; export { AuthInitializer } from "./auth-initializer"; export { defaultStorage } from "./storage"; export { createPersistStorage } from "./persist-storage"; -export { createWorkspaceAwareStorage, setCurrentWorkspaceId, getCurrentWorkspaceId } from "./workspace-storage"; +export { createWorkspaceAwareStorage, setCurrentWorkspaceId, getCurrentWorkspaceId, registerForWorkspaceRehydration, rehydrateAllWorkspaceStores } from "./workspace-storage"; +export { clearWorkspaceStorage } from "./storage-cleanup"; diff --git a/packages/core/platform/persist-storage.test.ts b/packages/core/platform/persist-storage.test.ts index 1dbfb6ac1..b895899f5 100644 --- a/packages/core/platform/persist-storage.test.ts +++ b/packages/core/platform/persist-storage.test.ts @@ -12,7 +12,7 @@ function mockAdapter(): StorageAdapter { } describe("createPersistStorage", () => { - it("delegates to StorageAdapter without namespace", () => { + it("delegates to StorageAdapter", () => { const adapter = mockAdapter(); const storage = createPersistStorage(adapter); @@ -27,28 +27,6 @@ describe("createPersistStorage", () => { expect(result).toEqual(JSON.stringify("value")); }); - it("namespaces keys when wsId is provided", () => { - const adapter = mockAdapter(); - const storage = createPersistStorage(adapter, "ws_123"); - - storage.setItem("draft", JSON.stringify({ title: "test" })); - expect(adapter.setItem).toHaveBeenCalledWith( - "draft:ws_123", - JSON.stringify({ title: "test" }), - ); - - storage.getItem("draft"); - expect(adapter.getItem).toHaveBeenCalledWith("draft:ws_123"); - }); - - it("removeItem namespaces correctly", () => { - const adapter = mockAdapter(); - const storage = createPersistStorage(adapter, "ws_abc"); - - storage.removeItem("draft"); - expect(adapter.removeItem).toHaveBeenCalledWith("draft:ws_abc"); - }); - it("returns null for missing keys", () => { const adapter = mockAdapter(); const storage = createPersistStorage(adapter); @@ -56,4 +34,12 @@ describe("createPersistStorage", () => { const result = storage.getItem("nonexistent"); expect(result).toBeNull(); }); + + it("removeItem delegates correctly", () => { + const adapter = mockAdapter(); + const storage = createPersistStorage(adapter); + + storage.removeItem("key"); + expect(adapter.removeItem).toHaveBeenCalledWith("key"); + }); }); diff --git a/packages/core/platform/persist-storage.ts b/packages/core/platform/persist-storage.ts index 6b7873e4a..77dc97a8e 100644 --- a/packages/core/platform/persist-storage.ts +++ b/packages/core/platform/persist-storage.ts @@ -1,14 +1,14 @@ import type { StateStorage } from "zustand/middleware"; import type { StorageAdapter } from "../types/storage"; -export function createPersistStorage( - adapter: StorageAdapter, - wsId?: string, -): StateStorage { - const resolve = (key: string) => (wsId ? `${key}:${wsId}` : key); +/** + * Bridge between Zustand persist middleware and our StorageAdapter DI system. + * For workspace-scoped stores, use createWorkspaceAwareStorage instead. + */ +export function createPersistStorage(adapter: StorageAdapter): StateStorage { return { - getItem: (key) => adapter.getItem(resolve(key)), - setItem: (key, value) => adapter.setItem(resolve(key), value), - removeItem: (key) => adapter.removeItem(resolve(key)), + getItem: (key) => adapter.getItem(key), + setItem: (key, value) => adapter.setItem(key, value), + removeItem: (key) => adapter.removeItem(key), }; } diff --git a/packages/core/platform/storage-cleanup.test.ts b/packages/core/platform/storage-cleanup.test.ts new file mode 100644 index 000000000..49cb8cb0a --- /dev/null +++ b/packages/core/platform/storage-cleanup.test.ts @@ -0,0 +1,22 @@ +import { describe, it, expect, vi } from "vitest"; +import { clearWorkspaceStorage } from "./storage-cleanup"; + +describe("clearWorkspaceStorage", () => { + it("removes all workspace-scoped keys for given wsId", () => { + const adapter = { + getItem: vi.fn(), + setItem: vi.fn(), + removeItem: vi.fn(), + }; + + clearWorkspaceStorage(adapter, "ws_123"); + + expect(adapter.removeItem).toHaveBeenCalledWith("multica_issue_draft:ws_123"); + expect(adapter.removeItem).toHaveBeenCalledWith("multica_issues_view:ws_123"); + expect(adapter.removeItem).toHaveBeenCalledWith("multica_issues_scope:ws_123"); + expect(adapter.removeItem).toHaveBeenCalledWith("multica_my_issues_view:ws_123"); + expect(adapter.removeItem).toHaveBeenCalledWith("multica:chat:selectedAgentId:ws_123"); + expect(adapter.removeItem).toHaveBeenCalledWith("multica:chat:activeSessionId:ws_123"); + expect(adapter.removeItem).toHaveBeenCalledTimes(6); + }); +}); diff --git a/packages/core/platform/storage-cleanup.ts b/packages/core/platform/storage-cleanup.ts new file mode 100644 index 000000000..a4ba0b780 --- /dev/null +++ b/packages/core/platform/storage-cleanup.ts @@ -0,0 +1,27 @@ +import type { StorageAdapter } from "../types/storage"; + +/** + * Keys that are namespaced per workspace (stored as `${key}:${wsId}`). + * + * IMPORTANT: When adding a new workspace-scoped persist store or storage key, + * add its key here so that workspace deletion and logout properly clean it up. + * Also ensure the store uses `createWorkspaceAwareStorage` for its persist config. + */ +const WORKSPACE_SCOPED_KEYS = [ + "multica_issue_draft", + "multica_issues_view", + "multica_issues_scope", + "multica_my_issues_view", + "multica:chat:selectedAgentId", + "multica:chat:activeSessionId", +]; + +/** Remove all workspace-scoped storage entries for the given workspace. */ +export function clearWorkspaceStorage( + adapter: StorageAdapter, + wsId: string, +) { + for (const key of WORKSPACE_SCOPED_KEYS) { + adapter.removeItem(`${key}:${wsId}`); + } +} diff --git a/packages/core/platform/workspace-storage.ts b/packages/core/platform/workspace-storage.ts index c1cb165d1..c148bb499 100644 --- a/packages/core/platform/workspace-storage.ts +++ b/packages/core/platform/workspace-storage.ts @@ -2,11 +2,24 @@ import type { StateStorage } from "zustand/middleware"; import type { StorageAdapter } from "../types/storage"; let _currentWsId: string | null = null; +const _rehydrateFns: Array<() => void> = []; export function setCurrentWorkspaceId(wsId: string | null) { _currentWsId = wsId; } +/** Register a persist store's rehydrate function to be called on workspace switch. */ +export function registerForWorkspaceRehydration(fn: () => void) { + _rehydrateFns.push(fn); +} + +/** Rehydrate all registered workspace-scoped persist stores from the new namespace. */ +export function rehydrateAllWorkspaceStores() { + for (const fn of _rehydrateFns) { + fn(); + } +} + export function getCurrentWorkspaceId(): string | null { return _currentWsId; } diff --git a/packages/core/realtime/use-realtime-sync.ts b/packages/core/realtime/use-realtime-sync.ts index c6b85bb55..f4c02aec7 100644 --- a/packages/core/realtime/use-realtime-sync.ts +++ b/packages/core/realtime/use-realtime-sync.ts @@ -7,6 +7,8 @@ import type { StoreApi, UseBoundStore } from "zustand"; import type { AuthState } from "../auth/store"; import type { WorkspaceStore } from "../workspace/store"; import { createLogger } from "../logger"; +import { clearWorkspaceStorage } from "../platform/storage-cleanup"; +import { defaultStorage } from "../platform/storage"; import { issueKeys } from "../issues/queries"; import { projectKeys } from "../projects/queries"; import { runtimeKeys } from "../runtimes/queries"; @@ -238,6 +240,7 @@ export function useRealtimeSync( const unsubWsDeleted = ws.on("workspace:deleted", (p) => { const { workspace_id } = p as WorkspaceDeletedPayload; + clearWorkspaceStorage(defaultStorage, workspace_id); const currentWs = workspaceStore.getState().workspace; if (currentWs?.id === workspace_id) { logger.warn("current workspace deleted, switching"); @@ -250,6 +253,8 @@ export function useRealtimeSync( const { user_id } = p as MemberRemovedPayload; const myUserId = authStore.getState().user?.id; if (user_id === myUserId) { + const wsId = workspaceStore.getState().workspace?.id; + if (wsId) clearWorkspaceStorage(defaultStorage, wsId); logger.warn("removed from workspace, switching"); onToast?.("You were removed from this workspace", "info"); workspaceStore.getState().refreshWorkspaces(); diff --git a/packages/core/workspace/store.ts b/packages/core/workspace/store.ts index 002f4b52f..7dfb3781b 100644 --- a/packages/core/workspace/store.ts +++ b/packages/core/workspace/store.ts @@ -2,7 +2,7 @@ import { create } from "zustand"; import type { Workspace, StorageAdapter } from "../types"; import type { ApiClient } from "../api/client"; import { createLogger } from "../logger"; -import { setCurrentWorkspaceId } from "../platform/workspace-storage"; +import { setCurrentWorkspaceId, rehydrateAllWorkspaceStores } from "../platform/workspace-storage"; const logger = createLogger("workspace-store"); @@ -59,6 +59,7 @@ export function createWorkspaceStore(api: ApiClient, options?: WorkspaceStoreOpt if (!nextWorkspace) { api.setWorkspaceId(null); setCurrentWorkspaceId(null); + rehydrateAllWorkspaceStores(); storage?.removeItem("multica_workspace_id"); set({ workspace: null }); return null; @@ -66,6 +67,7 @@ export function createWorkspaceStore(api: ApiClient, options?: WorkspaceStoreOpt api.setWorkspaceId(nextWorkspace.id); setCurrentWorkspaceId(nextWorkspace.id); + rehydrateAllWorkspaceStores(); storage?.setItem("multica_workspace_id", nextWorkspace.id); set({ workspace: nextWorkspace }); logger.debug("hydrate workspace", nextWorkspace.name, nextWorkspace.id); @@ -142,6 +144,7 @@ export function createWorkspaceStore(api: ApiClient, options?: WorkspaceStoreOpt clearWorkspace: () => { api.setWorkspaceId(null); setCurrentWorkspaceId(null); + rehydrateAllWorkspaceStores(); set({ workspace: null, workspaces: [] }); }, }));