Repository: tailwindlabs/headlessui
Pull request #3407 (2024-08-01 23:19:34)

Fix closing components using the `transition` prop, and after scrolling the page
This PR fixes an issue where components such as the `MenuItems` and the `ListboxOptions` won't close correctly if these are used in combination with the `transition` prop and _after_ you scrolled the page a little bit.

For some background information: the `Menu` and `Listbox` components can be used with the `transition` prop so that the `MenuItems` and `ListboxOptions` can be transitioned in and out based on the open/closed state.

To make the user experience better, we have some logic to check whether the `MenuButton` or `ListboxButton` moved from its initial position. If it did move when it's closing, then we immediately stop all transitions by unmounting the `MenuItems` or `ListboxOptions` immediately.

A situation this happens in:

1. Open the `Menu` (with the `transition` prop)
2. Press <kbd>tab</kbd> to focus the next element on the page
3. The browser scrolls the page to make the next element visible
4. The `MenuItems` will fade out while moving around (while the page scrolls)
   1. Bonus points, if you also use the `anchor` prop, then the `MenuItems` will reposition itself which makes the UX even worse.

While debugging this issue, I noticed a timing related issue because we are using the DOM element in a `useTransition` hook. Typically you store those DOM elements in a `useRef`. This is good practice because the DOM element will be filled in once you hit the `useEffect` hook, and you don't introduce unnecessary re-renders.

However, if the DOM element changes, then the `useRef` value will update, but the `useEffect` hook doesn't run again (because no state change happened).

There are other places in Headless UI where we also need to re-run effects when the DOM node changes. There are also places where we need access to the DOM element during render.

One of React's rules is to not read / write a ref during render.

So, this PR simplifies a few components and hooks by actually storing the DOM elements in state directly instead of in a ref. This way we can use the DOM element in dependency arrays of `useEffect` hooks, and we can use the DOM element during render.

Maybe a bit of an unorthodox solution, but it works quite well and it simplifies the code quite a bit.

---

It might be easier to review commit by commit. The reason why is that in some hooks I added temporary support for `HTMLElement | null` in places where we wanted `MutableRefObject<HTMLElement | null>`. These now temporary look like `MutableRefObject<HTMLElement | null> | HTMLElement | null`.

This way we can use both versions while refactoring. By the end of the commits those are simplified again such that we only handle `HTMLElement | null` instead of `MutableRefObject<HTMLElement | null>`.

Fixes: #3398
Fixes: #3403




vercel[bot]: [vc]: #+Mm17UcoxM8RU8d8J4eyAExXluvC0bPov0kEAzBahSc=:eyJpc01vbm9yZXBvIjp0cnVlLCJ0eXBlIjoiZ2l0aHViIiwicHJvamVjdHMiOlt7Im5hbWUiOiJoZWFkbGVzc3VpLXZ1ZSIsInJvb3REaXJlY3RvcnkiOiJwbGF5Z3JvdW5kcy92dWUiLCJpbnNwZWN0b3JVcmwiOiJodHRwczovL3ZlcmNlbC5jb20vdGFpbHdpbmRsYWJzL2hlYWRsZXNzdWktdnVlL0NyVDVkM2thbUVuejNMalZYVUhQQTI2YVluUWgiLCJwcmV2aWV3VXJsIjoiaGVhZGxlc3N1aS12dWUtZ2l0LWZpeC1pc3N1ZS0zMzk4LTM0MDMtdGFpbHdpbmRsYWJzLnZlcmNlbC5hcHAiLCJuZXh0Q29tbWl0U3RhdHVzIjoiUEVORElORyIsImxpdmVGZWVkYmFjayI6eyJyZXNvbHZlZCI6MCwidW5yZXNvbHZlZCI6MCwidG90YWwiOjAsImxpbmsiOiJoZWFkbGVzc3VpLXZ1ZS1naXQtZml4LWlzc3VlLTMzOTgtMzQwMy10YWlsd2luZGxhYnMudmVyY2VsLmFwcCJ9fV19
**The latest updates on your projects**. Learn more about [Vercel for Git ↗︎](https://vercel.link/github-learn-more)

| Name | Status | Preview | Comments | Updated (UTC) |
| :--- | :----- | :------ | :------- | :------ |
| **headlessui-vue** | 🔄 Building ([Inspect](https://vercel.com/tailwindlabs/headlessui-vue/CrT5d3kamEnz3LjVXUHPA26aYnQh)) | [Visit Preview](https://vercel.live/open-feedback/headlessui-vue-git-fix-issue-3398-3403-tailwindlabs.vercel.app?via=pr-comment-visit-preview-link&passThrough=1) | 💬 [**Add feedback**](https://vercel.live/open-feedback/headlessui-vue-git-fix-issue-3398-3403-tailwindlabs.vercel.app?via=pr-comment-feedback-link) | Aug 1, 2024 11:19pm |


thecrypticace: I like how much cleaner this makes some of the code 💯

Patch:
diff --git a/packages/@headlessui-react/CHANGELOG.md b/packages/@headlessui-react/CHANGELOG.md
index a8d6645..b470bb7 100644
--- a/packages/@headlessui-react/CHANGELOG.md
+++ b/packages/@headlessui-react/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Ensure `Transition` component state doesn't change when it becomes hidden ([#3372](https://github.com/tailwindlabs/headlessui/pull/3372))
+- Fix closing components using the `transition` prop, and after scrolling the page ([#3407](https://github.com/tailwindlabs/headlessui/pull/3407))
 
 ## [2.1.2] - 2024-07-05
 
diff --git a/packages/@headlessui-react/src/components/combobox/combobox.tsx b/packages/@headlessui-react/src/components/combobox/combobox.tsx
index 20b1652..131284a 100644
--- a/packages/@headlessui-react/src/components/combobox/combobox.tsx
+++ b/packages/@headlessui-react/src/components/combobox/combobox.tsx
@@ -116,6 +116,10 @@ interface StateDefinition<T> {
 
   isTyping: boolean
 
+  inputElement: HTMLInputElement | null
+  buttonElement: HTMLButtonElement | null
+  optionsElement: HTMLElement | null
+
   __demoMode: boolean
 }
 
@@ -132,6 +136,10 @@ enum ActionTypes {
   SetActivationTrigger,
 
   UpdateVirtualConfiguration,
+
+  SetInputElement,
+  SetButtonElement,
+  SetOptionsElement,
 }
 
 function adjustOrderedState<T>(
@@ -192,6 +200,9 @@ type Actions<T> =
       options: T[]
       disabled: ((value: any) => boolean) | null
     }
+  | { type: ActionTypes.SetInputElement; element: HTMLInputElement | null }
+  | { type: ActionTypes.SetButtonElement; element: HTMLButtonElement | null }
+  | { type: ActionTypes.SetOptionsElement; element: HTMLElement | null }
 
 let reducers: {
   [P in ActionTypes]: <T>(
@@ -245,7 +256,7 @@ let reducers: {
   [ActionTypes.GoToOption](state, action) {
     if (state.dataRef.current?.disabled) return state
     if (
-      state.dataRef.current?.optionsRef.current &&
+      state.optionsElement &&
       !state.dataRef.current?.optionsPropsRef.current.static &&
       state.comboboxState === ComboboxState.Closed
     ) {
@@ -419,6 +430,18 @@ let reducers: {
       virtual: { options: action.options, disabled: action.disabled ?? (() => false) },
     }
   },
+  [ActionTypes.SetInputElement]: (state, action) => {
+    if (state.inputElement === action.element) return state
+    return { ...state, inputElement: action.element }
+  },
+  [ActionTypes.SetButtonElement]: (state, action) => {
+    if (state.buttonElement === action.element) return state
+    return { ...state, buttonElement: action.element }
+  },
+  [ActionTypes.SetOptionsElement]: (state, action) => {
+    if (state.optionsElement === action.element) return state
+    return { ...state, optionsElement: action.element }
+  },
 }
 
 let ComboboxActionsContext = createContext<{
@@ -431,6 +454,10 @@ let ComboboxActionsContext = createContext<{
   selectActiveOption(): void
   setActivationTrigger(trigger: ActivationTrigger): void
   onChange(value: unknown): void
+
+  setInputElement(element: HTMLInputElement | null): void
+  setButtonElement(element: HTMLButtonElement | null): void
+  setOptionsElement(element: HTMLElement | null): void
 } | null>(null)
 ComboboxActionsContext.displayName = 'ComboboxActionsContext'
 
@@ -455,7 +482,7 @@ function VirtualProvider(props: {
   let { options } = data.virtual!
 
   let [paddingStart, paddingEnd] = useMemo(() => {
-    let el = data.optionsRef.current
+    let el = data.optionsElement
     if (!el) return [0, 0]
 
     let styles = window.getComputedStyle(el)
@@ -464,7 +491,7 @@ function VirtualProvider(props: {
       parseFloat(styles.paddingBlockStart || styles.paddingTop),
       parseFloat(styles.paddingBlockEnd || styles.paddingBottom),
     ]
-  }, [data.optionsRef.current])
+  }, [data.optionsElement])
 
   let virtualizer = useVirtualizer({
     enabled: options.length !== 0,
@@ -475,7 +502,7 @@ function VirtualProvider(props: {
       return 40
     },
     getScrollElement() {
-      return (data.optionsRef.current ?? null) as HTMLElement | null
+      return data.optionsElement
     },
     overscan: 12,
   })
@@ -573,10 +600,6 @@ let ComboboxDataContext = createContext<
         static: boolean
         hold: boolean
       }>
-
-      inputRef: MutableRefObject<HTMLInputElement | null>
-      buttonRef: MutableRefObject<HTMLButtonElement | null>
-      optionsRef: MutableRefObject<HTMLElement | null>
     } & Omit<StateDefinition<unknown>, 'dataRef'>)
   | null
 >(null)
@@ -688,6 +711,9 @@ function ComboboxFn<TValue, TTag extends ElementType = typeof DEFAULT_COMBOBOX_T
       : null,
     activeOptionIndex: null,
     activationTrigger: ActivationTrigger.Other,
+    inputElement: null,
+    buttonElement: null,
+    optionsElement: null,
     __demoMode,
   } as StateDefinition<TValue>)
 
@@ -695,10 +721,6 @@ function ComboboxFn<TValue, TTag extends ElementType = typeof DEFAULT_COMBOBOX_T
 
   let optionsPropsRef = useRef<_Data['optionsPropsRef']['current']>({ static: false, hold: false })
 
-  let inputRef = useRef<_Data['inputRef']['current']>(null)
-  let buttonRef = useRef<_Data['buttonRef']['current']>(null)
-  let optionsRef = useRef<_Data['optionsRef']['current']>(null)
-
   type TActualValue = true extends typeof multiple ? EnsureArray<TValue>[number] : TValue
   let compare = useByComparator<TActualValue>(by)
 
@@ -733,9 +755,6 @@ function ComboboxFn<TValue, TTag extends ElementType = typeof DEFAULT_COMBOBOX_T
       ...state,
       immediate,
       optionsPropsRef,
-      inputRef,
-      buttonRef,
-      optionsRef,
       value,
       defaultValue,
       disabled,
@@ -791,8 +810,10 @@ function ComboboxFn<TValue, TTag extends ElementType = typeof DEFAULT_COMBOBOX_T
 
   // Handle outside click
   let outsideClickEnabled = data.comboboxState === ComboboxState.Open
-  useOutsideClick(outsideClickEnabled, [data.buttonRef, data.inputRef, data.optionsRef], () =>
-    actions.closeCombobox()
+  useOutsideClick(
+    outsideClickEnabled,
+    [data.buttonElement, data.inputElement, data.optionsElement],
+    () => actions.closeCombobox()
   )
 
   let slot = useMemo(() => {
@@ -896,6 +917,18 @@ function ComboboxFn<TValue, TTag extends ElementType = typeof DEFAULT_COMBOBOX_T
     dispatch({ type: ActionTypes.SetActivationTrigger, trigger })
   })
 
+  let setInputElement = useEvent((element: HTMLInputElement | null) => {
+    dispatch({ type: ActionTypes.SetInputElement, element })
+  })
+
+  let setButtonElement = useEvent((element: HTMLButtonElement | null) => {
+    dispatch({ type: ActionTypes.SetButtonElement, element })
+  })
+
+  let setOptionsElement = useEvent((element: HTMLElement | null) => {
+    dispatch({ type: ActionTypes.SetOptionsElement, element })
+  })
+
   let actions = useMemo<_Actions>(
     () => ({
       onChange,
@@ -906,6 +939,9 @@ function ComboboxFn<TValue, TTag extends ElementType = typeof DEFAULT_COMBOBOX_T
       openCombobox,
       setActivationTrigger,
       selectActiveOption,
+      setInputElement,
+      setButtonElement,
+      setOptionsElement,
     }),
     []
   )
@@ -923,7 +959,7 @@ function ComboboxFn<TValue, TTag extends ElementType = typeof DEFAULT_COMBOBOX_T
     <LabelProvider
       value={labelledby}
       props={{
-        htmlFor: data.inputRef.current?.id,
+        htmlFor: data.inputElement?.id,
       }}
       slot={{
         open: data.comboboxState === ComboboxState.Open,
@@ -1019,15 +1055,16 @@ function InputFn<
     ...theirProps
   } = props
 
-  let inputRef = useSyncRefs(data.inputRef, ref, useFloatingReference())
-  let ownerDocument = useOwnerDocument(data.inputRef)
+  let internalInputRef = useRef<HTMLInputElement | null>(null)
+  let inputRef = useSyncRefs(internalInputRef, ref, useFloatingReference(), actions.setInputElement)
+  let ownerDocument = useOwnerDocument(data.inputElement)
 
   let d = useDisposables()
 
   let clear = useEvent(() => {
     actions.onChange(null)
-    if (data.optionsRef.current) {
-      data.optionsRef.current.scrollTop = 0
+    if (data.optionsElement) {
+      data.optionsElement.scrollTop = 0
     }
     actions.goToOption(Focus.Nothing)
   })
@@ -1071,7 +1108,7 @@ function InputFn<
       // using an IME, we don't want to mess with the input at all.
       if (data.isTyping) return
 
-      let input = data.inputRef.current
+      let input = internalInputRef.current
       if (!input) return
 
       if (oldState === ComboboxState.Open && state === ComboboxState.Closed) {
@@ -1121,7 +1158,7 @@ function InputFn<
         // using an IME, we don't want to mess with the input at all.
         if (data.isTyping) return
 
-        let input = data.inputRef.current
+        let input = internalInputRef.current
         if (!input) return
 
         // Capture current state
@@ -1232,7 +1269,7 @@ function InputFn<
       case Keys.Escape:
         if (data.comboboxState !== ComboboxState.Open) return
         event.preventDefault()
-        if (data.optionsRef.current && !data.optionsPropsRef.current.static) {
+        if (data.optionsElement && !data.optionsPropsRef.current.static) {
           event.stopPropagation()
         }
 
@@ -1286,10 +1323,10 @@ function InputFn<
       (event.relatedTarget as HTMLElement) ?? history.find((x) => x !== event.currentTarget)
 
     // Focus is moved into the list, we don't want to close yet.
-    if (data.optionsRef.current?.contains(relatedTarget)) return
+    if (data.optionsElement?.contains(relatedTarget)) return
 
     // Focus is moved to the button, we don't want to close yet.
-    if (data.buttonRef.current?.contains(relatedTarget)) return
+    if (data.buttonElement?.contains(relatedTarget)) return
 
     // Focus is moved, but the combobox is not open. This can mean two things:
     //
@@ -1316,8 +1353,8 @@ function InputFn<
   let handleFocus = useEvent((event: ReactFocusEvent) => {
     let relatedTarget =
       (event.relatedTarget as HTMLElement) ?? history.find((x) => x !== event.currentTarget)
-    if (data.buttonRef.current?.contains(relatedTarget)) return
-    if (data.optionsRef.current?.contains(relatedTarget)) return
+    if (data.buttonElement?.contains(relatedTarget)) return
+    if (data.optionsElement?.contains(relatedTarget)) return
     if (data.disabled) return
 
     if (!data.immediate) return
@@ -1378,7 +1415,7 @@ function InputFn<
       id,
       role: 'combobox',
       type,
-      'aria-controls': data.optionsRef.current?.id,
+      'aria-controls': data.optionsElement?.id,
       'aria-expanded': data.comboboxState === ComboboxState.Open,
       'aria-activedescendant':
         data.activeOptionIndex === null
@@ -1457,7 +1494,8 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
 ) {
   let data = useData('Combobox.Button')
   let actions = useActions('Combobox.Button')
-  let buttonRef = useSyncRefs(data.buttonRef, ref)
+  let buttonRef = useSyncRefs(ref, actions.setButtonElement)
+
   let internalId = useId()
   let {
     id = `headlessui-combobox-button-${internalId}`,
@@ -1466,7 +1504,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     ...theirProps
   } = props
 
-  let refocusInput = useRefocusableInput(data.inputRef)
+  let refocusInput = useRefocusableInput(data.inputElement)
 
   let handleKeyDown = useEvent((event: ReactKeyboardEvent<HTMLElement>) => {
     switch (event.key) {
@@ -1505,7 +1543,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
       case Keys.Escape:
         if (data.comboboxState !== ComboboxState.Open) return
         event.preventDefault()
-        if (data.optionsRef.current && !data.optionsPropsRef.current.static) {
+        if (data.optionsElement && !data.optionsPropsRef.current.static) {
           event.stopPropagation()
         }
         flushSync(() => actions.closeCombobox())
@@ -1561,10 +1599,10 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     {
       ref: buttonRef,
       id,
-      type: useResolveButtonType(props, data.buttonRef),
+      type: useResolveButtonType(props, data.buttonElement),
       tabIndex: -1,
       'aria-haspopup': 'listbox',
-      'aria-controls': data.optionsRef.current?.id,
+      'aria-controls': data.optionsElement?.id,
       'aria-expanded': data.comboboxState === ComboboxState.Open,
       'aria-labelledby': labelledBy,
       disabled: disabled || undefined,
@@ -1635,20 +1673,20 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
 
   let [floatingRef, style] = useFloatingPanel(anchor)
   let getFloatingPanelProps = useFloatingPanelProps()
-  let optionsRef = useSyncRefs(data.optionsRef, ref, anchor ? floatingRef : null)
-  let ownerDocument = useOwnerDocument(data.optionsRef)
+  let optionsRef = useSyncRefs(ref, anchor ? floatingRef : null, actions.setOptionsElement)
+  let ownerDocument = useOwnerDocument(data.optionsElement)
 
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    data.optionsRef,
+    data.optionsElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : data.comboboxState === ComboboxState.Open
   )
 
   // Ensure we close the combobox as soon as the input becomes hidden
-  useOnDisappear(visible, data.inputRef, actions.closeCombobox)
+  useOnDisappear(visible, data.inputElement, actions.closeCombobox)
 
   // Enable scroll locking when the combobox is visible, and `modal` is enabled
   let scrollLockEnabled = data.__demoMode
@@ -1661,11 +1699,10 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     ? false
     : modal && data.comboboxState === ComboboxState.Open
   useInertOthers(inertOthersEnabled, {
-    allowed: useEvent(() => [
-      data.inputRef.current,
-      data.buttonRef.current,
-      data.optionsRef.current,
-    ]),
+    allowed: useCallback(
+      () => [data.inputElement, data.buttonElement, data.optionsElement],
+      [data.inputElement, data.buttonElement, data.optionsElement]
+    ),
   })
 
   useIsoMorphicEffect(() => {
@@ -1676,7 +1713,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   }, [data.optionsPropsRef, hold])
 
   useTreeWalker(data.comboboxState === ComboboxState.Open, {
-    container: data.optionsRef.current,
+    container: data.optionsElement,
     accept(node) {
       if (node.getAttribute('role') === 'option') return NodeFilter.FILTER_REJECT
       if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
@@ -1687,7 +1724,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     },
   })
 
-  let labelledBy = useLabelledBy([data.buttonRef.current?.id])
+  let labelledBy = useLabelledBy([data.buttonElement?.id])
 
   let slot = useMemo(() => {
     return {
@@ -1728,8 +1765,8 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     style: {
       ...theirProps.style,
       ...style,
-      '--input-width': useElementSize(data.inputRef, true).width,
-      '--button-width': useElementSize(data.buttonRef, true).width,
+      '--input-width': useElementSize(data.inputElement, true).width,
+      '--button-width': useElementSize(data.buttonElement, true).width,
     } as CSSProperties,
     onWheel: data.activationTrigger === ActivationTrigger.Pointer ? undefined : handleWheel,
     onMouseDown: handleMouseDown,
@@ -1840,7 +1877,7 @@ function OptionFn<
     ...theirProps
   } = props
 
-  let refocusInput = useRefocusableInput(data.inputRef)
+  let refocusInput = useRefocusableInput(data.inputElement)
 
   let active = data.virtual
     ? data.activeOptionIndex === data.calculateIndex(value)
diff --git a/packages/@headlessui-react/src/components/disclosure/disclosure.tsx b/packages/@headlessui-react/src/components/disclosure/disclosure.tsx
index 66f5938..4a20bb2 100644
--- a/packages/@headlessui-react/src/components/disclosure/disclosure.tsx
+++ b/packages/@headlessui-react/src/components/disclosure/disclosure.tsx
@@ -57,10 +57,8 @@ enum DisclosureStates {
 interface StateDefinition {
   disclosureState: DisclosureStates
 
-  linkedPanel: boolean
-
-  buttonRef: MutableRefObject<HTMLButtonElement | null>
-  panelRef: MutableRefObject<HTMLElement | null>
+  buttonElement: HTMLButtonElement | null
+  panelElement: HTMLElement | null
 
   buttonId: string | null
   panelId: string | null
@@ -73,8 +71,8 @@ enum ActionTypes {
   SetButtonId,
   SetPanelId,
 
-  LinkPanel,
-  UnlinkPanel,
+  SetButtonElement,
+  SetPanelElement,
 }
 
 type Actions =
@@ -82,8 +80,8 @@ type Actions =
   | { type: ActionTypes.CloseDisclosure }
   | { type: ActionTypes.SetButtonId; buttonId: string | null }
   | { type: ActionTypes.SetPanelId; panelId: string | null }
-  | { type: ActionTypes.LinkPanel }
-  | { type: ActionTypes.UnlinkPanel }
+  | { type: ActionTypes.SetButtonElement; element: HTMLButtonElement | null }
+  | { type: ActionTypes.SetPanelElement; element: HTMLElement | null }
 
 let reducers: {
   [P in ActionTypes]: (
@@ -102,14 +100,6 @@ let reducers: {
     if (state.disclosureState === DisclosureStates.Closed) return state
     return { ...state, disclosureState: DisclosureStates.Closed }
   },
-  [ActionTypes.LinkPanel](state) {
-    if (state.linkedPanel === true) return state
-    return { ...state, linkedPanel: true }
-  },
-  [ActionTypes.UnlinkPanel](state) {
-    if (state.linkedPanel === false) return state
-    return { ...state, linkedPanel: false }
-  },
   [ActionTypes.SetButtonId](state, action) {
     if (state.buttonId === action.buttonId) return state
     return { ...state, buttonId: action.buttonId }
@@ -118,6 +108,14 @@ let reducers: {
     if (state.panelId === action.panelId) return state
     return { ...state, panelId: action.panelId }
   },
+  [ActionTypes.SetButtonElement](state, action) {
+    if (state.buttonElement === action.element) return state
+    return { ...state, buttonElement: action.element }
+  },
+  [ActionTypes.SetPanelElement](state, action) {
+    if (state.panelElement === action.element) return state
+    return { ...state, panelElement: action.element }
+  },
 }
 
 let DisclosureContext = createContext<[StateDefinition, Dispatch<Actions>] | null>(null)
@@ -195,14 +193,10 @@ function DisclosureFn<TTag extends ElementType = typeof DEFAULT_DISCLOSURE_TAG>(
     )
   )
 
-  let panelRef = useRef<StateDefinition['panelRef']['current']>(null)
-  let buttonRef = useRef<StateDefinition['buttonRef']['current']>(null)
-
   let reducerBag = useReducer(stateReducer, {
     disclosureState: defaultOpen ? DisclosureStates.Open : DisclosureStates.Closed,
-    linkedPanel: false,
-    buttonRef,
-    panelRef,
+    buttonElement: null,
+    panelElement: null,
     buttonId: null,
     panelId: null,
   } as StateDefinition)
@@ -301,7 +295,13 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   let isWithinPanel = panelContext === null ? false : panelContext === state.panelId
 
   let internalButtonRef = useRef<HTMLButtonElement | null>(null)
-  let buttonRef = useSyncRefs(internalButtonRef, ref, !isWithinPanel ? state.buttonRef : null)
+  let buttonRef = useSyncRefs(
+    internalButtonRef,
+    ref,
+    !isWithinPanel
+      ? useEvent((element) => dispatch({ type: ActionTypes.SetButtonElement, element }))
+      : null
+  )
   let mergeRefs = useMergeRefsFn()
 
   useEffect(() => {
@@ -323,7 +323,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
           event.preventDefault()
           event.stopPropagation()
           dispatch({ type: ActionTypes.ToggleDisclosure })
-          state.buttonRef.current?.focus()
+          state.buttonElement?.focus()
           break
       }
     } else {
@@ -355,7 +355,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
 
     if (isWithinPanel) {
       dispatch({ type: ActionTypes.ToggleDisclosure })
-      state.buttonRef.current?.focus()
+      state.buttonElement?.focus()
     } else {
       dispatch({ type: ActionTypes.ToggleDisclosure })
     }
@@ -376,7 +376,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     } satisfies ButtonRenderPropArg
   }, [state, hover, active, focus, disabled, autoFocus])
 
-  let type = useResolveButtonType(props, internalButtonRef)
+  let type = useResolveButtonType(props, state.buttonElement)
   let ourProps = isWithinPanel
     ? mergeProps(
         {
@@ -397,7 +397,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
           id,
           type,
           'aria-expanded': state.disclosureState === DisclosureStates.Open,
-          'aria-controls': state.linkedPanel ? state.panelId : undefined,
+          'aria-controls': state.panelElement ? state.panelId : undefined,
           disabled: disabled || undefined,
           autoFocus,
           onKeyDown: handleKeyDown,
@@ -451,9 +451,12 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   let { close } = useDisclosureAPIContext('Disclosure.Panel')
   let mergeRefs = useMergeRefsFn()
 
-  let panelRef = useSyncRefs(ref, state.panelRef, (el) => {
-    startTransition(() => dispatch({ type: el ? ActionTypes.LinkPanel : ActionTypes.UnlinkPanel }))
-  })
+  let panelRef = useSyncRefs(
+    ref,
+    useEvent((element) => {
+      startTransition(() => dispatch({ type: ActionTypes.SetPanelElement, element }))
+    })
+  )
 
   useEffect(() => {
     dispatch({ type: ActionTypes.SetPanelId, panelId: id })
@@ -465,7 +468,7 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    state.panelRef,
+    state.panelElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : state.disclosureState === DisclosureStates.Open
diff --git a/packages/@headlessui-react/src/components/listbox/listbox.tsx b/packages/@headlessui-react/src/components/listbox/listbox.tsx
index 9a7d166..e8c40e4 100644
--- a/packages/@headlessui-react/src/components/listbox/listbox.tsx
+++ b/packages/@headlessui-react/src/components/listbox/listbox.tsx
@@ -22,7 +22,6 @@ import React, {
 import { flushSync } from 'react-dom'
 import { useActivePress } from '../../hooks/use-active-press'
 import { useByComparator, type ByComparator } from '../../hooks/use-by-comparator'
-import { useComputed } from '../../hooks/use-computed'
 import { useControllable } from '../../hooks/use-controllable'
 import { useDefaultValue } from '../../hooks/use-default-value'
 import { useDidElementMove } from '../../hooks/use-did-element-move'
@@ -116,6 +115,9 @@ interface StateDefinition<T> {
   activeOptionIndex: number | null
   activationTrigger: ActivationTrigger
 
+  buttonElement: HTMLButtonElement | null
+  optionsElement: HTMLElement | null
+
   __demoMode: boolean
 }
 
@@ -129,6 +131,9 @@ enum ActionTypes {
 
   RegisterOption,
   UnregisterOption,
+
+  SetButtonElement,
+  SetOptionsElement,
 }
 
 function adjustOrderedState<T>(
@@ -173,6 +178,8 @@ type Actions<T> =
   | { type: ActionTypes.ClearSearch }
   | { type: ActionTypes.RegisterOption; id: string; dataRef: ListboxOptionDataRef<T> }
   | { type: ActionTypes.UnregisterOption; id: string }
+  | { type: ActionTypes.SetButtonElement; element: HTMLButtonElement | null }
+  | { type: ActionTypes.SetOptionsElement; element: HTMLElement | null }
 
 let reducers: {
   [P in ActionTypes]: <T>(
@@ -381,6 +388,14 @@ let reducers: {
       activationTrigger: ActivationTrigger.Other,
     }
   },
+  [ActionTypes.SetButtonElement]: (state, action) => {
+    if (state.buttonElement === action.element) return state
+    return { ...state, buttonElement: action.element }
+  },
+  [ActionTypes.SetOptionsElement]: (state, action) => {
+    if (state.optionsElement === action.element) return state
+    return { ...state, optionsElement: action.element }
+  },
 }
 
 let ListboxActionsContext = createContext<{
@@ -394,6 +409,8 @@ let ListboxActionsContext = createContext<{
   onChange(value: unknown): void
   search(query: string): void
   clearSearch(): void
+  setButtonElement(element: HTMLButtonElement | null): void
+  setOptionsElement(element: HTMLElement | null): void
 } | null>(null)
 ListboxActionsContext.displayName = 'ListboxActionsContext'
 
@@ -425,9 +442,6 @@ let ListboxDataContext = createContext<
       }>
 
       listRef: MutableRefObject<Map<string, HTMLElement | null>>
-
-      buttonRef: MutableRefObject<HTMLButtonElement | null>
-      optionsRef: MutableRefObject<HTMLElement | null>
     } & Omit<StateDefinition<unknown>, 'dataRef'>)
   | null
 >(null)
@@ -521,13 +535,13 @@ function ListboxFn<
     activeOptionIndex: null,
     activationTrigger: ActivationTrigger.Other,
     optionsVisible: false,
+    buttonElement: null,
+    optionsElement: null,
     __demoMode,
   } as StateDefinition<TType>)
 
   let optionsPropsRef = useRef<_Data['optionsPropsRef']['current']>({ static: false, hold: false })
 
-  let buttonRef = useRef<_Data['buttonRef']['current']>(null)
-  let optionsRef = useRef<_Data['optionsRef']['current']>(null)
   let listRef = useRef<_Data['listRef']['current']>(new Map())
 
   let compare = useByComparator(by)
@@ -556,8 +570,6 @@ function ListboxFn<
       compare,
       isSelected,
       optionsPropsRef,
-      buttonRef,
-      optionsRef,
       listRef,
     }),
     [value, disabled, invalid, multiple, state, listRef]
@@ -569,14 +581,18 @@ function ListboxFn<
 
   // Handle outside click
   let outsideClickEnabled = data.listboxState === ListboxStates.Open
-  useOutsideClick(outsideClickEnabled, [data.buttonRef, data.optionsRef], (event, target) => {
-    dispatch({ type: ActionTypes.CloseListbox })
+  useOutsideClick(
+    outsideClickEnabled,
+    [data.buttonElement, data.optionsElement],
+    (event, target) => {
+      dispatch({ type: ActionTypes.CloseListbox })
 
-    if (!isFocusableElement(target, FocusableMode.Loose)) {
-      event.preventDefault()
-      data.buttonRef.current?.focus()
+      if (!isFocusableElement(target, FocusableMode.Loose)) {
+        event.preventDefault()
+        data.buttonElement?.focus()
+      }
     }
-  })
+  )
 
   let slot = useMemo(() => {
     return {
@@ -647,6 +663,12 @@ function ListboxFn<
 
   let search = useEvent((value: string) => dispatch({ type: ActionTypes.Search, value }))
   let clearSearch = useEvent(() => dispatch({ type: ActionTypes.ClearSearch }))
+  let setButtonElement = useEvent((element: HTMLButtonElement | null) => {
+    dispatch({ type: ActionTypes.SetButtonElement, element })
+  })
+  let setOptionsElement = useEvent((element: HTMLElement | null) => {
+    dispatch({ type: ActionTypes.SetOptionsElement, element })
+  })
 
   let actions = useMemo<_Actions>(
     () => ({
@@ -659,6 +681,8 @@ function ListboxFn<
       selectOption,
       search,
       clearSearch,
+      setButtonElement,
+      setOptionsElement,
     }),
     []
   )
@@ -676,7 +700,7 @@ function ListboxFn<
     <LabelProvider
       value={labelledby}
       props={{
-        htmlFor: data.buttonRef.current?.id,
+        htmlFor: data.buttonElement?.id,
       }}
       slot={{
         open: data.listboxState === ListboxStates.Open,
@@ -760,7 +784,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     autoFocus = false,
     ...theirProps
   } = props
-  let buttonRef = useSyncRefs(data.buttonRef, ref, useFloatingReference())
+  let buttonRef = useSyncRefs(ref, useFloatingReference(), actions.setButtonElement)
   let getFloatingReferenceProps = useFloatingReferenceProps()
 
   let handleKeyDown = useEvent((event: ReactKeyboardEvent<HTMLButtonElement>) => {
@@ -801,7 +825,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     if (isDisabledReactIssue7711(event.currentTarget)) return event.preventDefault()
     if (data.listboxState === ListboxStates.Open) {
       flushSync(() => actions.closeListbox())
-      data.buttonRef.current?.focus({ preventScroll: true })
+      data.buttonElement?.focus({ preventScroll: true })
     } else {
       event.preventDefault()
       actions.openListbox()
@@ -836,9 +860,9 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     {
       ref: buttonRef,
       id,
-      type: useResolveButtonType(props, data.buttonRef),
+      type: useResolveButtonType(props, data.buttonElement),
       'aria-haspopup': 'listbox',
-      'aria-controls': data.optionsRef.current?.id,
+      'aria-controls': data.optionsElement?.id,
       'aria-expanded': data.listboxState === ListboxStates.Open,
       'aria-labelledby': labelledBy,
       'aria-describedby': describedBy,
@@ -916,19 +940,19 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   let data = useData('Listbox.Options')
   let actions = useActions('Listbox.Options')
 
-  let ownerDocument = useOwnerDocument(data.optionsRef)
+  let ownerDocument = useOwnerDocument(data.optionsElement)
 
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    data.optionsRef,
+    data.optionsElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : data.listboxState === ListboxStates.Open
   )
 
   // Ensure we close the listbox as soon as the button becomes hidden
-  useOnDisappear(visible, data.buttonRef, actions.closeListbox)
+  useOnDisappear(visible, data.buttonElement, actions.closeListbox)
 
   // Enable scroll locking when the listbox is visible, and `modal` is enabled
   let scrollLockEnabled = data.__demoMode
@@ -941,7 +965,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     ? false
     : modal && data.listboxState === ListboxStates.Open
   useInertOthers(inertOthersEnabled, {
-    allowed: useEvent(() => [data.buttonRef.current, data.optionsRef.current]),
+    allowed: useEvent(() => [data.buttonElement, data.optionsElement]),
   })
 
   // We keep track whether the button moved or not, we only check this when the menu state becomes
@@ -954,7 +978,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   // This can be solved by only transitioning the `opacity` instead of everything, but if you _do_
   // want to transition the y-axis for example you will run into the same issue again.
   let didElementMoveEnabled = data.listboxState !== ListboxStates.Open
-  let didButtonMove = useDidElementMove(didElementMoveEnabled, data.buttonRef)
+  let didButtonMove = useDidElementMove(didElementMoveEnabled, data.buttonElement)
 
   // Now that we know that the button did move or not, we can either disable the panel and all of
   // its transitions, or rely on the `visible` state to hide the panel whenever necessary.
@@ -999,18 +1023,18 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
 
   let [floatingRef, style] = useFloatingPanel(anchorOptions)
   let getFloatingPanelProps = useFloatingPanelProps()
-  let optionsRef = useSyncRefs(data.optionsRef, ref, anchor ? floatingRef : null)
+  let optionsRef = useSyncRefs(ref, anchor ? floatingRef : null, actions.setOptionsElement)
 
   let searchDisposables = useDisposables()
 
   useEffect(() => {
-    let container = data.optionsRef.current
+    let container = data.optionsElement
     if (!container) return
     if (data.listboxState !== ListboxStates.Open) return
     if (container === getOwnerDocument(container)?.activeElement) return
 
     container?.focus({ preventScroll: true })
-  }, [data.listboxState, data.optionsRef, data.optionsRef.current])
+  }, [data.listboxState, data.optionsElement])
 
   let handleKeyDown = useEvent((event: ReactKeyboardEvent<HTMLElement>) => {
     searchDisposables.dispose()
@@ -1036,7 +1060,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
         }
         if (data.mode === ValueMode.Single) {
           flushSync(() => actions.closeListbox())
-          data.buttonRef.current?.focus({ preventScroll: true })
+          data.buttonElement?.focus({ preventScroll: true })
         }
         break
 
@@ -1066,7 +1090,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
         event.preventDefault()
         event.stopPropagation()
         flushSync(() => actions.closeListbox())
-        data.buttonRef.current?.focus({ preventScroll: true })
+        data.buttonElement?.focus({ preventScroll: true })
         return
 
       case Keys.Tab:
@@ -1074,7 +1098,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
         event.stopPropagation()
         flushSync(() => actions.closeListbox())
         focusFrom(
-          data.buttonRef.current!,
+          data.buttonElement!,
           event.shiftKey ? FocusManagementFocus.Previous : FocusManagementFocus.Next
         )
         break
@@ -1088,7 +1112,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     }
   })
 
-  let labelledby = useComputed(() => data.buttonRef.current?.id, [data.buttonRef.current])
+  let labelledby = data.buttonElement?.id
   let slot = useMemo(() => {
     return {
       open: data.listboxState === ListboxStates.Open,
@@ -1112,7 +1136,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     style: {
       ...theirProps.style,
       ...style,
-      '--button-width': useElementSize(data.buttonRef, true).width,
+      '--button-width': useElementSize(data.buttonElement, true).width,
     } as CSSProperties,
     ...transitionDataAttributes(transitionData),
   })
@@ -1230,7 +1254,7 @@ function OptionFn<
     actions.onChange(value)
     if (data.mode === ValueMode.Single) {
       flushSync(() => actions.closeListbox())
-      data.buttonRef.current?.focus({ preventScroll: true })
+      data.buttonElement?.focus({ preventScroll: true })
     }
   })
 
diff --git a/packages/@headlessui-react/src/components/menu/menu.tsx b/packages/@headlessui-react/src/components/menu/menu.tsx
index 804e4e1..8da0075 100644
--- a/packages/@headlessui-react/src/components/menu/menu.tsx
+++ b/packages/@headlessui-react/src/components/menu/menu.tsx
@@ -6,7 +6,7 @@ import { useHover } from '@react-aria/interactions'
 import React, {
   Fragment,
   createContext,
-  createRef,
+  useCallback,
   useContext,
   useEffect,
   useMemo,
@@ -94,8 +94,8 @@ type MenuItemDataRef = MutableRefObject<{
 interface StateDefinition {
   __demoMode: boolean
   menuState: MenuStates
-  buttonRef: MutableRefObject<HTMLButtonElement | null>
-  itemsRef: MutableRefObject<HTMLElement | null>
+  buttonElement: HTMLButtonElement | null
+  itemsElement: HTMLElement | null
   items: { id: string; dataRef: MenuItemDataRef }[]
   searchQuery: string
   activeItemIndex: number | null
@@ -111,6 +111,9 @@ enum ActionTypes {
   ClearSearch,
   RegisterItem,
   UnregisterItem,
+
+  SetButtonElement,
+  SetItemsElement,
 }
 
 function adjustOrderedState(
@@ -152,6 +155,8 @@ type Actions =
   | { type: ActionTypes.ClearSearch }
   | { type: ActionTypes.RegisterItem; id: string; dataRef: MenuItemDataRef }
   | { type: ActionTypes.UnregisterItem; id: string }
+  | { type: ActionTypes.SetButtonElement; element: HTMLButtonElement | null }
+  | { type: ActionTypes.SetItemsElement; element: HTMLElement | null }
 
 let reducers: {
   [P in ActionTypes]: (
@@ -334,6 +339,14 @@ let reducers: {
       activationTrigger: ActivationTrigger.Other,
     }
   },
+  [ActionTypes.SetButtonElement]: (state, action) => {
+    if (state.buttonElement === action.element) return state
+    return { ...state, buttonElement: action.element }
+  },
+  [ActionTypes.SetItemsElement]: (state, action) => {
+    if (state.itemsElement === action.element) return state
+    return { ...state, itemsElement: action.element }
+  },
 }
 
 let MenuContext = createContext<[StateDefinition, Dispatch<Actions>] | null>(null)
@@ -379,24 +392,24 @@ function MenuFn<TTag extends ElementType = typeof DEFAULT_MENU_TAG>(
   let reducerBag = useReducer(stateReducer, {
     __demoMode,
     menuState: __demoMode ? MenuStates.Open : MenuStates.Closed,
-    buttonRef: createRef(),
-    itemsRef: createRef(),
+    buttonElement: null,
+    itemsElement: null,
     items: [],
     searchQuery: '',
     activeItemIndex: null,
     activationTrigger: ActivationTrigger.Other,
   } as StateDefinition)
-  let [{ menuState, itemsRef, buttonRef }, dispatch] = reducerBag
+  let [{ menuState, itemsElement, buttonElement }, dispatch] = reducerBag
   let menuRef = useSyncRefs(ref)
 
   // Handle outside click
   let outsideClickEnabled = menuState === MenuStates.Open
-  useOutsideClick(outsideClickEnabled, [buttonRef, itemsRef], (event, target) => {
+  useOutsideClick(outsideClickEnabled, [buttonElement, itemsElement], (event, target) => {
     dispatch({ type: ActionTypes.CloseMenu })
 
     if (!isFocusableElement(target, FocusableMode.Loose)) {
       event.preventDefault()
-      buttonRef.current?.focus()
+      buttonElement?.focus()
     }
   })
 
@@ -469,7 +482,11 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   } = props
   let [state, dispatch] = useMenuContext('Menu.Button')
   let getFloatingReferenceProps = useFloatingReferenceProps()
-  let buttonRef = useSyncRefs(state.buttonRef, ref, useFloatingReference())
+  let buttonRef = useSyncRefs(
+    ref,
+    useFloatingReference(),
+    useEvent((element) => dispatch({ type: ActionTypes.SetButtonElement, element }))
+  )
 
   let handleKeyDown = useEvent((event: ReactKeyboardEvent<HTMLButtonElement>) => {
     switch (event.key) {
@@ -509,7 +526,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     if (disabled) return
     if (state.menuState === MenuStates.Open) {
       flushSync(() => dispatch({ type: ActionTypes.CloseMenu }))
-      state.buttonRef.current?.focus({ preventScroll: true })
+      state.buttonElement?.focus({ preventScroll: true })
     } else {
       event.preventDefault()
       dispatch({ type: ActionTypes.OpenMenu })
@@ -536,9 +553,9 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     {
       ref: buttonRef,
       id,
-      type: useResolveButtonType(props, state.buttonRef),
+      type: useResolveButtonType(props, state.buttonElement),
       'aria-haspopup': 'menu',
-      'aria-controls': state.itemsRef.current?.id,
+      'aria-controls': state.itemsElement?.id,
       'aria-expanded': state.menuState === MenuStates.Open,
       disabled: disabled || undefined,
       autoFocus,
@@ -603,8 +620,12 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   let [state, dispatch] = useMenuContext('Menu.Items')
   let [floatingRef, style] = useFloatingPanel(anchor)
   let getFloatingPanelProps = useFloatingPanelProps()
-  let itemsRef = useSyncRefs(state.itemsRef, ref, anchor ? floatingRef : null)
-  let ownerDocument = useOwnerDocument(state.itemsRef)
+  let itemsRef = useSyncRefs(
+    ref,
+    anchor ? floatingRef : null,
+    useEvent((element) => dispatch({ type: ActionTypes.SetItemsElement, element }))
+  )
+  let ownerDocument = useOwnerDocument(state.itemsElement)
 
   // Always enable `portal` functionality, when `anchor` is enabled
   if (anchor) {
@@ -614,14 +635,14 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    state.itemsRef,
+    state.itemsElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : state.menuState === MenuStates.Open
   )
 
   // Ensure we close the menu as soon as the button becomes hidden
-  useOnDisappear(visible, state.buttonRef, () => {
+  useOnDisappear(visible, state.buttonElement, () => {
     dispatch({ type: ActionTypes.CloseMenu })
   })
 
@@ -632,7 +653,10 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   // Mark other elements as inert when the menu is visible, and `modal` is enabled
   let inertOthersEnabled = state.__demoMode ? false : modal && state.menuState === MenuStates.Open
   useInertOthers(inertOthersEnabled, {
-    allowed: useEvent(() => [state.buttonRef.current, state.itemsRef.current]),
+    allowed: useCallback(
+      () => [state.buttonElement, state.itemsElement],
+      [state.buttonElement, state.itemsElement]
+    ),
   })
 
   // We keep track whether the button moved or not, we only check this when the menu state becomes
@@ -645,23 +669,23 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   // This can be solved by only transitioning the `opacity` instead of everything, but if you _do_
   // want to transition the y-axis for example you will run into the same issue again.
   let didButtonMoveEnabled = state.menuState !== MenuStates.Open
-  let didButtonMove = useDidElementMove(didButtonMoveEnabled, state.buttonRef)
+  let didButtonMove = useDidElementMove(didButtonMoveEnabled, state.buttonElement)
 
   // Now that we know that the button did move or not, we can either disable the panel and all of
   // its transitions, or rely on the `visible` state to hide the panel whenever necessary.
   let panelEnabled = didButtonMove ? false : visible
 
   useEffect(() => {
-    let container = state.itemsRef.current
+    let container = state.itemsElement
     if (!container) return
     if (state.menuState !== MenuStates.Open) return
     if (container === ownerDocument?.activeElement) return
 
     container.focus({ preventScroll: true })
-  }, [state.menuState, state.itemsRef, ownerDocument, state.itemsRef.current])
+  }, [state.menuState, state.itemsElement, ownerDocument])
 
   useTreeWalker(state.menuState === MenuStates.Open, {
-    container: state.itemsRef.current,
+    container: state.itemsElement,
     accept(node) {
       if (node.getAttribute('role') === 'menuitem') return NodeFilter.FILTER_REJECT
       if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
@@ -695,7 +719,7 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
           let { dataRef } = state.items[state.activeItemIndex]
           dataRef.current?.domRef.current?.click()
         }
-        restoreFocusIfNecessary(state.buttonRef.current)
+        restoreFocusIfNecessary(state.buttonElement)
         break
 
       case Keys.ArrowDown:
@@ -724,7 +748,7 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
         event.preventDefault()
         event.stopPropagation()
         flushSync(() => dispatch({ type: ActionTypes.CloseMenu }))
-        state.buttonRef.current?.focus({ preventScroll: true })
+        state.buttonElement?.focus({ preventScroll: true })
         break
 
       case Keys.Tab:
@@ -732,7 +756,7 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
         event.stopPropagation()
         flushSync(() => dispatch({ type: ActionTypes.CloseMenu }))
         focusFrom(
-          state.buttonRef.current!,
+          state.buttonElement!,
           event.shiftKey ? FocusManagementFocus.Previous : FocusManagementFocus.Next
         )
         break
@@ -766,7 +790,7 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   let ourProps = mergeProps(anchor ? getFloatingPanelProps() : {}, {
     'aria-activedescendant':
       state.activeItemIndex === null ? undefined : state.items[state.activeItemIndex]?.id,
-    'aria-labelledby': state.buttonRef.current?.id,
+    'aria-labelledby': state.buttonElement?.id,
     id,
     onKeyDown: handleKeyDown,
     onKeyUp: handleKeyUp,
@@ -779,7 +803,7 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
     style: {
       ...theirProps.style,
       ...style,
-      '--button-width': useElementSize(state.buttonRef, true).width,
+      '--button-width': useElementSize(state.buttonElement, true).width,
     } as CSSProperties,
     ...transitionDataAttributes(transitionData),
   })
@@ -879,7 +903,7 @@ function ItemFn<TTag extends ElementType = typeof DEFAULT_ITEM_TAG>(
   let handleClick = useEvent((event: MouseEvent) => {
     if (disabled) return event.preventDefault()
     dispatch({ type: ActionTypes.CloseMenu })
-    restoreFocusIfNecessary(state.buttonRef.current)
+    restoreFocusIfNecessary(state.buttonElement)
   })
 
   let handleFocus = useEvent(() => {
diff --git a/packages/@headlessui-react/src/components/popover/popover.tsx b/packages/@headlessui-react/src/components/popover/popover.tsx
index 034e22c..018368d 100644
--- a/packages/@headlessui-react/src/components/popover/popover.tsx
+++ b/packages/@headlessui-react/src/components/popover/popover.tsx
@@ -537,7 +537,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     useFloatingReference(),
     isWithinPanel
       ? null
-      : (button) => {
+      : useEvent((button) => {
           if (button) {
             state.buttons.current.push(uniqueIdentifier)
           } else {
@@ -552,7 +552,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
           }
 
           button && dispatch({ type: ActionTypes.SetButton, button })
-        }
+        })
   )
   let withinPanelButtonRef = useSyncRefs(internalButtonRef, ref)
   let ownerDocument = useOwnerDocument(internalButtonRef)
@@ -643,7 +643,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     } satisfies ButtonRenderPropArg
   }, [visible, hover, focus, active, disabled, autoFocus])
 
-  let type = useResolveButtonType(props, internalButtonRef)
+  let type = useResolveButtonType(props, state.button)
   let ourProps = isWithinPanel
     ? mergeProps(
         {
@@ -763,13 +763,13 @@ function BackdropFn<TTag extends ElementType = typeof DEFAULT_BACKDROP_TAG>(
     ...theirProps
   } = props
   let [{ popoverState }, dispatch] = usePopoverContext('Popover.Backdrop')
-  let internalBackdropRef = useRef<HTMLElement | null>(null)
-  let backdropRef = useSyncRefs(ref, internalBackdropRef)
+  let [backdropElement, setBackdropElement] = useState<HTMLElement | null>(null)
+  let backdropRef = useSyncRefs(ref, setBackdropElement)
 
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    internalBackdropRef,
+    backdropElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : popoverState === PopoverStates.Open
@@ -865,9 +865,12 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
     portal = true
   }
 
-  let panelRef = useSyncRefs(internalPanelRef, ref, anchor ? floatingRef : null, (panel) => {
-    dispatch({ type: ActionTypes.SetPanel, panel })
-  })
+  let panelRef = useSyncRefs(
+    internalPanelRef,
+    ref,
+    anchor ? floatingRef : null,
+    useEvent((panel) => dispatch({ type: ActionTypes.SetPanel, panel }))
+  )
   let ownerDocument = useOwnerDocument(internalPanelRef)
   let mergeRefs = useMergeRefsFn()
 
@@ -881,7 +884,7 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    internalPanelRef,
+    state.panel,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : state.popoverState === PopoverStates.Open
diff --git a/packages/@headlessui-react/src/components/switch/switch.tsx b/packages/@headlessui-react/src/components/switch/switch.tsx
index 5a17ce0..f0f64bd 100644
--- a/packages/@headlessui-react/src/components/switch/switch.tsx
+++ b/packages/@headlessui-react/src/components/switch/switch.tsx
@@ -156,11 +156,13 @@ function SwitchFn<TTag extends ElementType = typeof DEFAULT_SWITCH_TAG>(
     ...theirProps
   } = props
   let groupContext = useContext(GroupContext)
+  let [switchElement, setSwitchElement] = useState<HTMLButtonElement | null>(null)
   let internalSwitchRef = useRef<HTMLButtonElement | null>(null)
   let switchRef = useSyncRefs(
     internalSwitchRef,
     ref,
-    groupContext === null ? null : groupContext.setSwitch
+    groupContext === null ? null : groupContext.setSwitch,
+    setSwitchElement
   )
 
   let defaultChecked = useDefaultValue(_defaultChecked)
@@ -221,7 +223,7 @@ function SwitchFn<TTag extends ElementType = typeof DEFAULT_SWITCH_TAG>(
       id,
       ref: switchRef,
       role: 'switch',
-      type: useResolveButtonType(props, internalSwitchRef),
+      type: useResolveButtonType(props, switchElement),
       tabIndex: props.tabIndex === -1 ? 0 : props.tabIndex ?? 0,
       'aria-checked': checked,
       'aria-labelledby': labelledBy,
diff --git a/packages/@headlessui-react/src/components/tabs/tabs.tsx b/packages/@headlessui-react/src/components/tabs/tabs.tsx
index e33f309..9969273 100644
--- a/packages/@headlessui-react/src/components/tabs/tabs.tsx
+++ b/packages/@headlessui-react/src/components/tabs/tabs.tsx
@@ -8,6 +8,7 @@ import React, {
   useMemo,
   useReducer,
   useRef,
+  useState,
   type ElementType,
   type MutableRefObject,
   type KeyboardEvent as ReactKeyboardEvent,
@@ -431,8 +432,9 @@ function TabFn<TTag extends ElementType = typeof DEFAULT_TAB_TAG>(
   let actions = useActions('Tab')
   let data = useData('Tab')
 
+  let [tabElement, setTabElement] = useState<HTMLElement | null>(null)
   let internalTabRef = useRef<HTMLElement | null>(null)
-  let tabRef = useSyncRefs(internalTabRef, ref)
+  let tabRef = useSyncRefs(internalTabRef, ref, setTabElement)
 
   useIsoMorphicEffect(() => actions.registerTab(internalTabRef), [actions, internalTabRef])
 
@@ -542,7 +544,7 @@ function TabFn<TTag extends ElementType = typeof DEFAULT_TAB_TAG>(
       onClick: handleSelection,
       id,
       role: 'tab',
-      type: useResolveButtonType(props, internalTabRef),
+      type: useResolveButtonType(props, tabElement),
       'aria-controls': panels[myIndex]?.current?.id,
       'aria-selected': selected,
       tabIndex: selected ? 0 : -1,
diff --git a/packages/@headlessui-react/src/components/transition/transition.tsx b/packages/@headlessui-react/src/components/transition/transition.tsx
index a61a7a0..31aaf07 100644
--- a/packages/@headlessui-react/src/components/transition/transition.tsx
+++ b/packages/@headlessui-react/src/components/transition/transition.tsx
@@ -319,10 +319,13 @@ function TransitionChildFn<TTag extends ElementType = typeof DEFAULT_TRANSITION_
 
     ...theirProps
   } = props as typeof props
+  let [containerElement, setContainerElement] = useState<HTMLElement | null>(null)
   let container = useRef<HTMLElement | null>(null)
   let requiresRef = shouldForwardRef(props)
 
-  let transitionRef = useSyncRefs(...(requiresRef ? [container, ref] : ref === null ? [] : [ref]))
+  let transitionRef = useSyncRefs(
+    ...(requiresRef ? [container, ref, setContainerElement] : ref === null ? [] : [ref])
+  )
   let strategy = theirProps.unmount ?? true ? RenderStrategy.Unmount : RenderStrategy.Hidden
 
   let { show, appear, initial } = useTransitionContext()
@@ -435,7 +438,7 @@ function TransitionChildFn<TTag extends ElementType = typeof DEFAULT_TRANSITION_
   // a leave transition on the `<Transition>` is done, but there is still a
   // child `<TransitionChild>` busy, then `visible` would be `false`, while
   // `state` would still be `TreeStates.Visible`.
-  let [, transitionData] = useTransition(enabled, container, show, { start, end })
+  let [, transitionData] = useTransition(enabled, containerElement, show, { start, end })
 
   let ourProps = compact({
     ref: transitionRef,
diff --git a/packages/@headlessui-react/src/hooks/use-did-element-move.ts b/packages/@headlessui-react/src/hooks/use-did-element-move.ts
index 9c9266f..aa32eca 100644
--- a/packages/@headlessui-react/src/hooks/use-did-element-move.ts
+++ b/packages/@headlessui-react/src/hooks/use-did-element-move.ts
@@ -1,21 +1,21 @@
-import { useRef, type MutableRefObject } from 'react'
+import { useRef } from 'react'
 import { useIsoMorphicEffect } from './use-iso-morphic-effect'
 
-export function useDidElementMove(enabled: boolean, element: MutableRefObject<HTMLElement | null>) {
+export function useDidElementMove(enabled: boolean, element: HTMLElement | null) {
   let elementPosition = useRef({ left: 0, top: 0 })
+
   useIsoMorphicEffect(() => {
-    let el = element.current
-    if (!el) return
+    if (!element) return
 
-    let DOMRect = el.getBoundingClientRect()
+    let DOMRect = element.getBoundingClientRect()
     if (DOMRect) elementPosition.current = DOMRect
-  }, [enabled])
+  }, [enabled, element])
 
-  if (element.current == null) return false
+  if (element == null) return false
   if (!enabled) return false
-  if (element.current === document.activeElement) return false
+  if (element === document.activeElement) return false
 
-  let buttonRect = element.current.getBoundingClientRect()
+  let buttonRect = element.getBoundingClientRect()
 
   let didElementMove =
     buttonRect.top !== elementPosition.current.top ||
diff --git a/packages/@headlessui-react/src/hooks/use-element-size.ts b/packages/@headlessui-react/src/hooks/use-element-size.ts
index 1928fae..1a22bcf 100644
--- a/packages/@headlessui-react/src/hooks/use-element-size.ts
+++ b/packages/@headlessui-react/src/hooks/use-element-size.ts
@@ -7,11 +7,7 @@ function computeSize(element: HTMLElement | null) {
   return { width, height }
 }
 
-export function useElementSize(
-  ref: React.MutableRefObject<HTMLElement | null> | HTMLElement | null,
-  unit = false
-) {
-  let element = ref === null ? null : 'current' in ref ? ref.current : ref
+export function useElementSize(element: HTMLElement | null, unit = false) {
   let [identity, forceRerender] = useReducer(() => ({}), {})
 
   // When the element changes during a re-render, we want to make sure we
diff --git a/packages/@headlessui-react/src/hooks/use-outside-click.ts b/packages/@headlessui-react/src/hooks/use-outside-click.ts
index 52b3062..43d47d0 100644
--- a/packages/@headlessui-react/src/hooks/use-outside-click.ts
+++ b/packages/@headlessui-react/src/hooks/use-outside-click.ts
@@ -1,4 +1,4 @@
-import { useCallback, useRef, type MutableRefObject } from 'react'
+import { useCallback, useRef } from 'react'
 import { FocusableMode, isFocusableElement } from '../utils/focus-management'
 import { isMobile } from '../utils/platform'
 import { useDocumentEvent } from './use-document-event'
@@ -6,7 +6,7 @@ import { useIsTopLayer } from './use-is-top-layer'
 import { useLatestValue } from './use-latest-value'
 import { useWindowEvent } from './use-window-event'
 
-type Container = MutableRefObject<HTMLElement | null> | HTMLElement | null
+type Container = HTMLElement | null
 type ContainerCollection = Container[] | Set<Container>
 type ContainerInput = Container | ContainerCollection
 
@@ -69,15 +69,14 @@ export function useOutsideClick(
       // Ignore if the target exists in one of the containers
       for (let container of _containers) {
         if (container === null) continue
-        let domNode = container instanceof HTMLElement ? container : container.current
-        if (domNode?.contains(target)) {
+        if (container.contains(target)) {
           return
         }
 
         // If the click crossed a shadow boundary, we need to check if the
         // container is inside the tree by using `composedPath` to "pierce" the
         // shadow boundary
-        if (event.composed && event.composedPath().includes(domNode as EventTarget)) {
+        if (event.composed && event.composedPath().includes(container as EventTarget)) {
           return
         }
       }
@@ -101,7 +100,7 @@ export function useOutsideClick(
 
       return cbRef.current(event, target)
     },
-    [cbRef]
+    [cbRef, containers]
   )
 
   let initialClickTarget = useRef<EventTarget | null>(null)
diff --git a/packages/@headlessui-react/src/hooks/use-refocusable-input.ts b/packages/@headlessui-react/src/hooks/use-refocusable-input.ts
index 2ed48df..e8e361b 100644
--- a/packages/@headlessui-react/src/hooks/use-refocusable-input.ts
+++ b/packages/@headlessui-react/src/hooks/use-refocusable-input.ts
@@ -1,4 +1,4 @@
-import { useRef, type MutableRefObject } from 'react'
+import { useRef } from 'react'
 import { useEvent } from './use-event'
 import { useEventListener } from './use-event-listener'
 
@@ -8,7 +8,7 @@ import { useEventListener } from './use-event-listener'
  * This hook will also keep the cursor position into account to make sure the
  * cursor is placed at the correct position as-if we didn't loose focus at all.
  */
-export function useRefocusableInput(ref: MutableRefObject<HTMLInputElement | null>) {
+export function useRefocusableInput(input: HTMLInputElement | null) {
   // Track the cursor position and the value of the input
   let info = useRef({
     value: '',
@@ -16,7 +16,7 @@ export function useRefocusableInput(ref: MutableRefObject<HTMLInputElement | nul
     selectionEnd: null as number | null,
   })
 
-  useEventListener(ref.current, 'blur', (event) => {
+  useEventListener(input, 'blur', (event) => {
     let target = event.target
     if (!(target instanceof HTMLInputElement)) return
 
@@ -28,8 +28,6 @@ export function useRefocusableInput(ref: MutableRefObject<HTMLInputElement | nul
   })
 
   return useEvent(() => {
-    let input = ref.current
-
     // If the input is already focused, we don't need to do anything
     if (document.activeElement === input) return
 
diff --git a/packages/@headlessui-react/src/hooks/use-resolve-button-type.ts b/packages/@headlessui-react/src/hooks/use-resolve-button-type.ts
index 163e5da..462dd55 100644
--- a/packages/@headlessui-react/src/hooks/use-resolve-button-type.ts
+++ b/packages/@headlessui-react/src/hooks/use-resolve-button-type.ts
@@ -1,33 +1,21 @@
-import { useState, type MutableRefObject } from 'react'
-import { useIsoMorphicEffect } from './use-iso-morphic-effect'
-
-function resolveType<TTag>(props: { type?: string; as?: TTag }) {
-  if (props.type) return props.type
-
-  let tag = props.as ?? 'button'
-  if (typeof tag === 'string' && tag.toLowerCase() === 'button') return 'button'
-
-  return undefined
-}
+import { useMemo } from 'react'
 
 export function useResolveButtonType<TTag>(
   props: { type?: string; as?: TTag },
-  ref: MutableRefObject<HTMLElement | null>
+  element: HTMLElement | null
 ) {
-  let [type, setType] = useState(() => resolveType(props))
-
-  useIsoMorphicEffect(() => {
-    setType(resolveType(props))
-  }, [props.type, props.as])
+  return useMemo(() => {
+    // A type was provided
+    if (props.type) return props.type
 
-  useIsoMorphicEffect(() => {
-    if (type) return
-    if (!ref.current) return
+    // Resolve the type based on the `as` prop
+    let tag = props.as ?? 'button'
+    if (typeof tag === 'string' && tag.toLowerCase() === 'button') return 'button'
 
-    if (ref.current instanceof HTMLButtonElement && !ref.current.hasAttribute('type')) {
-      setType('button')
-    }
-  }, [type, ref])
+    // Resolve the type based on the HTML element
+    if (element?.tagName === 'BUTTON' && !element.hasAttribute('type')) return 'button'
 
-  return type
+    // Could not resolve the type
+    return undefined
+  }, [props.type, props.as, element])
 }
diff --git a/packages/@headlessui-react/src/hooks/use-transition.ts b/packages/@headlessui-react/src/hooks/use-transition.ts
index 42b3f57..ae2602e 100644
--- a/packages/@headlessui-react/src/hooks/use-transition.ts
+++ b/packages/@headlessui-react/src/hooks/use-transition.ts
@@ -51,7 +51,7 @@ export function transitionDataAttributes(data: TransitionData) {
 
 export function useTransition(
   enabled: boolean,
-  elementRef: MutableRefObject<HTMLElement | null>,
+  element: HTMLElement | null,
   show: boolean,
   events?: {
     start?(show: boolean): void
@@ -68,98 +68,92 @@ export function useTransition(
 
   let d = useDisposables()
 
-  useIsoMorphicEffect(
-    function retry() {
-      if (!enabled) return
+  useIsoMorphicEffect(() => {
+    if (!enabled) return
 
+    if (show) {
+      setVisible(true)
+    }
+
+    if (!element) {
       if (show) {
-        setVisible(true)
+        addFlag(TransitionState.Enter | TransitionState.Closed)
       }
+      return
+    }
 
-      let node = elementRef.current
-      if (!node) {
-        // Retry if the DOM node isn't available yet
-        if (show) {
-          addFlag(TransitionState.Enter | TransitionState.Closed)
-          return d.nextFrame(() => retry())
+    events?.start?.(show)
+
+    return transition(element, {
+      inFlight,
+      prepare() {
+        if (cancelledRef.current) {
+          // Cancelled a cancellation, we're back to the original state.
+          cancelledRef.current = false
+        } else {
+          // If we were already in-flight, then we want to cancel the current
+          // transition.
+          cancelledRef.current = inFlight.current
         }
-        return
-      }
 
-      events?.start?.(show)
+        inFlight.current = true
 
-      return transition(node, {
-        inFlight,
-        prepare() {
-          if (cancelledRef.current) {
-            // Cancelled a cancellation, we're back to the original state.
-            cancelledRef.current = false
-          } else {
-            // If we were already in-flight, then we want to cancel the current
-            // transition.
-            cancelledRef.current = inFlight.current
-          }
-
-          inFlight.current = true
-
-          if (cancelledRef.current) return
+        if (cancelledRef.current) return
 
+        if (show) {
+          addFlag(TransitionState.Enter | TransitionState.Closed)
+          removeFlag(TransitionState.Leave)
+        } else {
+          addFlag(TransitionState.Leave)
+          removeFlag(TransitionState.Enter)
+        }
+      },
+      run() {
+        if (cancelledRef.current) {
+          // If we cancelled a transition, then the `show` state is going to
+          // be inverted already, but that doesn't mean we have to go to that
+          // new state.
+          //
+          // What we actually want is to revert to the "idle" state (the
+          // stable state where an `Enter` transitions to, and a `Leave`
+          // transitions from.)
+          //
+          // Because of this, it might look like we are swapping the flags in
+          // the following branches, but that's not the case.
           if (show) {
-            addFlag(TransitionState.Enter | TransitionState.Closed)
-            removeFlag(TransitionState.Leave)
-          } else {
+            removeFlag(TransitionState.Enter | TransitionState.Closed)
             addFlag(TransitionState.Leave)
-            removeFlag(TransitionState.Enter)
+          } else {
+            removeFlag(TransitionState.Leave)
+            addFlag(TransitionState.Enter | TransitionState.Closed)
           }
-        },
-        run() {
-          if (cancelledRef.current) {
-            // If we cancelled a transition, then the `show` state is going to
-            // be inverted already, but that doesn't mean we have to go to that
-            // new state.
-            //
-            // What we actually want is to revert to the "idle" state (the
-            // stable state where an `Enter` transitions to, and a `Leave`
-            // transitions from.)
-            //
-            // Because of this, it might look like we are swapping the flags in
-            // the following branches, but that's not the case.
-            if (show) {
-              removeFlag(TransitionState.Enter | TransitionState.Closed)
-              addFlag(TransitionState.Leave)
-            } else {
-              removeFlag(TransitionState.Leave)
-              addFlag(TransitionState.Enter | TransitionState.Closed)
-            }
+        } else {
+          if (show) {
+            removeFlag(TransitionState.Closed)
           } else {
-            if (show) {
-              removeFlag(TransitionState.Closed)
-            } else {
-              addFlag(TransitionState.Closed)
-            }
+            addFlag(TransitionState.Closed)
           }
-        },
-        done() {
-          if (cancelledRef.current) {
-            if (typeof node.getAnimations === 'function' && node.getAnimations().length > 0) {
-              return
-            }
+        }
+      },
+      done() {
+        if (cancelledRef.current) {
+          if (typeof element.getAnimations === 'function' && element.getAnimations().length > 0) {
+            return
           }
+        }
 
-          inFlight.current = false
+        inFlight.current = false
 
-          removeFlag(TransitionState.Enter | TransitionState.Leave | TransitionState.Closed)
+        removeFlag(TransitionState.Enter | TransitionState.Leave | TransitionState.Closed)
 
-          if (!show) {
-            setVisible(false)
-          }
+        if (!show) {
+          setVisible(false)
+        }
 
-          events?.end?.(show)
-        },
-      })
-    },
-    [enabled, show, elementRef, d]
-  )
+        events?.end?.(show)
+      },
+    })
+  }, [enabled, show, element, d])
 
   if (!enabled) {
     return [


Test patch:
diff --git a/packages/@headlessui-react/src/components/transition/__snapshots__/transition.test.tsx.snap b/packages/@headlessui-react/src/components/transition/__snapshots__/transition.test.tsx.snap
index 7d2fa9f..e2aa844 100644
--- a/packages/@headlessui-react/src/components/transition/__snapshots__/transition.test.tsx.snap
+++ b/packages/@headlessui-react/src/components/transition/__snapshots__/transition.test.tsx.snap
@@ -133,7 +133,6 @@ exports[`Setup API transition classes should be possible to passthrough the tran
   data-closed=""
   data-enter=""
   data-transition=""
-  style=""
 >
   Children
 </div>