feat: improve Studio navigation and namespace/organization switching

studio/src/components/dashboard/graph-command-group.tsx

@@ -120,15 +120,14 @@ function GraphCommandItem({
       )}
       value={value}
       onSelect={() => {
-        setNamespace(namespace.name);
         router.push({
           pathname,
           query: {
             organizationSlug,
             namespace: namespace.name,
             ...(isSubgraph ? { subgraphSlug: name } : { slug: name }),
           }
-        });
+        }).finally(() => setNamespace(namespace.name));
       }}
     >
       <span className="flex justify-between items-center gap-2 w-full">

studio/src/components/dashboard/namespace-selector.tsx

@@ -49,19 +49,22 @@ export function NamespaceSelector({ isViewingGraphOrSubgraph, truncateNamespace
   return (
     <div className="flex items-center justify-start">
       {isViewingGraphOrSubgraph && (
-        <Link
-          href={{
-            pathname,
-            query: { organizationSlug, namespace: namespace.name },
-          }}
-          className={cn(
-            "bg-primary/15 hover:bg-primary/30 text-primary transition-colors duration-150 pl-3 pr-2 py-1.5 rounded-l-lg text-sm flex-shrink-0",
-            truncateNamespace && "max-w-[180px] lg:max-w-xs truncate"
-          )}
-          onClick={() => setNamespace(namespace.name, false)}
-        >
-          {namespace.name}
-        </Link>
+        <>
+          <Link
+            href={{
+              pathname,
+              query: { organizationSlug, namespace: namespace.name },
+            }}
+            className={cn(
+              "bg-primary/15 hover:bg-primary/30 text-primary transition-colors duration-150 pl-3 pr-2 py-1.5 rounded-l-lg text-sm flex-shrink-0",
+              truncateNamespace && "max-w-[180px] lg:max-w-xs truncate"
+            )}
+            onClick={() => setNamespace(namespace.name, false)}
+          >
+            {namespace.name}
+          </Link>
+          <div className="w-[1px] h-8 bg-primary/30" />
+        </>
       )}
       <Popover
         modal
@@ -81,7 +84,7 @@ export function NamespaceSelector({ isViewingGraphOrSubgraph, truncateNamespace
               "bg-primary/15 hover:bg-primary/30 text-primary transition-colors duration-150 text-sm flex-shrink-0 border-none outline-none",
               isViewingGraphOrSubgraph
                 ? "rounded-r-lg pl-2 pr-3 py-2"
-                : "flex justify-start items-center gap-4 rounded-lg px-3 py-1.5"
+                : "flex justify-start items-center gap-4 rounded-lg px-3 py-1.5",
             )}
           >
             {!isViewingGraphOrSubgraph && (

studio/src/components/dashboard/workspace-command-wrapper.tsx

@@ -122,20 +122,18 @@ export function WorkspaceCommandWrapper({
                   No namespace, graph or subgraph matches your criteria.
                 </div>
               ) : filteredGraphs.map((wns, index) => (
-                <>
-                    <GraphCommandGroup
-                      key={`namespace-${index}`}
-                      isFiltering={isFiltering}
-                      namespace={wns}
-                      namespaceIndex={index}
-                      activeGraphId={activeGraph?.id}
-                      activeSubgraphId={activeSubgraph?.id}
-                      setNamespace={(ns) => {
-                        setNamespace(ns, false);
-                        close();
-                      }}
-                    />
-                </>
+                  <GraphCommandGroup
+                    key={`namespace-${index}`}
+                    isFiltering={isFiltering}
+                    namespace={wns}
+                    namespaceIndex={index}
+                    activeGraphId={activeGraph?.id}
+                    activeSubgraphId={activeSubgraph?.id}
+                    setNamespace={(ns) => {
+                      setNamespace(ns, false);
+                      close();
+                    }}
+                  />
               ))}
             </>
           ) : children}

studio/src/components/dashboard/workspace-provider.tsx

@@ -37,8 +37,9 @@ export function WorkspaceProvider({ children }: React.PropsWithChildren) {
       return;
     }
 
+    const actualNamespace = (router.query.namespace as string) || namespace;
     const currentNamespaces = data.namespaces.map((wns) => wns.name);
-    if (!currentNamespaces.some((ns) => ns.toLowerCase() === namespace.toLowerCase())) {
+    if (!currentNamespaces.some((ns) => ns.toLowerCase() === actualNamespace.toLowerCase())) {
       // The authenticated user doesn't have access to the namespace, pick between the `default` or the
       // first available namespace if the user doesn't have access to the `default` namespace
       const ns = currentNamespaces.find((n) => n === DEFAULT_NAMESPACE_NAME) || currentNamespaces[0];
@@ -50,12 +51,21 @@ export function WorkspaceProvider({ children }: React.PropsWithChildren) {
           namespace: ns,
         });
       }
-    } else if (!namespaceParam) {
-      applyParams({ namespace });
+    } else if (actualNamespace.toLowerCase() !== namespace.toLowerCase()) {
+      setNamespace(actualNamespace);
+      setStoredNamespace(actualNamespace);
     }
 
     setNamespaces(currentNamespaces);
-  }, [applyParams, data?.response?.code, data?.namespaces, namespace, namespaceParam, setStoredNamespace]);
+  }, [
+    applyParams,
+    data?.response?.code,
+    data?.namespaces,
+    router.query.namespace,
+    namespace,
+    namespaceParam,
+    setStoredNamespace,
+  ]);
 
   // Memoize context components
   const currentNamespace= useMemo(

studio/src/components/dashboard/workspace-selector.tsx

@@ -20,8 +20,8 @@ export function WorkspaceSelector({ children, truncateNamespace = true }: Worksp
 
   const [activeGraph, activeSubgraph] = useMemo(
     () => {
-      const routeSegment = router.pathname.split("/")[3]?.toLowerCase();
-      const currentSlug = router.query.slug as string;
+      const routeSegment = router.asPath.split("/")[3]?.toLowerCase();
+      const currentSlug = (router.query.slug as string)?.toLowerCase();
       return [
         routeSegment === "graph"
           ? namespace.graphs.find((graph) => graph.name.toLowerCase() === currentSlug)
@@ -33,13 +33,13 @@ export function WorkspaceSelector({ children, truncateNamespace = true }: Worksp
           : undefined,
       ];
     },
-    [namespace, router.pathname, router.query.slug, subgraphContext?.subgraph?.id],
+    [namespace, router.asPath, router.query.slug, subgraphContext?.subgraph?.id],
   );
 
   const isViewingGraphOrSubgraph = !!activeGraph || !!activeSubgraph;
   return (
     <div className={cn(
-      "flex justify-start items-center",
+      "flex justify-start items-center h-9 text-sm",
       isViewingGraphOrSubgraph && "gap-x-2",
     )}>
       <NamespaceSelector

studio/src/components/layout/dashboard-layout.tsx

@@ -34,6 +34,7 @@ import { UserGroupIcon } from "@heroicons/react/24/outline";
 import { useCheckUserAccess } from "@/hooks/use-check-user-access";
 import { useUser } from "@/hooks/use-user";
 import { useStarBannerDisabled } from "@/hooks/use-star-banner-disabled";
+import { useWorkspace } from "@/hooks/use-workspace";
 
 export const StarBanner = ({
   isDisabled,
@@ -115,6 +116,7 @@ export const DashboardLayout = ({ children }: LayoutProps) => {
   const organizationSlug = router.query.organizationSlug as string;
   const checkUserAccess = useCheckUserAccess();
   const [isStarBannerDisabled, setDisableStarBanner] = useStarBannerDisabled();
+  const { namespace } = useWorkspace();
 
   const isAdmin = checkUserAccess({ rolesToBe: ["organization-admin" ]});
   const isAdminOrDeveloper = checkUserAccess({ rolesToBe: ["organization-admin", "organization-developer"] });
@@ -134,21 +136,22 @@ export const DashboardLayout = ({ children }: LayoutProps) => {
 
   const links = useMemo(() => {
     const basePath = `/${user?.currentOrganization.slug || organizationSlug}`;
+    const nsQueryString = `?namespace=${encodeURIComponent(namespace.name)}`;
 
     const navigation: Partial<NavLink>[] = [
       {
         title: "Graphs",
-        href: basePath + "/graphs",
+        href: basePath + `/graphs${nsQueryString}`,
         icon: <PiGraphLight className="size-4" />,
       },
       {
         title: "Subgraphs",
-        href: basePath + "/subgraphs",
+        href: basePath + `/subgraphs${nsQueryString}`,
         icon: <Component2Icon className="size-4" />,
       },
       {
         title: "Feature Flags",
-        href: basePath + "/feature-flags",
+        href: basePath + `/feature-flags${nsQueryString}`,
         icon: <MdOutlineFeaturedPlayList className="size-4" />,
         matchExact: false,
         separator: !isAdminOrDeveloper,
@@ -159,17 +162,17 @@ export const DashboardLayout = ({ children }: LayoutProps) => {
       navigation.push(
         {
           title: "Policies",
-          href: basePath + "/policies",
+          href: basePath + `/policies${nsQueryString}`,
           icon: <MdOutlinePolicy className="size-4" />,
         },
         {
           title: "Check Extensions",
-          href: basePath + "/check-extensions",
+          href: basePath + `/check-extensions${nsQueryString}`,
           icon: <MdOutlineExtension className="size-4" />,
         },
         {
           title: "Cache Warmer",
-          href: basePath + "/cache-warmer",
+          href: basePath + `/cache-warmer${nsQueryString}`,
           icon: <FaGripfire className="size-4" />,
           separator: true,
         },
@@ -279,8 +282,10 @@ export const DashboardLayout = ({ children }: LayoutProps) => {
     organizationSlug,
     plans.data?.plans?.length,
     user?.currentOrganization.slug,
+    namespace,
     isAdmin,
     isAdminOrDeveloper,
+    isApiKeyManager,
     user?.invitations,
   ]);
 

studio/src/components/layout/sidenav.tsx

@@ -88,8 +88,8 @@ const MobileNav = () => {
 const Organizations = () => {
   const user = useContext(UserContext);
   const router = useRouter();
-  const currentPage = router.asPath.split("/")[2];
-  const isOrganizationRoot = router.asPath.split("/").length === 3;
+  const pathSegments = router.pathname.substring(1).split('/');
+  const isOrganizationRoot = pathSegments.length === 2;
 
   if (!user?.currentOrganization) return null;
 
@@ -112,18 +112,23 @@ const Organizations = () => {
               (org) => org.slug === orgSlug,
             );
             if (currentOrg) {
-              router.replace(
-                isOrganizationRoot && currentPage !== "invitations"
-                  ? `/${currentOrg.slug}/${currentPage}`
-                  : `/${currentOrg.slug}`,
-              );
+              router.replace({
+                pathname: isOrganizationRoot &&
+                  pathSegments[0]?.toLowerCase() !== "account" &&
+                  pathSegments[1]?.toLowerCase() !== "invitations"
+                  ? `/[organizationSlug]/${pathSegments[1]}`
+                  : `/[organizationSlug]`,
+                query: {
+                  organizationSlug: currentOrg.slug,
+                },
+              });
             }
           }}
         >
           {user?.organizations?.map(({ name, slug }) => {
             return (
-              <DropdownMenuRadioItem className="pl-2" key={slug} value={slug}>
-                {name}
+              <DropdownMenuRadioItem className="pl-2 gap-x-2" key={slug} value={slug}>
+                <span className="w-full">{name}</span>
               </DropdownMenuRadioItem>
             );
           })}
@@ -199,7 +204,7 @@ export const SideNav = (props: SideNavLayoutProps) => {
               const isCurrent =
                 item.href &&
                 isActive(
-                  encodeURI(item.href),
+                  encodeURI(item.href.split("?")[0]),
                   router.asPath.split("?")[0],
                   item.matchExact,
                 );