[Menu] Fixed Menu item not toggling highlighted prop when submenu popup is opened using keyboard

packages/react/src/menu/root/MenuRoot.tsx

@@ -61,6 +61,7 @@ const MenuRoot: React.FC<MenuRoot.Props> = function MenuRoot(props) {
       ...menuRoot,
       nested,
       parentContext,
+      setActiveIndex: menuRoot.setActiveIndex,
       disabled,
       allowMouseUpTriggerRef:
         parentContext?.allowMouseUpTriggerRef ?? menuRoot.allowMouseUpTriggerRef,

packages/react/src/menu/root/useMenuRoot.ts

@@ -250,6 +250,7 @@ export function useMenuRoot(parameters: useMenuRoot.Parameters): useMenuRoot.Ret
     () => ({
       activeIndex,
       allowMouseUpTriggerRef,
+      setActiveIndex,
       floatingRootContext,
       getItemProps,
       getPopupProps,
@@ -269,6 +270,7 @@ export function useMenuRoot(parameters: useMenuRoot.Parameters): useMenuRoot.Ret
     }),
     [
       activeIndex,
+      setActiveIndex,
       floatingRootContext,
       getItemProps,
       getPopupProps,
@@ -351,6 +353,7 @@ export namespace useMenuRoot {
 
   export interface ReturnValue {
     activeIndex: number | null;
+    setActiveIndex: (index: number | null) => void;
     floatingRootContext: FloatingRootContext;
     getItemProps: (externalProps?: GenericHTMLProps) => GenericHTMLProps;
     getPopupProps: (externalProps?: GenericHTMLProps) => GenericHTMLProps;

packages/react/src/menu/submenu-trigger/MenuSubmenuTrigger.test.tsx

@@ -0,0 +1,89 @@
+import * as React from 'react';
+import { expect } from 'chai';
+import { fireEvent, waitFor, screen } from '@mui/internal-test-utils';
+import { createRenderer } from '#test-utils';
+import { DirectionProvider } from '@base-ui-components/react/direction-provider';
+import { Menu } from '@base-ui-components/react/menu';
+
+type TextDirection = 'ltr' | 'rtl';
+
+describe('<Menu.SubmenuTrigger />', () => {
+  const { render } = createRenderer();
+
+  function TestComponent({ direction = 'ltr' }: { direction: TextDirection }) {
+    return (
+      <DirectionProvider direction={direction}>
+        <Menu.Root open>
+          <Menu.Portal>
+            <Menu.Positioner>
+              <Menu.Popup>
+                <Menu.Item>1</Menu.Item>
+                <Menu.Root>
+                  <Menu.SubmenuTrigger>2</Menu.SubmenuTrigger>
+                  <Menu.Portal>
+                    <Menu.Positioner>
+                      <Menu.Popup>
+                        <Menu.Item>2.1</Menu.Item>
+                        <Menu.Item>2.2</Menu.Item>
+                      </Menu.Popup>
+                    </Menu.Positioner>
+                  </Menu.Portal>
+                </Menu.Root>
+              </Menu.Popup>
+            </Menu.Positioner>
+          </Menu.Portal>
+        </Menu.Root>
+      </DirectionProvider>
+    );
+  }
+
+  const testCases = [
+    { direction: 'ltr', openKey: 'ArrowRight', closeKey: 'ArrowLeft' },
+    { direction: 'rtl', openKey: 'ArrowLeft', closeKey: 'ArrowRight' },
+  ];
+
+  testCases.forEach(({ direction, openKey }) => {
+    it(`opens the submenu with ${openKey} and highlights a single item in ${direction.toUpperCase()} direction`, async () => {
+      await render(<TestComponent direction={direction as TextDirection} />);
+      const submenuTrigger = screen.getByText('2');
+
+      fireEvent.focus(submenuTrigger);
+      fireEvent.keyDown(submenuTrigger, { key: openKey });
+
+      const submenuItems = await screen.findAllByRole('menuitem');
+      const submenuItem1 = submenuItems.find((item) => item.textContent === '2.1');
+
+      await waitFor(() => {
+        expect(submenuItem1).toHaveFocus();
+      });
+
+      submenuItems.forEach((item) => {
+        if (item === submenuItem1) {
+          expect(item).to.have.attribute('data-highlighted');
+        } else {
+          expect(item).not.to.have.attribute('data-highlighted');
+        }
+      });
+
+      // Check that parent menu items are not active
+      const parentMenuItems = screen
+        .getAllByRole('menuitem')
+        .filter((item) => item.textContent !== '2.1' && item.textContent !== '2.2');
+      parentMenuItems.forEach((item) => {
+        expect(item).not.to.have.attribute('data-highlighted');
+      });
+    });
+  });
+
+  it('sets tabIndex to 0 on the submenu trigger after opening the submenu with a keydown event', async () => {
+    await render(<TestComponent direction="ltr" />);
+    const submenuTrigger = screen.getByText('2');
+
+    fireEvent.focus(submenuTrigger);
+    fireEvent.keyDown(submenuTrigger, { key: 'ArrowRight' });
+
+    await waitFor(() => {
+      expect(submenuTrigger).to.have.attribute('tabIndex', '0');
+    });
+  });
+});

packages/react/src/menu/submenu-trigger/MenuSubmenuTrigger.tsx

@@ -37,15 +37,14 @@ const MenuSubmenuTrigger = React.forwardRef(function SubmenuTriggerComponent(
     throw new Error('Base UI: ItemTrigger must be placed in a nested Menu.');
   }
 
-  const { activeIndex, getItemProps } = parentContext;
+  const { activeIndex, getItemProps, setActiveIndex } = parentContext;
   const item = useCompositeListItem();
 
   const highlighted = activeIndex === item.index;
 
   const mergedRef = useForkRef(forwardedRef, item.ref);
 
   const { events: menuEvents } = useFloatingTree()!;
-
   const { getRootProps } = useMenuSubmenuTrigger({
     id,
     highlighted,
@@ -55,6 +54,7 @@ const MenuSubmenuTrigger = React.forwardRef(function SubmenuTriggerComponent(
     setTriggerElement,
     allowMouseUpTriggerRef,
     typingRef,
+    setActiveIndex,
   });
 
   const state: MenuSubmenuTrigger.State = React.useMemo(

packages/react/src/menu/submenu-trigger/useMenuSubmenuTrigger.ts

@@ -4,6 +4,14 @@ import { FloatingEvents } from '@floating-ui/react';
 import { useMenuItem } from '../item/useMenuItem';
 import { useForkRef } from '../../utils/useForkRef';
 import { GenericHTMLProps } from '../../utils/types';
+import { mergeReactProps } from '../../utils/mergeReactProps';
+import { useDirection } from '../../direction-provider/DirectionContext';
+
+const { useState } = React;
+
+type MenuKeyboardEvent = {
+  key: 'ArrowRight' | 'ArrowLeft' | 'ArrowUp' | 'ArrowDown' | 'Tab';
+} & React.KeyboardEvent;
 
 export function useMenuSubmenuTrigger(
   parameters: useSubmenuTrigger.Parameters,
@@ -17,6 +25,7 @@ export function useMenuSubmenuTrigger(
     setTriggerElement,
     allowMouseUpTriggerRef,
     typingRef,
+    setActiveIndex,
   } = parameters;
 
   const { getRootProps: getMenuItemProps, rootRef: menuItemRef } = useMenuItem({
@@ -32,17 +41,37 @@ export function useMenuSubmenuTrigger(
 
   const menuTriggerRef = useForkRef(menuItemRef, setTriggerElement);
 
+  const direction = useDirection();
+
+  const [isSubmenuOpen, setIsSubmenuOpen] = useState(false);
+
   const getRootProps = React.useCallback(
     (externalProps?: GenericHTMLProps) => {
-      return {
+      const openKey = direction === 'rtl' ? 'ArrowLeft' : 'ArrowRight';
+      const handleOpenSubmenu = () => {
+        if (highlighted) {
+          setActiveIndex(null);
+          setIsSubmenuOpen(true);
+        }
+      };
+      return mergeReactProps(externalProps, {
         ...getMenuItemProps({
+          // Once the submenu is opened, retain the tab index of the trigger element
+          tabIndex: highlighted || isSubmenuOpen ? 0 : -1,
           'aria-haspopup': 'menu' as const,
-          ...externalProps,
+          onKeyDown: (event: MenuKeyboardEvent) => {
+            if (event.key === openKey) {
+              handleOpenSubmenu();
+            } else if (event.key === 'Tab') {
+              setActiveIndex(null);
+            }
+          },
+          onClick: handleOpenSubmenu,
         }),
         ref: menuTriggerRef,
-      };
+      });
     },
-    [getMenuItemProps, menuTriggerRef],
+    [getMenuItemProps, menuTriggerRef, highlighted, setActiveIndex, direction, isSubmenuOpen],
   );
 
   return React.useMemo(
@@ -82,6 +111,11 @@ export namespace useSubmenuTrigger {
      * A ref that is set to `true` when the user is using the typeahead feature.
      */
     typingRef: React.RefObject<boolean>;
+    /**
+     * Callback to update the active (highlighted) item index.
+     * Set to null to remove highlighting from all items.
+     */
+    setActiveIndex: (index: number | null) => void;
   }
 
   export interface ReturnValue {