Fix components not properly closing when using the transition prop (#3448)

This PR fixes a bug where the components don't always properly close
when using the `transition` prop on those components.

The issue here is that the internal `useTransition(…)` hook relies on a
DOM node. Whenever the DOM node changes, we need to re-run the
`useTransition(…)`. This is why we store the DOM element in state
instead of relying on a `useRef(…)`.

Let's say you have a `Popover` component, then the structure looks like
this:
```ts
<Popover>
  <PopoverButton>Show</PopoverButton>
  <PopoverPanel>Contents</PopoverPanel>
</Popover>
```

We store a DOM reference to the button and the panel in state, and the
state lives in the `Popover` component. The reason we do that is so that
the button can reference the panel and the panel can reference the
button. This is needed for some `aria-*` attributes for example:
```ts
<PopoverButton aria-controls={panelElement.id}>
```

For the transitions, we set some state to make sure that the panel is
visible or hidden, then we wait for transitions to finish by listening
to transition related events on the DOM node directly.

If you now say, "hey panel, please re-render because you have to become
visible/hidden" then the component re-renders, the panel DOM node
(stored in the `Popover` component) eventually updates and then the
`useTransition(…)` hooks receives the new value (either the DOM node or
null when the leave transition is complete).

The problem here is the round trip that it first has to go to the root
`<Popover/>` component, re-render everything and provide the new DOM
node to the `useTransition(…)` hook.

The solution? Local state so that the panel can re-render on its own and
doesn't require the round trip via the parent.

Fixes: #3438
Fixes: #3437
Fixes: tailwindlabs/tailwind-plus-issues#1625

---------

Co-authored-by: Jonathan Reinink <jonathan@reinink.ca>

Author: RobinMalfait

packages/@headlessui-react/CHANGELOG.md

@@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-- Nothing yet!
+### Fixed
+
+- Fix components not closing properly when using the `transition` prop ([#3448](https://github.com/tailwindlabs/headlessui/pull/3448))
 
 ## [2.1.3] - 2024-08-23
 

packages/@headlessui-react/src/components/combobox/combobox.test.tsx

@@ -1,4 +1,4 @@
-import { render } from '@testing-library/react'
+import { render, waitFor } from '@testing-library/react'
 import React, { Fragment, createElement, useEffect, useState } from 'react'
 import {
   ComboboxMode,
@@ -42,7 +42,13 @@ import {
 } from '../../test-utils/interactions'
 import { mockingConsoleLogs, suppressConsoleLogs } from '../../test-utils/suppress-console-logs'
 import { Transition } from '../transition/transition'
-import { Combobox } from './combobox'
+import {
+  Combobox,
+  ComboboxButton,
+  ComboboxInput,
+  ComboboxOption,
+  ComboboxOptions,
+} from './combobox'
 
 let NOOP = () => {}
 
@@ -6060,3 +6066,39 @@ describe('Form compatibility', () => {
     ])
   })
 })
+
+describe('transitions', () => {
+  it(
+    'should be possible to close the Combobox when using the `transition` prop',
+    suppressConsoleLogs(async () => {
+      render(
+        <Combobox>
+          <ComboboxButton>Toggle</ComboboxButton>
+          <ComboboxInput />
+          <ComboboxOptions transition>
+            <ComboboxOption value="alice">Alice</ComboboxOption>
+            <ComboboxOption value="bob">Bob</ComboboxOption>
+            <ComboboxOption value="charlie">Charlie</ComboboxOption>
+          </ComboboxOptions>
+        </Combobox>
+      )
+
+      // Open the combobox
+      await click(getComboboxButton())
+
+      // Ensure the combobox is visible
+      assertCombobox({ state: ComboboxState.Visible })
+
+      // Close the combobox
+      await click(getComboboxButton())
+
+      // Wait for the transition to finish, and the combobox to close
+      await waitFor(() => {
+        assertComboboxList({ state: ComboboxState.InvisibleUnmounted })
+      })
+
+      // Ensure the input got the restored focus
+      assertActiveElement(getComboboxInput())
+    })
+  )
+})

packages/@headlessui-react/src/components/combobox/combobox.tsx

@@ -1672,14 +1672,26 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   }
 
   let [floatingRef, style] = useFloatingPanel(anchor)
+
+  // To improve the correctness of transitions (timing related race conditions),
+  // we track the element locally to this component, instead of relying on the
+  // context value. This way, the component can re-render independently of the
+  // parent component when the `useTransition(…)` hook performs a state change.
+  let [localOptionsElement, setLocalOptionsElement] = useState<HTMLElement | null>(null)
+
   let getFloatingPanelProps = useFloatingPanelProps()
-  let optionsRef = useSyncRefs(ref, anchor ? floatingRef : null, actions.setOptionsElement)
+  let optionsRef = useSyncRefs(
+    ref,
+    anchor ? floatingRef : null,
+    actions.setOptionsElement,
+    setLocalOptionsElement
+  )
   let ownerDocument = useOwnerDocument(data.optionsElement)
 
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    data.optionsElement,
+    localOptionsElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : data.comboboxState === ComboboxState.Open

packages/@headlessui-react/src/components/disclosure/disclosure.test.tsx

@@ -1,18 +1,18 @@
-import { render } from '@testing-library/react'
-import React, { createElement, Suspense, useEffect, useRef } from 'react'
+import { render, waitFor } from '@testing-library/react'
+import React, { Suspense, createElement, useEffect, useRef } from 'react'
 import {
+  DisclosureState,
   assertActiveElement,
   assertDisclosureButton,
   assertDisclosurePanel,
-  DisclosureState,
   getByText,
   getDisclosureButton,
   getDisclosurePanel,
 } from '../../test-utils/accessibility-assertions'
-import { click, focus, Keys, MouseButton, press } from '../../test-utils/interactions'
+import { Keys, MouseButton, click, focus, press } from '../../test-utils/interactions'
 import { suppressConsoleLogs } from '../../test-utils/suppress-console-logs'
 import { Transition } from '../transition/transition'
-import { Disclosure } from './disclosure'
+import { Disclosure, DisclosureButton, DisclosurePanel } from './disclosure'
 
 jest.mock('../../hooks/use-id')
 
@@ -985,3 +985,40 @@ describe('Mouse interactions', () => {
     })
   )
 })
+
+describe('transitions', () => {
+  it(
+    'should be possible to close the Disclosure when using the `transition` prop',
+    suppressConsoleLogs(async () => {
+      render(
+        <Disclosure>
+          <DisclosureButton>Toggle</DisclosureButton>
+          <DisclosurePanel transition>Contents</DisclosurePanel>
+        </Disclosure>
+      )
+
+      // Focus the button
+      await focus(getDisclosureButton())
+
+      // Ensure the button is focused
+      assertActiveElement(getDisclosureButton())
+
+      // Open the disclosure
+      await click(getDisclosureButton())
+
+      // Ensure the disclosure is visible
+      assertDisclosurePanel({ state: DisclosureState.Visible })
+
+      // Close the disclosure
+      await click(getDisclosureButton())
+
+      // Wait for the transition to finish, and the disclosure to close
+      await waitFor(() => {
+        assertDisclosurePanel({ state: DisclosureState.InvisibleUnmounted })
+      })
+
+      // Ensure the button got the restored focus
+      assertActiveElement(getDisclosureButton())
+    })
+  )
+})

packages/@headlessui-react/src/components/disclosure/disclosure.tsx

@@ -11,6 +11,7 @@ import React, {
   useMemo,
   useReducer,
   useRef,
+  useState,
   type ContextType,
   type Dispatch,
   type ElementType,
@@ -451,11 +452,18 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   let { close } = useDisclosureAPIContext('Disclosure.Panel')
   let mergeRefs = useMergeRefsFn()
 
+  // To improve the correctness of transitions (timing related race conditions),
+  // we track the element locally to this component, instead of relying on the
+  // context value. This way, the component can re-render independently of the
+  // parent component when the `useTransition(…)` hook performs a state change.
+  let [localPanelElement, setLocalPanelElement] = useState<HTMLElement | null>(null)
+
   let panelRef = useSyncRefs(
     ref,
     useEvent((element) => {
       startTransition(() => dispatch({ type: ActionTypes.SetPanelElement, element }))
-    })
+    }),
+    setLocalPanelElement
   )
 
   useEffect(() => {
@@ -468,7 +476,7 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    state.panelElement,
+    localPanelElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : state.disclosureState === DisclosureStates.Open

packages/@headlessui-react/src/components/listbox/listbox.test.tsx

@@ -1,4 +1,4 @@
-import { render } from '@testing-library/react'
+import { render, waitFor } from '@testing-library/react'
 import React, { createElement, useEffect, useState } from 'react'
 import {
   ListboxMode,
@@ -35,7 +35,7 @@ import {
 } from '../../test-utils/interactions'
 import { suppressConsoleLogs } from '../../test-utils/suppress-console-logs'
 import { Transition } from '../transition/transition'
-import { Listbox } from './listbox'
+import { Listbox, ListboxButton, ListboxOption, ListboxOptions } from './listbox'
 
 jest.mock('../../hooks/use-id')
 
@@ -4811,3 +4811,44 @@ describe('Form compatibility', () => {
     ])
   })
 })
+
+describe('transitions', () => {
+  it(
+    'should be possible to close the Listbox when using the `transition` prop',
+    suppressConsoleLogs(async () => {
+      render(
+        <Listbox>
+          <ListboxButton>Toggle</ListboxButton>
+          <ListboxOptions transition>
+            <ListboxOption value="alice">Alice</ListboxOption>
+            <ListboxOption value="bob">Bob</ListboxOption>
+            <ListboxOption value="charlie">Charlie</ListboxOption>
+          </ListboxOptions>
+        </Listbox>
+      )
+
+      // Focus the button
+      await focus(getListboxButton())
+
+      // Ensure the button is focused
+      assertActiveElement(getListboxButton())
+
+      // Open the listbox
+      await click(getListboxButton())
+
+      // Ensure the listbox is visible
+      assertListbox({ state: ListboxState.Visible })
+
+      // Close the listbox
+      await click(getListboxButton())
+
+      // Wait for the transition to finish, and the listbox to close
+      await waitFor(() => {
+        assertListbox({ state: ListboxState.InvisibleUnmounted })
+      })
+
+      // Ensure the button got the restored focus
+      assertActiveElement(getListboxButton())
+    })
+  )
+})

packages/@headlessui-react/src/components/listbox/listbox.tsx

@@ -12,6 +12,7 @@ import React, {
   useMemo,
   useReducer,
   useRef,
+  useState,
   type CSSProperties,
   type ElementType,
   type MutableRefObject,
@@ -932,6 +933,12 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   } = props
   let anchor = useResolvedAnchor(rawAnchor)
 
+  // To improve the correctness of transitions (timing related race conditions),
+  // we track the element locally to this component, instead of relying on the
+  // context value. This way, the component can re-render independently of the
+  // parent component when the `useTransition(…)` hook performs a state change.
+  let [localOptionsElement, setLocalOptionsElement] = useState<HTMLElement | null>(null)
+
   // Always enable `portal` functionality, when `anchor` is enabled
   if (anchor) {
     portal = true
@@ -945,7 +952,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    data.optionsElement,
+    localOptionsElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : data.listboxState === ListboxStates.Open
@@ -1023,7 +1030,12 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
 
   let [floatingRef, style] = useFloatingPanel(anchorOptions)
   let getFloatingPanelProps = useFloatingPanelProps()
-  let optionsRef = useSyncRefs(ref, anchor ? floatingRef : null, actions.setOptionsElement)
+  let optionsRef = useSyncRefs(
+    ref,
+    anchor ? floatingRef : null,
+    actions.setOptionsElement,
+    setLocalOptionsElement
+  )
 
   let searchDisposables = useDisposables()
 

packages/@headlessui-react/src/components/menu/menu.test.tsx

@@ -1,4 +1,4 @@
-import { render } from '@testing-library/react'
+import { render, waitFor } from '@testing-library/react'
 import React, { createElement, useEffect } from 'react'
 import {
   MenuState,
@@ -31,7 +31,7 @@ import {
 } from '../../test-utils/interactions'
 import { suppressConsoleLogs } from '../../test-utils/suppress-console-logs'
 import { Transition } from '../transition/transition'
-import { Menu } from './menu'
+import { Menu, MenuButton, MenuItem, MenuItems } from './menu'
 
 jest.mock('../../hooks/use-id')
 
@@ -3531,3 +3531,44 @@ describe('Mouse interactions', () => {
     })
   )
 })
+
+describe('transitions', () => {
+  it(
+    'should be possible to close the Menu when using the `transition` prop',
+    suppressConsoleLogs(async () => {
+      render(
+        <Menu>
+          <MenuButton>Toggle</MenuButton>
+          <MenuItems transition>
+            <MenuItem as="a">Alice</MenuItem>
+            <MenuItem as="a">Bob</MenuItem>
+            <MenuItem as="a">Charlie</MenuItem>
+          </MenuItems>
+        </Menu>
+      )
+
+      // Focus the button
+      await focus(getMenuButton())
+
+      // Ensure the button is focused
+      assertActiveElement(getMenuButton())
+
+      // Open the menu
+      await click(getMenuButton())
+
+      // Ensure the menu is visible
+      assertMenu({ state: MenuState.Visible })
+
+      // Close the menu
+      await click(getMenuButton())
+
+      // Wait for the transition to finish, and the menu to close
+      await waitFor(() => {
+        assertMenu({ state: MenuState.InvisibleUnmounted })
+      })
+
+      // Ensure the button got the restored focus
+      assertActiveElement(getMenuButton())
+    })
+  )
+})