feat: enable toggle workspace, invalidate on workspace update (#130)

* feat: enable toggle workspace, invalidate on workspace update

* fix: failing tests

* fix(workspaces selection): checkbox & bg color

Author: alex-mcgovern

src/components/react-query-provider.tsx

@@ -0,0 +1,85 @@
+import { V1ListActiveWorkspacesResponse } from "@/api/generated";
+import { v1ListActiveWorkspacesQueryKey } from "@/api/generated/@tanstack/react-query.gen";
+import { toast } from "@stacklok/ui-kit";
+import {
+  QueryCacheNotifyEvent,
+  QueryClient,
+  QueryClientProvider as VendorQueryClientProvider,
+} from "@tanstack/react-query";
+import { ReactNode, useState, useEffect } from "react";
+
+/**
+ * Responsible for determining whether a queryKey attached to a queryCache event
+ * is for the "list active workspaces" query.
+ */
+function isActiveWorkspacesQueryKey(queryKey: unknown): boolean {
+  return (
+    Array.isArray(queryKey) &&
+    queryKey[0]._id === v1ListActiveWorkspacesQueryKey()[0]?._id
+  );
+}
+
+/**
+ * Responsible for extracting the incoming active workspace name from the deeply
+ * nested payload attached to a queryCache event.
+ */
+function getWorkspaceName(event: QueryCacheNotifyEvent): string | null {
+  if ("action" in event === false || "data" in event.action === false)
+    return null;
+  return (
+    (event.action.data as V1ListActiveWorkspacesResponse | undefined | null)
+      ?.workspaces[0]?.name ?? null
+  );
+}
+
+export function QueryClientProvider({ children }: { children: ReactNode }) {
+  const [activeWorkspaceName, setActiveWorkspaceName] = useState<string | null>(
+    null,
+  );
+
+  const [queryClient] = useState(() => new QueryClient());
+
+  useEffect(() => {
+    const queryCache = queryClient.getQueryCache();
+    const unsubscribe = queryCache.subscribe((event) => {
+      if (
+        event.type === "updated" &&
+        event.action.type === "success" &&
+        isActiveWorkspacesQueryKey(event.query.options.queryKey)
+      ) {
+        const newWorkspaceName: string | null = getWorkspaceName(event);
+        if (
+          newWorkspaceName === activeWorkspaceName ||
+          newWorkspaceName === null
+        )
+          return;
+
+        setActiveWorkspaceName(newWorkspaceName);
+        toast.info(
+          <span className="block whitespace-nowrap">
+            Activated workspace:{" "}
+            <span className="font-semibold">"{newWorkspaceName}"</span>
+          </span>,
+        );
+
+        void queryClient.invalidateQueries({
+          refetchType: "all",
+          // Avoid a continuous loop
+          predicate(query) {
+            return !isActiveWorkspacesQueryKey(query.queryKey);
+          },
+        });
+      }
+    });
+
+    return () => {
+      return unsubscribe();
+    };
+  }, [activeWorkspaceName, queryClient]);
+
+  return (
+    <VendorQueryClientProvider client={queryClient}>
+      {children}
+    </VendorQueryClientProvider>
+  );
+}

src/features/workspace/components/workspaces-selection.tsx

@@ -1,40 +1,50 @@
-import { useWorkspacesData } from "@/hooks/useWorkspacesData";
+import { useListWorkspaces } from "@/features/workspace/hooks/use-list-workspaces";
 import {
   Button,
   DialogTrigger,
   Input,
+  LinkButton,
   ListBox,
   ListBoxItem,
   Popover,
   SearchField,
   Separator,
 } from "@stacklok/ui-kit";
 import { useQueryClient } from "@tanstack/react-query";
-import clsx from "clsx";
 import { ChevronDown, Search, Settings } from "lucide-react";
 import { useState } from "react";
-import { Link } from "react-router-dom";
+import { useActiveWorkspaces } from "../hooks/use-active-workspaces";
+import { useActivateWorkspace } from "../hooks/use-activate-workspace";
+import clsx from "clsx";
 
 export function WorkspacesSelection() {
   const queryClient = useQueryClient();
-  const { data } = useWorkspacesData();
+
+  const { data: workspacesResponse } = useListWorkspaces();
+  const { data: activeWorkspacesResponse } = useActiveWorkspaces();
+  const { mutateAsync: activateWorkspace } = useActivateWorkspace();
+
+  const activeWorkspaceName: string | null =
+    activeWorkspacesResponse?.workspaces[0]?.name ?? null;
+
   const [isOpen, setIsOpen] = useState(false);
   const [searchWorkspace, setSearchWorkspace] = useState("");
-  const workspaces = data?.workspaces ?? [];
+  const workspaces = workspacesResponse?.workspaces ?? [];
   const filteredWorkspaces = workspaces.filter((workspace) =>
     workspace.name.toLowerCase().includes(searchWorkspace.toLowerCase()),
   );
-  const activeWorkspace = workspaces.find((workspace) => workspace.is_active);
 
-  const handleWorkspaceClick = () => {
-    queryClient.invalidateQueries({ refetchType: "all" });
-    setIsOpen(false);
+  const handleWorkspaceClick = (name: string) => {
+    activateWorkspace({ body: { name } }).then(() => {
+      queryClient.invalidateQueries({ refetchType: "all" });
+      setIsOpen(false);
+    });
   };
 
   return (
     <DialogTrigger isOpen={isOpen} onOpenChange={(test) => setIsOpen(test)}>
       <Button variant="tertiary" className="flex cursor-pointer">
-        Workspace {activeWorkspace?.name ?? "default"}
+        Workspace {activeWorkspaceName ?? "default"}
         <ChevronDown />
       </Button>
 
@@ -51,40 +61,44 @@ export function WorkspacesSelection() {
           </div>
 
           <ListBox
-            className="pb-2 pt-3"
             aria-label="Workspaces"
             items={filteredWorkspaces}
-            selectedKeys={activeWorkspace?.name ?? []}
+            selectedKeys={activeWorkspaceName ? [activeWorkspaceName] : []}
+            onAction={(v) => {
+              handleWorkspaceClick(v?.toString());
+            }}
+            className="py-2 pt-3"
             renderEmptyState={() => (
               <p className="text-center">No workspaces found</p>
             )}
           >
             {(item) => (
               <ListBoxItem
                 id={item.name}
-                onAction={() => handleWorkspaceClick()}
+                key={item.name}
+                data-is-selected={item.name === activeWorkspaceName}
                 className={clsx(
                   "cursor-pointer py-2 m-1 text-base hover:bg-gray-300",
                   {
-                    "bg-gray-900 text-white hover:text-secondary":
+                    "!bg-gray-900 hover:bg-gray-900 !text-gray-25 hover:!text-gray-25":
                       item.is_active,
                   },
                 )}
-                key={item.name}
               >
                 {item.name}
               </ListBoxItem>
             )}
           </ListBox>
           <Separator className="" />
-          <Link
-            to="/workspaces"
-            onClick={() => setIsOpen(false)}
-            className="text-secondary pt-3 px-2 gap-2 flex"
+          <LinkButton
+            href="/workspaces"
+            onPress={() => setIsOpen(false)}
+            variant="tertiary"
+            className="text-secondary h-8 pl-2 gap-2 flex mt-2 justify-start"
           >
             <Settings />
             Manage Workspaces
-          </Link>
+          </LinkButton>
         </div>
       </Popover>
     </DialogTrigger>

src/features/workspace/hooks/use-activate-workspace.ts

@@ -0,0 +1,8 @@
+import { v1ActivateWorkspaceMutation } from "@/api/generated/@tanstack/react-query.gen";
+import { useMutation } from "@tanstack/react-query";
+
+export function useActivateWorkspace() {
+  return useMutation({
+    ...v1ActivateWorkspaceMutation(),
+  });
+}

src/features/workspace/hooks/use-active-workspaces.ts

@@ -0,0 +1,14 @@
+import { v1ListActiveWorkspacesOptions } from "@/api/generated/@tanstack/react-query.gen";
+import { useQuery } from "@tanstack/react-query";
+
+export function useActiveWorkspaces() {
+  return useQuery({
+    ...v1ListActiveWorkspacesOptions(),
+    refetchInterval: 5_000,
+    refetchIntervalInBackground: true,
+    refetchOnMount: true,
+    refetchOnReconnect: true,
+    refetchOnWindowFocus: true,
+    retry: false,
+  });
+}

src/features/workspace/hooks/use-list-workspaces.ts

@@ -0,0 +1,14 @@
+import { useQuery } from "@tanstack/react-query";
+import { v1ListWorkspacesOptions } from "@/api/generated/@tanstack/react-query.gen";
+
+export const useListWorkspaces = () => {
+  return useQuery({
+    ...v1ListWorkspacesOptions(),
+    refetchInterval: 5_000,
+    refetchIntervalInBackground: true,
+    refetchOnMount: true,
+    refetchOnReconnect: true,
+    refetchOnWindowFocus: true,
+    retry: false,
+  });
+};

src/hooks/useWorkspacesData.ts

@@ -5,11 +5,11 @@ import "@stacklok/ui-kit/style";
 import App from "./App.tsx";
 import { BrowserRouter } from "react-router-dom";
 import { SidebarProvider } from "./components/ui/sidebar.tsx";
-import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import ErrorBoundary from "./components/ErrorBoundary.tsx";
 import { Error } from "./components/Error.tsx";
-import { DarkModeProvider } from "@stacklok/ui-kit";
+import { DarkModeProvider, Toaster } from "@stacklok/ui-kit";
 import { client } from "./api/generated/index.ts";
+import { QueryClientProvider } from "./components/react-query-provider.tsx";
 
 // Initialize the API client
 client.setConfig({
@@ -22,7 +22,8 @@ createRoot(document.getElementById("root")!).render(
       <DarkModeProvider>
         <SidebarProvider>
           <ErrorBoundary fallback={<Error />}>
-            <QueryClientProvider client={new QueryClient()}>
+            <QueryClientProvider>
+              <Toaster />
               <App />
             </QueryClientProvider>
           </ErrorBoundary>

src/main.tsx

@@ -12,16 +12,25 @@ export const handlers = [
       error: null,
     }),
   ),
-  http.get("*/dashboard/version", () =>
+  http.get("*/api/v1/dashboard/version", () =>
     HttpResponse.json({ status: "healthy" }),
   ),
-  http.get("*/dashboard/messages", () => {
+  http.get("*/api/v1/workspaces/active", () =>
+    HttpResponse.json([
+      {
+        name: "my-awesome-workspace",
+        is_active: true,
+        last_updated: new Date(Date.now()).toISOString(),
+      },
+    ]),
+  ),
+  http.get("*/api/v1/dashboard/messages", () => {
     return HttpResponse.json(mockedPrompts);
   }),
-  http.get("*/dashboard/alerts", () => {
+  http.get("*/api/v1/dashboard/alerts", () => {
     return HttpResponse.json(mockedAlerts);
   }),
-  http.get("*/workspaces", () => {
+  http.get("*/api/v1/workspaces", () => {
     return HttpResponse.json(mockedWorkspaces);
   }),
 ];

src/mocks/msw/handlers.ts

@@ -1,4 +1,4 @@
-import { useWorkspacesData } from "@/hooks/useWorkspacesData";
+import { useListWorkspaces } from "@/features/workspace/hooks/use-list-workspaces";
 import {
   Cell,
   Column,
@@ -12,7 +12,7 @@ import {
 import { Settings } from "lucide-react";
 
 export function RouteWorkspaces() {
-  const result = useWorkspacesData();
+  const result = useListWorkspaces();
   const workspaces = result.data?.workspaces ?? [];
 
   return (