fix(container-menu)!: improves logic for returning focus to trigger

packages/menu/src/MenuContainer.spec.tsx

@@ -371,6 +371,20 @@ describe('MenuContainer', () => {
 
           expect(trigger).toHaveFocus();
         });
+
+        it('does not return focus to trigger when Escape key pressed in expanded menu if `restoreFocus` is false', async () => {
+          const { getByTestId } = render(<TestMenu items={ITEMS} restoreFocus={false} />);
+          const trigger = getByTestId('trigger');
+
+          trigger.focus();
+
+          await waitFor(async () => {
+            await user.keyboard('{ArrowDown}');
+            await user.keyboard('{Escape}');
+          });
+
+          expect(trigger).not.toHaveFocus();
+        });
       });
 
       describe('menu items', () => {

packages/menu/src/MenuContainer.tsx

@@ -29,5 +29,10 @@ MenuContainer.propTypes = {
   defaultExpanded: PropTypes.bool,
   selectedItems: PropTypes.arrayOf(PropTypes.any),
   focusedValue: PropTypes.oneOfType([PropTypes.string]),
-  defaultFocusedValue: PropTypes.oneOfType([PropTypes.string])
+  defaultFocusedValue: PropTypes.oneOfType([PropTypes.string]),
+  restoreFocus: PropTypes.bool
+};
+
+MenuContainer.defaultProps = {
+  restoreFocus: true
 };

packages/menu/src/types.ts

@@ -66,6 +66,8 @@ export interface IUseMenuProps<T = HTMLButtonElement, M = HTMLElement> {
   focusedValue?: string | null;
   /** Determines focused value on menu initialization */
   defaultFocusedValue?: string;
+  /** Returns keyboard focus to the element that triggered the menu */
+  restoreFocus?: boolean;
   /** Sets the selected values in a controlled menu */
   selectedItems?: ISelectedItem[];
   /**

packages/menu/src/useMenu.ts

@@ -12,7 +12,6 @@ import React, {
   useEffect,
   useMemo,
   useReducer,
-  useRef,
   useState
 } from 'react';
 import { useSelection } from '@zendeskgarden/container-selection';
@@ -49,6 +48,7 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
   onChange = () => undefined,
   isExpanded,
   defaultExpanded = false,
+  restoreFocus = true,
   selectedItems,
   focusedValue,
   defaultFocusedValue
@@ -94,8 +94,6 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
    */
 
   const [menuVisible, setMenuVisible] = useState<boolean>(false);
-  const focusTriggerRef = useRef<boolean>(false);
-
   const [state, dispatch] = useReducer(stateReducer, {
     focusedValue,
     isExpanded: isExpanded || defaultExpanded,
@@ -135,6 +133,15 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
 
   // Internal
 
+  const returnFocusToTrigger = useCallback(
+    (skip?: boolean) => {
+      if (!skip && restoreFocus && triggerRef.current) {
+        triggerRef.current.focus();
+      }
+    },
+    [triggerRef, restoreFocus]
+  );
+
   const closeMenu = useCallback(
     (changeType: string) => {
       dispatch({
@@ -281,13 +288,22 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
         }
       });
 
+      // Skip focus return when isExpanded === true
+      returnFocusToTrigger(!controlledIsExpanded);
+
       onChange({
         type: changeType,
         focusedValue: null,
         isExpanded: !controlledIsExpanded
       });
     },
-    [controlledIsExpanded, isFocusedValueControlled, isExpandedControlled, onChange]
+    [
+      isFocusedValueControlled,
+      isExpandedControlled,
+      controlledIsExpanded,
+      returnFocusToTrigger,
+      onChange
+    ]
   );
 
   const handleTriggerKeyDown = useCallback(
@@ -321,14 +337,23 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
           }
         });
 
+        returnFocusToTrigger();
+
         onChange({
           type: changeType,
           focusedValue: defaultFocusedValue || nextFocusedValue,
           isExpanded: true
         });
       }
     },
-    [isExpandedControlled, isFocusedValueControlled, defaultFocusedValue, onChange, values]
+    [
+      values,
+      isFocusedValueControlled,
+      defaultFocusedValue,
+      isExpandedControlled,
+      returnFocusToTrigger,
+      onChange
+    ]
   );
 
   const handleMenuKeyDown = useCallback(
@@ -341,25 +366,25 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
 
         const type = StateChangeTypes[key === KEYS.ESCAPE ? 'MenuKeyDownEscape' : 'MenuKeyDownTab'];
 
-        if (KEYS.ESCAPE === key) {
-          focusTriggerRef.current = true;
-        }
+        // TODO: Investigate why focus goes to body instead of next interactive element on TAB keydown. Meanwhile, returning focus to trigger.
+        returnFocusToTrigger();
 
         closeMenu(type);
       }
     },
-    [closeMenu, focusTriggerRef]
+    [closeMenu, returnFocusToTrigger]
   );
 
   const handleMenuBlur = useCallback(
     (event: MouseEvent) => {
       const path = event.composedPath();
 
       if (!path.includes(menuRef.current!) && !path.includes(triggerRef.current!)) {
+        returnFocusToTrigger();
         closeMenu(StateChangeTypes.MenuBlur);
       }
     },
-    [closeMenu, menuRef, triggerRef]
+    [closeMenu, menuRef, returnFocusToTrigger, triggerRef]
   );
 
   const handleMenuMouseLeave = useCallback(() => {
@@ -395,6 +420,8 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
         }
       });
 
+      returnFocusToTrigger(isTransitionItem);
+
       onChange({
         type: changeType,
         value: item.value,
@@ -403,10 +430,11 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
       });
     },
     [
+      getSelectedItems,
       state.nestedPathIds,
       isExpandedControlled,
       isSelectedItemsControlled,
-      getSelectedItems,
+      returnFocusToTrigger,
       onChange
     ]
   );
@@ -448,19 +476,13 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
           ...(nextSelection && { selectedItems: nextSelection })
         };
 
-        if (item.href) {
-          if (key === KEYS.SPACE) {
-            event.preventDefault();
+        event.preventDefault();
 
-            triggerLink(event.target as HTMLAnchorElement, environment || window);
-          }
-        } else {
-          event.preventDefault();
+        if (item.href) {
+          triggerLink(event.target as HTMLAnchorElement, environment || window);
         }
 
-        if (!isTransitionItem) {
-          focusTriggerRef.current = true;
-        }
+        returnFocusToTrigger(isTransitionItem);
       } else if (key === KEYS.RIGHT) {
         if (rtl && isPrevious) {
           changeType = StateChangeTypes.MenuItemKeyDownPrevious;
@@ -529,15 +551,15 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
       }
     },
     [
-      rtl,
-      state.nestedPathIds,
+      getSelectedItems,
       isExpandedControlled,
-      isFocusedValueControlled,
       isSelectedItemsControlled,
-      focusTriggerRef,
+      returnFocusToTrigger,
       environment,
+      rtl,
       getNextFocusedValue,
-      getSelectedItems,
+      isFocusedValueControlled,
+      state.nestedPathIds,
       onChange
     ]
   );
@@ -594,12 +616,8 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
   }, [controlledIsExpanded, handleMenuBlur, handleMenuKeyDown, environment]);
 
   /**
-   * Handles focus depending on menu state:
-   * - When opened, focus the menu via `focusedValue`
-   * - When closed, focus the trigger via `triggerRef`
-   *
-   * This effect is intended to prevent focusing the trigger or menu
-   * unless the menu is in the right expansion state.
+   * When the menu is opened, this effect sets focus on the current menu item using `focusedValue`
+   * or on the first menu item.
    */
   useEffect(() => {
     if (state.focusOnOpen && menuVisible && controlledFocusedValue && controlledIsExpanded) {
@@ -614,13 +632,7 @@ export const useMenu = <T extends HTMLElement = HTMLElement, M extends HTMLEleme
 
       ref && ref.focus();
     }
-
-    if (!menuVisible && !controlledIsExpanded && focusTriggerRef.current) {
-      triggerRef?.current?.focus();
-      focusTriggerRef.current = false;
-    }
   }, [
-    focusTriggerRef,
     values,
     menuVisible,
     itemRefs,