fix: tree component focus loss issue

src/components/Tree/Tree.constants.ts

@@ -7,6 +7,7 @@ const DEFAULTS = {
   SELECTABLE_NODES: 'leafOnly' as const,
   IS_REQUIRED: false,
   NODE_ID_PREFIX: 'md-tree-node',
+  SHOULD_NODE_FOCUS_BE_INSET: true,
 };
 
 const STYLE = {

src/components/Tree/Tree.hooks.tsx

@@ -9,13 +9,18 @@ import { NODE_ID_ATTRIBUTE_NAME, NODE_ID_DATA_NAME } from '../TreeNodeBase/TreeN
 /**
  * Handle DOM changes for virtual tree
  *
+ * This hook manage 2 use cases:
+ * 1) When Virtual Tree removes active node, the hook adds a cloned node to not lose focus.
+ * 2) When the active node added back to the DOM, the hook invoke focus on the element and trigger
+ *    key event if it happened on the cloned element.
+ *
  * @param props
  * @internal
  */
 export const useVirtualTreeNavigation = ({
   virtualTreeConnector,
   treeRef,
-  activeNodeId,
+  activeNodeIdRef,
 }: UseVirtualTreeNavigationProps): void => {
   // Handle DOM changes for virtual tree
   useEffect(() => {
@@ -36,12 +41,13 @@ export const useVirtualTreeNavigation = ({
       );
 
     // Filter element by active node id
-    const getByNodeId = (node: HTMLElement) => node.dataset[NODE_ID_DATA_NAME] === activeNodeId;
+    const getByNodeId = (node: HTMLElement) =>
+      node.dataset[NODE_ID_DATA_NAME] === activeNodeIdRef.current;
 
     // Mutation observer to handle the focus change
     const mutationHandler: MutationCallback = (mutationList) => {
       for (const mutation of mutationList) {
-        // Handle removed active node
+        // Active node moved out of view and removed
         const removed = Array.from(mutation.removedNodes).find(getByNodeId) as HTMLElement;
         if (removed) {
           // Clone the active node and make some modifications
@@ -58,11 +64,11 @@ export const useVirtualTreeNavigation = ({
 
           // add focus and keydown event listeners
           clonedNode.addEventListener('focus', () => {
-            virtualTreeConnector.scrollToNode(activeNodeId);
+            virtualTreeConnector.scrollToNode(activeNodeIdRef.current);
           });
           clonedNode.addEventListener('keydown', (evt) => {
             if (TREE_NAVIGATION_KEYS.includes(evt.key) || (evt.key === 'Tab' && !evt.shiftKey)) {
-              virtualTreeConnector.scrollToNode(activeNodeId);
+              virtualTreeConnector.scrollToNode(activeNodeIdRef.current);
               keyDownEvent = new KeyboardEvent('keydown', evt);
               evt.stopPropagation();
               evt.preventDefault();
@@ -71,7 +77,7 @@ export const useVirtualTreeNavigation = ({
           return;
         }
 
-        // Handle adding back the focused node
+        // Active node moved back into the view and it added back to the DOM
         const added = Array.from(mutation.addedNodes).find(getByNodeId) as HTMLElement;
         if (added) {
           cleanUp();
@@ -103,5 +109,5 @@ export const useVirtualTreeNavigation = ({
       observer.disconnect();
       cleanUp();
     };
-  }, [virtualTreeConnector, activeNodeId, treeRef.current]);
+  }, [virtualTreeConnector, treeRef.current]);
 };

src/components/Tree/Tree.test.tsx

@@ -485,9 +485,9 @@ describe('<Tree />', () => {
       node.remove();
 
       await waitFor(() => {
-        const clonedNode = getByText('1');
+        const clonedNode = getByTestId('1');
         expect(clonedNode).toHaveFocus();
-        // get a dedicated clas name
+        // get a dedicated class name
         expect(clonedNode).toHaveClass(TREE_CONSTANTS.STYLE.clonedVirtualTreeNode);
         // nodeid unsetted
         expect(clonedNode.dataset.nodeid).toBe(undefined);
@@ -533,7 +533,7 @@ describe('<Tree />', () => {
       expect(scrollToNode).toHaveBeenNthCalledWith(1, '1');
     });
 
-    it('should not call scrollToNode when cloned node exists and user press Shift+Tab', async () => {
+    it('should call scrollToNode when the focus moves back to the tree and the active node is not in the DOM', async () => {
       expect.assertions(2);
       const scrollToNode = jest.fn();
       const { getByText } = await renderTreeAndRemoveNode({
@@ -548,6 +548,30 @@ describe('<Tree />', () => {
       expect(scrollToNode).toHaveBeenCalledTimes(0);
     });
 
+    it.each`
+      keyPressed
+      ${'ArrowUp'}
+      ${'ArrowDown'}
+      ${'ArrowLeft'}
+      ${'ArrowRight'}
+    `(
+      'should call scrollToNode when active node is not in the DOM and $kepPressed key pressed',
+      async ({ keyPressed }) => {
+        expect.assertions(2);
+        const scrollToNode = jest.fn();
+        const { getByText } = await renderTreeAndRemoveNode({
+          scrollToNode,
+          setNodeOpen: jest.fn(),
+        });
+
+        await userEvent.keyboard(`{${keyPressed}}`);
+
+        // Only adding back Node 1 will remove the cloned node
+        expect(getByText('1')).not.toBeNull();
+        expect(scrollToNode).toHaveBeenCalledTimes(1);
+      }
+    );
+
     it('should remove the cloned node when the real node added back', async () => {
       expect.assertions(5);
       const scrollToNode = jest.fn();

src/components/Tree/Tree.tsx

@@ -14,11 +14,11 @@ import { TreeIdNodeMap, Props, TreeContextValue, TreeNodeId, TreeRefObject } fro
 import './Tree.style.scss';
 import {
   convertNestedTree2MappedTree,
+  isActiveNodeInDOM,
   getFistActiveNode,
   getNextActiveNode,
   getNodeDOMId,
   getTreeRootId,
-  isActiveNodeInDOM,
   migrateTreeState,
   toggleTreeNodeRecord,
   TreeContext,
@@ -35,7 +35,7 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
     id,
     style,
     children,
-    shouldNodeFocusBeInset,
+    shouldNodeFocusBeInset = DEFAULTS.SHOULD_NODE_FOCUS_BE_INSET,
     treeStructure,
     isRenderedFlat = DEFAULTS.IS_RENDERED_FLAT,
     excludeTreeRoot = DEFAULTS.EXCLUDE_TREE_ROOT,
@@ -60,10 +60,11 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
     isRequired,
   });
   const [tree, setTree] = useState<TreeIdNodeMap>(convertNestedTree2MappedTree(treeStructure));
-  const [activeNodeId, setActiveNodeId] = useState<TreeNodeId | undefined>(
+  const [activeNode, setActiveNode] = useState<TreeNodeId | undefined>(
     getFistActiveNode(tree, excludeTreeRoot)
   );
   const [isFocusWithin, setIsFocusWithin] = useState(false);
+  const activeNodeIdRef = useRef<TreeNodeId>(activeNode);
 
   const previousTree = usePrevious(tree);
 
@@ -72,7 +73,7 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
     migrateTreeState(previousTree, newTree);
     setTree(newTree);
     // Find the closest node to the last active node in the new tree
-    let newActiveNodeId = activeNodeId;
+    let newActiveNodeId = activeNode;
     while (newActiveNodeId) {
       if (newTree.has(newActiveNodeId) && !newTree.get(newActiveNodeId).isHidden) break;
       newActiveNodeId = previousTree.get(newActiveNodeId)?.parent;
@@ -86,17 +87,35 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
 
   const isVirtualTree = virtualTreeConnector !== undefined;
 
+  const scrollToVTreeNode = useCallback(
+    (nodeId: TreeNodeId) => {
+      if (isVirtualTree && virtualTreeConnector && !isActiveNodeInDOM(treeRef, nodeId)) {
+        virtualTreeConnector.scrollToNode?.(nodeId);
+      }
+    },
+    [isVirtualTree, virtualTreeConnector]
+  );
+
   // Handle DOM changes for virtual tree
-  useVirtualTreeNavigation({ virtualTreeConnector, treeRef: treeRef, activeNodeId });
+  useVirtualTreeNavigation({ virtualTreeConnector, treeRef: treeRef, activeNodeIdRef });
+
+  const setActiveNodeId = useCallback(
+    (newNodeId) => {
+      if (activeNodeIdRef.current !== newNodeId) {
+        activeNodeIdRef.current = newNodeId;
+        setActiveNode(newNodeId);
+        scrollToVTreeNode(newNodeId);
+      }
+    },
+    [isVirtualTree, setActiveNode]
+  );
 
   const toggleTreeNode = useCallback(
     async (id: TreeNodeId, isOpen?: boolean): Promise<void> => {
       const newOpenState = isOpen !== undefined ? isOpen : !tree.get(id).isOpen;
 
       if (isVirtualTree) {
-        if (!isActiveNodeInDOM(treeRef, activeNodeId)) {
-          virtualTreeConnector.scrollToNode?.(id);
-        }
+        scrollToVTreeNode(id);
         await virtualTreeConnector.setNodeOpen?.(id, newOpenState);
       }
 
@@ -115,10 +134,10 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
 
       const parent = node.parent && tree.get(node.parent);
       const isRoot = parent === undefined;
-
       // tabindex depends on the treeNodeBase params as well
       const nodeProps = {
         id: getNodeDOMId(node.id),
+
         'aria-setsize': isRoot ? 1 : parent.children.length,
         'aria-level': node.level + (excludeTreeRoot ? 0 : 1),
         'aria-posinset': node.index + 1,
@@ -132,16 +151,16 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
       if (!node.isLeaf) {
         nodeProps['aria-expanded'] = (!!node.isOpen).toString();
       }
+
+      const contentId = `${getNodeDOMId(node.id)}-content`;
+      const nodeContentProps = { id: contentId };
       const groupProps = {
         role: 'group',
         'aria-owns': node.children.map(getNodeDOMId).join(' '),
-        'aria-labelledby': getNodeDOMId(node.id),
+        'aria-labelledby': contentId,
       };
 
-      return {
-        nodeProps,
-        groupProps,
-      };
+      return { nodeProps, nodeContentProps, groupProps };
     },
     [tree]
   );
@@ -152,15 +171,15 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
       getNodeDetails,
       isRenderedFlat,
       shouldNodeFocusBeInset,
-      activeNodeId,
+      activeNodeId: activeNode,
       setActiveNodeId,
       toggleTreeNode,
       selectableNodes,
       itemSelection,
-      isFocusWithin: isFocusWithin,
+      isFocusWithin,
     }),
     [
-      activeNodeId,
+      activeNode,
       getNodeAriaProps,
       getNodeDetails,
       isRenderedFlat,
@@ -186,7 +205,12 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
   );
 
   const { focusWithinProps } = useFocusWithin({
-    onFocusWithinChange: setIsFocusWithin,
+    onFocusWithinChange: (val) => {
+      setIsFocusWithin(val);
+      if (val) {
+        scrollToVTreeNode(activeNode);
+      }
+    },
   });
 
   const { keyboardProps } = useKeyboard({
@@ -198,26 +222,20 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
         case 'ArrowRight':
         case 'ArrowLeft': {
           evt.preventDefault();
-          if (activeNodeId) {
-            const next = getNextActiveNode(
-              tree,
-              activeNodeId,
-              key,
-              excludeTreeRoot,
-              toggleTreeNode
-            );
+          if (activeNode) {
+            const next = getNextActiveNode(tree, activeNode, key, excludeTreeRoot, toggleTreeNode);
             setActiveNodeId(next);
           }
           break;
         }
         case 'Space': // Space
           if (
             selectionMode !== 'none' &&
-            activeNodeId &&
-            (selectableNodes === 'any' || tree.get(activeNodeId)?.isLeaf)
+            activeNode &&
+            (selectableNodes === 'any' || tree.get(activeNode)?.isLeaf)
           ) {
             evt.preventDefault();
-            itemSelection.toggle(activeNodeId);
+            itemSelection.toggle(activeNode);
           }
           break;
 

src/components/Tree/Tree.types.ts

@@ -204,9 +204,12 @@ export interface Props
  */
 export interface UseVirtualTreeNavigationProps extends Pick<Props, 'virtualTreeConnector'> {
   /**
-   * The active node id in the tree.
+   * Reference to the active node id in the tree.
+   *
+   * @remarks This prevent to destroy and re-create MutationObserver every time when active node
+   * changes. Also, this solves the problem of missed mutations.
    */
-  activeNodeId: TreeNodeId;
+  activeNodeIdRef: MutableRefObject<TreeNodeId>;
   /**
    * The reference of the tree DOM element.
    */
@@ -221,6 +224,10 @@ export interface NodeAriaProps {
    * attributes to re-build the semantic structure of the tree.
    */
   nodeProps: Partial<HTMLAttributes<HTMLElement>>;
+  /**
+   * Additional attributes for the tree node's content.
+   */
+  nodeContentProps: Partial<HTMLAttributes<HTMLElement>>;
   /**
    * Additional attributes for the node connection group.
    *

src/components/Tree/Tree.utils.test.tsx

@@ -21,7 +21,6 @@ import {
 } from './Tree.types';
 import { createTreeNode as tNode } from './test.utils';
 import { renderHook } from '@testing-library/react-hooks';
-import tree from './Tree';
 
 const createSingleLevelTree = () =>
   // prettier-ignore

src/components/TreeNodeBase/TreeNodeBase.constants.ts

@@ -24,6 +24,7 @@ const DEFAULTS = {
 const STYLE = {
   wrapper: `${CLASS_PREFIX}-wrapper`,
   contextMenuWrapper: `${CLASS_PREFIX}-context-menu-wrapper`,
+  content: `${CLASS_PREFIX}-content`,
   group: `${CLASS_PREFIX}-group`,
 };