🦾 Fix search overlay focus restoration on close (#60299)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Evan Bonsignori <evanabonsignori@gmail.com>

Author: 3 people

src/frame/components/page-header/Header.tsx

@@ -43,16 +43,23 @@ export const Header = () => {
   const isEarlyAccessPage = currentProduct && currentProduct.id === 'early-access'
   const { width } = useInnerWindowWidth()
   const returnFocusRef = useRef(null)
-  const searchButtonRef = useRef<HTMLButtonElement>(null)
+  const searchButtonRefLarge = useRef<HTMLButtonElement>(null)
+  const searchButtonRefSmall = useRef<HTMLButtonElement>(null)
   const { initializeCTA } = useCTAPopoverContext()
   const { isSearchOpen, setIsSearchOpen } = useSearchOverlayContext()
 
+  // The lg breakpoint (1012px) determines which search button is visible.
+  // Pass the correct ref to SearchOverlayContainer so Primer's Overlay
+  // restores focus to the visible trigger element on close.
+  const isLargeViewport = width !== null && width >= 1012
+  const searchButtonRef = isLargeViewport ? searchButtonRefLarge : searchButtonRefSmall
+
   const SearchButtonLarge: JSX.Element = (
     <SearchBarButton
       isSearchOpen={isSearchOpen}
       setIsSearchOpen={setIsSearchOpen}
       params={params}
-      searchButtonRef={searchButtonRef}
+      searchButtonRef={searchButtonRefLarge}
       instanceId="large"
     />
   )
@@ -62,7 +69,7 @@ export const Header = () => {
       isSearchOpen={isSearchOpen}
       setIsSearchOpen={setIsSearchOpen}
       params={params}
-      searchButtonRef={searchButtonRef}
+      searchButtonRef={searchButtonRefSmall}
       instanceId="small"
     />
   )
@@ -90,17 +97,18 @@ export const Header = () => {
     return () => window.removeEventListener('keydown', close)
   }, [])
 
-  // For the UI in smaller browser widths, and focus the picker menu button when the search
-  // input is closed.
+  // For the UI in smaller browser widths, focus the picker menu button when
+  // the sidebar is closed (not when the search overlay is closed — the
+  // overlay's returnFocusRef handles focus restoration to the search button).
   useEffect(() => {
-    if (!isSearchOpen && isMounted.current && menuButtonRef.current) {
+    if (!isSidebarOpen && isMounted.current && menuButtonRef.current) {
       menuButtonRef.current.focus()
     }
 
     if (!isMounted.current) {
       isMounted.current = true
     }
-  }, [isSearchOpen])
+  }, [isSidebarOpen])
 
   // When the sidebar overlay is opened, prevent the main content from being
   // scrollable.

src/search/components/input/SearchBarButton.tsx

@@ -52,55 +52,56 @@ export function SearchBarButton({
 
   return (
     <>
-      {/* We don't want to show the input when overlay is open */}
-      {!isSearchOpen ? (
-        <>
-          {/* On mobile only the IconButton is shown */}
-          <IconButton
-            data-testid="mobile-search-button"
-            data-instance={instanceId}
-            ref={searchButtonRef}
-            className={styles.searchIconButton}
-            onClick={handleClick}
-            tabIndex={0}
-            aria-label={t('search.input.placeholder_no_icon')}
-            icon={SearchIcon}
-          />
-          {/* On large and up the SearchBarButton is shown */}
-          <button
-            data-testid="search"
-            data-instance={instanceId}
-            tabIndex={0}
-            aria-label={t('search.input.placeholder_no_icon')}
-            className={styles.searchInputButton}
-            onKeyDown={handleKeyDown}
-            onClick={handleClick}
-            ref={searchButtonRef}
+      {/* Keep buttons in DOM even when overlay is open so returnFocusRef
+          can restore focus to the trigger element on overlay close. */}
+      <div style={isSearchOpen ? { visibility: 'hidden', position: 'absolute' } : undefined}>
+        {/* On mobile only the IconButton is shown. The "small" instance is
+            the mobile one, so it gets the ref on the icon button. */}
+        <IconButton
+          data-testid="mobile-search-button"
+          data-instance={instanceId}
+          ref={instanceId === 'small' ? searchButtonRef : undefined}
+          className={styles.searchIconButton}
+          onClick={handleClick}
+          tabIndex={isSearchOpen ? -1 : 0}
+          aria-label={t('search.input.placeholder_no_icon')}
+          icon={SearchIcon}
+        />
+        {/* On large and up the SearchBarButton is shown. The "large" instance
+            is the desktop one, so it gets the ref on the desktop button. */}
+        <button
+          data-testid="search"
+          data-instance={instanceId}
+          tabIndex={isSearchOpen ? -1 : 0}
+          aria-label={t('search.input.placeholder_no_icon')}
+          className={styles.searchInputButton}
+          onKeyDown={handleKeyDown}
+          onClick={handleClick}
+          ref={instanceId === 'large' ? searchButtonRef : undefined}
+        >
+          {/* Styled to look like an input */}
+          <div
+            className={cx('d-flex align-items-center flex-grow-1', styles.searchInputContainer)}
+            aria-hidden
+            tabIndex={-1}
           >
-            {/* Styled to look like an input */}
-            <div
-              className={cx('d-flex align-items-center flex-grow-1', styles.searchInputContainer)}
-              aria-hidden
-              tabIndex={-1}
+            <span
+              className={cx(styles.queryText, !urlSearchInputQuery ? styles.placeholder : null)}
             >
-              <span
-                className={cx(styles.queryText, !urlSearchInputQuery ? styles.placeholder : null)}
-              >
-                {urlSearchInputQuery ? (
-                  urlSearchInputQuery
-                ) : (
-                  <>
-                    <span className={styles.placeholderText}>{placeHolderElements}</span>
-                  </>
-                )}
-              </span>
-            </div>
-            <span className={styles.searchIconContainer} aria-hidden tabIndex={-1}>
-              <SearchIcon />
+              {urlSearchInputQuery ? (
+                urlSearchInputQuery
+              ) : (
+                <>
+                  <span className={styles.placeholderText}>{placeHolderElements}</span>
+                </>
+              )}
             </span>
-          </button>
-        </>
-      ) : null}
+          </div>
+          <span className={styles.searchIconContainer} aria-hidden tabIndex={-1}>
+            <SearchIcon />
+          </span>
+        </button>
+      </div>
     </>
   )
 }