Fix button block focus trap after a URL has been added (WordPress#34314)

* Rework button block link UI to match RichText format implementation

* Refine some more, determine visibility by selection and url state

* Add e2e test

* Also focus rich text when unlinking using a keyboard shortcut

Author: talldan

packages/block-library/src/button/edit.js

@@ -7,7 +7,7 @@ import classnames from 'classnames';
  * WordPress dependencies
  */
 import { __ } from '@wordpress/i18n';
-import { useCallback, useState, useRef } from '@wordpress/element';
+import { useCallback, useEffect, useState, useRef } from '@wordpress/element';
 import {
 	Button,
 	ButtonGroup,
@@ -73,27 +73,42 @@ function URLPicker( {
 	opensInNewTab,
 	onToggleOpenInNewTab,
 	anchorRef,
+	richTextRef,
 } ) {
-	const [ isURLPickerOpen, setIsURLPickerOpen ] = useState( false );
-	const urlIsSet = !! url;
-	const urlIsSetandSelected = urlIsSet && isSelected;
-	const openLinkControl = () => {
-		setIsURLPickerOpen( true );
-		return false; // prevents default behaviour for event
+	const [ isEditingURL, setIsEditingURL ] = useState( false );
+	const isURLSet = !! url;
+
+	const startEditing = ( event ) => {
+		event.preventDefault();
+		setIsEditingURL( true );
 	};
-	const unlinkButton = () => {
+
+	const unlink = () => {
 		setAttributes( {
 			url: undefined,
 			linkTarget: undefined,
 			rel: undefined,
 		} );
-		setIsURLPickerOpen( false );
+		setIsEditingURL( false );
 	};
-	const linkControl = ( isURLPickerOpen || urlIsSetandSelected ) && (
+
+	useEffect( () => {
+		if ( ! isSelected ) {
+			setIsEditingURL( false );
+		}
+	}, [ isSelected ] );
+
+	const isLinkControlVisible = isSelected && ( isEditingURL || isURLSet );
+
+	const linkControl = isLinkControlVisible && (
 		<Popover
 			position="bottom center"
-			onClose={ () => setIsURLPickerOpen( false ) }
+			onClose={ () => {
+				setIsEditingURL( false );
+				richTextRef.current?.focus();
+			} }
 			anchorRef={ anchorRef?.current }
+			focusOnMount={ isEditingURL ? 'firstElement' : false }
 		>
 			<LinkControl
 				className="wp-block-navigation-link__inline-link-input"
@@ -108,28 +123,34 @@ function URLPicker( {
 						onToggleOpenInNewTab( newOpensInNewTab );
 					}
 				} }
+				onRemove={ () => {
+					unlink();
+					richTextRef.current?.focus();
+				} }
+				forceIsEditingLink={ isEditingURL }
 			/>
 		</Popover>
 	);
+
 	return (
 		<>
 			<BlockControls group="block">
-				{ ! urlIsSet && (
+				{ ! isURLSet && (
 					<ToolbarButton
 						name="link"
 						icon={ link }
 						title={ __( 'Link' ) }
 						shortcut={ displayShortcut.primary( 'k' ) }
-						onClick={ openLinkControl }
+						onClick={ startEditing }
 					/>
 				) }
-				{ urlIsSetandSelected && (
+				{ isURLSet && (
 					<ToolbarButton
 						name="link"
 						icon={ linkOff }
 						title={ __( 'Unlink' ) }
 						shortcut={ displayShortcut.primaryShift( 'k' ) }
-						onClick={ unlinkButton }
+						onClick={ unlink }
 						isActive={ true }
 					/>
 				) }
@@ -138,8 +159,11 @@ function URLPicker( {
 				<KeyboardShortcuts
 					bindGlobal
 					shortcuts={ {
-						[ rawShortcut.primary( 'k' ) ]: openLinkControl,
-						[ rawShortcut.primaryShift( 'k' ) ]: unlinkButton,
+						[ rawShortcut.primary( 'k' ) ]: startEditing,
+						[ rawShortcut.primaryShift( 'k' ) ]: () => {
+							unlink();
+							richTextRef.current?.focus();
+						},
 					} }
 				/>
 			) }
@@ -201,6 +225,7 @@ function ButtonEdit( props ) {
 	const colorProps = useColorProps( attributes );
 	const spacingProps = useSpacingProps( attributes );
 	const ref = useRef();
+	const richTextRef = useRef();
 	const blockProps = useBlockProps( { ref } );
 
 	return (
@@ -213,6 +238,7 @@ function ButtonEdit( props ) {
 				} ) }
 			>
 				<RichText
+					ref={ richTextRef }
 					aria-label={ __( 'Button text' ) }
 					placeholder={ placeholder || __( 'Add text…' ) }
 					value={ text }
@@ -252,6 +278,7 @@ function ButtonEdit( props ) {
 				opensInNewTab={ linkTarget === '_blank' }
 				onToggleOpenInNewTab={ onToggleOpenInNewTab }
 				anchorRef={ ref }
+				richTextRef={ richTextRef }
 			/>
 			<InspectorControls>
 				<WidthPanel

packages/e2e-tests/specs/editor/blocks/buttons.test.js

@@ -38,6 +38,36 @@ describe( 'Buttons', () => {
 		expect( await getEditedPostContent() ).toMatchSnapshot();
 	} );
 
+	it( 'moves focus from the link editor back to the button when escape is pressed after the URL has been submitted', async () => {
+		// Regression: https://github.com/WordPress/gutenberg/issues/34307
+		await insertBlock( 'Buttons' );
+		await pressKeyWithModifier( 'primary', 'k' );
+		await page.waitForFunction(
+			() => !! document.activeElement.closest( '.block-editor-url-input' )
+		);
+		await page.keyboard.type( 'https://example.com' );
+		await page.keyboard.press( 'Enter' );
+		await page.waitForFunction(
+			() =>
+				document.activeElement ===
+				document.querySelector(
+					'.block-editor-link-control a[href="https://example.com"]'
+				)
+		);
+		await page.keyboard.press( 'Escape' );
+
+		// Focus should move from the link control to the button block's text.
+		await page.waitForFunction(
+			() =>
+				document.activeElement ===
+				document.querySelector( '[aria-label="Button text"]' )
+		);
+
+		// The link control should still be visible when a URL is set.
+		const linkControl = await page.$( '.block-editor-link-control' );
+		expect( linkControl ).toBeTruthy();
+	} );
+
 	it( 'can jump to the link editor using the keyboard shortcut', async () => {
 		await insertBlock( 'Buttons' );
 		await page.keyboard.type( 'WordPress' );