fix(block-theme): ensure Overlay Menu block panel opens in editor after template part switch

src/blocks/overlay-menu/edit.js

@@ -2,7 +2,6 @@
  * WordPress dependencies
  */
 import { useEffect, useState } from '@wordpress/element';
-import { useSelect } from '@wordpress/data';
 import { InnerBlocks, useBlockProps } from '@wordpress/block-editor';
 
 /**
@@ -37,23 +36,15 @@ export default function OverlayMenuEdit( { attributes, setAttributes, clientId }
 		}
 	}, [ clientId ] ); // eslint-disable-line react-hooks/exhaustive-deps
 
-	// Find the child panel block to key the preview state Maps.
-	const panelClientId = useSelect( select => {
-		const block = select( 'core/block-editor' ).getBlock( clientId );
-		return block?.innerBlocks?.find( b => b.name === 'newspack/overlay-menu-panel' )?.clientId;
-	} );
-
 	// Mirror the panel's open state so the toolbar button label and isPressed stay correct.
 	const [ isPreviewOpen, setIsPreviewOpen ] = useState( false );
 	useEffect( () => {
-		if ( ! panelClientId ) {
-			return;
-		}
-		return subscribeToPanel( panelClientId, setIsPreviewOpen );
-	}, [ panelClientId ] );
+		return subscribeToPanel( clientId, setIsPreviewOpen );
+	}, [ clientId ] );
 
 	// Delegate the actual toggle to the panel via its registered ref function.
-	const togglePreview = () => panelToggles.get( panelClientId )?.();
+	// The panel keys its entry by parentClientId === this block's clientId.
+	const togglePreview = () => panelToggles.get( clientId )?.();
 
 	const blockProps = useBlockProps( {
 		className: 'is-layout-flex',

src/blocks/overlay-menu/panel/edit.js

@@ -3,7 +3,8 @@
  */
 import { __ } from '@wordpress/i18n';
 import { close as closeIcon } from '@wordpress/icons';
-import { useEffect, useRef, useState } from '@wordpress/element';
+import { useEffect, useLayoutEffect, useRef, useState } from '@wordpress/element';
+import { useSelect } from '@wordpress/data';
 import {
 	// eslint-disable-next-line @wordpress/no-unsafe-wp-apis
 	__experimentalColorGradientSettingsDropdown as ColorGradientSettingsDropdown,
@@ -48,26 +49,57 @@ export default function OverlayMenuPanelEdit( { attributes, clientId, setAttribu
 
 	const [ isPreviewOpen, setIsPreviewOpen ] = useState( false );
 
-	// Keep a ref to the current open state so the toggle registered in panelToggles
-	// never has a stale closure over isPreviewOpen.
+	// Keep a ref to the current open state so the toggle never has a stale closure.
 	const isOpenRef = useRef( false );
 	isOpenRef.current = isPreviewOpen;
 
-	// Register a toggle function keyed by clientId so the parent toolbar button
-	// can open/close the panel without sharing block attributes.
+	// Key everything by the parent's clientId so the parent/trigger can look up the toggle using their own clientId or getBlockRootClientId.
+	const parentClientId = useSelect( select => select( 'core/block-editor' ).getBlockRootClientId( clientId ), [ clientId ] );
+
+	// Keep a stable ref to the toggle function so the Map entry is always current.
+	const toggleFnRef = useRef( null );
+	toggleFnRef.current = () => {
+		const next = ! isOpenRef.current;
+		setIsPreviewOpen( next );
+		if ( parentClientId ) {
+			notifySubscribers( parentClientId, next );
+		}
+	};
+
+	// Render-phase registration: runs even when the component renders inside a
+	// React transition that hasn't committed yet (e.g. the site editor wrapping a
+	// template-part switch in startTransition). This is the fallback that makes the
+	// toggle available before the commit phase fires.
+	if ( parentClientId ) {
+		panelToggles.set( parentClientId, () => toggleFnRef.current?.() );
+	}
+
+	// Authoritative registration and cleanup in the commit phase. useLayoutEffect
+	// overwrites the render-phase entry once the component commits, and its cleanup
+	// reliably removes the entry on unmount or parentClientId change — preventing
+	// stale Map entries from aborted renders lingering after the component unmounts.
+	useLayoutEffect( () => {
+		if ( ! parentClientId ) {
+			return;
+		}
+		panelToggles.set( parentClientId, () => toggleFnRef.current?.() );
+		return () => panelToggles.delete( parentClientId );
+	}, [ parentClientId ] ); // eslint-disable-line react-hooks/exhaustive-deps
+
 	useEffect( () => {
-		panelToggles.set( clientId, () => {
-			const next = ! isOpenRef.current;
-			setIsPreviewOpen( next );
-			notifySubscribers( clientId, next );
-		} );
-		return () => panelToggles.delete( clientId );
-	}, [ clientId ] ); // eslint-disable-line react-hooks/exhaustive-deps
+		return () => {
+			if ( parentClientId ) {
+				panelToggles.delete( parentClientId );
+			}
+		};
+	}, [ parentClientId ] );
 
 	// Update local state and notify all subscribers (parent + trigger toolbar buttons).
 	const togglePreview = open => {
 		setIsPreviewOpen( open );
-		notifySubscribers( clientId, open );
+		if ( parentClientId ) {
+			notifySubscribers( parentClientId, open );
+		}
 	};
 
 	const { positionClass } = DIRECTION_CONFIG[ slideDirection ] ?? DIRECTION_CONFIG.left;

src/blocks/overlay-menu/preview-refs.js

@@ -1,14 +1,16 @@
 /**
  * Module-level Maps for ephemeral editor-only preview state.
  *
- * panelToggles: the panel registers a toggle function here so sibling/parent
- * blocks can call it without needing a shared reactive store.
+ * Both Maps are keyed by the PARENT overlay-menu block's clientId, not the
+ * panel's clientId. This lets the parent and trigger look up entries using
+ * their own clientId or a single getBlockRootClientId() call, avoiding the
+ * innerBlocks traversal that can return empty during a template-part switch.
  *
- * subscribers: any block that needs to mirror the panel's open state registers
- * a React state setter here (keyed by panel clientId). Multiple blocks can
- * subscribe to the same panel.
+ * panelToggles: the panel registers a toggle function here so the parent and
+ * trigger toolbar buttons can open/close it without a shared reactive store.
  *
- * Both Maps are keyed by panel clientId.
+ * subscribers: any block that needs to mirror the panel's open state registers
+ * a React state setter here. Multiple blocks can subscribe to the same panel.
  */
 
 /** @type {Map<string, function(): void>} */
@@ -21,24 +23,24 @@ const subscribers = new Map();
  * Subscribe a React state setter to open-state changes for a panel.
  * Returns an unsubscribe function suitable for useEffect cleanup.
  *
- * @param {string}                  panelClientId Panel clientId.
- * @param {function(boolean): void} setter        React state setter.
+ * @param {string}                  parentClientId Parent overlay-menu block clientId.
+ * @param {function(boolean): void} setter         React state setter.
  * @return {function(): void} Unsubscribe function.
  */
-export function subscribeToPanel( panelClientId, setter ) {
-	if ( ! subscribers.has( panelClientId ) ) {
-		subscribers.set( panelClientId, new Set() );
+export function subscribeToPanel( parentClientId, setter ) {
+	if ( ! subscribers.has( parentClientId ) ) {
+		subscribers.set( parentClientId, new Set() );
 	}
-	subscribers.get( panelClientId ).add( setter );
-	return () => subscribers.get( panelClientId )?.delete( setter );
+	subscribers.get( parentClientId ).add( setter );
+	return () => subscribers.get( parentClientId )?.delete( setter );
 }
 
 /**
  * Notify all subscribers of a new open state.
  *
- * @param {string}  panelClientId Panel clientId.
- * @param {boolean} isOpen        New open state.
+ * @param {string}  parentClientId Parent overlay-menu block clientId.
+ * @param {boolean} isOpen         New open state.
  */
-export function notifySubscribers( panelClientId, isOpen ) {
-	subscribers.get( panelClientId )?.forEach( fn => fn( isOpen ) );
+export function notifySubscribers( parentClientId, isOpen ) {
+	subscribers.get( parentClientId )?.forEach( fn => fn( isOpen ) );
 }

src/blocks/overlay-menu/trigger/edit.js

@@ -35,32 +35,25 @@ export default function OverlayMenuTriggerEdit( { attributes, setAttributes, cli
 	const isTextOnly = classes.includes( 'is-style-text-only' );
 	const showTriggerIcon = ! isTextOnly;
 
-	// Find the panel sibling block to key the preview state Maps.
-	const panelClientId = useSelect(
-		select => {
-			const { getBlockRootClientId, getBlocks } = select( 'core/block-editor' );
-			const parentClientId = getBlockRootClientId( clientId );
-			return getBlocks( parentClientId ).find( b => b.name === 'newspack/overlay-menu-panel' )?.clientId;
-		},
-		[ clientId ]
-	);
+	// The panel registers its toggle under the parent's clientId.
+	const parentClientId = useSelect( select => select( 'core/block-editor' ).getBlockRootClientId( clientId ), [ clientId ] );
 
 	// Mirror the panel's open state so the toolbar button label and isPressed stay correct.
 	const [ isPanelOpen, setIsPanelOpen ] = useState( false );
 	useEffect( () => {
-		if ( ! panelClientId ) {
+		if ( ! parentClientId ) {
 			return;
 		}
-		return subscribeToPanel( panelClientId, setIsPanelOpen );
-	}, [ panelClientId ] );
+		return subscribeToPanel( parentClientId, setIsPanelOpen );
+	}, [ parentClientId ] );
 
 	const blockProps = useBlockProps( {
 		className: 'overlay-menu__trigger wp-block-button__link wp-element-button',
 	} );
 
 	return (
 		<>
-			<PanelPreviewToggle isOpen={ isPanelOpen } onToggle={ () => panelToggles.get( panelClientId )?.() } />
+			<PanelPreviewToggle isOpen={ isPanelOpen } onToggle={ () => panelToggles.get( parentClientId )?.() } />
 
 			<button
 				{ ...blockProps }