From 7fae6fdd9dcc6944052a1ec2ff04076853d390f9 Mon Sep 17 00:00:00 2001
From: James Singleton <jsingleton@wandb.com>
Date: Mon, 27 Apr 2026 16:28:57 -0500
Subject: [PATCH] Fix workspace state

---
 apps/web/app/routes/ws/_components/WorkspaceSelector.tsx | 8 +++-----
 apps/web/app/routes/ws/_layout.tsx                       | 6 +-----
 2 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/apps/web/app/routes/ws/_components/WorkspaceSelector.tsx b/apps/web/app/routes/ws/_components/WorkspaceSelector.tsx
index bace15de23..47cfc536d8 100644
--- a/apps/web/app/routes/ws/_components/WorkspaceSelector.tsx
+++ b/apps/web/app/routes/ws/_components/WorkspaceSelector.tsx
@@ -1,6 +1,7 @@
 import { Check, ChevronDown } from "lucide-react";
 import { NavLink } from "react-router";
 
+import { useWorkspace } from "~/components/WorkspaceProvider";
 import { Button } from "~/components/ui/button";
 import {
   DropdownMenu,
@@ -18,12 +19,9 @@ import {
 
 export const WorkspaceSelector: React.FC<{
   viewer: { email: string };
-  activeWorkspaceId?: string | null;
   workspaces: Array<{ id: string; slug: string; name: string }>;
-}> = ({ viewer, activeWorkspaceId, workspaces }) => {
-  const workspace =
-    workspaces.find((w) => w.id === activeWorkspaceId) ?? workspaces.at(0);
-  if (workspace == null) return;
+}> = ({ viewer, workspaces }) => {
+  const { workspace } = useWorkspace();
   return (
     <DropdownMenu>
       <DropdownMenuTrigger asChild>
diff --git a/apps/web/app/routes/ws/_layout.tsx b/apps/web/app/routes/ws/_layout.tsx
index 6067b4409c..b37d23bd9e 100644
--- a/apps/web/app/routes/ws/_layout.tsx
+++ b/apps/web/app/routes/ws/_layout.tsx
@@ -196,11 +196,7 @@ export default function WorkspaceLayout() {
           <SidebarMenu>
             <SidebarMenuItem>
               {viewer != null && (
-                <WorkspaceSelector
-                  viewer={viewer}
-                  activeWorkspaceId={viewer.activeWorkspaceId}
-                  workspaces={workspaces}
-                />
+                <WorkspaceSelector viewer={viewer} workspaces={workspaces} />
               )}
             </SidebarMenuItem>
           </SidebarMenu>