prevent focus from being lost to the body when submenutrigger is virtually focused

typically submenus dont have focus restore turned on since it would move focus manually back to the trigger when keyboard closing the menu. However, we cant move focus to virtually focused triggers so enable focus restore on the submenu in these cases

Author: LFDanLu

packages/@react-aria/menu/src/useSubmenuTrigger.ts

@@ -38,7 +38,9 @@ export interface AriaSubmenuTriggerProps {
    * The delay time in milliseconds for the submenu to appear after hovering over the trigger.
    * @default 200
    */
-  delay?: number
+  delay?: number,
+  /** Whether the submenu trigger uses virtual focus. */
+  isVirtualFocus?: boolean
 }
 
 interface SubmenuTriggerProps extends Omit<AriaMenuItemProps, 'key'> {
@@ -67,7 +69,7 @@ export interface SubmenuTriggerAria<T> {
  * @param ref - Ref to the submenu trigger element.
  */
 export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: SubmenuTriggerState, ref: RefObject<FocusableElement | null>): SubmenuTriggerAria<T> {
-  let {parentMenuRef, submenuRef, type = 'menu', isDisabled, delay = 200} = props;
+  let {parentMenuRef, submenuRef, type = 'menu', isDisabled, delay = 200, isVirtualFocus} = props;
   let submenuTriggerId = useId();
   let overlayId = useId();
   let {direction} = useLocale();
@@ -101,14 +103,18 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
         if (direction === 'ltr' && e.currentTarget.contains(e.target as Element)) {
           e.stopPropagation();
           onSubmenuClose();
-          ref.current?.focus();
+          if (!isVirtualFocus) {
+            ref.current?.focus();
+          }
         }
         break;
       case 'ArrowRight':
         if (direction === 'rtl' && e.currentTarget.contains(e.target as Element)) {
           e.stopPropagation();
           onSubmenuClose();
-          ref.current?.focus();
+          if (!isVirtualFocus) {
+            ref.current?.focus();
+          }
         }
         break;
       case 'Escape':
@@ -121,9 +127,10 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
   let subDialogKeyDown = (e: KeyboardEvent) => {
     switch (e.key) {
       case 'Escape':
-        e.stopPropagation();
         onSubmenuClose();
-        ref.current?.focus();
+        if (!isVirtualFocus) {
+          ref.current?.focus();
+        }
         break;
     }
   };
@@ -247,8 +254,7 @@ export function useSubmenuTrigger<T>(props: AriaSubmenuTriggerProps, state: Subm
     submenuProps,
     popoverProps: {
       isNonModal: true,
-      // TODO: does this break anything in RSP implementation?
-      disableFocusManagement: type === 'menu',
+      disableFocusManagement: type === 'menu' && !isVirtualFocus,
       shouldCloseOnInteractOutside
     }
   };

packages/react-aria-components/src/Menu.tsx

@@ -118,7 +118,7 @@ export interface SubmenuTriggerProps {
   delay?: number
 }
 
-const SubmenuTriggerContext = createContext<{parentMenuRef: RefObject<HTMLElement | null>} | null>(null);
+const SubmenuTriggerContext = createContext<{parentMenuRef: RefObject<HTMLElement | null>, isVirtualFocus?: boolean} | null>(null);
 
 /**
  * A submenu trigger is used to wrap a submenu's trigger item and the submenu itself.
@@ -132,11 +132,12 @@ export const SubmenuTrigger =  /*#__PURE__*/ createBranchComponent('submenutrigg
   let submenuTriggerState = useSubmenuTriggerState({triggerKey: item.key}, rootMenuTriggerState);
   let submenuRef = useRef<HTMLDivElement>(null);
   let itemRef = useObjectRef(ref);
-  let {parentMenuRef} = useContext(SubmenuTriggerContext)!;
+  let {parentMenuRef, isVirtualFocus} = useContext(SubmenuTriggerContext)!;
   let {submenuTriggerProps, submenuProps, popoverProps} = useSubmenuTrigger({
     parentMenuRef,
     submenuRef,
-    delay: props.delay
+    delay: props.delay,
+    isVirtualFocus
   }, submenuTriggerState, itemRef);
 
   return (
@@ -187,7 +188,6 @@ export const SubdialogTrigger =  /*#__PURE__*/ createBranchComponent('subdialogt
   let submenuTriggerState = useSubmenuTriggerState({triggerKey: item.key}, rootMenuTriggerState);
   let subdialogRef = useRef<HTMLDivElement>(null);
   let itemRef = useObjectRef(ref);
-  // TODO: We will probably support nested subdialogs so test that use case
   let {parentMenuRef} = useContext(SubmenuTriggerContext)!;
   let {submenuTriggerProps, submenuProps, popoverProps} = useSubmenuTrigger({
     parentMenuRef,
@@ -302,8 +302,9 @@ function MenuInner<T extends object>({props, collection, menuRef: ref}: MenuInne
             [MenuStateContext, state],
             [SeparatorContext, {elementType: 'div'}],
             [SectionContext, {name: 'MenuSection', render: MenuSectionInner}],
-            [SubmenuTriggerContext, {parentMenuRef: ref}],
+            [SubmenuTriggerContext, {parentMenuRef: ref, isVirtualFocus: autocompleteMenuProps?.shouldUseVirtualFocus}],
             [MenuItemContext, null],
+            [UNSTABLE_InternalAutocompleteContext, null],
             [SelectionManagerContext, state.selectionManager]
           ]}>
           <CollectionRoot

packages/react-aria-components/stories/Autocomplete.stories.tsx

@@ -671,7 +671,6 @@ export const AutocompleteInPopoverDialogTrigger = {
   }
 };
 
-// TODO: hitting escape sometimes closes both the root menu and the leaf menu, seems to happen if you arrow key to an option in the submenu's options and then hits escape...
 export const AutocompleteMenuInPopoverDialogTrigger = {
   render: (args) => {
     let {onAction, onSelectionChange, selectionMode} = args;