fix(desktop): key-based React remount for workspace switching (#415)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

desktop/src/features/workspaces/useWorkspaceInit.ts

@@ -1,12 +1,25 @@
-import { useEffect, useState } from "react";
+import { useEffect, useRef, useState } from "react";
 
+import { relayClient } from "@/shared/api/relayClient";
 import { applyWorkspace, getDefaultRelayUrl } from "@/shared/api/tauri";
+import { resetMediaCaches } from "@/shared/lib/mediaUrl";
+import { clearSearchHitEventCache } from "@/app/navigation/searchHitEventCache";
+import { clearAllDrafts } from "@/features/messages/lib/useDrafts";
 
-import {
-  loadActiveWorkspaceId,
-  loadWorkspaces,
-  saveActiveWorkspaceId,
-} from "./workspaceStorage";
+import type { Workspace } from "./types";
+
+/**
+ * Tear down all workspace-scoped module singletons so the new
+ * workspace starts with a clean slate. If you add a new module-level
+ * cache or singleton that holds workspace data, add its reset here.
+ * See AGENTS.md "Workspace Switching" for the full contract.
+ */
+function resetWorkspaceState(): void {
+  relayClient.disconnect();
+  resetMediaCaches();
+  clearSearchHitEventCache();
+  clearAllDrafts();
+}
 
 type WorkspaceInitResult =
   | { isReady: true; needsSetup: false }
@@ -14,36 +27,38 @@ type WorkspaceInitResult =
   | { isReady: false; needsSetup: false };
 
 /**
- * Runs once on mount. Loads the active workspace from localStorage
- * and calls the Tauri backend to apply the workspace config
- * (keys, relay URL, token).
+ * Applies the active workspace config to the Tauri backend and resets
+ * all workspace-scoped module singletons when the workspace changes.
  *
  * Returns a discriminated union — only render the app after the
  * workspace is applied. When `needsSetup` is true, the caller
  * should show a first-run welcome screen.
  */
-export function useWorkspaceInit(): WorkspaceInitResult {
+export function useWorkspaceInit(
+  activeWorkspace: Workspace | null,
+): WorkspaceInitResult {
   const [result, setResult] = useState<WorkspaceInitResult>({
     isReady: false,
     needsSetup: false,
   });
 
+  // Track whether this is the initial mount or a workspace switch.
+  // On the initial mount we skip resetting singletons (they're fresh).
+  const hasInitializedRef = useRef(false);
+
+  // biome-ignore lint/correctness/useExhaustiveDependencies: we intentionally depend on specific properties (id/relayUrl/token) — depending on the whole object would trigger resets on name-only changes
   useEffect(() => {
     let cancelled = false;
 
     async function init() {
-      const workspaces = loadWorkspaces();
-
-      if (workspaces.length === 0) {
-        // No workspaces at all — fetch the build default relay URL
-        // so the welcome screen can pre-fill it.
+      if (!activeWorkspace) {
+        // No workspace — need setup
         try {
           const defaultRelayUrl = await getDefaultRelayUrl();
           if (!cancelled) {
             setResult({ isReady: false, needsSetup: true, defaultRelayUrl });
           }
         } catch {
-          // If we can't get the default, fall back to localhost
           if (!cancelled) {
             setResult({
               isReady: false,
@@ -55,24 +70,23 @@ export function useWorkspaceInit(): WorkspaceInitResult {
         return;
       }
 
-      // Determine active workspace
-      let activeId = loadActiveWorkspaceId();
-      if (!activeId || !workspaces.find((w) => w.id === activeId)) {
-        activeId = workspaces[0].id;
-        saveActiveWorkspaceId(activeId);
+      // On workspace switch (not initial mount), reset module singletons
+      // so the new tree starts with a clean slate.
+      if (hasInitializedRef.current) {
+        resetWorkspaceState();
       }
+      hasInitializedRef.current = true;
 
-      const active = workspaces.find((w) => w.id === activeId);
-      if (!active) {
-        if (!cancelled) {
-          setResult({ isReady: true, needsSetup: false });
-        }
-        return;
-      }
+      // Show loading gate while we apply the new workspace config
+      setResult({ isReady: false, needsSetup: false });
 
       // Apply workspace config to the Tauri backend
       try {
-        await applyWorkspace(active.relayUrl, active.nsec, active.token);
+        await applyWorkspace(
+          activeWorkspace.relayUrl,
+          activeWorkspace.nsec,
+          activeWorkspace.token,
+        );
       } catch (error) {
         console.error("Failed to apply workspace to backend:", error);
       }
@@ -87,7 +101,7 @@ export function useWorkspaceInit(): WorkspaceInitResult {
     return () => {
       cancelled = true;
     };
-  }, []);
+  }, [activeWorkspace?.id, activeWorkspace?.relayUrl, activeWorkspace?.token]);
 
   return result;
 }

desktop/src/features/workspaces/useWorkspaces.tsx

@@ -1,4 +1,12 @@
-import { useCallback, useMemo, useRef, useState } from "react";
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
+import type { ReactNode } from "react";
 
 import type { Workspace } from "./types";
 import {
@@ -11,6 +19,8 @@ import {
 export type UseWorkspacesReturn = {
   workspaces: Workspace[];
   activeWorkspace: Workspace | null;
+  /** Counter bumped when the active workspace's config changes (relayUrl/token). */
+  reinitKey: number;
   /** Add a workspace, deduplicating by relayUrl. Returns the final ID in the list. */
   addWorkspace: (workspace: Workspace) => string;
   removeWorkspace: (id: string) => void;
@@ -21,12 +31,32 @@ export type UseWorkspacesReturn = {
   ) => void;
 };
 
+const WorkspacesContext = createContext<UseWorkspacesReturn | null>(null);
+
+export function WorkspacesProvider({ children }: { children: ReactNode }) {
+  const value = useWorkspacesInternal();
+  return (
+    <WorkspacesContext.Provider value={value}>
+      {children}
+    </WorkspacesContext.Provider>
+  );
+}
+
 export function useWorkspaces(): UseWorkspacesReturn {
+  const ctx = useContext(WorkspacesContext);
+  if (!ctx) {
+    throw new Error("useWorkspaces must be used within a WorkspacesProvider");
+  }
+  return ctx;
+}
+
+function useWorkspacesInternal(): UseWorkspacesReturn {
   const [workspaces, setWorkspacesState] =
     useState<Workspace[]>(loadWorkspaces);
   const [activeId, setActiveId] = useState<string | null>(
     loadActiveWorkspaceId,
   );
+  const [reinitKey, setReinitKey] = useState(0);
   const workspacesRef = useRef(workspaces);
   workspacesRef.current = workspaces;
 
@@ -76,9 +106,8 @@ export function useWorkspaces(): UseWorkspacesReturn {
 
         // If removing the active workspace, switch to first remaining
         if (activeId === id && next.length > 0) {
-          setActiveId(next[0].id);
           saveActiveWorkspaceId(next[0].id);
-          window.location.reload();
+          setActiveId(next[0].id);
         }
 
         return next;
@@ -93,7 +122,7 @@ export function useWorkspaces(): UseWorkspacesReturn {
         return;
       }
       saveActiveWorkspaceId(id);
-      window.location.reload();
+      setActiveId(id);
     },
     [activeId],
   );
@@ -115,12 +144,13 @@ export function useWorkspaces(): UseWorkspacesReturn {
         saveWorkspaces(next);
         return next;
       });
-      // If the active workspace's relay URL or token changed, reload to reconnect
+      // If the active workspace's relay URL or token changed, bump reinitKey
+      // so the React tree remounts with the new config.
       if (
         id === activeId &&
         (updates.relayUrl || updates.token !== undefined)
       ) {
-        window.location.reload();
+        setReinitKey((k) => k + 1);
       }
     },
     [activeId],
@@ -129,6 +159,7 @@ export function useWorkspaces(): UseWorkspacesReturn {
   return {
     workspaces,
     activeWorkspace,
+    reinitKey,
     addWorkspace,
     removeWorkspace,
     switchWorkspace,