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

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

AGENTS.md

@@ -127,6 +127,35 @@ just desktop-dev   # web-only dev server (faster iteration)
 just desktop-app   # full Tauri app with native shell
 ```
 
+### Workspace Switching
+
+The desktop app supports multiple workspaces (each backed by a different relay).
+Switching workspaces does **not** reload the page — it uses React key-based
+remounting. `<AppReady key={workspaceKey} />` in `App.tsx` forces the entire
+workspace-scoped subtree to unmount and remount with fresh state.
+
+**Module-level singletons must be explicitly reset.** React remounting only
+clears React state (useState, useRef, context). Module-level variables (Maps,
+class instances, cached promises) survive across remounts. Every workspace-scoped
+singleton needs a reset function wired into `resetWorkspaceState()` in
+`desktop/src/features/workspaces/useWorkspaceInit.ts`.
+
+Current singletons that are reset on workspace switch:
+- `relayClient.disconnect()` — WebSocket teardown + promise rejection
+- `resetMediaCaches()` — proxy port and relay origin caches
+- `clearSearchHitEventCache()` — search result event cache
+- `clearAllDrafts()` — message draft cache
+
+**If you add a new module-level cache, Map, or class instance that holds
+workspace-scoped data, you must add its reset to `resetWorkspaceState()`.**
+Failure to do so causes data from the old workspace to leak into the new one.
+
+Key files:
+- `desktop/src/app/App.tsx` — workspace key, init gate, remount boundary
+- `desktop/src/features/workspaces/useWorkspaceInit.ts` — `resetWorkspaceState()`, applies config to Tauri backend
+- `desktop/src/features/workspaces/useWorkspaces.tsx` — `WorkspacesProvider` context (shared state for App + AppShell)
+- `desktop/src/main.tsx` — provider hierarchy (`QueryClientProvider` > `WorkspacesProvider` > `App`)
+
 ---
 
 ## Mobile App (Flutter)

desktop/scripts/check-file-sizes.mjs

@@ -40,7 +40,7 @@ const overrides = new Map([
   ["src/features/settings/ui/SettingsView.tsx", 600],
   ["src/features/sidebar/ui/AppSidebar.tsx", 860], // channels + forums creation forms + Pulse nav
   ["src/features/tokens/ui/TokenSettingsCard.tsx", 800],
-  ["src/shared/api/relayClientSession.ts", 835], // durable websocket session manager with reconnect/replay/recovery state + sendTypingIndicator + fetchChannelHistoryBefore + subscribeToChannelLive (huddle TTS) + subscribeToHuddleEvents (huddle indicator)
+  ["src/shared/api/relayClientSession.ts", 890], // durable websocket session manager with reconnect/replay/recovery state + sendTypingIndicator + fetchChannelHistoryBefore + subscribeToChannelLive (huddle TTS) + subscribeToHuddleEvents (huddle indicator) + disconnect() for workspace switch teardown
   ["src/shared/api/tauri.ts", 1100], // remote agent provider API bindings + canvas API functions
   ["src-tauri/src/lib.rs", 710], // sprout-media:// proxy + Range headers + Sprout nest init (ensure_nest) in setup() + huddle command registration + PTT global shortcut handler + persona pack commands + app_handle storage for event emission
   ["src-tauri/src/commands/media.rs", 720], // ffmpeg video transcode + poster frame extraction + run_ffmpeg_with_timeout (find_ffmpeg, is_video_file, transcode_to_mp4, extract_poster_frame, transcode_and_extract_poster) + spawn_blocking wrappers + tests

desktop/src/app/App.tsx

@@ -1,4 +1,5 @@
 import { getCurrentWindow } from "@tauri-apps/api/window";
+import { useQueryClient } from "@tanstack/react-query";
 import { RouterProvider } from "@tanstack/react-router";
 import { useCallback, useLayoutEffect } from "react";
 
@@ -6,6 +7,7 @@ import { router } from "@/app/router";
 import { useAppOnboardingState } from "@/features/onboarding/hooks";
 import { OnboardingFlow } from "@/features/onboarding/ui/OnboardingFlow";
 import { useWorkspaceInit } from "@/features/workspaces/useWorkspaceInit";
+import { useWorkspaces } from "@/features/workspaces/useWorkspaces";
 import { WelcomeSetup } from "@/features/workspaces/ui/WelcomeSetup";
 
 function AppLoadingGate() {
@@ -52,11 +54,25 @@ export function App() {
     void getCurrentWindow().show();
   }, []);
 
-  const workspace = useWorkspaceInit();
+  const queryClient = useQueryClient();
+  const { activeWorkspace, reinitKey } = useWorkspaces();
+  const workspace = useWorkspaceInit(activeWorkspace);
+
+  // Composite key: changes when workspace ID changes OR when
+  // the active workspace's config is updated (relayUrl/token).
+  const workspaceKey = `${activeWorkspace?.id ?? "none"}-${reinitKey}`;
+
+  // Clear stale React Query cache synchronously when workspace changes.
+  // useLayoutEffect fires before child useEffect hooks, preventing stale
+  // data from being served to the new workspace's components.
+  // biome-ignore lint/correctness/useExhaustiveDependencies: workspaceKey drives the re-run intentionally
+  useLayoutEffect(() => {
+    queryClient.clear();
+  }, [workspaceKey, queryClient]);
 
   const handleSetupComplete = useCallback(() => {
-    // Force a full reload so useWorkspaceInit re-runs and picks up
-    // the newly-created workspace from localStorage.
+    // Force a full reload so useWorkspaces re-initializes from localStorage.
+    // This only runs once — during first-run setup when no workspace existed.
     window.location.reload();
   }, []);
 
@@ -76,5 +92,5 @@ export function App() {
     return <AppLoadingGate />;
   }
 
-  return <AppReady />;
+  return <AppReady key={workspaceKey} />;
 }

desktop/src/app/navigation/searchHitEventCache.ts

@@ -38,6 +38,10 @@ export function cacheSearchHitEvent(hit: SearchHit): RelayEvent {
   return event;
 }
 
+export function clearSearchHitEventCache(): void {
+  searchHitEventCache.clear();
+}
+
 export function getCachedSearchHitEvent(
   eventId: string | null | undefined,
 ): RelayEvent | null {

desktop/src/features/agents/ui/ManagedAgentRow.tsx

@@ -460,6 +460,9 @@ function AgentOriginBadge({ agent }: { agent: ManagedAgent }) {
   );
 }
 
+/** Grace period after mount before treating "running + no presence" as "Starting…" */
+const PRESENCE_GRACE_MS = 15_000;
+
 function AgentStatusBadge({
   presenceLoaded,
   presenceStatus,
@@ -469,8 +472,16 @@ function AgentStatusBadge({
   presenceStatus: PresenceStatus | undefined;
   status: ManagedAgent["status"];
 }) {
+  const [inGracePeriod, setInGracePeriod] = React.useState(true);
+
+  React.useEffect(() => {
+    const timer = setTimeout(() => setInGracePeriod(false), PRESENCE_GRACE_MS);
+    return () => clearTimeout(timer);
+  }, []);
+
   const isActive = status === "running" || status === "deployed";
   const isStarting =
+    !inGracePeriod &&
     presenceLoaded &&
     status === "running" &&
     (!presenceStatus || presenceStatus === "offline");

desktop/src/features/channels/ui/ChannelPane.tsx

@@ -282,10 +282,14 @@ export const ChannelPane = React.memo(function ChannelPane({
           channelName={activeChannel?.name ?? "channel"}
           currentPubkey={currentPubkey}
           disabled={isComposerDisabled}
+          editTarget={editTarget}
           isSending={isSending}
+          onCancelEdit={onCancelEdit}
           onCancelReply={onCancelThreadReply}
           onClose={onCloseThread}
           onDelete={onDelete}
+          onEdit={onEdit}
+          onEditSave={onEditSave}
           onExpandReplies={onExpandThreadReplies}
           onSelectReplyTarget={onSelectThreadReplyTarget}
           onSend={onSendThreadReply}

desktop/src/features/messages/lib/useDrafts.ts

@@ -10,6 +10,10 @@ export type DraftState = {
 
 const sharedDrafts = new Map<string, DraftState>();
 
+export function clearAllDrafts(): void {
+  sharedDrafts.clear();
+}
+
 export function useDrafts() {
   const saveDraft = React.useCallback(
     (channelId: string, draft: DraftState) => {

desktop/src/features/messages/ui/MessageThreadPanel.tsx

@@ -19,10 +19,14 @@ type MessageThreadPanelProps = {
   channelName: string;
   currentPubkey?: string;
   disabled?: boolean;
+  editTarget?: { author: string; body: string; id: string } | null;
   isSending: boolean;
+  onCancelEdit?: () => void;
   onCancelReply: () => void;
   onClose: () => void;
   onDelete?: (message: TimelineMessage) => void;
+  onEdit?: (message: TimelineMessage) => void;
+  onEditSave?: (content: string) => Promise<void>;
   onExpandReplies: (message: TimelineMessage) => void;
   onResetWidth: () => void;
   onResizeStart: (event: React.PointerEvent<HTMLButtonElement>) => void;
@@ -66,10 +70,14 @@ export function MessageThreadPanel({
   channelName,
   currentPubkey,
   disabled = false,
+  editTarget,
   isSending,
+  onCancelEdit,
   onCancelReply,
   onClose,
   onDelete,
+  onEdit,
+  onEditSave,
   onExpandReplies,
   onResetWidth,
   onResizeStart,
@@ -179,6 +187,11 @@ export function MessageThreadPanel({
                     ? onDelete
                     : undefined
                 }
+                onEdit={
+                  onEdit && canManageMessage(threadHead, currentPubkey)
+                    ? onEdit
+                    : undefined
+                }
                 onToggleReaction={onToggleReaction}
                 profiles={profiles}
               />
@@ -205,6 +218,12 @@ export function MessageThreadPanel({
                             ? onDelete
                             : undefined
                         }
+                        onEdit={
+                          onEdit &&
+                          canManageMessage(entry.message, currentPubkey)
+                            ? onEdit
+                            : undefined
+                        }
                         onReply={onSelectReplyTarget}
                         onToggleReaction={onToggleReaction}
                         profiles={profiles}
@@ -259,8 +278,11 @@ export function MessageThreadPanel({
           channelName={channelName}
           disabled={disabled || isSending || !channelId}
           draftKey={`thread:${threadHead.id}`}
+          editTarget={editTarget}
           isSending={isSending}
+          onCancelEdit={onCancelEdit}
           onCancelReply={composerReplyTarget ? onCancelReply : undefined}
+          onEditSave={onEditSave}
           onSend={onSend}
           placeholder={`Reply in thread to ${threadHead.author}`}
           replyTarget={composerReplyTarget}

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,

desktop/src/main.tsx

@@ -3,6 +3,7 @@ import ReactDOM from "react-dom/client";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { App } from "@/app/App";
 import "@/shared/styles/globals.css";
+import { WorkspacesProvider } from "@/features/workspaces/useWorkspaces";
 import { ThemeProvider } from "@/shared/theme/ThemeProvider";
 import { Toaster } from "@/shared/ui/sonner";
 import { TooltipProvider } from "@/shared/ui/tooltip";
@@ -29,12 +30,14 @@ function renderApp() {
   ReactDOM.createRoot(document.getElementById("root") as HTMLElement).render(
     <React.StrictMode>
       <QueryClientProvider client={queryClient}>
-        <ThemeProvider defaultTheme="houston">
-          <TooltipProvider delayDuration={300}>
-            <App />
-            <Toaster />
-          </TooltipProvider>
-        </ThemeProvider>
+        <WorkspacesProvider>
+          <ThemeProvider defaultTheme="houston">
+            <TooltipProvider delayDuration={300}>
+              <App />
+              <Toaster />
+            </TooltipProvider>
+          </ThemeProvider>
+        </WorkspacesProvider>
       </QueryClientProvider>
     </React.StrictMode>,
   );

desktop/src/shared/api/relayClientSession.ts

@@ -48,6 +48,62 @@ export class RelayClient {
   private notifyReconnectListeners = false;
   private onMessageChannel: Channel<unknown> | null = null;
 
+  /**
+   * Cleanly tear down the connection without scheduling a reconnect.
+   * Used during workspace switches to reset the singleton before the
+   * new workspace applies.
+   */
+  disconnect() {
+    const error = new Error("Relay disconnected for workspace switch.");
+
+    if (this.reconnectTimeout) {
+      window.clearTimeout(this.reconnectTimeout);
+      this.reconnectTimeout = null;
+    }
+    this.keepAliveRequested = false;
+    this.relayUrl = null;
+    this.hasConnectedOnce = false;
+    this.notifyReconnectListeners = false;
+
+    if (this.wsId !== null) {
+      void invoke("plugin:websocket|disconnect", { id: this.wsId }).catch(
+        () => {},
+      );
+      this.wsId = null;
+    }
+
+    this.connectPromise = null;
+
+    if (this.authRequest) {
+      window.clearTimeout(this.authRequest.timeout);
+      this.authRequest.reject(error);
+      this.authRequest = null;
+    }
+
+    for (const [subId, sub] of this.subscriptions) {
+      if (sub.mode === "history") {
+        window.clearTimeout(sub.timeout);
+        sub.reject(error);
+      }
+      this.subscriptions.delete(subId);
+    }
+
+    for (const [eventId, pending] of this.pendingEvents) {
+      window.clearTimeout(pending.timeout);
+      pending.reject(error);
+      this.pendingEvents.delete(eventId);
+    }
+
+    if (this.flushTimeout !== null) {
+      window.clearTimeout(this.flushTimeout);
+      this.flushTimeout = null;
+    }
+    this.eventBuffer = [];
+    this.reconnectListeners.clear();
+    this.onMessageChannel = null;
+    this.reconnectDelayMs = RECONNECT_BASE_DELAY_MS;
+  }
+
   async fetchChannelHistory(channelId: string, limit = 50) {
     return this.fetchHistory(this.buildChannelFilter(channelId, limit));
   }

desktop/src/shared/lib/mediaUrl.ts

@@ -70,6 +70,16 @@ if (typeof window !== "undefined") {
   portPromise = fetchProxyPort();
 }
 
+/**
+ * Reset module-level caches so the next render re-fetches the proxy port
+ * and relay origin for the new workspace.
+ */
+export function resetMediaCaches(): void {
+  cachedPort = null;
+  portPromise = null;
+  cachedRelayOrigin = null;
+}
+
 /**
  * If `url` is a Blossom media URL hosted on the Sprout relay, rewrite it
  * to go through the localhost streaming proxy. External Blossom URLs and

desktop/src/shared/ui/avatar.tsx

@@ -37,7 +37,7 @@ const AvatarFallback = React.forwardRef<
   <AvatarPrimitive.Fallback
     ref={ref}
     className={cn(
-      "flex h-full w-full items-center justify-center rounded-full bg-muted",
+      "flex h-full w-full items-center justify-center rounded-[inherit] bg-muted",
       className,
     )}
     {...props}