MUI tree view with React-dnd

I have a material ui v4 treeview that is working fine with react-dnd using the code below.

When I upgrade to mui v5 treeview, the drag does not work anymore, the item is no longer draggable.

I had a look between the 2 TreeItem implementation but there is a lot of change and I am lost.

https://github.com/mui-org/material-ui/blob/v4.x/packages/material-ui-lab/src/TreeItem/TreeItem.js

https://github.com/mui-org/material-ui/blob/master/packages/mui-lab/src/TreeItem/TreeItem.js

What am I missing? Many thanks for your answers :)

//recursive function to generate TreeItem tree with Drag embedded
function Box({ treeItem }) {
  const [{ isDragging }, drag, preview] = useDrag(() => ({
    type: "TREEVIEW",
    item: !treeItem.children.length //if the resource has no child
      ? treeItem.data // we provide only the resource data
      : [
          treeItem.data,
          ...flatten(extractChildren(treeItem), extractChildren).map(
            //other wise we provide resource + child resources
            (x) => delete x.children && x
          ),
        ],
    collect: (monitor) => ({
      isDragging: monitor.isDragging(),
    }),
  }));

  return (
    <>
      <DragPreviewImage connect={preview} src={knightImage} />
      <TreeItem
        nodeId={treeItem.data.TreeID}
        label={treeItem.data.TreeName}
        ref={drag}
        style={{ isDragging }}
      >
        {treeItem.children &&
          treeItem.children.map((treeItem) => <Box treeItem={treeItem} />)}
      </TreeItem>
    </>
  );
}

const renderedListItems = tree.map((treeItem) => (
  <Box treeItem={treeItem} />
));

Probably, this happens because of the tree item getting selected while clicking on it. Even though disableSelection is set as true in TreeView, somehow the clicked child is being focused through aria-activedescendant property at the root level.

We can skip this focus event by adding onFocusCapture={e => e.stopPropagation()} at TreeItem which effectively makes the item draggable.