fix(menu): refactor to useClickOutsideEvent

src/core/components/menu/menu.tsx

@@ -1,6 +1,6 @@
 import {forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef} from 'react'
 import {styled} from 'styled-components'
-import {useClickOutside, useGlobalKeyDown} from '../../hooks'
+import {useClickOutsideEvent, useGlobalKeyDown} from '../../hooks'
 import {Box, Stack} from '../../primitives'
 import {ResponsivePaddingProps} from '../../primitives/types'
 import {useLayer} from '../../utils'
@@ -75,16 +75,29 @@ export const Menu = forwardRef(function Menu(
     handleItemMouseLeave,
     handleKeyDown,
     mount,
-    rootElement,
-    setRootElement,
-  } = useMenuController({onKeyDown, originElement, shouldFocus})
+  } = useMenuController({onKeyDown, originElement, shouldFocus, rootElementRef: ref})
 
+  const unregisterElementRef = useRef<(() => void) | null>(null)
   const handleRefChange = useCallback(
     (el: HTMLDivElement | null) => {
-      setRootElement(el)
+      // Run cleanup of previously registered elements
+      if (unregisterElementRef.current) {
+        // The `registerElement` callback were originally used in a `useEffect`, so it returns a cleanup function that is a bit gnarly to handle in a ref callback.
+        // Since we can't change the `registerElement` implementation itself without making breaking change,
+        // that is explained in the code comments for createGlobalScopedContext.tsx,
+        // we need to handle with a ref that holds on to the cleanup function last returned when the ref callback is called.
+        unregisterElementRef.current()
+        unregisterElementRef.current = null
+      }
+
       ref.current = el
+
+      // Register root element (for nested menus)
+      if (ref.current && registerElement) {
+        unregisterElementRef.current = registerElement(ref.current)
+      }
     },
-    [setRootElement],
+    [registerElement],
   )
 
   // Trigger `onItemSelect` when active index changes
@@ -93,13 +106,7 @@ export const Menu = forwardRef(function Menu(
   }, [activeIndex, onItemSelect])
 
   // Close menu when clicking outside
-  useClickOutside(
-    useCallback(
-      (event) => isTopLayer && onClickOutside && onClickOutside(event),
-      [isTopLayer, onClickOutside],
-    ),
-    [rootElement],
-  )
+  useClickOutsideEvent(isTopLayer && onClickOutside, () => [ref.current])
 
   // Close menu when pressing Escape
   useGlobalKeyDown(
@@ -116,13 +123,6 @@ export const Menu = forwardRef(function Menu(
     ),
   )
 
-  // Register root element (for nested menus)
-  useEffect(() => {
-    if (!rootElement || !registerElement) return
-
-    return registerElement(rootElement)
-  }, [registerElement, rootElement])
-
   const value: MenuContextValue = useMemo(
     () => ({
       version: 0.0,

src/core/components/menu/menuButton.tsx

@@ -196,11 +196,7 @@ export const MenuButton = forwardRef(function MenuButton(
   const registerElement = useCallback((el: HTMLElement) => {
     setChildMenuElements((els) => els.concat([el]))
 
-    return () => {
-      setChildMenuElements((els) => {
-        return els.filter((_el) => _el !== el)
-      })
-    }
+    return () => setChildMenuElements((els) => els.filter((_el) => _el !== el))
   }, [])
 
   const menuProps: MenuProps = useMemo(

src/core/components/menu/useMenuController.ts

@@ -1,4 +1,4 @@
-import {useCallback, useEffect, useRef, useState} from 'react'
+import {useCallback, useEffect, useMemo, useRef, useState} from 'react'
 import {_getFocusableElements, _sortElements} from './helpers'
 
 /**
@@ -11,8 +11,6 @@ export interface MenuController {
   handleItemMouseLeave: () => void
   handleKeyDown: (event: React.KeyboardEvent<HTMLDivElement>) => void
   mount: (element: HTMLElement | null, selected?: boolean) => () => void
-  rootElement: HTMLDivElement | null
-  setRootElement: (el: HTMLDivElement | null) => void
 }
 
 /**
@@ -24,14 +22,14 @@ export function useMenuController(props: {
   onKeyDown?: React.KeyboardEventHandler
   originElement?: HTMLElement | null
   shouldFocus: 'first' | 'last' | null
+  rootElementRef: React.MutableRefObject<HTMLDivElement | null>
 }): MenuController {
-  const {onKeyDown, originElement, shouldFocus} = props
+  const {onKeyDown, originElement, shouldFocus, rootElementRef} = props
   const elementsRef = useRef<HTMLElement[]>([])
-  const [rootElement, setRootElement] = useState<HTMLDivElement | null>(null)
   const [activeIndex, _setActiveIndex] = useState(-1)
   const activeIndexRef = useRef(activeIndex)
-  const activeElement = elementsRef.current[activeIndex] || null
-  const mounted = Boolean(rootElement)
+  const activeElement = useMemo(() => elementsRef.current[activeIndex] || null, [activeIndex])
+  const mounted = Boolean(rootElementRef.current)
 
   const setActiveIndex = useCallback((nextActiveIndex: number) => {
     _setActiveIndex(nextActiveIndex)
@@ -44,7 +42,7 @@ export function useMenuController(props: {
 
       if (elementsRef.current.indexOf(element) === -1) {
         elementsRef.current.push(element)
-        _sortElements(rootElement, elementsRef.current)
+        _sortElements(rootElementRef.current, elementsRef.current)
       }
 
       if (selected) {
@@ -61,7 +59,7 @@ export function useMenuController(props: {
         }
       }
     },
-    [rootElement, setActiveIndex],
+    [rootElementRef, setActiveIndex],
   )
 
   const handleKeyDown = useCallback(
@@ -179,17 +177,15 @@ export function useMenuController(props: {
     // which would be incorrect when the user hovers over a gap
     // between two menu items or a menu divider.
     setActiveIndex(-2)
-    rootElement?.focus()
-  }, [setActiveIndex, rootElement])
+    rootElementRef.current?.focus()
+  }, [rootElementRef, setActiveIndex])
 
   // Set focus on the currently active element
   useEffect(() => {
     if (!mounted) return
 
     const rafId = window.requestAnimationFrame(() => {
-      const _activeIndex = activeIndexRef.current
-
-      if (_activeIndex === -1) {
+      if (activeIndex === -1) {
         if (shouldFocus === 'first') {
           const focusableElements = _getFocusableElements(elementsRef.current)
           const el = focusableElements[0]
@@ -217,7 +213,7 @@ export function useMenuController(props: {
         return
       }
 
-      const element = elementsRef.current[_activeIndex] || null
+      const element = elementsRef.current[activeIndex] || null
 
       element?.focus()
     })
@@ -234,7 +230,5 @@ export function useMenuController(props: {
     handleItemMouseLeave,
     handleKeyDown,
     mount,
-    rootElement,
-    setRootElement,
   }
 }

stories/components/Menu.stories.tsx

@@ -1,12 +1,18 @@
 import {BinaryDocumentIcon, CircleIcon, RestoreIcon} from '@sanity/icons'
 import type {Meta, StoryObj} from '@storybook/react'
+import {fn} from '@storybook/test'
 import {Menu, MenuDivider, MenuGroup, MenuItem} from '../../src/core/components'
 import {Badge, Card, Container, Flex, Stack, Text} from '../../src/core/primitives'
 import {LayerProvider} from '../../src/core/utils'
 import {getSpaceControls} from '../controls'
 
 const meta: Meta<typeof Menu> = {
-  args: {},
+  args: {
+    onClickOutside: fn(),
+    onEscape: fn(),
+    onItemClick: fn(),
+    onItemSelect: fn(),
+  },
   argTypes: {
     padding: getSpaceControls(),
     space: getSpaceControls(),

stories/components/MenuButton.stories.tsx

@@ -1,12 +1,14 @@
 import {ClockIcon, CommentIcon, ExpandIcon, SearchIcon} from '@sanity/icons'
 import type {Meta, StoryObj} from '@storybook/react'
-import {expect} from '@storybook/test'
+import {expect, fn} from '@storybook/test'
 import {userEvent, within} from '@storybook/test'
 import {Menu, MenuButton, MenuDivider, MenuGroup, MenuItem} from '../../src/core/components'
 import {Button, Flex} from '../../src/core/primitives'
 
 const meta: Meta<typeof MenuButton> = {
   args: {
+    onOpen: fn(),
+    onClose: fn(),
     button: <Button text="Open" />,
     menu: (
       <Menu>