fix(menu-button): handle child refs safely

src/core/components/menu/menuButton.tsx

@@ -1,5 +1,14 @@
 import {ThemeColorSchemeKey} from '@sanity/ui/theme'
-import {cloneElement, forwardRef, useCallback, useMemo, useState, useEffect, useRef} from 'react'
+import {
+  cloneElement,
+  forwardRef,
+  useCallback,
+  useMemo,
+  useState,
+  useEffect,
+  useRef,
+  useImperativeHandle,
+} from 'react'
 import {Popover, PopoverProps} from '../../primitives'
 import {Placement, Radius} from '../../types'
 import {MenuProps} from './menu'
@@ -52,7 +61,7 @@ export interface MenuButtonProps {
  */
 export const MenuButton = forwardRef(function MenuButton(
   props: MenuButtonProps,
-  ref: React.ForwardedRef<HTMLButtonElement | null>,
+  forwardedRef: React.ForwardedRef<HTMLButtonElement | null>,
 ) {
   const {
     __unstable_disableRestoreFocusOnClose: disableRestoreFocusOnClose = false,
@@ -219,19 +228,7 @@ export const MenuButton = forwardRef(function MenuButton(
 
   const menu = menuProp && cloneElement(menuProp, menuProps)
 
-  const setButtonRef = useCallback(
-    (el: HTMLButtonElement | null) => {
-      if (typeof ref === 'function') {
-        ref(el)
-      } else if (ref) {
-        ref.current = el
-      }
-
-      setButtonElement(el)
-    },
-    [ref],
-  )
-
+  const ref = useRef<HTMLButtonElement | null>(null)
   const button = useMemo(
     () =>
       buttonProp &&
@@ -243,12 +240,28 @@ export const MenuButton = forwardRef(function MenuButton(
         onMouseDown: handleMouseDown,
         'aria-haspopup': true,
         'aria-expanded': open,
-        ref: setButtonRef,
+        ref: ref,
         selected: open,
       }),
-    [buttonProp, handleButtonClick, handleButtonKeyDown, handleMouseDown, id, open, setButtonRef],
+    [buttonProp, handleButtonClick, handleButtonKeyDown, handleMouseDown, id, open],
   )
 
+  // Forward button ref to parent
+  useImperativeHandle<HTMLButtonElement | null, HTMLButtonElement | null>(
+    forwardedRef,
+    () => ref.current,
+  )
+
+  // If there's a button then we need to set the reference element to the cloned button ref
+  // and if button changes we make sure to update or remove the reference element.
+  useEffect(() => {
+    if (!button) return undefined
+
+    setButtonElement(ref.current)
+
+    return () => setButtonElement(null)
+  }, [button])
+
   const popoverProps: MenuButtonProps['popover'] = useMemo(
     () => ({
       boundaryElement: deprecated_boundaryElement,