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

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

@@ -6,6 +6,7 @@ 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 useForkRef from '@mui/utils/useForkRef';
 import { unstable_composeClasses as composeClasses } from '@mui/base';
 import { TreeItemContent } from './TreeItemContent';
 import { treeItemClasses, getTreeItemUtilityClass } from './treeItemClasses';
@@ -177,9 +178,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 +279,24 @@ 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;
 
   const item = (
     <TreeItemRoot
@@ -290,11 +306,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 +334,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,36 +51,44 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
     [state.focusedNodeId],
   );
 
-  const focusNode = useEventCallback((event: React.SyntheticEvent, nodeId: string | null) => {
-    if (nodeId) {
-      setFocusedNodeId(nodeId);
+  const isTreeViewFocused = () =>
+    !!rootRef.current && rootRef.current.contains(getActiveElement(ownerDocument(rootRef.current)));
+
+  const focusNode = useEventCallback(
+    (event: React.SyntheticEvent | null, nodeId: string | null) => {
+      if (nodeId) {
+        const node = instance.getNode(nodeId);
+        const nodeElement = document.getElementById(
+          instance.getTreeItemId(nodeId, node.idAttribute),
+        );
+        if (nodeElement) {
+          nodeElement.focus({ preventScroll: true });
+        }
+
+        setFocusedNodeId(nodeId);
 
-      if (params.onNodeFocus) {
-        params.onNodeFocus(event, nodeId);
+        if (params.onNodeFocus) {
+          params.onNodeFocus(event, nodeId);
+        }
+      } else {
+        setFocusedNodeId(null);
       }
-    }
-  });
+    },
+  );
 
-  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
-      ) {
-        return instance.getChildrenIds(null)[0];
-      }
-      return oldFocusedNodeId;
-    });
+    if (state.focusedNodeId === id) {
+      instance.focusNode(null, instance.getChildrenIds(null)[0]);
+    }
   });
 
   const createHandleFocus =
@@ -86,12 +120,6 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
       }
     };
 
-  const createHandleBlur =
-    (otherHandlers: EventHandlers) => (event: React.FocusEvent<HTMLUListElement>) => {
-      otherHandlers.onBlur?.(event);
-      setFocusedNodeId(null);
-    };
-
   const focusedNode = instance.getNode(state.focusedNodeId!);
   const activeDescendant = focusedNode
     ? instance.getTreeItemId(focusedNode.id, focusedNode.idAttribute)
@@ -100,7 +128,6 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
   return {
     getRootProps: (otherHandlers) => ({
       onFocus: createHandleFocus(otherHandlers),
-      onBlur: createHandleBlur(otherHandlers),
       'aria-activedescendant': activeDescendant ?? undefined,
     }),
   };

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;
-  focusNode: (event: React.SyntheticEvent, nodeId: string | null) => void;
-  focusRoot: () => void;
+  isTreeViewFocused: () => boolean;
+  canNodeBeTabbed: (nodeId: string) => boolean;
+  focusNode: (event: React.SyntheticEvent | null, nodeId: string | null) => void;
 }
 
 export interface UseTreeViewFocusParameters {
@@ -18,7 +19,7 @@ export interface UseTreeViewFocusParameters {
    * @param {string} nodeId The id of the node focused.
    * @param {string} value of the focused node.
    */
-  onNodeFocus?: (event: React.SyntheticEvent, nodeId: string) => void;
+  onNodeFocus?: (event: React.SyntheticEvent | null, nodeId: string) => void;
 }
 
 export type UseTreeViewFocusDefaultizedParameters = UseTreeViewFocusParameters;