Add workspace UI by Gkrumbach07 · Pull Request #19468 · mlflow/mlflow

Gkrumbach07 · 2025-12-17T17:47:58Z

Related Issues/PRs

#1464
#5844
Design document: https://docs.google.com/document/d/1IbFfceCmWV3knJfc0ninn58EXLVbHUh1meV_pknOM40/edit

What changes are proposed in this pull request?

This PR adds the frontend UI components for multi-workspace support in MLflow.

⚠️ Important: Backward Compatibility
The new workspace UI behavior is only activated when the server's /ajax-api/2.0/mlflow/server-features endpoint returns { "workspacesEnabled": true }. When workspaces are not enabled (default behavior), the UI operates exactly as before with no visible changes.

New Components:

WorkspaceSelector - Dropdown component with search and tooltips for selecting workspaces
WorkspacePermissionError - Error page displayed when accessing an unauthorized workspace
useWorkspaces hook - Custom hook for fetching available workspaces from the API
ServerFeaturesContext - Context provider for detecting server feature flags via /ajax-api/2.0/mlflow/server-features
WorkspaceUtils - Utilities for workspace state management and URL handling
WorkspaceRouteUtils - Utilities for prefixing routes with workspace paths

Core Changes:

MlflowRouter - Added workspace-aware routing with WorkspaceRouterSync component (only active when workspaces enabled)
app.tsx - Integrated ServerFeaturesProvider for feature detection
FetchUtils - Added X-MLFLOW-WORKSPACE header injection for all API requests (only when workspace is set)
graphql/client - Added workspace header support for GraphQL requests

Behavior when workspaces are enabled:

Frontend URLs include workspace: /#/workspaces/team-a/experiments
Backend API calls pass workspace via X-MLFLOW-WORKSPACE header (not URL path)
Workspace selector appears in header
Automatic redirect to default workspace when none is specified
Workspace state persisted to localStorage

Behavior when workspaces are disabled (default):

No workspace selector shown
No workspace prefix in URLs
No X-MLFLOW-WORKSPACE header sent
UI operates exactly as before

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

New tests added for:

ServerFeaturesContext.test.tsx - Tests feature detection and workspace enabled/disabled states
WorkspaceRouteUtils.test.ts - Tests route prefixing logic
ArtifactUtils.test.ts - Tests workspace header injection
graphql/client.test.ts - Tests GraphQL workspace header support
Updated existing tests to mock workspace context

Does this PR require documentation update?

No. You can skip the rest of this section.

Documentation will be added in a follow-up PR once the full workspace feature is complete.

Release Notes

Is this a user-facing change?

Yes. Give a description of this change to be included in the release notes for MLflow users.

Added multi-workspace UI support allowing users to switch between different workspaces via a dropdown selector in the MLflow header. This feature is only enabled when the MLflow server is configured with workspace support (determined by the /ajax-api/2.0/mlflow/server-features endpoint). When workspaces are not enabled, the UI behaves exactly as before.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

github-actions · 2025-12-17T17:48:11Z

@Gkrumbach07 Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

Gkrumbach07 · 2025-12-17T17:49:05Z

Will add screenshots / video in a sec

Gkrumbach07 · 2026-01-20T04:38:05Z

Here is the new landing page

note:

i removed the notion of navigating to a default/last used workspace automatically otherwise this workspace page would never get seen

mprahl · 2026-01-20T15:17:55Z

@Gkrumbach07 a few UX things:

In the "Create Experiment" modal, could you hide the "Artifact Location" form field in workspace mode? In workspace mode, the artifact location can only be set at server level or the workspace level.
Once I select a workspace, there's no easy way to navigate back to the workspace landing page. Should we include a "Workspaces" in the left-nav to go back to the landing page?
The "Create Workspace" modal should have a form field to set default_artifact_root (replaces the "Artifact Location" customization on experiments)
On the workspace landing page, I think the "workspace" selector in the navigation bar shouldn't show a select workspace. Right now, it shows the "last used"
The / redirects to /#/workspaces. Would it be better to just have / be the workspace landing page to not have quick refresh.
Should we have a way to edit the workspace description and default_artifact_root similar to being able to edit an experiment?

mprahl · 2026-01-20T18:56:51Z

@Gkrumbach07 one other thing I noticed is if there is an API error when creating a workspace through the UI (e.g. if the workspace provider doesn't support workspace creation), there is no error notification.

Gkrumbach07 · 2026-01-20T22:32:15Z

Fixed create experiment not showing error
hide artifact location in experiment create modal when in workspace mode
artifact root added to workspace create modal and table
able to edit description and artifact root via workspace page table edit buttons inline
route now uses root instead of /workspaces for workspace landing page
can navigate back to workspace landing page via side nav

mprahl · 2026-01-22T14:58:00Z

@Gkrumbach07 manually testing seems good. I had GPT 5.2 Codex review it:

Findings

High — Server feature detection bypasses default headers (cookie-derived headers + Authorization), so OAuth/K8s deployments that rely on these headers can get a 401/403 on /server-features and silently disable workspaces for the session. The fetch is raw and never uses getDefaultHeaders, which is where those headers are injected.

async function fetchServerFeatures(): Promise<ServerFeatures> {
  try {
    const response = await fetch(getAjaxUrl(SERVER_FEATURES_ENDPOINT));

    if (response.status === 404) {
      return { workspacesEnabled: Boolean(process.env['MLFLOW_ENABLE_WORKSPACES']) };
    }

export const getDefaultHeaders = (cookieStr: any) => {
  const cookieHeaders = getDefaultHeadersFromCookies(cookieStr);

  // Forward Authorization header for OAuth/Kubernetes integration
  const authHeader = typeof sessionStorage !== 'undefined' ? sessionStorage.getItem('mlflow-auth-header') : null;

  // getActiveWorkspace() returns null if workspaces feature is not enabled
  const workspace = getActiveWorkspace();

  return {
    ...cookieHeaders,
    ...(authHeader ? { Authorization: authHeader } : {}),
    ...(workspace ? { 'X-MLFLOW-WORKSPACE': workspace } : {}),
  };
};

Medium — Two different redirect policies run for “workspace not found”: WorkspaceRouterSync forces /, while WorkspaceSelector redirects to a fallback workspace. Depending on effect ordering, users can land on the selector or a workspace with a flicker. This is non-deterministic and hard to reason about.

  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;
      }

  useEffect(() => {
    if (!workspacesEnabled || isLoading || workspaces.length === 0 || !currentWorkspace) {
      return;
    }

    const currentWorkspaceExists = workspaces.some((ws) => ws.name === currentWorkspace);

    if (!currentWorkspaceExists) {
      // Get first available workspace as fallback (or 'default' if available)
      const fallbackWorkspace = workspaces.find((ws) => ws.name === 'default') || workspaces[0];
      if (fallbackWorkspace) {
        setActiveWorkspace(fallbackWorkspace.name);
        const currentSection = getNavigationSection(location.pathname);
        navigate(`/workspaces/${encodeURIComponent(fallbackWorkspace.name)}${currentSection}`);
      }
    }
  }, [workspacesEnabled, isLoading, workspaces, currentWorkspace, location.pathname, navigate]);

Low — extractWorkspaceFromPathname can throw on malformed percent-encoding (decodeURIComponent). That would crash rendering for a bad URL instead of gracefully treating it as “no workspace.”

export const extractWorkspaceFromPathname = (pathname: string): string | null => {
  if (!pathname || !pathname.startsWith(WORKSPACE_PREFIX)) {
    return null;
  }
  const segments = pathname.split('/');
  if (segments.length < 3 || !segments[2]) {
    return null;
  }
  const workspaceName = decodeURIComponent(segments[2]);

  // Validate workspace name format
  if (!WORKSPACE_NAME_PATTERN.test(workspaceName)) {
    return null;
  }

Gkrumbach07 · 2026-01-22T15:37:15Z

@mprahl I have addressed your review

mlflow/server/js/src/common/utils/ServerFeaturesContext.tsx

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

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

mlflow/server/js/src/common/utils/FetchUtils.ts

mprahl · 2026-01-22T22:08:37Z

/lgtm

github-actions · 2026-01-23T02:14:39Z

Documentation preview for f883189 is available at:

https://pr-19468--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

mlflow/server/js/src/workspaces/components/WorkspaceSelector.tsx

hubertzub-db · 2026-02-02T07:31:31Z

mlflow/server/js/src/workspaces/components/WorkspaceSelector.tsx

+  const currentWorkspace = extractWorkspaceFromSearchParams(searchParams);
+
+  // Fetch workspaces using custom hook
+  const { workspaces, isLoading, isError, refetch } = useWorkspaces(workspacesEnabled);


If the workspaces are persisted in local storage anyway, should we use that as a cache and load workspaces only after the dropdown is opened? This could save some API calls

useWorkspaces(workspacesEnabled) uses react query and therefore it is cached there. I am hesitant to lean on local storage for cache as that has no internal invalidation so we have to then deal with

when do we refresh the dropdown

what if a user looses access, one is deleted outside their browser, ...

mlflow/server/js/src/experiment-tracking/components/ExperimentListView.tsx

mlflow/server/js/src/experiment-tracking/pages/experiment-page-tabs/ExperimentPageTabs.test.tsx

hubertzub-db · 2026-02-02T07:40:03Z

mlflow/server/js/src/home/components/LogTracesDrawer.tsx


-const CONFIGURE_EXPERIMENT_SNIPPET = `import mlflow
+const getConfigureExperimentSnippet = (workspace?: string | null) => {
+  const workspaceLine = workspace ? `mlflow.set_workspace("${workspace}")\n` : '';


@B-Step62 can you confirm if it's ok to introduce those changes to snippets already?

mlflow/server/js/src/shared/web-shared/genai-traces-table/utils/FetchUtils.ts

...ing/components/experiment-logged-models/ExperimentLoggedModelListPageColumnSelector.test.tsx

mlflow/server/js/src/common/utils/WorkspaceUtils.ts

mlflow/server/js/src/common/utils/FetchUtils.ts

...rver/js/src/shared/web-shared/model-trace-explorer/hooks/useUpdateTraceTagsMutation.test.tsx

Gkrumbach07 · 2026-02-03T19:52:07Z

I addressed comments

DaoDaoNoCode · 2026-02-04T21:21:01Z

Hi @Gkrumbach07
On https://github.com/mlflow/mlflow/blob/master/mlflow/server/js/src/experiment-tracking/components/experiment-logged-models/hooks/useExperimentLoggedModelListPageMode.tsx#L18-L20. You might want to change it to

setParams((prevParams) => {
  const newParams = new URLSearchParams(prevParams);
  newParams.set(VIEW_MODE_QUERY_PARAM, mode);
  return newParams;
});

in case the workspace is removed from the query params when users switch to the Chart View on the models tab.

hubertzub-db

Took another pass, this looks good! Thanks a lot for redoing everything to use query params instead of a prefix!

Left a few minor questions
A major one: what's the exact issue/error that made you add <TooltipProvider> to many tests in the app? In general, <TooltipProvider> should be automatically added by <DesignSystemProvider> - if it doesn't, it indicates some deeper problem in design system we'd like to investigate
CC @daniellok-db would you like to take a quick look at this PR as well?

hubertzub-db · 2026-02-07T06:09:06Z

mlflow/server/js/src/common/utils/FetchUtils.ts

-  // eslint-disable-next-line no-restricted-globals
-  const fetchFn = fetch;


can we keep this fetchFn abstraction? this helps sync with the other codebase

hubertzub-db · 2026-02-07T06:13:33Z

...xperiment-tracking/components/experiment-page/hooks/useGetExperimentPageActiveTabByRoute.tsx

 const getTabNameFromRoutePath = (pathname: string) =>
  ExperimentPageRoutePathToTabNameMap
    // Find the first route path that matches the given pathname
+    // Note: With query param-based workspace routing, pathname no longer contains workspace prefix


I think this comment is not needed now

hubertzub-db · 2026-02-07T06:32:28Z

mlflow/server/js/src/workspaces/utils/WorkspaceUtils.ts

+
+export const setActiveWorkspace = (workspace: string | null) => {
+  activeWorkspace = workspace;
+  listeners.forEach((listener) => listener(activeWorkspace));


Is it absolutely necessary for workspace switching to be a dynamic action that updates stuff using this pubsub pattern, instead of just hard reloading the entire UI app with a new query parameter? It may sound barbaric, but it could make it a way simpler and less prone to bugs. WDYT?

So the ONLY reason for the pub sub is because the react query cache is not under the router. I can move things around so that is the case and no need for pub sub.

Otherwise I guess we could hard refresh, but it makes a very jumpy UI every time a workflow changes

i did the refresh pattern

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

- Remove redundant TooltipProvider from all test files (DesignSystemProvider already includes RadixTooltipProvider internally) - Refactor workspace switching from pub/sub to hard reload for simplicity and to eliminate redirect race conditions - Restore fetchFn abstraction in FetchUtils.ts for codebase sync - Remove stale workspace routing comment - Fix upstream bug in useMlflowTraces.test.tsx (fetchFn -> fetchAPI) Signed-off-by: Gage Krumbach <gkrumbach@gmail.com>

Gkrumbach07 · 2026-02-09T23:09:24Z

Addressed all feedback from the latest review:

TooltipProvider in tests: Investigated and confirmed DesignSystemProvider already includes RadixTooltipProvider internally (they resolve to the same @radix-ui/react-tooltip). The extra wrappers were unnecessary so i removed them all
fetchFn abstraction — Restored const fetchFn = fetch in FetchUtils.ts to keep in sync with the internal codebase.
Stale comment in useGetExperimentPageActiveTabByRoute.tsx -> Removed.
Pub/sub → hard reload -> Agreed this is the better trade off. Removed the listeners/subscribeToWorkspaceChanges pattern entirely. Workspace switching now does window.location.hash + window.location.reload(), which naturally handles cache invalidation

hubertzub-db

@Gkrumbach07 thanks for all the changes and fixes!
Stamping as it looks good in the current form, I'll defer to @daniellok-db for further review

daniellok-db

Overall LGTM but there seems to be some bug with links to the chat sessions page, e.g.

mlflow/mlflow/server/js/src/shared/web-shared/model-trace-explorer/ModelTraceHeaderSessionIdTag.tsx

Line 76 in c551c52

My suspicion is that this is due to the packages in src/shared/web-shared having their own individual copies of RoutingUtils. Unfortunately this is a quirk of our setup where the UI has some upstream dependencies (we cannot import from @mlflow/mlflow from within that folder as it will cause conflicts). See:

mlflow/mlflow/server/js/src/shared/web-shared/genai-traces-table/utils/RoutingUtils.tsx

Line 4 in c551c52

    
            * Duplicate of mlflow/server/js/src/common/utils/RoutingUtils.tsx, because unfortunately

mlflow/mlflow/server/js/src/shared/web-shared/model-trace-explorer/RoutingUtils.tsx

Line 4 in c551c52

/* eslint-disable no-restricted-imports */

Could you copy paste the changes you made over to those files too?

Doesn't have to block the merge, I'll leave it up to @B-Step62 but we should definitely follow up at least in a future PR

github-actions bot added area/uiux Front-end, user experience, plotting, JavaScript, JavaScript dev server rn/feature Mention under Features in Changelogs. labels Dec 17, 2025

Gkrumbach07 force-pushed the workspace-ui-pr branch from f9479a0 to 78b80dd Compare December 17, 2025 17:53

mprahl mentioned this pull request Jan 9, 2026

Support for Multi-Tenancy #5844

Open

Gkrumbach07 force-pushed the workspace-ui-pr branch 2 times, most recently from 32988a0 to 59886b1 Compare January 12, 2026 16:43

B-Step62 force-pushed the orgnization-support branch from 0730f36 to 6cf60c4 Compare January 19, 2026 06:02

Gkrumbach07 force-pushed the workspace-ui-pr branch 3 times, most recently from b5247eb to 6ccc89e Compare January 20, 2026 04:37

Gkrumbach07 force-pushed the workspace-ui-pr branch from 6ccc89e to 47219d4 Compare January 20, 2026 15:09

Gkrumbach07 force-pushed the workspace-ui-pr branch from 47219d4 to 948728e Compare January 20, 2026 15:58

Gkrumbach07 force-pushed the workspace-ui-pr branch from 948728e to 3be539c Compare January 20, 2026 22:26

Gkrumbach07 force-pushed the workspace-ui-pr branch 2 times, most recently from 46bb816 to 56b2e5b Compare January 21, 2026 17:21

Gkrumbach07 force-pushed the workspace-ui-pr branch 2 times, most recently from d7be8ea to 1fded7b Compare January 22, 2026 15:36