From 4c4cf2a541ffa28f529a104aa66064cfbef2f8c2 Mon Sep 17 00:00:00 2001
From: Bradley Axen <baxen@squareup.com>
Date: Fri, 8 May 2026 17:36:16 -0700
Subject: [PATCH] Fix workspace switch cache isolation

---
 desktop/src/app/App.tsx                       | 44 +++++++++++--------
 .../features/workspaces/useWorkspaceInit.ts   | 32 +++++++++++---
 desktop/src/main.tsx                          | 34 +++++---------
 desktop/src/shared/api/queryClient.ts         | 17 +++++++
 4 files changed, 77 insertions(+), 50 deletions(-)
 create mode 100644 desktop/src/shared/api/queryClient.ts

diff --git a/desktop/src/app/App.tsx b/desktop/src/app/App.tsx
index 28790de36..62c9bdd08 100644
--- a/desktop/src/app/App.tsx
+++ b/desktop/src/app/App.tsx
@@ -1,15 +1,21 @@
 import { getCurrentWindow } from "@tauri-apps/api/window";
-import { useQueryClient } from "@tanstack/react-query";
+import { QueryClientProvider } from "@tanstack/react-query";
 import { RouterProvider } from "@tanstack/react-router";
-import { useCallback, useEffect, useLayoutEffect } from "react";
+import {
+  type ReactNode,
+  useCallback,
+  useEffect,
+  useLayoutEffect,
+  useState,
+} from "react";
 
 import { router } from "@/app/router";
-import { UpdaterProvider } from "@/features/settings/hooks/UpdaterProvider";
 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";
+import { createSproutQueryClient } from "@/shared/api/queryClient";
 import { listenForDeepLinks } from "@/shared/deep-link";
 
 function AppLoadingGate() {
@@ -30,6 +36,14 @@ function AppLoadingGate() {
   );
 }
 
+function WorkspaceQueryProvider({ children }: { children: ReactNode }) {
+  const [queryClient] = useState(createSproutQueryClient);
+
+  return (
+    <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+  );
+}
+
 function AppReady() {
   const onboarding = useAppOnboardingState();
 
@@ -55,7 +69,6 @@ export function App() {
     void getCurrentWindow().show();
   }, []);
 
-  const queryClient = useQueryClient();
   const {
     activeWorkspace,
     reinitKey,
@@ -74,19 +87,10 @@ export function App() {
       void unlisten.then((fn) => fn());
     };
   }, [addWorkspace, switchWorkspace, reconnectWorkspace]);
-  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 workspace = useWorkspaceInit(activeWorkspace, workspaceKey);
 
   const handleSetupComplete = useCallback(() => {
     // Force a full reload so useWorkspaces re-initializes from localStorage.
@@ -104,15 +108,17 @@ export function App() {
     );
   }
 
-  // Wait for workspace config to be applied to the backend before
-  // rendering anything that connects to the relay.
-  if (!workspace.isReady) {
+  // Wait for this exact workspace config to be applied to the backend before
+  // rendering anything that connects to the relay. The appliedKey check avoids
+  // a one-render race where React sees the new active workspace while the Tauri
+  // backend is still configured for the previous one.
+  if (!workspace.isReady || workspace.appliedKey !== workspaceKey) {
     return <AppLoadingGate />;
   }
 
   return (
-    <UpdaterProvider>
+    <WorkspaceQueryProvider key={workspaceKey}>
       <AppReady key={workspaceKey} />
-    </UpdaterProvider>
+    </WorkspaceQueryProvider>
   );
 }
diff --git a/desktop/src/features/workspaces/useWorkspaceInit.ts b/desktop/src/features/workspaces/useWorkspaceInit.ts
index f15f7169d..265004485 100644
--- a/desktop/src/features/workspaces/useWorkspaceInit.ts
+++ b/desktop/src/features/workspaces/useWorkspaceInit.ts
@@ -24,13 +24,13 @@ function resetWorkspaceState(): void {
 }
 
 type WorkspaceInitResult =
-  | { isReady: true; needsSetup: false }
+  | { isReady: true; needsSetup: false; appliedKey: string }
   | {
       isReady: false;
       needsSetup: true;
       defaultRelayUrl: string;
     }
-  | { isReady: false; needsSetup: false };
+  | { isReady: false; needsSetup: false; appliedKey: string | null };
 
 /**
  * Applies the active workspace config to the Tauri backend and resets
@@ -42,10 +42,12 @@ type WorkspaceInitResult =
  */
 export function useWorkspaceInit(
   activeWorkspace: Workspace | null,
+  workspaceKey: string,
 ): WorkspaceInitResult {
   const [result, setResult] = useState<WorkspaceInitResult>({
     isReady: false,
     needsSetup: false,
+    appliedKey: null,
   });
 
   // Track whether this is the initial mount or a workspace switch.
@@ -80,6 +82,16 @@ export function useWorkspaceInit(
         return;
       }
 
+      // Mark this workspace config as pending while it is applied to the
+      // backend. App.tsx also checks appliedKey against the active workspaceKey,
+      // which prevents rendering workspace-scoped UI for a new workspace until
+      // that exact config has finished applying.
+      setResult({
+        isReady: false,
+        needsSetup: false,
+        appliedKey: workspaceKey,
+      });
+
       // On workspace switch (not initial mount), reset module singletons
       // so the new tree starts with a clean slate.
       if (hasInitializedRef.current) {
@@ -87,9 +99,6 @@ export function useWorkspaceInit(
       }
       hasInitializedRef.current = true;
 
-      // Show loading gate while we apply the new workspace config
-      setResult({ isReady: false, needsSetup: false });
-
       // Apply workspace config to the Tauri backend.
       //
       // Note: we deliberately do NOT pass an nsec here. The persisted
@@ -110,7 +119,11 @@ export function useWorkspaceInit(
       }
 
       if (!cancelled) {
-        setResult({ isReady: true, needsSetup: false });
+        setResult({
+          isReady: true,
+          needsSetup: false,
+          appliedKey: workspaceKey,
+        });
       }
     }
 
@@ -119,7 +132,12 @@ export function useWorkspaceInit(
     return () => {
       cancelled = true;
     };
-  }, [activeWorkspace?.id, activeWorkspace?.relayUrl, activeWorkspace?.token]);
+  }, [
+    activeWorkspace?.id,
+    activeWorkspace?.relayUrl,
+    activeWorkspace?.token,
+    workspaceKey,
+  ]);
 
   return result;
 }
diff --git a/desktop/src/main.tsx b/desktop/src/main.tsx
index ff8988551..e706404ed 100644
--- a/desktop/src/main.tsx
+++ b/desktop/src/main.tsx
@@ -1,8 +1,8 @@
 import React from "react";
 import ReactDOM from "react-dom/client";
-import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { App } from "@/app/App";
 import "@/shared/styles/globals.css";
+import { UpdaterProvider } from "@/features/settings/hooks/UpdaterProvider";
 import { WorkspacesProvider } from "@/features/workspaces/useWorkspaces";
 import { ThemeProvider } from "@/shared/theme/ThemeProvider";
 import { Toaster } from "@/shared/ui/sonner";
@@ -12,33 +12,19 @@ type E2eWindow = Window & {
   __SPROUT_E2E__?: unknown;
 };
 
-const queryClient = new QueryClient({
-  defaultOptions: {
-    queries: {
-      retry: 1,
-      refetchOnWindowFocus: false,
-      networkMode: "always",
-      gcTime: 5 * 60 * 1_000,
-    },
-    mutations: {
-      networkMode: "always",
-    },
-  },
-});
-
 function renderApp() {
   ReactDOM.createRoot(document.getElementById("root") as HTMLElement).render(
     <React.StrictMode>
-      <QueryClientProvider client={queryClient}>
-        <WorkspacesProvider>
-          <ThemeProvider defaultTheme="houston">
-            <TooltipProvider delayDuration={300}>
+      <WorkspacesProvider>
+        <ThemeProvider defaultTheme="houston">
+          <TooltipProvider delayDuration={300}>
+            <UpdaterProvider>
               <App />
-              <Toaster />
-            </TooltipProvider>
-          </ThemeProvider>
-        </WorkspacesProvider>
-      </QueryClientProvider>
+            </UpdaterProvider>
+            <Toaster />
+          </TooltipProvider>
+        </ThemeProvider>
+      </WorkspacesProvider>
     </React.StrictMode>,
   );
 }
diff --git a/desktop/src/shared/api/queryClient.ts b/desktop/src/shared/api/queryClient.ts
new file mode 100644
index 000000000..196998cca
--- /dev/null
+++ b/desktop/src/shared/api/queryClient.ts
@@ -0,0 +1,17 @@
+import { QueryClient } from "@tanstack/react-query";
+
+export function createSproutQueryClient() {
+  return new QueryClient({
+    defaultOptions: {
+      queries: {
+        retry: 1,
+        refetchOnWindowFocus: false,
+        networkMode: "always",
+        gcTime: 5 * 60 * 1_000,
+      },
+      mutations: {
+        networkMode: "always",
+      },
+    },
+  });
+}
