SingleTopNav: Fix losing focus when moving between menu states (#95069)

* handle button focus between states

* fix focus when closing the overlay menu from the close button

* move into a separate hook

* properly feature toggle the hook

* use constants

public/app/core/components/AppChrome/AppChrome.tsx

@@ -15,6 +15,7 @@ import { KioskMode } from 'app/types';
 import { AppChromeMenu } from './AppChromeMenu';
 import { DOCKED_LOCAL_STORAGE_KEY, DOCKED_MENU_OPEN_LOCAL_STORAGE_KEY } from './AppChromeService';
 import { MegaMenu, MENU_WIDTH } from './MegaMenu/MegaMenu';
+import { useMegaMenuFocusHelper } from './MegaMenu/utils';
 import { NavToolbar } from './NavToolbar/NavToolbar';
 import { ReturnToPrevious } from './ReturnToPrevious/ReturnToPrevious';
 import { SingleTopBar } from './TopBar/SingleTopBar';
@@ -50,6 +51,7 @@ export function AppChrome({ children }: Props) {
       }
     },
   });
+  useMegaMenuFocusHelper(state.megaMenuOpen, state.megaMenuDocked);
 
   const contentClass = cx({
     [styles.content]: true,

public/app/core/components/AppChrome/AppChromeMenu.tsx

@@ -58,9 +58,13 @@ export function AppChromeMenu({}: Props) {
           classNames={animationStyles.overlay}
           timeout={{ enter: animationSpeed, exit: 0 }}
         >
-          <FocusScope contain autoFocus restoreFocus>
-            <MegaMenu className={styles.menu} onClose={onClose} ref={ref} {...overlayProps} {...dialogProps} />
-          </FocusScope>
+          <>
+            {isOpen && (
+              <FocusScope contain autoFocus restoreFocus>
+                <MegaMenu className={styles.menu} onClose={onClose} ref={ref} {...overlayProps} {...dialogProps} />
+              </FocusScope>
+            )}
+          </>
         </CSSTransition>
         <CSSTransition
           nodeRef={backdropRef}

public/app/core/components/AppChrome/MegaMenu/MegaMenu.tsx

@@ -13,9 +13,10 @@ import { setBookmark } from 'app/core/reducers/navBarTree';
 import { usePatchUserPreferencesMutation } from 'app/features/preferences/api/index';
 import { useDispatch, useSelector } from 'app/types';
 
+import { MEGA_MENU_TOGGLE_ID } from '../TopBar/SingleTopBar';
 import { TOP_BAR_LEVEL_HEIGHT } from '../types';
 
-import { MegaMenuHeader } from './MegaMenuHeader';
+import { DOCK_MENU_BUTTON_ID, MegaMenuHeader } from './MegaMenuHeader';
 import { MegaMenuItem } from './MegaMenuItem';
 import { usePinnedItems } from './hooks';
 import { enrichWithInteractionTracking, findByUrl, getActiveItem } from './utils';
@@ -76,9 +77,11 @@ export const MegaMenu = memo(
       }
 
       // refocus on undock/menu open button when changing state
-      setTimeout(() => {
-        document.getElementById(state.megaMenuDocked ? 'mega-menu-toggle' : 'dock-menu-button')?.focus();
-      });
+      if (!config.featureToggles.singleTopNav) {
+        setTimeout(() => {
+          document.getElementById(state.megaMenuDocked ? MEGA_MENU_TOGGLE_ID : DOCK_MENU_BUTTON_ID)?.focus();
+        });
+      }
     };
 
     const isPinned = useCallback(

public/app/core/components/AppChrome/MegaMenu/MegaMenuHeader.tsx

@@ -15,6 +15,9 @@ export interface Props {
   onClose: () => void;
 }
 
+export const DOCK_MENU_BUTTON_ID = 'dock-menu-button';
+export const MEGA_MENU_HEADER_TOGGLE_ID = 'mega-menu-header-toggle';
+
 export function MegaMenuHeader({ handleMegaMenu, handleDockedMenu, onClose }: Props) {
   const theme = useTheme2();
   const { chrome } = useGrafana();
@@ -24,13 +27,18 @@ export function MegaMenuHeader({ handleMegaMenu, handleDockedMenu, onClose }: Pr
   return (
     <div className={styles.header}>
       <Stack alignItems="center" minWidth={0} gap={0.25}>
-        <ToolbarButton narrow onClick={handleMegaMenu} tooltip={t('navigation.megamenu.close', 'Close menu')}>
+        <ToolbarButton
+          narrow
+          id={MEGA_MENU_HEADER_TOGGLE_ID}
+          onClick={handleMegaMenu}
+          tooltip={t('navigation.megamenu.close', 'Close menu')}
+        >
           <Branding.MenuLogo className={styles.img} />
         </ToolbarButton>
         <OrganizationSwitcher />
       </Stack>
       <IconButton
-        id="dock-menu-button"
+        id={DOCK_MENU_BUTTON_ID}
         className={styles.dockMenuButton}
         tooltip={
           state.megaMenuDocked

public/app/core/components/AppChrome/MegaMenu/utils.ts

@@ -1,3 +1,5 @@
+import { useEffect } from 'react';
+
 import { NavModelItem } from '@grafana/data';
 import { config, reportInteraction } from '@grafana/runtime';
 import { t } from 'app/core/internationalization';
@@ -7,6 +9,9 @@ import { ShowModalReactEvent } from '../../../../types/events';
 import appEvents from '../../../app_events';
 import { getFooterLinks } from '../../Footer/Footer';
 import { HelpModal } from '../../help/HelpModal';
+import { MEGA_MENU_TOGGLE_ID } from '../TopBar/SingleTopBar';
+
+import { DOCK_MENU_BUTTON_ID, MEGA_MENU_HEADER_TOGGLE_ID } from './MegaMenuHeader';
 
 export const enrichHelpItem = (helpItem: NavModelItem) => {
   let menuItems = helpItem.children || [];
@@ -146,3 +151,33 @@ export function findByUrl(nodes: NavModelItem[], url: string): NavModelItem | nu
   }
   return null;
 }
+
+/**
+ * helper to manage focus when opening/closing and docking/undocking the mega menu
+ * @param isOpen whether the mega menu is open
+ * @param isDocked whether mega menu is docked
+ */
+export function useMegaMenuFocusHelper(isOpen: boolean, isDocked: boolean) {
+  const isSingleTopNav = config.featureToggles.singleTopNav;
+  // manage focus when opening/closing
+  useEffect(() => {
+    if (isSingleTopNav) {
+      if (isOpen) {
+        document.getElementById(MEGA_MENU_HEADER_TOGGLE_ID)?.focus();
+      } else {
+        document.getElementById(MEGA_MENU_TOGGLE_ID)?.focus();
+      }
+    }
+  }, [isOpen, isSingleTopNav]);
+
+  // manage focus when docking/undocking
+  useEffect(() => {
+    if (isSingleTopNav) {
+      if (isDocked) {
+        document.getElementById(DOCK_MENU_BUTTON_ID)?.focus();
+      } else {
+        document.getElementById(MEGA_MENU_TOGGLE_ID)?.focus();
+      }
+    }
+  }, [isDocked, isSingleTopNav]);
+}

public/app/core/components/AppChrome/TopBar/SingleTopBar.tsx

@@ -23,6 +23,8 @@ import { SignInLink } from './SignInLink';
 import { TopNavBarMenu } from './TopNavBarMenu';
 import { TopSearchBarCommandPaletteTrigger } from './TopSearchBarCommandPaletteTrigger';
 
+export const MEGA_MENU_TOGGLE_ID = 'mega-menu-toggle';
+
 interface Props {
   sectionNav: NavModelItem;
   pageNav?: NavModelItem;
@@ -52,7 +54,12 @@ export const SingleTopBar = memo(function SingleTopBar({
     <div className={styles.layout}>
       <Stack minWidth={0} gap={0.5} alignItems="center">
         {!menuDockedAndOpen && (
-          <ToolbarButton narrow onClick={onToggleMegaMenu} tooltip={t('navigation.megamenu.open', 'Open menu')}>
+          <ToolbarButton
+            narrow
+            id={MEGA_MENU_TOGGLE_ID}
+            onClick={onToggleMegaMenu}
+            tooltip={t('navigation.megamenu.open', 'Open menu')}
+          >
             <Stack gap={0} alignItems="center">
               <Branding.MenuLogo className={styles.img} />
               <Icon size="sm" name="angle-down" />