Persist focus when tabbing back to the block toolbar (#25760)

* Initial implementation for persisting block toolbar focus

* Update code with failing test

* Fix failing test

* Reset the initial toolbar item index when navigating blocks

packages/block-editor/src/components/block-list/block-popover.js

@@ -7,7 +7,13 @@ import classnames from 'classnames';
 /**
  * WordPress dependencies
  */
-import { useState, useCallback, useContext } from '@wordpress/element';
+import {
+	useState,
+	useCallback,
+	useContext,
+	useRef,
+	useEffect,
+} from '@wordpress/element';
 import { isUnmodifiedDefaultBlock } from '@wordpress/blocks';
 import { Popover } from '@wordpress/components';
 import { useSelect } from '@wordpress/data';
@@ -92,6 +98,16 @@ function BlockPopover( {
 		}
 	);
 
+	// Stores the active toolbar item index so the block toolbar can return focus
+	// to it when re-mounting.
+	const initialToolbarItemIndexRef = useRef();
+
+	useEffect( () => {
+		// Resets the index whenever the active block changes so this is not
+		// persisted. See https://github.com/WordPress/gutenberg/pull/25760#issuecomment-717906169
+		initialToolbarItemIndexRef.current = undefined;
+	}, [ clientId ] );
+
 	if (
 		! shouldShowBreadcrumb &&
 		! shouldShowContextualToolbar &&
@@ -190,6 +206,12 @@ function BlockPopover( {
 					// If the toolbar is being shown because of being forced
 					// it should focus the toolbar right after the mount.
 					focusOnMount={ isToolbarForced }
+					__experimentalInitialIndex={
+						initialToolbarItemIndexRef.current
+					}
+					__experimentalOnIndexChange={ ( index ) => {
+						initialToolbarItemIndexRef.current = index;
+					} }
 				/>
 			) }
 			{ shouldShowBreadcrumb && (

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

@@ -18,6 +18,14 @@ function hasOnlyToolbarItem( elements ) {
 	return ! elements.some( ( element ) => ! ( dataProp in element.dataset ) );
 }
 
+function getAllToolbarItemsIn( container ) {
+	return Array.from( container.querySelectorAll( '[data-toolbar-item]' ) );
+}
+
+function hasFocusWithin( container ) {
+	return container.contains( container.ownerDocument.activeElement );
+}
+
 function focusFirstTabbableIn( container ) {
 	const [ firstTabbable ] = focus.tabbable.find( container );
 	if ( firstTabbable ) {
@@ -74,14 +82,22 @@ function useIsAccessibleToolbar( ref ) {
 	return isAccessibleToolbar;
 }
 
-function useToolbarFocus( ref, focusOnMount, isAccessibleToolbar ) {
+function useToolbarFocus(
+	ref,
+	focusOnMount,
+	isAccessibleToolbar,
+	defaultIndex,
+	onIndexChange
+) {
 	// Make sure we don't use modified versions of this prop
 	const [ initialFocusOnMount ] = useState( focusOnMount );
+	const [ initialIndex ] = useState( defaultIndex );
 
 	const focusToolbar = useCallback( () => {
 		focusFirstTabbableIn( ref.current );
 	}, [] );
 
+	// Focus on toolbar when pressing alt+F10 when the toolbar is visible
 	useShortcut( 'core/block-editor/focus-toolbar', focusToolbar, {
 		bindGlobal: true,
 		eventName: 'keydown',
@@ -92,21 +108,55 @@ function useToolbarFocus( ref, focusOnMount, isAccessibleToolbar ) {
 			focusToolbar();
 		}
 	}, [ isAccessibleToolbar, initialFocusOnMount, focusToolbar ] );
-}
 
-function NavigableToolbar( { children, focusOnMount, ...props } ) {
-	const wrapper = useRef();
-	const isAccessibleToolbar = useIsAccessibleToolbar( wrapper );
+	useEffect( () => {
+		// If initialIndex is passed, we focus on that toolbar item when the
+		// toolbar gets mounted and initial focus is not forced.
+		// We have to wait for the next browser paint because block controls aren't
+		// rendered right away when the toolbar gets mounted.
+		let raf = 0;
+		if ( initialIndex && ! initialFocusOnMount ) {
+			raf = window.requestAnimationFrame( () => {
+				const items = getAllToolbarItemsIn( ref.current );
+				const index = initialIndex || 0;
+				if ( items[ index ] && hasFocusWithin( ref.current ) ) {
+					items[ index ].focus();
+				}
+			} );
+		}
+		return () => {
+			window.cancelAnimationFrame( raf );
+			if ( ! onIndexChange ) return;
+			// When the toolbar element is unmounted and onIndexChange is passed, we
+			// pass the focused toolbar item index so it can be hydrated later.
+			const items = getAllToolbarItemsIn( ref.current );
+			const index = items.findIndex( ( item ) => item.tabIndex === 0 );
+			onIndexChange( index );
+		};
+	}, [ initialIndex, initialFocusOnMount ] );
+}
 
-	useToolbarFocus( wrapper, focusOnMount, isAccessibleToolbar );
+function NavigableToolbar( {
+	children,
+	focusOnMount,
+	__experimentalInitialIndex: initialIndex,
+	__experimentalOnIndexChange: onIndexChange,
+	...props
+} ) {
+	const ref = useRef();
+	const isAccessibleToolbar = useIsAccessibleToolbar( ref );
+
+	useToolbarFocus(
+		ref,
+		focusOnMount,
+		isAccessibleToolbar,
+		initialIndex,
+		onIndexChange
+	);
 
 	if ( isAccessibleToolbar ) {
 		return (
-			<Toolbar
-				label={ props[ 'aria-label' ] }
-				ref={ wrapper }
-				{ ...props }
-			>
+			<Toolbar label={ props[ 'aria-label' ] } ref={ ref } { ...props }>
 				{ children }
 			</Toolbar>
 		);
@@ -116,7 +166,7 @@ function NavigableToolbar( { children, focusOnMount, ...props } ) {
 		<NavigableMenu
 			orientation="horizontal"
 			role="toolbar"
-			ref={ wrapper }
+			ref={ ref }
 			{ ...props }
 		>
 			{ children }

packages/e2e-tests/specs/editor/various/toolbar-roving-tabindex.test.js

@@ -4,6 +4,7 @@
 import {
 	createNewPost,
 	pressKeyWithModifier,
+	clickBlockToolbarButton,
 	insertBlock,
 } from '@wordpress/e2e-test-utils';
 
@@ -103,4 +104,14 @@ describe( 'Toolbar roving tabindex', () => {
 		await wrapCurrentBlockWithGroup();
 		await testGroupKeyboardNavigation( 'Block: Custom HTML' );
 	} );
+
+	it( 'ensures block toolbar remembers the last focused item', async () => {
+		await insertBlock( 'Paragraph' );
+		await page.keyboard.type( 'Paragraph' );
+		await focusBlockToolbar();
+		await clickBlockToolbarButton( 'Bold' );
+		await page.keyboard.type( 'a' );
+		await pressKeyWithModifier( 'shift', 'Tab' );
+		await expectLabelToHaveFocus( 'Bold' );
+	} );
 } );