BlockLockModal: restore focus on fallback toolbar button when original button is not rendered

packages/block-editor/src/components/block-lock/toolbar.js

@@ -3,7 +3,8 @@
  */
 import { __ } from '@wordpress/i18n';
 import { ToolbarButton, ToolbarGroup } from '@wordpress/components';
-import { useReducer } from '@wordpress/element';
+import { focus } from '@wordpress/dom';
+import { useReducer, useRef, useEffect } from '@wordpress/element';
 import { lock } from '@wordpress/icons';
 
 /**
@@ -12,19 +13,45 @@ import { lock } from '@wordpress/icons';
 import BlockLockModal from './modal';
 import useBlockLock from './use-block-lock';
 
-export default function BlockLockToolbar( { clientId } ) {
+export default function BlockLockToolbar( { clientId, wrapperRef } ) {
 	const { canEdit, canMove, canRemove, canLock } = useBlockLock( clientId );
 
 	const [ isModalOpen, toggleModal ] = useReducer(
 		( isActive ) => ! isActive,
 		false
 	);
 
-	if ( ! canLock ) {
-		return null;
-	}
+	const lockButtonRef = useRef( null );
+	const isFirstRender = useRef( true );
+
+	const shouldHideBlockLockUI =
+		! canLock || ( canEdit && canMove && canRemove );
+
+	// Restore focus manually on the first focusable element in the toolbar
+	// when the block lock modal is closed and the block is not locked anymore.
+	// See https://github.com/WordPress/gutenberg/issues/51447
+	useEffect( () => {
+		if ( isFirstRender.current ) {
+			isFirstRender.current = false;
+			return;
+		}
+
+		if ( ! isModalOpen && shouldHideBlockLockUI ) {
+			focus.focusable
+				.find( wrapperRef.current, {
+					sequential: false,
+				} )
+				.find(
+					( element ) =>
+						element.tagName === 'BUTTON' &&
+						element !== lockButtonRef.current
+				)
+				?.focus();
+		}
+		// wrapperRef is a reference object and should be stable
+	}, [ isModalOpen, shouldHideBlockLockUI, wrapperRef ] );
 
-	if ( canEdit && canMove && canRemove ) {
+	if ( shouldHideBlockLockUI ) {
 		return null;
 	}
 
@@ -35,6 +62,7 @@ export default function BlockLockToolbar( { clientId } ) {
 					icon={ lock }
 					label={ __( 'Unlock' ) }
 					onClick={ toggleModal }
+					ref={ lockButtonRef }
 				/>
 			</ToolbarGroup>
 			{ isModalOpen && (

packages/block-editor/src/components/block-toolbar/index.js

@@ -78,6 +78,8 @@ const BlockToolbar = ( { hideDragHandle } ) => {
 		};
 	}, [] );
 
+	const toolbarWrapperRef = useRef( null );
+
 	// Handles highlighting the current block outline on hover or focus of the
 	// block type toolbar area.
 	const { toggleBlockHighlight } = useDispatch( blockEditorStore );
@@ -123,7 +125,7 @@ const BlockToolbar = ( { hideDragHandle } ) => {
 	} );
 
 	return (
-		<div className={ classes }>
+		<div className={ classes } ref={ toolbarWrapperRef }>
 			{ ! isMultiToolbar &&
 				isLargeViewport &&
 				blockEditingMode === 'default' && <BlockParentSelector /> }
@@ -135,6 +137,7 @@ const BlockToolbar = ( { hideDragHandle } ) => {
 							{ ! isMultiToolbar && (
 								<BlockLockToolbar
 									clientId={ blockClientIds[ 0 ] }
+									wrapperRef={ toolbarWrapperRef }
 								/>
 							) }
 							<BlockMover