Writing flow: Fix focus loss on toolbar return when last focus element isn't found

docs/reference-guides/data/data-core-block-editor.md

@@ -588,6 +588,18 @@ _Properties_
 -   _isDisabled_ `boolean`: Whether or not the user should be prevented from inserting this item.
 -   _frecency_ `number`: Heuristic that combines frequency and recency.
 
+### getLastFocus
+
+Returns the element of the last element that had focus when focus left the editor canvas.
+
+_Parameters_
+
+-   _state_ `Object`: Block editor state.
+
+_Returns_
+
+-   `Object`: Element.
+
 ### getLastMultiSelectedBlockClientId
 
 Returns the client ID of the last block in the multi-selection set, or null if there is no multi-selection.
@@ -1651,6 +1663,18 @@ _Parameters_
 -   _clientId_ `string`: The block's clientId.
 -   _hasControlledInnerBlocks_ `boolean`: True if the block's inner blocks are controlled.
 
+### setLastFocus
+
+Action that sets the element that had focus when focus leaves the editor canvas.
+
+_Parameters_
+
+-   _lastFocus_ `Object`: The last focused element.
+
+_Returns_
+
+-   `Object`: Action object.
+
 ### setNavigationMode
 
 Action that enables or disables the navigation mode.

packages/block-editor/src/components/writing-flow/use-tab-nav.js

@@ -15,17 +15,23 @@ import { isInSameBlock, isInsideRootBlock } from '../../utils/dom';
 
 export default function useTabNav() {
 	const container = useRef();
+	const lastBlock = useRef();
 	const focusCaptureBeforeRef = useRef();
 	const focusCaptureAfterRef = useRef();
-	const lastFocus = useRef();
+
 	const { hasMultiSelection, getSelectedBlockClientId, getBlockCount } =
 		useSelect( blockEditorStore );
-	const { setNavigationMode } = useDispatch( blockEditorStore );
+	const { setNavigationMode, setLastFocus } = useDispatch( blockEditorStore );
 	const isNavigationMode = useSelect(
 		( select ) => select( blockEditorStore ).isNavigationMode(),
 		[]
 	);
 
+	const lastFocus = useSelect(
+		( select ) => select( blockEditorStore ).getLastFocus(),
+		[]
+	);
+
 	// Don't allow tabbing to this element in Navigation mode.
 	const focusCaptureTabIndex = ! isNavigationMode ? '0' : undefined;
 
@@ -40,7 +46,31 @@ export default function useTabNav() {
 		} else if ( hasMultiSelection() ) {
 			container.current.focus();
 		} else if ( getSelectedBlockClientId() ) {
-			lastFocus.current.focus();
+			// If last focus position matches last block, this likely means we're on the last block wrapper. Let's try to find a better place to focus before defaulting to the wrapper.
+			if ( lastBlock.current !== lastFocus.current ) {
+				// Try to focus the last element which had focus.
+				lastFocus.current.focus();
+				// Check to see if focus worked.
+				if (
+					lastFocus.current.ownerDocument.activeElement ===
+					lastFocus.current
+				) {
+					// Looks like yes, return early.
+					return;
+				}
+			}
+			// Last element focus did not work. Now try to find the first tabbable in the last block to focus.
+			const firstBlockTabbable = focus.tabbable.findNext(
+				lastBlock.current
+			);
+			// Check to ensure tabbable is inside the last block and that it is a form element.
+			if ( isInsideRootBlock( lastBlock.current, firstBlockTabbable ) ) {
+				// Focus the found tabbable in the last block.
+				firstBlockTabbable.focus();
+			} else {
+				// Focus the last block wrapper if no tabbable was found.
+				lastBlock.current.focus();
+			}
 		} else {
 			setNavigationMode( true );
 
@@ -158,7 +188,10 @@ export default function useTabNav() {
 		}
 
 		function onFocusOut( event ) {
-			lastFocus.current = event.target;
+			// Capture the last element with focus.
+			setLastFocus( { ...lastFocus, current: event.target } );
+			// Capture the last known block before focus leaves writing flow.
+			lastBlock.current = event.target.closest( '[data-block]' );
 
 			const { ownerDocument } = node;
 

packages/block-editor/src/store/actions.js

@@ -1967,3 +1967,18 @@ export function unsetBlockEditingMode( clientId = '' ) {
 		clientId,
 	};
 }
+
+/**
+ * Action that sets the element that had focus when focus leaves the editor canvas.
+ *
+ * @param {Object} lastFocus The last focused element.
+ *
+ *
+ * @return {Object} Action object.
+ */
+export function setLastFocus( lastFocus = null ) {
+	return {
+		type: 'LAST_FOCUS',
+		lastFocus,
+	};
+}

packages/block-editor/src/store/reducer.js

@@ -1946,6 +1946,22 @@ export function styleOverrides( state = new Map(), action ) {
 			return newState;
 		}
 	}
+}
+
+/**
+ * Reducer setting last focused element
+ *
+ * @param {boolean} state  Current state.
+ * @param {Object}  action Dispatched action.
+ *
+ * @return {boolean} Updated state.
+ */
+export function lastFocus( state = false, action ) {
+	switch ( action.type ) {
+		case 'LAST_FOCUS':
+			return action.lastFocus;
+	}
+
 	return state;
 }
 
@@ -1965,6 +1981,7 @@ const combinedReducers = combineReducers( {
 	settings,
 	preferences,
 	lastBlockAttributesChange,
+	lastFocus,
 	editorMode,
 	hasBlockMovingClientId,
 	highlightedBlock,

packages/block-editor/src/store/selectors.js

@@ -3018,3 +3018,14 @@ export const isGroupable = createRegistrySelector(
 			);
 		}
 );
+
+/*
+ * Returns the element of the last element that had focus when focus left the editor canvas.
+ *
+ * @param {Object} state Block editor state.
+ *
+ * @return {Object} Element.
+ */
+export function getLastFocus( state ) {
+	return state.lastFocus;
+}