fix(ui): preserve refs across hidden subtrees

.changeset/calm-refs-rest.md

@@ -0,0 +1,5 @@
+---
+"@sanity/ui": patch
+---
+
+fix: prevent stateful element refs from updating during Suspense and Activity detach cycles

apps/docs/src/lib/arcade/ArcadeFrame.tsx

@@ -22,6 +22,9 @@ export function ArcadeFrame({
   const [frame, setFrame] = useState<HTMLIFrameElement | null>(null)
   const [ready, setReady] = useState(false)
   const msgQueueRef = useRef<any[]>([])
+  const handleFrameRef = useCallback((element: HTMLIFrameElement | null) => {
+    if (element) setFrame(element)
+  }, [])
 
   // Handle messages from frame
   useEffect(() => {
@@ -73,5 +76,5 @@ export function ArcadeFrame({
     [hookCode, jsxCode, postMessage],
   )
 
-  return <Root ref={setFrame} src={`${basePath}/arcade/frame`} />
+  return <Root ref={handleFrameRef} src={`${basePath}/arcade/frame`} />
 }

packages/ui/src/core/components/autocomplete/autocomplete.tsx

@@ -23,6 +23,7 @@ import {
 import {EMPTY_ARRAY, EMPTY_RECORD} from '../../constants'
 import {_raf} from '../../helpers/animation'
 import {_hasFocus, focusFirstDescendant} from '../../helpers/focus'
+import {useConnectedRef} from '../../hooks/useConnectedRef'
 import {Box, BoxProps} from '../../primitives/box/box'
 import {Button} from '../../primitives/button/button'
 import {Card} from '../../primitives/card/card'
@@ -202,9 +203,10 @@ export function Autocomplete<Option extends BaseAutocompleteOption>(
    * This doesn't happen on React 19 due to automatic batching of all state updates, the startTransition wrapper here gives a type of batching for 18 users in a way that still works with 19.
    * NOTE: The startTransition wrapper is not needed in UI v4, since the baseline there is React 19.
    */
-  const setInputElement = useCallback((node: HTMLInputElement | null) => {
+  const handleInputElementChange = useCallback((node: HTMLInputElement | null) => {
     startTransition(() => _setInputElement(node))
   }, [])
+  const setInputElement = useConnectedRef(handleInputElementChange)
 
   // Value refs
   const listFocusedRef = useRef(false)

packages/ui/src/core/components/menu/menu.test.tsx

@@ -1,13 +1,45 @@
 /** @vitest-environment jsdom */
 
-import React, {useCallback, useMemo} from 'react'
+// oxlint-disable-next-line no-unassigned-import
+import '../../../../test/mocks/matchMedia.mock'
+import React, {act, Activity, useCallback, useMemo} from 'react'
 import {describe, expect, it, vi} from 'vitest'
 
 import {render} from '../../../../test/utils'
+import {LayerProvider} from '../../utils/layer/layerProvider'
+import {Menu} from './menu'
 import {MenuContext, MenuContextValue} from './menuContext'
 import {useMenu} from './useMenu'
 
 describe('components/menu', () => {
+  it('preserves element registration across Activity hide and reveal', async () => {
+    const unregisterElement = vi.fn()
+    const registerElement = vi.fn(() => unregisterElement)
+    const renderMenu = (mode: 'hidden' | 'visible') => (
+      <LayerProvider>
+        <Activity mode={mode}>
+          <Menu registerElement={registerElement}>Item</Menu>
+        </Activity>
+      </LayerProvider>
+    )
+    const {rerender, unmount} = render(renderMenu('visible'))
+
+    expect(registerElement).toHaveBeenCalledOnce()
+
+    rerender(renderMenu('hidden'))
+    await act(async () => undefined)
+
+    expect(unregisterElement).not.toHaveBeenCalled()
+
+    rerender(renderMenu('visible'))
+    expect(registerElement).toHaveBeenCalledOnce()
+
+    unmount()
+    await act(async () => undefined)
+
+    expect(unregisterElement).toHaveBeenCalledOnce()
+  })
+
   describe('useMenu', () => {
     it('should get context value', async () => {
       const log = vi.fn()

packages/ui/src/core/components/menu/menu.tsx

@@ -2,6 +2,7 @@ import {useCallback, useEffect, useImperativeHandle, useMemo, useRef} from 'reac
 import {styled} from 'styled-components'
 
 import {useClickOutsideEvent} from '../../hooks/useClickOutsideEvent'
+import {useConnectedRef} from '../../hooks/useConnectedRef'
 import {useGlobalKeyDown} from '../../hooks/useGlobalKeyDown'
 import {Box} from '../../primitives/box/box'
 import {Stack} from '../../primitives/stack/stack'
@@ -83,7 +84,7 @@ export function Menu(
   } = useMenuController({onKeyDown, originElement, shouldFocus, rootElementRef: ref})
 
   const unregisterElementRef = useRef<(() => void) | null>(null)
-  const handleRefChange = useCallback(
+  const handleRootElementChange = useCallback(
     (el: HTMLDivElement | null) => {
       // Run cleanup of previously registered elements
       if (unregisterElementRef.current) {
@@ -104,6 +105,7 @@ export function Menu(
     },
     [registerElement],
   )
+  const handleRefChange = useConnectedRef(handleRootElementChange)
 
   // Trigger `onItemSelect` when active index changes
   useEffect(() => {

packages/ui/src/core/components/menu/menuButton.tsx

@@ -8,6 +8,7 @@ import {
   useState,
 } from 'react'
 
+import {useConnectedRef} from '../../hooks/useConnectedRef'
 import {Popover, PopoverProps} from '../../primitives/popover/popover'
 import {MenuProps} from './menu'
 
@@ -72,6 +73,7 @@ export function MenuButton(props: MenuButtonProps) {
   const [open, setOpen] = useState(false)
   const [shouldFocus, setShouldFocus] = useState<'first' | 'last' | null>(null)
   const [buttonElement, setButtonElement] = useState<HTMLButtonElement | null>(null)
+  const buttonElementRef = useConnectedRef<HTMLButtonElement>(setButtonElement)
   const [menuElements, setChildMenuElements] = useState<HTMLElement[]>([])
   const openRef = useRef<boolean>(open)
 
@@ -214,10 +216,18 @@ export function MenuButton(props: MenuButtonProps) {
         'onMouseDown': handleMouseDown,
         'aria-haspopup': true,
         'aria-expanded': open,
-        'ref': setButtonElement,
+        'ref': buttonElementRef,
         'selected': buttonProp.props.selected ?? open,
       }),
-    [buttonProp, handleButtonClick, handleButtonKeyDown, handleMouseDown, id, open],
+    [
+      buttonElementRef,
+      buttonProp,
+      handleButtonClick,
+      handleButtonKeyDown,
+      handleMouseDown,
+      id,
+      open,
+    ],
   )
 
   // Forward button ref to parent

packages/ui/src/core/components/menu/menuGroup.tsx

@@ -2,6 +2,7 @@ import {ChevronRightIcon} from '@sanity/icons/ChevronRight'
 import {isValidElement, useCallback, useEffect, useState} from 'react'
 import {isValidElementType} from 'react-is'
 
+import {useConnectedRef} from '../../hooks/useConnectedRef'
 import {Selectable} from '../../primitives/_selectable/selectable'
 import {Box} from '../../primitives/box/box'
 import {Flex} from '../../primitives/flex/flex'
@@ -82,6 +83,7 @@ const MenuGroupComponent = function MenuGroup(
   } = menu
   const onItemMouseEnter = _onItemMouseEnter ?? menu.onItemMouseEnter
   const [rootElement, setRootElement] = useState<HTMLButtonElement | HTMLDivElement | null>(null)
+  const rootElementRef = useConnectedRef<HTMLButtonElement | HTMLDivElement>(setRootElement)
   const [open, setOpen] = useState(false)
   const [shouldFocus, setShouldFocus] = useState<'first' | 'last' | null>(null)
   const active = Boolean(activeElement) && activeElement === rootElement
@@ -201,7 +203,7 @@ const MenuGroupComponent = function MenuGroup(
         onClick={handleClick}
         onKeyDown={handleKeyDown}
         onMouseEnter={handleMouseEnter}
-        ref={setRootElement}
+        ref={rootElementRef}
         tabIndex={-1}
         type={as === 'button' ? 'button' : undefined}
       >

packages/ui/src/core/components/menu/menuItem.tsx

@@ -9,6 +9,7 @@ import {
 } from 'react'
 import {isValidElementType} from 'react-is'
 
+import {useConnectedRef} from '../../hooks/useConnectedRef'
 import {Selectable} from '../../primitives/_selectable/selectable'
 import {Box} from '../../primitives/box/box'
 import {Flex} from '../../primitives/flex/flex'
@@ -121,10 +122,11 @@ const MenuItemComponent = function MenuItem(
 
   const hotkeysFontSize = _getArrayProp(fontSize).map((s) => s - 1)
 
-  const setRef = useCallback((el: HTMLDivElement | null) => {
+  const handleRootElementChange = useCallback((el: HTMLDivElement | null) => {
     ref.current = el
     setRootElement(el)
   }, [])
+  const setRef = useConnectedRef(handleRootElementChange)
 
   return (
     <Selectable

packages/ui/src/core/components/tree/treeItem.tsx

@@ -3,6 +3,7 @@ import {startTransition, useCallback, useEffect, useId, useMemo, useRef, useStat
 import {styled} from 'styled-components'
 
 import {ThemeFontWeightKey} from '../../../theme/system/font'
+import {useConnectedRef} from '../../hooks/useConnectedRef'
 import {Box} from '../../primitives/box/box'
 import {Flex} from '../../primitives/flex/flex'
 import {Text} from '../../primitives/text/text'
@@ -92,9 +93,10 @@ export function TreeItem(
    * This doesn't happen on React 19 due to automatic batching of all state updates, the startTransition wrapper here gives a type of batching for 18 users in a way that still works with 19.
    * NOTE: The startTransition wrapper is not needed in UI v4, since the baseline there is React 19.
    */
-  const setRootElement = useCallback((node: HTMLLIElement | null) => {
+  const handleRootElementChange = useCallback((node: HTMLLIElement | null) => {
     startTransition(() => _setRootElement(node))
   }, [])
+  const setRootElement = useConnectedRef(handleRootElementChange)
 
   const treeitemRef = useRef<HTMLAnchorElement | null>(null)
   const tree = useTree()

packages/ui/src/core/hooks/useConnectedRef.test.tsx

@@ -0,0 +1,121 @@
+/** @vitest-environment jsdom */
+
+import {act, render} from '@testing-library/react'
+import {Activity, type ReactNode, Suspense} from 'react'
+import {describe, expect, it, vi} from 'vitest'
+
+import {useConnectedRef} from './useConnectedRef'
+
+function Fixture({
+  children = 'content',
+  onChange,
+  visible = true,
+}: {
+  children?: ReactNode
+  onChange: (node: HTMLDivElement | null) => void
+  visible?: boolean
+}) {
+  const ref = useConnectedRef(onChange)
+
+  return visible ? <div ref={ref}>{children}</div> : null
+}
+
+function createSuspender() {
+  let suspended = false
+  let resolve: (() => void) | undefined
+  let promise: Promise<void> | undefined
+
+  const Component = ({children}: {children: ReactNode}) => {
+    if (suspended) throw promise
+    return children
+  }
+
+  return {
+    Component,
+    suspend() {
+      suspended = true
+      promise = new Promise<void>((nextResolve) => {
+        resolve = nextResolve
+      })
+    },
+    resume() {
+      suspended = false
+      resolve?.()
+    },
+  }
+}
+
+describe('useConnectedRef', () => {
+  it('ignores Activity detach and reveal callbacks for connected nodes', async () => {
+    const onChange = vi.fn()
+    const {rerender, unmount} = render(
+      <Activity mode="visible">
+        <Fixture onChange={onChange} />
+      </Activity>,
+    )
+
+    expect(onChange).toHaveBeenCalledOnce()
+    expect(onChange).toHaveBeenLastCalledWith(expect.any(HTMLDivElement))
+    onChange.mockClear()
+
+    rerender(
+      <Activity mode="hidden">
+        <Fixture onChange={onChange} />
+      </Activity>,
+    )
+    await act(async () => undefined)
+
+    rerender(
+      <Activity mode="visible">
+        <Fixture onChange={onChange} />
+      </Activity>,
+    )
+
+    expect(onChange).not.toHaveBeenCalled()
+
+    unmount()
+    await act(async () => undefined)
+  })
+
+  it('ignores Suspense detach and reveal callbacks for connected nodes', async () => {
+    const onChange = vi.fn()
+    const suspender = createSuspender()
+    const renderFixture = () => (
+      <Suspense fallback={<div>loading</div>}>
+        <suspender.Component>
+          <Fixture onChange={onChange} />
+        </suspender.Component>
+      </Suspense>
+    )
+    const {rerender, unmount} = render(renderFixture())
+
+    expect(onChange).toHaveBeenCalledOnce()
+    onChange.mockClear()
+
+    suspender.suspend()
+    rerender(renderFixture())
+    await act(async () => undefined)
+
+    await act(async () => {
+      suspender.resume()
+      rerender(renderFixture())
+    })
+
+    expect(onChange).not.toHaveBeenCalled()
+
+    unmount()
+    await act(async () => undefined)
+  })
+
+  it('forwards null after a node is removed from the document', async () => {
+    const onChange = vi.fn()
+    const {rerender} = render(<Fixture onChange={onChange} />)
+    onChange.mockClear()
+
+    rerender(<Fixture onChange={onChange} visible={false} />)
+    await act(async () => undefined)
+
+    expect(onChange).toHaveBeenCalledOnce()
+    expect(onChange).toHaveBeenCalledWith(null)
+  })
+})

packages/ui/src/core/hooks/useConnectedRef.ts

@@ -0,0 +1,54 @@
+import {type RefCallback, useCallback, useRef} from 'react'
+
+/**
+ * Creates a callback ref that preserves connected elements across React's
+ * temporary Suspense and Activity detach cycles.
+ *
+ * React calls refs with `null` when hiding a subtree even though its DOM nodes
+ * remain connected. Forwarding that transient `null` to a state setter can
+ * cause a hide/reveal update loop. Real removals are forwarded after the
+ * commit, when the detached node is no longer connected.
+ *
+ * @internal
+ */
+export function useConnectedRef<T extends Node>(
+  onChange: (node: T | null) => void,
+): RefCallback<T> {
+  const currentNodeRef = useRef<T | null>(null)
+  const currentOnChangeRef = useRef(onChange)
+  const detachVersionRef = useRef(0)
+
+  return useCallback(
+    (node: T | null) => {
+      const detachVersion = ++detachVersionRef.current
+
+      if (node) {
+        const changed = currentNodeRef.current !== node || currentOnChangeRef.current !== onChange
+
+        currentNodeRef.current = node
+        currentOnChangeRef.current = onChange
+
+        if (changed) onChange(node)
+        return
+      }
+
+      const detachedNode = currentNodeRef.current
+      if (!detachedNode) return
+
+      queueMicrotask(() => {
+        if (
+          detachVersionRef.current !== detachVersion ||
+          currentNodeRef.current !== detachedNode ||
+          detachedNode.isConnected
+        ) {
+          return
+        }
+
+        currentNodeRef.current = null
+        currentOnChangeRef.current = onChange
+        onChange(null)
+      })
+    },
+    [onChange],
+  )
+}