[TreeView] Set focus on the focused TreeItem instead of the Tree View

packages/x-tree-view/src/TreeItem/TreeItem.tsx

@@ -6,13 +6,16 @@ import { resolveComponentProps, useSlotProps } from '@mui/base/utils';
 import { alpha, styled, useThemeProps } from '@mui/material/styles';
 import unsupportedProp from '@mui/utils/unsupportedProp';
 import elementTypeAcceptingRef from '@mui/utils/elementTypeAcceptingRef';
+import ownerDocument from '@mui/utils/ownerDocument';
+import useForkRef from '@mui/utils/useForkRef';
 import { unstable_composeClasses as composeClasses } from '@mui/base';
 import { TreeItemContent } from './TreeItemContent';
 import { treeItemClasses, getTreeItemUtilityClass } from './treeItemClasses';
 import { TreeItemOwnerState, TreeItemProps } from './TreeItem.types';
 import { useTreeViewContext } from '../internals/TreeViewProvider/useTreeViewContext';
 import { DefaultTreeViewPlugins } from '../internals/plugins';
 import { TreeViewCollapseIcon, TreeViewExpandIcon } from '../icons';
+import { getActiveElement } from '../internals/utils/utils';
 
 const useUtilityClasses = (ownerState: TreeItemOwnerState) => {
   const { classes } = ownerState;
@@ -77,6 +80,7 @@ const StyledTreeItemContent = styled(TreeItemContent, {
   [`&.${treeItemClasses.disabled}`]: {
     opacity: (theme.vars || theme).palette.action.disabledOpacity,
     backgroundColor: 'transparent',
+    pointerEvents: 'none',
   },
   [`&.${treeItemClasses.focused}`]: {
     backgroundColor: (theme.vars || theme).palette.action.focus,
@@ -177,9 +181,18 @@ export const TreeItem = React.forwardRef(function TreeItem(
     onMouseDown,
     TransitionComponent = Collapse,
     TransitionProps,
+    onFocus,
+    onBlur,
+    onKeyDown,
     ...other
   } = props;
 
+  const handleContentRef = React.useRef<HTMLDivElement>(null);
+  const contentRef = useForkRef(ContentProps?.ref, handleContentRef);
+
+  const handleRootRef = React.useRef<HTMLLIElement>(null);
+  const rootRef = useForkRef(ref, handleRootRef);
+
   const slots = {
     expandIcon: inSlots?.expandIcon ?? contextIcons.slots.expandIcon ?? TreeViewExpandIcon,
     collapseIcon: inSlots?.collapseIcon ?? contextIcons.slots.collapseIcon ?? TreeViewCollapseIcon,
@@ -269,18 +282,45 @@ export const TreeItem = React.forwardRef(function TreeItem(
   }
 
   function handleFocus(event: React.FocusEvent<HTMLLIElement>) {
-    // DOM focus stays on the tree which manages focus with aria-activedescendant
-    if (event.target === event.currentTarget) {
-      instance.focusRoot();
-    }
-
     const canBeFocused = !disabled || disabledItemsFocusable;
     if (!focused && canBeFocused && event.currentTarget === event.target) {
       instance.focusNode(event, nodeId);
     }
   }
 
+  function handleBlur(event: React.FocusEvent<HTMLLIElement>) {
+    onBlur?.(event);
+    instance.focusNode(event, null);
+  }
+
+  const handleKeyDown = (event: React.KeyboardEvent<HTMLLIElement>) => {
+    onKeyDown?.(event);
+    instance.handleItemKeyDown(event, nodeId);
+  };
+
   const idAttribute = instance.getTreeItemId(nodeId, id);
+  const tabIndex = instance.canNodeBeTabbed(nodeId) ? 0 : -1;
+
+  React.useEffect(() => {
+    if (
+      !focused ||
+      !handleContentRef.current ||
+      !handleRootRef.current ||
+      !instance.isTreeViewFocused()
+    ) {
+      return;
+    }
+
+    const activeElement = getActiveElement(ownerDocument(handleContentRef.current));
+    if (!handleContentRef.current.contains(activeElement)) {
+      handleRootRef.current.focus({ preventScroll: true });
+    }
+  }, [focused]); // eslint-disable-line react-hooks/exhaustive-deps
+
+  const focusedRef = React.useRef(focused);
+  React.useEffect(() => {
+    focusedRef.current = focused;
+  });
 
   const item = (
     <TreeItemRoot
@@ -290,11 +330,13 @@ export const TreeItem = React.forwardRef(function TreeItem(
       aria-selected={ariaSelected}
       aria-disabled={disabled || undefined}
       id={idAttribute}
-      tabIndex={-1}
       {...other}
       ownerState={ownerState}
+      tabIndex={tabIndex}
       onFocus={handleFocus}
-      ref={ref}
+      onBlur={handleBlur}
+      onKeyDown={handleKeyDown}
+      ref={rootRef}
     >
       <StyledTreeItemContent
         as={ContentComponent}
@@ -316,6 +358,7 @@ export const TreeItem = React.forwardRef(function TreeItem(
         displayIcon={displayIcon}
         ownerState={ownerState}
         {...ContentProps}
+        ref={contentRef}
       />
       {children && (
         <TreeItemGroup

packages/x-tree-view/src/TreeItem/TreeItem.types.ts

@@ -61,7 +61,7 @@ export interface TreeItemProps extends Omit<React.HTMLAttributes<HTMLLIElement>,
   /**
    * Props applied to ContentComponent.
    */
-  ContentProps?: React.HTMLAttributes<HTMLElement>;
+  ContentProps?: React.HTMLAttributes<HTMLElement> & { ref?: React.Ref<HTMLDivElement> };
   /**
    * If `true`, the node is disabled.
    * @default false

packages/x-tree-view/src/internals/plugins/useTreeViewExpansion/useTreeViewExpansion.ts

@@ -51,7 +51,7 @@ export const useTreeViewExpansion: TreeViewPlugin<UseTreeViewExpansionSignature>
     },
   );
 
-  const expandAllSiblings = (event: React.KeyboardEvent<HTMLUListElement>, nodeId: string) => {
+  const expandAllSiblings = (event: React.KeyboardEvent, nodeId: string) => {
     const node = instance.getNode(nodeId);
     const siblings = instance.getChildrenIds(node.parentId);
 

packages/x-tree-view/src/internals/plugins/useTreeViewExpansion/useTreeViewExpansion.types.ts

@@ -6,7 +6,7 @@ export interface UseTreeViewExpansionInstance {
   isNodeExpanded: (nodeId: string) => boolean;
   isNodeExpandable: (nodeId: string) => boolean;
   toggleNodeExpansion: (event: React.SyntheticEvent, value: string) => void;
-  expandAllSiblings: (event: React.KeyboardEvent<HTMLUListElement>, nodeId: string) => void;
+  expandAllSiblings: (event: React.KeyboardEvent, nodeId: string) => void;
 }
 
 export interface UseTreeViewExpansionParameters {

packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.ts

@@ -2,10 +2,34 @@ import * as React from 'react';
 import useEventCallback from '@mui/utils/useEventCallback';
 import { EventHandlers } from '@mui/base/utils';
 import ownerDocument from '@mui/utils/ownerDocument';
-import { TreeViewPlugin } from '../../models';
+import { TreeViewPlugin, TreeViewUsedInstance } from '../../models';
 import { populateInstance } from '../../useTreeView/useTreeView.utils';
 import { UseTreeViewFocusSignature } from './useTreeViewFocus.types';
 import { useInstanceEventHandler } from '../../hooks/useInstanceEventHandler';
+import { getActiveElement } from '../../utils/utils';
+
+const useTabbableNodeId = (
+  instance: TreeViewUsedInstance<UseTreeViewFocusSignature>,
+  selectedNodes: string | string[] | null,
+) => {
+  const isNodeVisible = (nodeId: string) => {
+    const node = instance.getNode(nodeId);
+    return node && (node.parentId == null || instance.isNodeExpanded(node.parentId));
+  };
+
+  let tabbableNodeId: string | null | undefined;
+  if (Array.isArray(selectedNodes)) {
+    tabbableNodeId = selectedNodes.find(isNodeVisible);
+  } else if (selectedNodes != null && isNodeVisible(selectedNodes)) {
+    tabbableNodeId = selectedNodes;
+  }
+
+  if (tabbableNodeId == null) {
+    tabbableNodeId = instance.getNavigableChildrenIds(null)[0];
+  }
+
+  return tabbableNodeId;
+};
 
 export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
   instance,
@@ -15,6 +39,8 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
   models,
   rootRef,
 }) => {
+  const tabbableNodeId = useTabbableNodeId(instance, models.selectedNodes.value);
+
   const setFocusedNodeId = useEventCallback((nodeId: React.SetStateAction<string | null>) => {
     const cleanNodeId = typeof nodeId === 'function' ? nodeId(state.focusedNodeId) : nodeId;
     setState((prevState) => ({ ...prevState, focusedNodeId: cleanNodeId }));
@@ -25,6 +51,9 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
     [state.focusedNodeId],
   );
 
+  const isTreeViewFocused = () =>
+    !!rootRef.current && rootRef.current.contains(getActiveElement(ownerDocument(rootRef.current)));
+
   const focusNode = useEventCallback((event: React.SyntheticEvent, nodeId: string | null) => {
     if (nodeId) {
       setFocusedNodeId(nodeId);
@@ -35,22 +64,18 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
     }
   });
 
-  const focusRoot = useEventCallback(() => {
-    rootRef.current?.focus({ preventScroll: true });
-  });
+  const canNodeBeTabbed = (nodeId: string) => nodeId === tabbableNodeId;
 
   populateInstance<UseTreeViewFocusSignature>(instance, {
     isNodeFocused,
+    isTreeViewFocused,
+    canNodeBeTabbed,
     focusNode,
-    focusRoot,
   });
 
   useInstanceEventHandler(instance, 'removeNode', ({ id }) => {
     setFocusedNodeId((oldFocusedNodeId) => {
-      if (
-        oldFocusedNodeId === id &&
-        rootRef.current === ownerDocument(rootRef.current).activeElement
-      ) {
+      if (oldFocusedNodeId === id) {
         return instance.getChildrenIds(null)[0];
       }
       return oldFocusedNodeId;
@@ -88,8 +113,13 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
 
   const createHandleBlur =
     (otherHandlers: EventHandlers) => (event: React.FocusEvent<HTMLUListElement>) => {
+      window.setTimeout(() => {
+        const activeElement = getActiveElement(ownerDocument(rootRef.current));
+        if (rootRef.current && !rootRef.current.contains(activeElement)) {
+          setFocusedNodeId(null);
+        }
+      });
       otherHandlers.onBlur?.(event);
-      setFocusedNodeId(null);
     };
 
   const focusedNode = instance.getNode(state.focusedNodeId!);

packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.types.ts

@@ -7,8 +7,9 @@ import { UseTreeViewExpansionSignature } from '../useTreeViewExpansion';
 
 export interface UseTreeViewFocusInstance {
   isNodeFocused: (nodeId: string) => boolean;
+  isTreeViewFocused: () => boolean;
+  canNodeBeTabbed: (nodeId: string) => boolean;
   focusNode: (event: React.SyntheticEvent, nodeId: string | null) => void;
-  focusRoot: () => void;
 }
 
 export interface UseTreeViewFocusParameters {