[Bug] Fix useFocusTrap scrolling issue (#7242)

Co-authored-by: Siddharth Kshetrapal <[email protected]>

Author: pksjce

.changeset/unlucky-icons-speak.md

@@ -0,0 +1,5 @@
+---
+"@primer/react": patch
+---
+
+useFocusTrap - Fix [bug related to restoring focus on scrolling](https://github.com/github/primer/issues/5200)

packages/react/src/hooks/useFocusTrap.ts

@@ -1,6 +1,7 @@
 import React from 'react'
 import {focusTrap} from '@primer/behaviors'
 import {useProvidedRefOrCreate} from './useProvidedRefOrCreate'
+import {useOnOutsideClick} from './useOnOutsideClick'
 
 export interface FocusTrapHookSettings {
   /**
@@ -34,6 +35,12 @@ export interface FocusTrapHookSettings {
    * Overrides restoreFocusOnCleanUp
    */
   returnFocusRef?: React.RefObject<HTMLElement | null>
+  /**
+   * If true, it should allow focus to escape the trap when clicking outside of the trap container and mark it as disabled.
+   *
+   * Overrides restoreFocusOnCleanUp and returnFocusRef
+   */
+  allowOutsideClick?: boolean
 }
 
 /**
@@ -45,6 +52,7 @@ export function useFocusTrap(
   settings?: FocusTrapHookSettings,
   dependencies: React.DependencyList = [],
 ): {containerRef: React.RefObject<HTMLElement | null>; initialFocusRef: React.RefObject<HTMLElement | null>} {
+  const [outsideClicked, setOutsideClicked] = React.useState(false)
   const containerRef = useProvidedRefOrCreate(settings?.containerRef)
   const initialFocusRef = useProvidedRefOrCreate(settings?.initialFocusRef)
   const disabled = settings?.disabled
@@ -53,14 +61,17 @@ export function useFocusTrap(
 
   // If we are enabling a focus trap and haven't already stored the previously focused element
   // go ahead an do that so we can restore later when the trap is disabled.
-  if (!previousFocusedElement.current && !settings?.disabled) {
+  if (!previousFocusedElement.current && !disabled) {
     previousFocusedElement.current = document.activeElement
   }
 
   // This function removes the event listeners that enable the focus trap and restores focus
   // to the previously-focused element (if necessary).
   function disableTrap() {
     abortController.current?.abort()
+    if (settings?.allowOutsideClick && outsideClicked) {
+      return
+    }
     if (settings?.returnFocusRef && settings.returnFocusRef.current instanceof HTMLElement) {
       settings.returnFocusRef.current.focus()
     } else if (settings?.restoreFocusOnCleanUp && previousFocusedElement.current instanceof HTMLElement) {
@@ -85,6 +96,17 @@ export function useFocusTrap(
     // eslint-disable-next-line react-hooks/exhaustive-deps
     [containerRef, initialFocusRef, disabled, ...dependencies],
   )
+  useOnOutsideClick({
+    containerRef: containerRef as React.RefObject<HTMLDivElement>,
+    onClickOutside: () => {
+      setOutsideClicked(true)
+      if (settings?.allowOutsideClick) {
+        if (settings.returnFocusRef) settings.returnFocusRef = undefined
+        settings.restoreFocusOnCleanUp = false
+        abortController.current?.abort()
+      }
+    },
+  })
 
   return {containerRef, initialFocusRef}
 }

packages/react/src/stories/useFocusTrap.stories.tsx

@@ -3,6 +3,7 @@ import type {Meta} from '@storybook/react-vite'
 
 import {Button, Flash, Stack, Text} from '..'
 import {useFocusTrap} from '../hooks/useFocusTrap'
+import {useOnEscapePress} from '../hooks/useOnEscapePress'
 import classes from './FocusTrap.stories.module.css'
 
 export default {
@@ -118,6 +119,95 @@ export const RestoreFocus = () => {
   )
 }
 
+export const RestoreFocusMinimal = () => {
+  const [enabled, setEnabled] = React.useState(false)
+  const toggleButtonRef = React.useRef<HTMLButtonElement>(null)
+  const {containerRef} = useFocusTrap({
+    disabled: !enabled,
+    restoreFocusOnCleanUp: true,
+    returnFocusRef: toggleButtonRef,
+    allowOutsideClick: true,
+  })
+
+  useOnEscapePress(
+    React.useCallback(
+      e => {
+        if (!enabled) return
+        e.preventDefault()
+        setEnabled(false)
+      },
+      [enabled, setEnabled],
+    ),
+    [enabled, setEnabled],
+  )
+
+  return (
+    <>
+      <HelperGlobalStyling />
+      <Stack direction="vertical" gap="normal">
+        <Flash style={{marginBottom: 'var(--base-size-12)'}}>
+          Minimal focus trap example. Click to toggle focus trap to toggle. While enabled, focus stays inside the green
+          zone. Disabling restores focus to the toggle button.
+        </Flash>
+        <Button
+          ref={toggleButtonRef}
+          onClick={() => {
+            if (enabled) {
+              setEnabled(false)
+            } else {
+              setEnabled(true)
+            }
+          }}
+        >
+          {enabled ? 'Disable' : 'Enable'} focus trap
+        </Button>
+        <div
+          style={{
+            height: '900px',
+            overflow: 'auto',
+            border: '1px dashed var(--borderColor-default)',
+            padding: 'var(--base-size-16)',
+            background: 'var(--bgColor-muted)',
+          }}
+          aria-hidden="true"
+        >
+          <Text
+            as="p"
+            style={{
+              fontSize: '12px',
+              lineHeight: '1.25',
+              margin: 0,
+            }}
+          >
+            Scroll down to reach the trap zone. This spacer exists so that when the trap zone becomes active you can
+            scroll such that the original toggle button is no longer visible. When you press Escape or the Close trap
+            button, focus will still restore to the toggle button and the browser will scroll it back into view.
+          </Text>
+          <Text
+            as="p"
+            style={{
+              fontSize: '12px',
+              lineHeight: '1.25',
+              margin: 0,
+            }}
+          >
+            (Content intentionally verbose to create vertical space.)
+          </Text>
+        </div>
+        <div className={classes.TrapZone} ref={containerRef as React.RefObject<HTMLDivElement>}>
+          <Stack direction="vertical" gap="normal">
+            <MarginButton>First</MarginButton>
+            <MarginButton>Second</MarginButton>
+            <MarginButton>Third</MarginButton>
+            <Button onClick={() => setEnabled(false)}>Close trap</Button>
+          </Stack>
+        </div>
+        <Button>Click here to escape trap</Button>
+      </Stack>
+    </>
+  )
+}
+
 export const CustomInitialFocus = () => {
   const [trapEnabled, setTrapEnabled] = React.useState(false)
   const {containerRef, initialFocusRef} = useFocusTrap({disabled: !trapEnabled})