I actually just solved this myself using the React Flow controls with a custom delete button. In the <ReactFlow> component I added the onElementClick prop that uses a React hook to store the selected element in state. In the onClick prop for the button I have another hook that gets the selected element from state, all edges for the tree, and call the getConnectedEdges() which will return all edges connected to the selected element. Then just pass an array with the selected element and connected edges to your onElementsRemove().

Here is the documentation for onElementClick. https://reactflow.dev/docs/api/component-props/

<ReactFlow 
    className={flowStyles}
    elements={elements}
    onElementClick={onClickElement}
>
  <Controls showInteractive={false}>
    <ControlButton onClick={onClickElementDelete}>
      <DeleteIcon />
    </ControlButton>
  </Controls>
</ReactFlow>

const onClickElement = useCallback((event: ReactMouseEvent, element: Node | Edge) => {
  // Set the clicked element in local state
  setState({
    clickedElement: [element]
  })
}, [])

Here is the documentation for getConnectedEdges. https://reactflow.dev/docs/api/helper-functions/

  const onClickElementDelete = useCallback(() => {
    // Get all edges for the flow
    const edges = elements.filter((element: Node | Edge) => isEdge(element))
    // Get edges connected to selected node
    const edgesToRemove = getConnectedEdges(state.clickedElement, edges)

    onElementsRemove([...state.clickedElement, ...edgesToRemove])
    }
  }, [elements, onElementsRemove, state.clickedElement])

You have to add delete function in the data property of your custom node, and on click of the delete button of your custom node you have to call that function

While creating custom node =>

const newNode = {
  id: nodeID,
  type,
  position,
  data: {
      onDelete : yourFunctionToDelete
  }}

Custom Node Code in the delete button => onClick={props.data.onDelete}