feat(ui): Add workspace landing page and multi-workspace support

Signed-off-by: Gage Krumbach <gkrumbach@gmail.com>
Co-Authored-By: Matt Prahl <mprahl@users.noreply.github.com>

Author: Gkrumbach07

mlflow/server/js/src/MlflowRouter.tsx

@@ -9,11 +9,14 @@ import {
   RouterProvider,
   Outlet,
   createLazyRouteElement,
+  useLocation,
+  useNavigate,
   useParams,
   usePageTitle,
 } from './common/utils/RoutingUtils';
 import { MlflowHeader } from './common/components/MlflowHeader';
 import { useDarkThemeContext } from './common/contexts/DarkThemeContext';
+import { useWorkspacesEnabled } from './common/utils/ServerFeaturesContext';
 
 // Route definition imports:
 import { getRouteDefs as getExperimentTrackingRouteDefs } from './experiment-tracking/route-defs';
@@ -24,6 +27,17 @@ import { useInitializeExperimentRunColors } from './experiment-tracking/componen
 import { MlflowSidebar } from './common/components/MlflowSidebar';
 import { AssistantProvider, AssistantRouteContextProvider } from './assistant';
 import { RootAssistantLayout } from './common/components/RootAssistantLayout';
+import { extractWorkspaceFromPathname, setActiveWorkspace, getActiveWorkspace } from './common/utils/WorkspaceUtils';
+import { prefixRoutePathWithWorkspace } from './common/utils/WorkspaceRouteUtils';
+import { useWorkspaces } from './common/hooks/useWorkspaces';
+
+type MlflowRouteDef = {
+  path?: string;
+  element?: React.ReactNode;
+  pageId?: string;
+  children?: MlflowRouteDef[];
+  [key: string]: unknown;
+};
 
 /**
  * This is root element for MLflow routes, containing app header.
@@ -89,9 +103,77 @@ const MlflowRootRoute = () => {
     </AssistantProvider>
   );
 };
+
+const WorkspaceRouterSync = ({ workspacesEnabled }: { workspacesEnabled: boolean }) => {
+  const location = useLocation();
+  const navigate = useNavigate({ bypassWorkspacePrefix: true });
+  const { workspaces, isLoading } = useWorkspaces(workspacesEnabled);
+
+  useEffect(() => {
+    if (!workspacesEnabled) {
+      setActiveWorkspace(null);
+      return;
+    }
+
+    const workspaceFromPath = extractWorkspaceFromPathname(location.pathname);
+    const activeWorkspace = getActiveWorkspace();
+
+    if (isLoading) {
+      return;
+    }
+
+    // Validate workspace from path
+    if (workspaceFromPath) {
+      const isValid = workspaces.some((w) => w.name === workspaceFromPath);
+      if (!isValid) {
+        setActiveWorkspace(null);
+        navigate('/', { replace: true });
+        return;
+      }
+      if (activeWorkspace !== workspaceFromPath) {
+        setActiveWorkspace(workspaceFromPath);
+      }
+      return;
+    }
+
+    // If not in a workspace path and not on root, redirect to selector
+    const isRootPath = location.pathname === '/' || location.pathname === '';
+    if (!isRootPath) {
+      setActiveWorkspace(null);
+      navigate('/', { replace: true });
+      return;
+    }
+  }, [location, navigate, workspacesEnabled, workspaces, isLoading]);
+
+  return null;
+};
+
+const WorkspaceAwareRootRoute = ({ workspacesEnabled }: { workspacesEnabled: boolean }) => (
+  <>
+    <WorkspaceRouterSync workspacesEnabled={workspacesEnabled} />
+    <MlflowRootRoute />
+  </>
+);
+
+const prependWorkspaceToRoutes = (routeDefs: MlflowRouteDef[], isChild = false): MlflowRouteDef[] =>
+  routeDefs.map((route) => {
+    // Only prepend workspace to child routes if they're absolute paths starting with /
+    // Otherwise keep them relative to preserve React Router's nested routing
+    const shouldPrependToPath = !isChild || (route.path && route.path.startsWith('/'));
+    const children = route.children ? prependWorkspaceToRoutes(route.children, true) : undefined;
+
+    return {
+      ...route,
+      path: shouldPrependToPath ? prefixRoutePathWithWorkspace(route.path) : route.path,
+      ...(children ? { children } : {}),
+    };
+  });
+
 export const MlflowRouter = () => {
+  const { workspacesEnabled, loading: featuresLoading } = useWorkspacesEnabled();
+
   // eslint-disable-next-line react-hooks/rules-of-hooks
-  const routes = useMemo(
+  const routes = useMemo<MlflowRouteDef[]>(
     () => [
       ...getExperimentTrackingRouteDefs(),
       ...getModelRegistryRouteDefs(),
@@ -100,19 +182,37 @@ export const MlflowRouter = () => {
     ],
     [],
   );
+
+  const workspaceRoutes = useMemo(
+    () => (workspacesEnabled ? prependWorkspaceToRoutes(routes) : []),
+    [routes, workspacesEnabled],
+  );
+  const combinedRoutes = useMemo(
+    () => (workspacesEnabled ? [...routes, ...workspaceRoutes] : routes),
+    [routes, workspaceRoutes, workspacesEnabled],
+  );
+
   // eslint-disable-next-line react-hooks/rules-of-hooks
   const hashRouter = useMemo(
     () =>
-      createHashRouter([
-        {
-          path: '/',
-          element: <MlflowRootRoute />,
-          children: routes,
-        },
-      ]),
-    [routes],
+      // Don't create router while still loading features
+      featuresLoading
+        ? null
+        : createHashRouter([
+            {
+              path: '/',
+              element: <WorkspaceAwareRootRoute workspacesEnabled={workspacesEnabled} />,
+              children: combinedRoutes,
+            },
+          ]),
+    [combinedRoutes, workspacesEnabled, featuresLoading],
   );
 
+  // Show loading skeleton while determining if workspaces are enabled
+  if (featuresLoading || !hashRouter) {
+    return <LegacySkeleton />;
+  }
+
   return (
     <React.Suspense fallback={<LegacySkeleton />}>
       <RouterProvider router={hashRouter} />

mlflow/server/js/src/app.tsx

@@ -20,15 +20,37 @@ import { MlflowRouter as MlflowRouter } from './MlflowRouter';
 import { useMLflowDarkTheme } from './common/hooks/useMLflowDarkTheme';
 import { DarkThemeProvider } from './common/contexts/DarkThemeContext';
 import { telemetryClient } from './telemetry';
+import { subscribeToWorkspaceChanges } from './common/utils/WorkspaceUtils';
+import { ServerFeaturesProvider, SERVER_FEATURES_QUERY_KEY } from './common/utils/ServerFeaturesContext';
 
 export function MLFlowRoot() {
   // eslint-disable-next-line react-hooks/rules-of-hooks
   const intl = useI18nInit();
+
   // eslint-disable-next-line react-hooks/rules-of-hooks
+  // Create clients once - we'll clear caches on workspace changes instead of recreating
   const apolloClient = useMemo(() => createApolloClient(), []);
   // eslint-disable-next-line react-hooks/rules-of-hooks
   const queryClient = useMemo(() => new QueryClient(), []);
 
+  // Clear caches when workspace changes instead of recreating clients
+  useEffect(() => {
+    const unsubscribe = subscribeToWorkspaceChanges(() => {
+      // Clear React Query cache except server features (they don't change with workspace)
+      queryClient.removeQueries({
+        predicate: (query) => {
+          // Keep server features cached
+          return query.queryKey[0] !== SERVER_FEATURES_QUERY_KEY[0];
+        },
+      });
+
+      // Clear Apollo cache
+      apolloClient.clearStore();
+    });
+
+    return unsubscribe;
+  }, [queryClient, apolloClient]);
+
   // eslint-disable-next-line react-hooks/rules-of-hooks
   const [isDarkTheme, setIsDarkTheme, MlflowThemeGlobalStyles] = useMLflowDarkTheme();
 
@@ -54,7 +76,9 @@ export function MLFlowRoot() {
               <MlflowThemeGlobalStyles />
               <DarkThemeProvider setIsDarkTheme={setIsDarkTheme}>
                 <QueryClientProvider client={queryClient}>
-                  <MlflowRouter />
+                  <ServerFeaturesProvider>
+                    <MlflowRouter />
+                  </ServerFeaturesProvider>
                 </QueryClientProvider>
               </DarkThemeProvider>
             </DesignSystemContainer>

mlflow/server/js/src/assistant/AssistantService.ts

@@ -3,7 +3,7 @@
  */
 
 import type { MessageRequest } from './types';
-import { getAjaxUrl } from '@mlflow/mlflow/src/common/utils/FetchUtils';
+import { getAjaxUrl, getDefaultHeaders } from '@mlflow/mlflow/src/common/utils/FetchUtils';
 
 const API_BASE = getAjaxUrl('ajax-api/3.0/mlflow/assistant');
 
@@ -28,10 +28,12 @@ export const sendMessageStream = async (
 ): Promise<void> => {
   try {
     // Step 1: POST the message to initiate processing
+    // eslint-disable-next-line no-restricted-globals -- See go/spog-fetch
     const response = await fetch(`${API_BASE}/message`, {
       method: 'POST',
       headers: {
         'Content-Type': 'application/json',
+        ...getDefaultHeaders(document.cookie),
       },
       body: JSON.stringify(request),
     });

mlflow/server/js/src/common/components/ErrorView.test.tsx

@@ -1,10 +1,33 @@
-import { describe, test, expect, it } from '@jest/globals';
+import { describe, test, expect, it, beforeAll, afterAll, jest } from '@jest/globals';
 import React from 'react';
 import { ErrorView } from './ErrorView';
 import { renderWithIntl, screen } from '@mlflow/mlflow/src/common/utils/TestUtils.react18';
 import { MemoryRouter } from '../utils/RoutingUtils';
+import { setActiveWorkspace } from '../utils/WorkspaceUtils';
+import { getWorkspacesEnabledSync } from '../utils/ServerFeaturesContext';
+
+jest.mock('../utils/ServerFeaturesContext', () => ({
+  ...jest.requireActual<typeof import('../utils/ServerFeaturesContext')>('../utils/ServerFeaturesContext'),
+  getWorkspacesEnabledSync: jest.fn(),
+}));
+
+const getWorkspacesEnabledSyncMock = jest.mocked(getWorkspacesEnabledSync);
+
+const TEST_WORKSPACE = 'test-workspace';
 
 describe('ErrorView', () => {
+  const workspacePrefixed = (path: string) => `/workspaces/${TEST_WORKSPACE}${path}`;
+
+  beforeAll(() => {
+    getWorkspacesEnabledSyncMock.mockReturnValue(true);
+    setActiveWorkspace(TEST_WORKSPACE);
+  });
+
+  afterAll(() => {
+    jest.restoreAllMocks();
+    setActiveWorkspace(null);
+  });
+
   test('should render 400', () => {
     renderWithIntl(
       <MemoryRouter>
@@ -26,7 +49,7 @@ describe('ErrorView', () => {
 
     const link = screen.getByRole('link');
     expect(link).toBeInTheDocument();
-    expect(link).toHaveAttribute('href', '/path/to');
+    expect(link).toHaveAttribute('href', workspacePrefixed('/path/to'));
   });
 
   it('should render 404', () => {
@@ -50,7 +73,7 @@ describe('ErrorView', () => {
 
     const link = screen.getByRole('link');
     expect(link).toBeInTheDocument();
-    expect(link).toHaveAttribute('href', '/path/to');
+    expect(link).toHaveAttribute('href', workspacePrefixed('/path/to'));
   });
 
   test('should render 404 with sub message', () => {
@@ -74,6 +97,6 @@ describe('ErrorView', () => {
 
     const link = screen.getByRole('link');
     expect(link).toBeInTheDocument();
-    expect(link).toHaveAttribute('href', '/path/to');
+    expect(link).toHaveAttribute('href', workspacePrefixed('/path/to'));
   });
 });

mlflow/server/js/src/common/components/MlflowHeader.tsx

@@ -4,6 +4,7 @@ import { HomePageDocsUrl, Version } from '../constants';
 import { DarkThemeSwitch } from '@mlflow/mlflow/src/common/components/DarkThemeSwitch';
 import { Button, MenuIcon, useDesignSystemTheme } from '@databricks/design-system';
 import { MlflowLogo } from './MlflowLogo';
+import { WorkspaceSelector } from './WorkspaceSelector';
 
 export const MlflowHeader = ({
   isDarkTheme = false,
@@ -66,6 +67,7 @@ export const MlflowHeader = ({
       </div>
       <div css={{ flex: 1 }} />
       <div css={{ display: 'flex', gap: theme.spacing.lg, alignItems: 'center' }}>
+        <WorkspaceSelector />
         <DarkThemeSwitch isDarkTheme={isDarkTheme} setIsDarkTheme={setIsDarkTheme} />
         <a href="https://github.com/mlflow/mlflow">GitHub</a>
         <a href={HomePageDocsUrl}>Docs</a>