[menu] Fold focusItem into list navigation

atomiks · atomiks · commit 2fbd46480cde · 2026-05-18T15:05:15.000+10:00
diff --git a/packages/react/src/floating-ui-react/hooks/useListNavigation.ts b/packages/react/src/floating-ui-react/hooks/useListNavigation.ts
@@ -134,6 +134,15 @@ export interface UseListNavigationProps {
    * @default true
    */
   focusItemOnHover?: boolean | undefined;
+  /**
+   * A pending imperative item focus request to resolve against the list.
+   * @default null
+   */
+  pendingFocusItem?: UseListNavigationFocusItem | null | undefined;
+  /**
+   * Callback fired when a pending imperative item focus request has been consumed.
+   */
+  onPendingFocusItemChange?: ((pendingFocusItem: null) => void) | undefined;
   /**
    * Whether pressing an arrow key on the navigation's main axis opens the
    * floating element.
@@ -223,6 +232,8 @@ export interface UseListNavigationProps {
   externalTree?: FloatingTreeStore | undefined;
 }
 
+export type UseListNavigationFocusItem = 'first' | 'last' | 'none';
+
 /**
  * Adds arrow key-based navigation of a list of items, either using real DOM
  * focus or virtual focus.
@@ -245,6 +256,8 @@ export function useListNavigation(
     virtual = false,
     focusItemOnOpen = 'auto',
     focusItemOnHover = true,
+    pendingFocusItem = null,
+    onPendingFocusItemChange,
     openOnArrowKeyDown = true,
     disabledIndices = undefined,
     orientation = 'vertical',
@@ -298,6 +311,10 @@ export function useListNavigation(
     onNavigateProp(indexRef.current === -1 ? null : indexRef.current, event);
   });
 
+  const clearPendingFocusItem = useStableCallback(() => {
+    onPendingFocusItemChange?.(null);
+  });
+
   const previousOnNavigateRef = React.useRef(onNavigate);
   const previousMountedRef = React.useRef(!!floatingElement);
   const previousOpenRef = React.useRef(open);
@@ -311,7 +328,6 @@ export function useListNavigation(
   const resetOnPointerLeaveRef = useValueAsRef(resetOnPointerLeave);
 
   const focusFrame = useAnimationFrame();
-  const waitForListPopulatedFrame = useAnimationFrame();
 
   const focusItem = useStableCallback(() => {
     function runFocus(item: HTMLElement) {
@@ -391,21 +407,70 @@ export function useListNavigation(
   // open.
   useIsoLayoutEffect(() => {
     if (!enabled) {
-      return;
+      return undefined;
     }
     if (!open) {
       forceSyncFocusRef.current = false;
-      return;
+      return undefined;
     }
     if (!floatingElement) {
-      return;
+      return undefined;
+    }
+
+    const resolveFocusItem = (
+      target: Exclude<UseListNavigationFocusItem, 'none'>,
+      onDone?: () => void,
+    ) => {
+      let cancelled = false;
+
+      const waitForListPopulated = () => {
+        if (cancelled) {
+          return;
+        }
+
+        if (listRef.current[0] == null) {
+          onDone?.();
+          return;
+        }
+
+        indexRef.current = target === 'last' ? getMaxListIndex(listRef) : getMinListIndex(listRef);
+        keyRef.current = null;
+        onNavigate();
+        onDone?.();
+      };
+
+      if (listRef.current[0] == null) {
+        // Some composed items register after their indexes are assigned by
+        // CompositeList, which can land one layout-effect flush after the
+        // popup opens. Retry once before clearing the pending request.
+        queueMicrotask(waitForListPopulated);
+      } else {
+        waitForListPopulated();
+      }
+
+      return () => {
+        cancelled = true;
+      };
+    };
+
+    if (pendingFocusItem != null) {
+      forceSyncFocusRef.current = false;
+
+      if (pendingFocusItem === 'none') {
+        indexRef.current = -1;
+        onNavigate();
+        clearPendingFocusItem();
+        return undefined;
+      }
+
+      return resolveFocusItem(pendingFocusItem, clearPendingFocusItem);
     }
 
     if (activeIndex == null) {
       forceSyncFocusRef.current = false;
 
       if (selectedIndexRef.current != null) {
-        return;
+        return undefined;
       }
 
       // Reset while the floating element was open (e.g. the list changed).
@@ -420,52 +485,36 @@ export function useListNavigation(
         focusItemOnOpenRef.current &&
         (keyRef.current != null || (focusItemOnOpenRef.current === true && keyRef.current == null))
       ) {
-        let runs = 0;
-        const waitForListPopulated = () => {
-          if (listRef.current[0] == null) {
-            // Avoid letting the browser paint if possible on the first try,
-            // otherwise use rAF. Don't try more than twice, since something
-            // is wrong otherwise.
-            if (runs < 2) {
-              const scheduler = runs
-                ? (callback: () => void) => waitForListPopulatedFrame.request(callback)
-                : queueMicrotask;
-              scheduler(waitForListPopulated);
-            }
-            runs += 1;
-          } else {
-            // initially focus the first non-disabled item
-            indexRef.current =
-              keyRef.current == null ||
-              isMainOrientationToEndKey(keyRef.current, orientation, rtl) ||
-              nested
-                ? getMinListIndex(listRef)
-                : getMaxListIndex(listRef);
-            keyRef.current = null;
-            onNavigate();
-          }
-        };
-
-        waitForListPopulated();
+        const target =
+          keyRef.current == null ||
+          isMainOrientationToEndKey(keyRef.current, orientation, rtl) ||
+          nested
+            ? 'first'
+            : 'last';
+
+        return resolveFocusItem(target);
       }
     } else if (!isIndexOutOfListBounds(listRef.current, activeIndex)) {
       indexRef.current = activeIndex;
       focusItem();
       forceScrollIntoViewRef.current = false;
     }
+
+    return undefined;
   }, [
     enabled,
     open,
     floatingElement,
     activeIndex,
+    pendingFocusItem,
     selectedIndexRef,
     nested,
     listRef,
     orientation,
     rtl,
     onNavigate,
+    clearPendingFocusItem,
     focusItem,
-    waitForListPopulatedFrame,
   ]);
 
   // Ensure the parent floating element has focus when a nested child closes
diff --git a/packages/react/src/menu/root/MenuRoot.test.tsx b/packages/react/src/menu/root/MenuRoot.test.tsx
@@ -10,6 +10,7 @@ import {
 } from '@mui/internal-test-utils';
 import { DirectionProvider } from '@base-ui/react/direction-provider';
 import { useRefWithInit } from '@base-ui/utils/useRefWithInit';
+import { useIsoLayoutEffect } from '@base-ui/utils/useIsoLayoutEffect';
 import { Menu } from '@base-ui/react/menu';
 import { Dialog } from '@base-ui/react/dialog';
 import { AlertDialog } from '@base-ui/react/alert-dialog';
@@ -1453,6 +1454,60 @@ describe('<Menu.Root />', () => {
             expect(item).toHaveAttribute('tabindex', '-1');
           });
         });
+
+        it('waits for composed items that register after a layout effect', async () => {
+          const actionsRef: React.RefObject<Menu.Root.Actions | null> = { current: null };
+
+          function DelayedItem(props: { children: React.ReactNode; testId: string }) {
+            const [mounted, setMounted] = React.useState(false);
+            useIsoLayoutEffect(() => {
+              setMounted(true);
+            }, []);
+
+            if (!mounted) {
+              return null;
+            }
+
+            return <Menu.Item data-testid={props.testId}>{props.children}</Menu.Item>;
+          }
+
+          function App() {
+            const [open, setOpen] = React.useState(false);
+            return (
+              <div>
+                <button
+                  type="button"
+                  onClick={() => {
+                    setOpen(true);
+                    actionsRef.current?.focusItem('first');
+                  }}
+                >
+                  external
+                </button>
+                <TestMenu
+                  rootProps={{ open, onOpenChange: setOpen, actionsRef }}
+                  popupProps={{
+                    children: (
+                      <React.Fragment>
+                        <DelayedItem testId="delayed-1">One</DelayedItem>
+                        <DelayedItem testId="delayed-2">Two</DelayedItem>
+                      </React.Fragment>
+                    ),
+                  }}
+                />
+              </div>
+            );
+          }
+
+          const { user } = await render(<App />);
+
+          await user.click(screen.getByRole('button', { name: 'external' }));
+
+          const firstItem = await screen.findByTestId('delayed-1');
+          await waitFor(() => {
+            expect(firstItem).toHaveFocus();
+          });
+        });
       });
     });
 
diff --git a/packages/react/src/menu/root/MenuRoot.tsx b/packages/react/src/menu/root/MenuRoot.tsx
@@ -5,7 +5,6 @@ import { useStableCallback } from '@base-ui/utils/useStableCallback';
 import { useId } from '@base-ui/utils/useId';
 import { useIsoLayoutEffect } from '@base-ui/utils/useIsoLayoutEffect';
 import { useOnFirstRender } from '@base-ui/utils/useOnFirstRender';
-import { useAnimationFrame } from '@base-ui/utils/useAnimationFrame';
 import { EMPTY_ARRAY, EMPTY_OBJECT } from '@base-ui/utils/empty';
 import { fastComponent } from '@base-ui/utils/fastHooks';
 import {
@@ -17,7 +16,6 @@ import {
   useTypeahead,
   useSyncedFloatingRootContext,
 } from '../../floating-ui-react';
-import { getMaxListIndex, getMinListIndex } from '../../floating-ui-react/utils/composite';
 import { MenuRootContext, useMenuRootContext } from './MenuRootContext';
 import { MenubarContext, useMenubarContext } from '../../menubar/MenubarContext';
 import { TYPEAHEAD_RESET_MS } from '../../internals/constants';
@@ -421,6 +419,10 @@ export const MenuRoot = fastComponent(function MenuRoot<Payload>(props: MenuRoot
   });
 
   const direction = useDirection();
+  const pendingFocusItem = store.useState('pendingFocusItem');
+  const clearPendingFocusItem = useStableCallback(() => {
+    store.set('pendingFocusItem', null);
+  });
 
   const listNavigation = useListNavigation(floatingRootContext, {
     enabled: !disabled,
@@ -436,49 +438,10 @@ export const MenuRoot = fastComponent(function MenuRoot<Payload>(props: MenuRoot
     openOnArrowKeyDown: parent.type !== 'context-menu',
     externalTree: nested ? floatingTreeRoot : undefined,
     focusItemOnHover: highlightItemOnHover,
+    pendingFocusItem,
+    onPendingFocusItemChange: clearPendingFocusItem,
   });
 
-  const pendingFocusItem = store.useState('pendingFocusItem');
-  const focusItemFrame = useAnimationFrame();
-
-  useIsoLayoutEffect(() => {
-    if (pendingFocusItem == null || !open) {
-      return undefined;
-    }
-
-    if (pendingFocusItem === 'none') {
-      store.update({ activeIndex: null, pendingFocusItem: null });
-      return undefined;
-    }
-
-    let cancelled = false;
-    let runs = 0;
-    const elements = store.context.itemDomElements;
-
-    const resolve = () => {
-      if (cancelled) {
-        return;
-      }
-      if (elements.current[0] == null) {
-        if (runs < 2) {
-          runs += 1;
-          focusItemFrame.request(resolve);
-        }
-        return;
-      }
-      const index =
-        pendingFocusItem === 'last' ? getMaxListIndex(elements) : getMinListIndex(elements);
-      store.update({ activeIndex: index, pendingFocusItem: null });
-    };
-
-    resolve();
-
-    return () => {
-      cancelled = true;
-      focusItemFrame.cancel();
-    };
-  }, [open, positionerElement, pendingFocusItem, store, focusItemFrame]);
-
   const onTyping = React.useCallback(
     (nextTyping: boolean) => {
       store.context.typingRef.current = nextTyping;
