From 1a4a95e350f578f8296815771baaebb30ddaf642 Mon Sep 17 00:00:00 2001
From: ling1726 <lingfangao@hotmail.com>
Date: Mon, 23 Sep 2024 14:13:46 +0200
Subject: [PATCH] fix(Menu): Do not rely on tabster restorer for focus restore
 (#32840)

Co-authored-by: Juraj Kapsiar <jurokapsiar@gmail.com>
---
 ...-6534522e-919d-4ebb-973b-b93815bd1abe.json |  7 +++++++
 .../Menu/__snapshots__/Menu.test.tsx.snap     |  1 -
 .../library/src/components/Menu/useMenu.tsx   | 20 ++++++++++++++++++-
 .../MenuTrigger/MenuTrigger.test.tsx          |  2 --
 .../__snapshots__/MenuTrigger.test.tsx.snap   |  1 -
 .../components/MenuTrigger/useMenuTrigger.ts  |  4 +---
 6 files changed, 27 insertions(+), 8 deletions(-)
 create mode 100644 change/@fluentui-react-menu-6534522e-919d-4ebb-973b-b93815bd1abe.json

diff --git a/change/@fluentui-react-menu-6534522e-919d-4ebb-973b-b93815bd1abe.json b/change/@fluentui-react-menu-6534522e-919d-4ebb-973b-b93815bd1abe.json
new file mode 100644
index 0000000000000..4ba90098a0734
--- /dev/null
+++ b/change/@fluentui-react-menu-6534522e-919d-4ebb-973b-b93815bd1abe.json
@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "fix(Menu): Do not rely on tabster restorer for focus restore",
+  "packageName": "@fluentui/react-menu",
+  "email": "lingfangao@hotmail.com",
+  "dependentChangeType": "patch"
+}
diff --git a/packages/react-components/react-menu/library/src/components/Menu/__snapshots__/Menu.test.tsx.snap b/packages/react-components/react-menu/library/src/components/Menu/__snapshots__/Menu.test.tsx.snap
index c604de7dd8c5e..fc7e36d350dc0 100644
--- a/packages/react-components/react-menu/library/src/components/Menu/__snapshots__/Menu.test.tsx.snap
+++ b/packages/react-components/react-menu/library/src/components/Menu/__snapshots__/Menu.test.tsx.snap
@@ -4,7 +4,6 @@ exports[`Menu renders a default state 1`] = `
 <div>
   <button
     aria-haspopup="menu"
-    data-tabster="{\\"restorer\\":{\\"type\\":1}}"
     id="menu1"
   >
     Menu trigger
diff --git a/packages/react-components/react-menu/library/src/components/Menu/useMenu.tsx b/packages/react-components/react-menu/library/src/components/Menu/useMenu.tsx
index 507b8f072ccad..507bec8bf9b28 100644
--- a/packages/react-components/react-menu/library/src/components/Menu/useMenu.tsx
+++ b/packages/react-components/react-menu/library/src/components/Menu/useMenu.tsx
@@ -13,6 +13,7 @@ import {
   useOnScrollOutside,
   elementContains,
   useTimeout,
+  useFirstMount,
 } from '@fluentui/react-utilities';
 import { useFluent_unstable as useFluent } from '@fluentui/react-shared-contexts';
 import { useFocusFinders } from '@fluentui/react-tabster';
@@ -173,6 +174,8 @@ const useMenuOpenState = (
   > &
     Pick<MenuProps, 'open' | 'defaultOpen' | 'onOpenChange'>,
 ) => {
+  'use no memo';
+
   const { targetDocument } = useFluent();
   const parentSetOpen = useMenuContext_unstable(context => context.setOpen);
   const onOpenChange: MenuProps['onOpenChange'] = useEventCallback((e, data) => state.onOpenChange?.(e, data));
@@ -265,11 +268,26 @@ const useMenuOpenState = (
     firstFocusable?.focus();
   }, [findFirstFocusable, state.menuPopoverRef]);
 
+  const firstMount = useFirstMount();
   React.useEffect(() => {
     if (open) {
       focusFirst();
+    } else {
+      if (!firstMount) {
+        if (targetDocument?.activeElement === targetDocument?.body) {
+          // We know that React effects are sync so we focus the trigger here
+          // after any event handler (event handlers will update state and re-render).
+          // Since the browser only performs the default behaviour for the Tab key once
+          // keyboard events have fully bubbled up the window, the browser will move
+          // focus to the next tabbable element before/after the trigger if needed.
+          // If the Tab key was not pressed, focus will remain on the trigger as expected.
+          state.triggerRef.current?.focus();
+        }
+      }
     }
-  }, [open, focusFirst]);
+    // firstMount change should not re-run this effect
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [state.triggerRef, state.isSubmenu, open, focusFirst, targetDocument, state.menuPopoverRef]);
 
   return [open, setOpen] as const;
 };
diff --git a/packages/react-components/react-menu/library/src/components/MenuTrigger/MenuTrigger.test.tsx b/packages/react-components/react-menu/library/src/components/MenuTrigger/MenuTrigger.test.tsx
index 240ae1f22a9a7..4a55ade8fc0de 100644
--- a/packages/react-components/react-menu/library/src/components/MenuTrigger/MenuTrigger.test.tsx
+++ b/packages/react-components/react-menu/library/src/components/MenuTrigger/MenuTrigger.test.tsx
@@ -57,7 +57,6 @@ describe('MenuTrigger', () => {
       Array [
         <button
           aria-haspopup="menu"
-          data-tabster="{\\"restorer\\":{\\"type\\":1}}"
           id="id"
         >
           Trigger
@@ -88,7 +87,6 @@ describe('MenuTrigger', () => {
       Array [
         <button
           aria-haspopup="menu"
-          data-tabster="{\\"restorer\\":{\\"type\\":1}}"
           id="id"
         >
           Trigger
diff --git a/packages/react-components/react-menu/library/src/components/MenuTrigger/__snapshots__/MenuTrigger.test.tsx.snap b/packages/react-components/react-menu/library/src/components/MenuTrigger/__snapshots__/MenuTrigger.test.tsx.snap
index 8d08c72a08737..7826501c5a172 100644
--- a/packages/react-components/react-menu/library/src/components/MenuTrigger/__snapshots__/MenuTrigger.test.tsx.snap
+++ b/packages/react-components/react-menu/library/src/components/MenuTrigger/__snapshots__/MenuTrigger.test.tsx.snap
@@ -3,7 +3,6 @@
 exports[`MenuTrigger renders a default state 1`] = `
 <button
   aria-haspopup="menu"
-  data-tabster="{\\"restorer\\":{\\"type\\":1}}"
   id="id"
   onClick={[Function]}
   onContextMenu={[Function]}
diff --git a/packages/react-components/react-menu/library/src/components/MenuTrigger/useMenuTrigger.ts b/packages/react-components/react-menu/library/src/components/MenuTrigger/useMenuTrigger.ts
index 14a983e3450d6..ac27b978354a5 100644
--- a/packages/react-components/react-menu/library/src/components/MenuTrigger/useMenuTrigger.ts
+++ b/packages/react-components/react-menu/library/src/components/MenuTrigger/useMenuTrigger.ts
@@ -2,7 +2,7 @@ import * as React from 'react';
 import { MenuTriggerProps, MenuTriggerState } from './MenuTrigger.types';
 import { useMenuContext_unstable } from '../../contexts/menuContext';
 import { useIsSubmenu } from '../../utils/useIsSubmenu';
-import { useFocusFinders, useRestoreFocusTarget } from '@fluentui/react-tabster';
+import { useFocusFinders } from '@fluentui/react-tabster';
 import { useFluent_unstable as useFluent } from '@fluentui/react-shared-contexts';
 import { ArrowRight, ArrowLeft, Escape, ArrowDown } from '@fluentui/keyboard-keys';
 import {
@@ -31,7 +31,6 @@ export const useMenuTrigger_unstable = (props: MenuTriggerProps): MenuTriggerSta
   const triggerId = useMenuContext_unstable(context => context.triggerId);
   const openOnHover = useMenuContext_unstable(context => context.openOnHover);
   const openOnContext = useMenuContext_unstable(context => context.openOnContext);
-  const restoreFocusTargetAttribute = useRestoreFocusTarget();
 
   const isSubmenu = useIsSubmenu();
 
@@ -125,7 +124,6 @@ export const useMenuTrigger_unstable = (props: MenuTriggerProps): MenuTriggerSta
 
   const contextMenuProps = {
     id: triggerId,
-    ...restoreFocusTargetAttribute,
     ...child?.props,
     ref: useMergedRefs(triggerRef, child?.ref),
     onMouseEnter: useEventCallback(mergeCallbacks(child?.props.onMouseEnter, onMouseEnter)),