diff --git a/packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js b/packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js index 297aed456df7f..056ade045d165 100644 --- a/packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js +++ b/packages/block-editor/src/components/block-list/use-block-props/use-block-refs.js @@ -1,8 +1,14 @@ /** * WordPress dependencies */ -import { useContext, useMemo, useRef } from '@wordpress/element'; -import { useRefEffect, useObservableValue } from '@wordpress/compose'; +import { + useContext, + useMemo, + useRef, + useState, + useLayoutEffect, +} from '@wordpress/element'; +import { useRefEffect } from '@wordpress/compose'; /** * Internal dependencies @@ -67,7 +73,17 @@ function useBlockRef( clientId ) { */ function useBlockElement( clientId ) { const { refsMap } = useContext( BlockRefs ); - return useObservableValue( refsMap, clientId ) ?? null; + const [ blockElement, setBlockElement ] = useState( null ); + // Delay setting the resulting `blockElement` until an effect. If the block element + // changes (i.e., the block is unmounted and re-mounted), this allows enough time + // for the ref callbacks to clean up the old element and set the new one. + useLayoutEffect( () => { + setBlockElement( refsMap.get( clientId ) ); + return refsMap.subscribe( clientId, () => + setBlockElement( refsMap.get( clientId ) ) + ); + }, [ refsMap, clientId ] ); + return blockElement; } export { useBlockRef as __unstableUseBlockRef };