Add preventFocusOnClose to ActionMenu

.changeset/actionmenu-prevent-focus-on-close.md

@@ -0,0 +1,5 @@
+---
+'@primer/react': minor
+---
+
+ActionMenu, Overlay: Add `preventFocusOnClose` to let consumers manage focus when an overlay closes.

packages/react/src/ActionMenu/ActionMenu.docs.json

@@ -117,6 +117,12 @@
           "defaultValue": "",
           "description": "Callback that is called when the position of the overlay changes"
         },
+        {
+          "name": "preventFocusOnClose",
+          "type": "boolean",
+          "defaultValue": "false",
+          "description": "If true, focus will not be returned to the anchor or returnFocusRef when the menu closes. Consumers must manage focus themselves."
+        },
         {
           "name": "data-test-id",
           "type": "unknown",

packages/react/src/ActionMenu/ActionMenu.examples.stories.tsx

@@ -1,4 +1,7 @@
 import React, {useState, useCallback} from 'react'
+import {flushSync} from 'react-dom'
+import type {StoryObj} from '@storybook/react-vite'
+import {expect, userEvent, within} from 'storybook/test'
 import {ActionMenu} from '.'
 import {ActionList} from '../ActionList'
 import {Button, IconButton} from '../Button'
@@ -200,6 +203,66 @@ export const ControlledMenu = () => {
   )
 }
 
+export const PreventFocusOnClose: StoryObj = {
+  render: function PreventFocusOnCloseExample() {
+    const [open, setOpen] = React.useState(false)
+    const [preview, setPreview] = React.useState(false)
+    const editorRef = React.useRef<HTMLInputElement>(null)
+    const selectionRef = React.useRef({start: 0, end: 0})
+
+    const onOpenChange = (nextOpen: boolean) => {
+      const editor = editorRef.current
+      if (nextOpen) {
+        selectionRef.current = {start: editor?.selectionStart ?? 0, end: editor?.selectionEnd ?? 0}
+        setOpen(true)
+      } else {
+        // Unmount the menu and its focus trap before restoring the editor selection.
+        flushSync(() => setOpen(false))
+        editor?.focus()
+        editor?.setSelectionRange(selectionRef.current.start, selectionRef.current.end)
+      }
+    }
+
+    return (
+      <>
+        <FormControl>
+          <FormControl.Label>Document</FormControl.Label>
+          <TextInput ref={editorRef} defaultValue="Edit this document, then change the link presentation." />
+        </FormControl>
+        <ActionMenu open={open} onOpenChange={onOpenChange}>
+          <ActionMenu.Button key={preview ? 'preview' : 'link'}>
+            {preview ? 'Preview actions' : 'Link actions'}
+          </ActionMenu.Button>
+          <ActionMenu.Overlay preventFocusOnClose>
+            <ActionList>
+              <ActionList.Item onSelect={() => setPreview(true)}>Convert to preview</ActionList.Item>
+              <ActionList.Item onSelect={() => setPreview(false)}>Convert to link</ActionList.Item>
+            </ActionList>
+          </ActionMenu.Overlay>
+        </ActionMenu>
+      </>
+    )
+  },
+  play: async ({canvasElement}) => {
+    const canvas = within(canvasElement)
+    const editor = canvas.getByRole<HTMLInputElement>('textbox', {name: 'Document'})
+    await userEvent.click(editor)
+    await userEvent.keyboard('{Home}{ArrowRight}{ArrowRight}{ArrowRight}')
+    const selectionStart = editor.selectionStart
+    const selectionEnd = editor.selectionEnd
+    const anchor = canvas.getByRole('button', {name: 'Link actions'})
+
+    await userEvent.click(anchor)
+    await userEvent.click(within(canvasElement.ownerDocument.body).getByRole('menuitem', {name: 'Convert to preview'}))
+
+    await expect(anchor).not.toBeInTheDocument()
+    await expect(canvas.getByRole('button', {name: 'Preview actions'})).not.toHaveFocus()
+    await expect(editor).toHaveFocus()
+    await expect(editor.selectionStart).toBe(selectionStart)
+    await expect(editor.selectionEnd).toBe(selectionEnd)
+  },
+}
+
 export const ShortcutMenu = () => {
   const [open, setOpen] = React.useState(false)
   const triggerRef = React.useRef<HTMLButtonElement>(null)

packages/react/src/ActionMenu/ActionMenu.test.tsx

@@ -41,12 +41,12 @@ vi.mock('@primer/behaviors', async () => {
   }
 })
 
-function Example(): JSX.Element {
+function Example({preventFocusOnClose}: {preventFocusOnClose?: boolean} = {}): JSX.Element {
   return (
     <BaseStyles>
       <ActionMenu>
         <ActionMenu.Button>Toggle Menu</ActionMenu.Button>
-        <ActionMenu.Overlay>
+        <ActionMenu.Overlay preventFocusOnClose={preventFocusOnClose}>
           <ActionList>
             <ActionList.Item>New file</ActionList.Item>
             <ActionList.Divider />
@@ -244,6 +244,147 @@ describe('ActionMenu', () => {
     expect(component.queryByRole('menu')).toBeNull()
   })
 
+  describe('preventFocusOnClose', () => {
+    it.each([
+      {gesture: 'item-click', preventFocusOnClose: undefined},
+      {gesture: 'item-click', preventFocusOnClose: false},
+      {gesture: 'item-click', preventFocusOnClose: true},
+      {gesture: 'Enter', preventFocusOnClose: true},
+      {gesture: 'Space', preventFocusOnClose: true},
+      {gesture: 'Escape', preventFocusOnClose: undefined},
+      {gesture: 'Escape', preventFocusOnClose: false},
+      {gesture: 'Escape', preventFocusOnClose: true},
+      {gesture: 'outside-click', preventFocusOnClose: true},
+    ])('handles $gesture with preventFocusOnClose=$preventFocusOnClose', async ({gesture, preventFocusOnClose}) => {
+      const component = HTMLRender(
+        <>
+          <Example preventFocusOnClose={preventFocusOnClose} />
+          <Button>Outside</Button>
+        </>,
+      )
+      const user = userEvent.setup()
+      const trigger = component.getByRole('button', {name: 'Toggle Menu'})
+      await user.click(trigger)
+      const item = component.getByRole('menuitem', {name: 'New file'})
+      item.focus()
+
+      const focus = vi.spyOn(trigger, 'focus')
+      if (gesture === 'item-click') {
+        await user.click(item)
+      } else if (gesture === 'outside-click') {
+        await user.click(component.getByRole('button', {name: 'Outside'}))
+      } else {
+        await user.keyboard(gesture === 'Space' ? ' ' : `{${gesture}}`)
+      }
+
+      expect(component.queryByRole('menu')).not.toBeInTheDocument()
+      if (preventFocusOnClose) {
+        expect(focus).not.toHaveBeenCalled()
+        expect(trigger).not.toHaveFocus()
+      } else {
+        expect(focus).toHaveBeenCalled()
+        expect(trigger).toHaveFocus()
+      }
+      if (gesture === 'outside-click') {
+        expect(component.getByRole('button', {name: 'Outside'})).toHaveFocus()
+      }
+    })
+
+    it.each([undefined, false, true])(
+      'handles anchor-click closure with preventFocusOnClose=%s',
+      async preventFocusOnClose => {
+        const returnFocusRef = {current: null as HTMLButtonElement | null}
+        const component = HTMLRender(
+          <BaseStyles>
+            <Button ref={returnFocusRef}>Return focus</Button>
+            <ActionMenu>
+              <ActionMenu.Button>Toggle Menu</ActionMenu.Button>
+              <ActionMenu.Overlay preventFocusOnClose={preventFocusOnClose} returnFocusRef={returnFocusRef}>
+                <ActionList>
+                  <ActionList.Item>New file</ActionList.Item>
+                </ActionList>
+              </ActionMenu.Overlay>
+            </ActionMenu>
+          </BaseStyles>,
+        )
+        const user = userEvent.setup()
+        const trigger = component.getByRole('button', {name: 'Toggle Menu'})
+        const returnFocusTarget = component.getByRole('button', {name: 'Return focus'})
+
+        await user.click(trigger)
+        const focus = vi.spyOn(returnFocusTarget, 'focus')
+        await user.click(trigger)
+
+        expect(component.queryByRole('menu')).not.toBeInTheDocument()
+        if (preventFocusOnClose) {
+          expect(focus).not.toHaveBeenCalled()
+          expect(returnFocusTarget).not.toHaveFocus()
+        } else {
+          expect(focus).toHaveBeenCalled()
+          expect(returnFocusTarget).toHaveFocus()
+        }
+      },
+    )
+
+    it('skips restoration to a custom returnFocusRef on controlled closure', async () => {
+      const returnFocusRef = {current: document.createElement('button')}
+      const example = (open: boolean) => (
+        <BaseStyles>
+          <ActionMenu open={open} onOpenChange={() => {}}>
+            <ActionMenu.Button>Toggle Menu</ActionMenu.Button>
+            <ActionMenu.Overlay preventFocusOnClose returnFocusRef={returnFocusRef}>
+              <ActionList>
+                <ActionList.Item>New file</ActionList.Item>
+              </ActionList>
+            </ActionMenu.Overlay>
+          </ActionMenu>
+        </BaseStyles>
+      )
+      const component = HTMLRender(example(true))
+      const focus = vi.spyOn(returnFocusRef.current, 'focus')
+
+      await waitFor(() => expect(component.getByRole('menuitem', {name: 'New file'})).toHaveFocus())
+      component.rerender(example(false))
+
+      expect(component.queryByRole('menu')).not.toBeInTheDocument()
+      expect(focus).not.toHaveBeenCalled()
+      expect(component.getByRole('button', {name: 'Toggle Menu'})).not.toHaveFocus()
+    })
+
+    it('does not focus a captured anchor that an action replaces', async () => {
+      function ReplaceableAnchor() {
+        const [preview, setPreview] = useState(false)
+
+        return (
+          <BaseStyles>
+            <ActionMenu>
+              <ActionMenu.Anchor>
+                <Button key={preview ? 'preview' : 'link'}>{preview ? 'Preview' : 'Link'}</Button>
+              </ActionMenu.Anchor>
+              <ActionMenu.Overlay preventFocusOnClose>
+                <ActionList>
+                  <ActionList.Item onSelect={() => setPreview(true)}>Convert to preview</ActionList.Item>
+                </ActionList>
+              </ActionMenu.Overlay>
+            </ActionMenu>
+          </BaseStyles>
+        )
+      }
+      const component = HTMLRender(<ReplaceableAnchor />)
+      const user = userEvent.setup()
+      const trigger = component.getByRole('button', {name: 'Link'})
+      await user.click(trigger)
+      const focus = vi.spyOn(trigger, 'focus')
+
+      await user.click(component.getByRole('menuitem', {name: 'Convert to preview'}))
+
+      expect(component.queryByRole('menu')).not.toBeInTheDocument()
+      expect(trigger).not.toBeInTheDocument()
+      expect(component.getByRole('button', {name: 'Preview'})).not.toHaveFocus()
+      expect(focus).not.toHaveBeenCalled()
+    })
+  })
+
   it('should not close Menu if event is prevented', async () => {
     const component = HTMLRender(<Example />)
     const button = component.getByRole('button')

packages/react/src/ActionMenu/ActionMenu.types.test.tsx

@@ -16,3 +16,12 @@ export function actionButtonWithInvalidSize() {
   //@ts-expect-error size must be one of the valid values
   return <ActionMenu.Button size="some-unknownsize">Click me!</ActionMenu.Button>
 }
+
+export function actionOverlayWithPreventFocusOnClose() {
+  return <ActionMenu.Overlay preventFocusOnClose>Menu contents</ActionMenu.Overlay>
+}
+
+export function actionOverlayWithInvalidPreventFocusOnClose() {
+  // @ts-expect-error preventFocusOnClose must be a boolean
+  return <ActionMenu.Overlay preventFocusOnClose="true">Menu contents</ActionMenu.Overlay>
+}

packages/react/src/Overlay/Overlay.docs.json

@@ -16,6 +16,12 @@
   ],
   "importPath": "@primer/react",
   "props": [
+    {
+      "name": "preventFocusOnClose",
+      "type": "boolean",
+      "defaultValue": "false",
+      "description": "If true, focus will not be returned to returnFocusRef when the Overlay is unmounted. Consumers must manage focus themselves."
+    },
     {
       "name": "returnFocusRef",
       "type": "React.RefObject<HTMLElement>",
@@ -155,4 +161,4 @@
     }
   ],
   "subcomponents": []
-}
+}

packages/react/src/Overlay/Overlay.test.tsx

@@ -15,11 +15,13 @@ type TestComponentSettings = {
   width?: 'small' | 'medium' | 'large' | 'auto' | 'xlarge' | 'xxlarge'
   callback?: () => void
   preventFocusOnOpen?: boolean
+  preventFocusOnClose?: boolean
 }
 const TestComponent = ({
   initialFocus,
   width = 'small',
   preventFocusOnOpen = undefined,
+  preventFocusOnClose,
   callback,
 }: TestComponentSettings) => {
   const [isOpen, setIsOpen] = useState(false)
@@ -57,6 +59,7 @@ const TestComponent = ({
             onClickOutside={closeOverlay}
             width={width}
             preventFocusOnOpen={preventFocusOnOpen}
+            preventFocusOnClose={preventFocusOnClose}
             role="dialog"
           >
             <div
@@ -159,6 +162,46 @@ describe('Overlay', () => {
     expect(document.activeElement).toEqual(openButton)
   })
 
+  it.each([undefined, false, true])(
+    'controls focus restoration with preventFocusOnClose set to %s',
+    async preventFocusOnClose => {
+      const user = userEvent.setup()
+      const component = render(<TestComponent preventFocusOnClose={preventFocusOnClose} />)
+      const trigger = component.getByRole('button', {name: 'open overlay'})
+      await user.click(trigger)
+
+      expect(component.getByRole('button', {name: 'Cancel'})).toHaveFocus()
+      expect(component.getByRole('dialog')).not.toHaveAttribute('preventFocusOnClose')
+
+      await user.click(component.getByRole('button', {name: 'Cancel'}))
+      expect(component.queryByRole('dialog')).not.toBeInTheDocument()
+
+      if (preventFocusOnClose) {
+        expect(trigger).not.toHaveFocus()
+      } else {
+        expect(trigger).toHaveFocus()
+      }
+    },
+  )
+
+  it('preserves consumer-managed focus when automatic restoration is prevented', async () => {
+    const user = userEvent.setup()
+    const component = render(
+      <TestComponent
+        preventFocusOnClose
+        callback={() => {
+          component.getByRole('button', {name: 'outside'}).focus()
+        }}
+      />,
+    )
+
+    await user.click(component.getByRole('button', {name: 'open overlay'}))
+    await user.click(component.getByRole('button', {name: 'Cancel'}))
+
+    expect(component.queryByRole('dialog')).not.toBeInTheDocument()
+    expect(component.getByRole('button', {name: 'outside'})).toHaveFocus()
+  })
+
   it('should call function when user clicks outside container', async () => {
     const user = userEvent.setup()
     const mockFunction = vi.fn()

packages/react/src/Overlay/Overlay.tsx

@@ -127,6 +127,7 @@ type ContainerProps = {
   portalContainerName?: string
   preventOverflow?: boolean
   preventFocusOnOpen?: boolean
+  preventFocusOnClose?: boolean
   returnFocusRef: React.RefObject<HTMLElement | null>
 }
 
@@ -143,6 +144,7 @@ type internalOverlayProps = Merge<OwnOverlayProps, ContainerProps>
  * @param portalContainerName Optional. The name of the portal container to render the Overlay into.
  * @param preventOverflow Optional. The Overlay width will be adjusted responsively if there is not enough space to display the Overlay. If `preventOverflow` is `true`, the width of the `Overlay` will not be adjusted.
  * @param preventFocusOnOpen Optional. If 'true', focus will not be applied when the component is first mounted, even if initialFocusRef prop is given.
+ * @param preventFocusOnClose Optional. If true, focus will not be returned when the component is unmounted. Consumers must manage focus themselves.
  * @param returnFocusRef Required. Ref for the element to focus when the `Overlay` is closed.
  * @param right Optional. Horizontal right position of the overlay, relative to its closest positioned ancestor (often its `Portal`).
  * @param width Sets the width of the `Overlay`, pick from our set list of widths, or pass `auto` to automatically set the width based on the content of the `Overlay`. `small` corresponds to `256px`, `medium` corresponds to `320px`, `large` corresponds to `480px`, `xlarge` corresponds to `640px`, `xxlarge` corresponds to `960px`.
@@ -162,6 +164,7 @@ const Overlay = React.forwardRef<HTMLDivElement, internalOverlayProps>(
       portalContainerName,
       preventOverflow = true,
       preventFocusOnOpen,
+      preventFocusOnClose,
       returnFocusRef,
       right,
       role = 'none',
@@ -187,6 +190,7 @@ const Overlay = React.forwardRef<HTMLDivElement, internalOverlayProps>(
       onClickOutside,
       initialFocusRef,
       preventFocusOnOpen,
+      preventFocusOnClose,
     })
 
     useEffect(() => {