Multi select: set focus back after attempt

packages/block-editor/src/components/block-list/use-multi-selection.js

@@ -55,6 +55,18 @@ function selector( select ) {
 	};
 }
 
+function toggleRichText( container, toggle ) {
+	Array
+		.from( container.querySelectorAll( '.rich-text' ) )
+		.forEach( ( node ) => {
+			if ( toggle ) {
+				node.setAttribute( 'contenteditable', true );
+			} else {
+				node.removeAttribute( 'contenteditable' );
+			}
+		} );
+}
+
 export default function useMultiSelection( ref ) {
 	const {
 		isSelectionEnabled,
@@ -72,14 +84,15 @@ export default function useMultiSelection( ref ) {
 	} = useDispatch( 'core/block-editor' );
 	const rafId = useRef();
 	const startClientId = useRef();
+	const anchorElement = useRef();
 
 	/**
 	 * When the component updates, and there is multi selection, we need to
 	 * select the entire block contents.
 	 */
 	useEffect( () => {
 		if ( ! hasMultiSelection || isMultiSelecting ) {
-			if ( ! selectedBlockClientId ) {
+			if ( ! selectedBlockClientId || isMultiSelecting ) {
 				return;
 			}
 
@@ -167,6 +180,19 @@ export default function useMultiSelection( ref ) {
 		rafId.current = window.requestAnimationFrame( () => {
 			onSelectionChange();
 			stopMultiSelect();
+			toggleRichText( ref.current, true );
+
+			const selection = window.getSelection();
+
+			// If the anchor element contains the selection, set focus back to
+			// the anchor element.
+			if ( selection.rangeCount ) {
+				const { commonAncestorContainer } = selection.getRangeAt( 0 );
+
+				if ( anchorElement.current.contains( commonAncestorContainer ) ) {
+					anchorElement.current.focus();
+				}
+			}
 		} );
 	}, [ onSelectionChange, stopMultiSelect ] );
 
@@ -187,6 +213,7 @@ export default function useMultiSelection( ref ) {
 		}
 
 		startClientId.current = clientId;
+		anchorElement.current = document.activeElement;
 		startMultiSelect();
 
 		// `onSelectionStart` is called after `mousedown` and `mouseleave`
@@ -203,7 +230,6 @@ export default function useMultiSelection( ref ) {
 		// especially in Safari for the blocks that are asynchonously rendered.
 		// To ensure the browser instantly removes the selection boundaries, we
 		// remove the contenteditable attributes manually.
-		Array.from( ref.current.querySelectorAll( '.rich-text' ) )
-			.forEach( ( node ) => node.removeAttribute( 'contenteditable' ) );
+		toggleRichText( ref.current, false );
 	}, [ isSelectionEnabled, startMultiSelect, onSelectionEnd ] );
 }

packages/e2e-tests/specs/editor/various/__snapshots__/multi-block-selection.test.js.snap

@@ -52,6 +52,12 @@ exports[`Multi-block selection should only trigger multi-selection when at the e
 <!-- /wp:paragraph -->"
 `;
 
+exports[`Multi-block selection should return original focus after failed multi selection attempt 1`] = `
+"<!-- wp:paragraph -->
+<p>2</p>
+<!-- /wp:paragraph -->"
+`;
+
 exports[`Multi-block selection should use selection direction to determine vertical edge 1`] = `
 "<!-- wp:paragraph -->
 <p>1<br>2.</p>

packages/e2e-tests/specs/editor/various/multi-block-selection.test.js

@@ -363,4 +363,49 @@ describe( 'Multi-block selection', () => {
 
 		expect( await getEditedPostContent() ).toMatchSnapshot();
 	} );
+
+	it( 'should return original focus after failed multi selection attempt', async () => {
+		await clickBlockAppender();
+		await page.keyboard.type( '1' );
+		await page.keyboard.type( '2' );
+		await page.keyboard.press( 'ArrowLeft' );
+
+		const [ coord1, coord2 ] = await page.evaluate( () => {
+			const selection = window.getSelection();
+
+			if ( ! selection.rangeCount ) {
+				return;
+			}
+
+			const range = selection.getRangeAt( 0 );
+			const rect1 = range.getClientRects()[ 0 ];
+			const element = document.querySelector( '.wp-block-paragraph' );
+			const rect2 = element.getBoundingClientRect();
+
+			return [
+				{
+					x: rect1.x,
+					y: rect1.y + ( rect1.height / 2 ),
+				},
+				{
+					// Move a bit outside the paragraph.
+					x: rect2.x - 10,
+					y: rect2.y + ( rect2.height / 2 ),
+				},
+			];
+		} );
+
+		await page.mouse.move( coord1.x, coord1.y );
+		await page.mouse.down();
+		await page.mouse.move( coord2.x, coord2.y, { steps: 10 } );
+		await page.mouse.up();
+
+		// Wait for the selection to update.
+		await page.evaluate( () => new Promise( window.requestAnimationFrame ) );
+
+		// Only "1" should be deleted.
+		await page.keyboard.press( 'Backspace' );
+
+		expect( await getEditedPostContent() ).toMatchSnapshot();
+	} );
 } );