Repository: tailwindlabs/headlessui
Pull request #3448 (2024-09-03 10:53:58)

Fix components not properly closing when using the `transition` prop
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: https://github.com/tailwindlabs/headlessui/issues/3438
Fixes: https://github.com/tailwindlabs/headlessui/issues/3437
Fixes: https://github.com/tailwindlabs/tailwindui-issues/issues/1625



vercel[bot]: [vc]: #fDmvtAPFvxGvThj9QNku1zQxvripjcJ2BacbIicefxY=:eyJpc01vbm9yZXBvIjp0cnVlLCJ0eXBlIjoiZ2l0aHViIiwicHJvamVjdHMiOlt7Im5hbWUiOiJoZWFkbGVzc3VpLXJlYWN0Iiwicm9vdERpcmVjdG9yeSI6InBsYXlncm91bmRzL3JlYWN0IiwiaW5zcGVjdG9yVXJsIjoiaHR0cHM6Ly92ZXJjZWwuY29tL3RhaWx3aW5kbGFicy9oZWFkbGVzc3VpLXJlYWN0L0RIeWVKYkhvZFp4TFpNQWpvcVNBeEJZS3dySHQiLCJwcmV2aWV3VXJsIjoiaGVhZGxlc3N1aS1yZWFjdC1naXQtZml4LWlzc3VlLTM0MzctdGFpbHdpbmRsYWJzLnZlcmNlbC5hcHAiLCJuZXh0Q29tbWl0U3RhdHVzIjoiUEVORElORyIsImxpdmVGZWVkYmFjayI6eyJyZXNvbHZlZCI6MCwidW5yZXNvbHZlZCI6MCwidG90YWwiOjAsImxpbmsiOiJoZWFkbGVzc3VpLXJlYWN0LWdpdC1maXgtaXNzdWUtMzQzNy10YWlsd2luZGxhYnMudmVyY2VsLmFwcCJ9fSx7Im5hbWUiOiJoZWFkbGVzc3VpLXZ1ZSIsInJvb3REaXJlY3RvcnkiOiJwbGF5Z3JvdW5kcy92dWUiLCJpbnNwZWN0b3JVcmwiOiJodHRwczovL3ZlcmNlbC5jb20vdGFpbHdpbmRsYWJzL2hlYWRsZXNzdWktdnVlL0FSbzhQYXhQdTZaa1R2eENCUDN3aFM2NHE0Tm4iLCJwcmV2aWV3VXJsIjoiaGVhZGxlc3N1aS12dWUtZ2l0LWZpeC1pc3N1ZS0zNDM3LXRhaWx3aW5kbGFicy52ZXJjZWwuYXBwIiwibmV4dENvbW1pdFN0YXR1cyI6IlBFTkRJTkciLCJsaXZlRmVlZGJhY2siOnsicmVzb2x2ZWQiOjAsInVucmVzb2x2ZWQiOjAsInRvdGFsIjowLCJsaW5rIjoiaGVhZGxlc3N1aS12dWUtZ2l0LWZpeC1pc3N1ZS0zNDM3LXRhaWx3aW5kbGFicy52ZXJjZWwuYXBwIn19XX0=
**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-react** | 🔄 Building ([Inspect](https://vercel.com/tailwindlabs/headlessui-react/DHyeJbHodZxLZMAjoqSAxBYKwrHt)) | [Visit Preview](https://vercel.live/open-feedback/headlessui-react-git-fix-issue-3437-tailwindlabs.vercel.app?via=pr-comment-visit-preview-link&passThrough=1) | 💬 [**Add feedback**](https://vercel.live/open-feedback/headlessui-react-git-fix-issue-3437-tailwindlabs.vercel.app?via=pr-comment-feedback-link) | Sep 3, 2024 10:54am |
| **headlessui-vue** | 🔄 Building ([Inspect](https://vercel.com/tailwindlabs/headlessui-vue/ARo8PaxPu6ZkTvxCBP3whS64q4Nn)) | [Visit Preview](https://vercel.live/open-feedback/headlessui-vue-git-fix-issue-3437-tailwindlabs.vercel.app?via=pr-comment-visit-preview-link&passThrough=1) | 💬 [**Add feedback**](https://vercel.live/open-feedback/headlessui-vue-git-fix-issue-3437-tailwindlabs.vercel.app?via=pr-comment-feedback-link) | Sep 3, 2024 10:54am |




Patch:
diff --git a/packages/@headlessui-react/CHANGELOG.md b/packages/@headlessui-react/CHANGELOG.md
index 89fb6bc..a8af42d 100644
--- a/packages/@headlessui-react/CHANGELOG.md
+++ b/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
 
diff --git a/packages/@headlessui-react/src/components/combobox/combobox.tsx b/packages/@headlessui-react/src/components/combobox/combobox.tsx
index 131284a..8e582a0 100644
--- a/packages/@headlessui-react/src/components/combobox/combobox.tsx
+++ b/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
diff --git a/packages/@headlessui-react/src/components/disclosure/disclosure.tsx b/packages/@headlessui-react/src/components/disclosure/disclosure.tsx
index 4a20bb2..2f3e522 100644
--- a/packages/@headlessui-react/src/components/disclosure/disclosure.tsx
+++ b/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
diff --git a/packages/@headlessui-react/src/components/listbox/listbox.tsx b/packages/@headlessui-react/src/components/listbox/listbox.tsx
index e8c40e4..e82cad3 100644
--- a/packages/@headlessui-react/src/components/listbox/listbox.tsx
+++ b/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()
 
diff --git a/packages/@headlessui-react/src/components/menu/menu.tsx b/packages/@headlessui-react/src/components/menu/menu.tsx
index 8da0075..88b0074 100644
--- a/packages/@headlessui-react/src/components/menu/menu.tsx
+++ b/packages/@headlessui-react/src/components/menu/menu.tsx
@@ -12,6 +12,7 @@ import React, {
   useMemo,
   useReducer,
   useRef,
+  useState,
   type CSSProperties,
   type Dispatch,
   type ElementType,
@@ -620,10 +621,18 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   let [state, dispatch] = useMenuContext('Menu.Items')
   let [floatingRef, style] = useFloatingPanel(anchor)
   let getFloatingPanelProps = useFloatingPanelProps()
+
+  // 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 [localItemsElement, setLocalItemsElement] = useState<HTMLElement | null>(null)
+
   let itemsRef = useSyncRefs(
     ref,
     anchor ? floatingRef : null,
-    useEvent((element) => dispatch({ type: ActionTypes.SetItemsElement, element }))
+    useEvent((element) => dispatch({ type: ActionTypes.SetItemsElement, element })),
+    setLocalItemsElement
   )
   let ownerDocument = useOwnerDocument(state.itemsElement)
 
@@ -635,7 +644,7 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    state.itemsElement,
+    localItemsElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : state.menuState === MenuStates.Open
diff --git a/packages/@headlessui-react/src/components/popover/popover.tsx b/packages/@headlessui-react/src/components/popover/popover.tsx
index 018368d..b1efa7e 100644
--- a/packages/@headlessui-react/src/components/popover/popover.tsx
+++ b/packages/@headlessui-react/src/components/popover/popover.tsx
@@ -763,13 +763,19 @@ function BackdropFn<TTag extends ElementType = typeof DEFAULT_BACKDROP_TAG>(
     ...theirProps
   } = props
   let [{ popoverState }, dispatch] = usePopoverContext('Popover.Backdrop')
-  let [backdropElement, setBackdropElement] = useState<HTMLElement | null>(null)
-  let backdropRef = useSyncRefs(ref, setBackdropElement)
+
+  // 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 [localBackdropElement, setLocalBackdropElement] = useState<HTMLElement | null>(null)
+
+  let backdropRef = useSyncRefs(ref, setLocalBackdropElement)
 
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    backdropElement,
+    localBackdropElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : popoverState === PopoverStates.Open
@@ -865,11 +871,18 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
     portal = true
   }
 
+  // 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(
     internalPanelRef,
     ref,
     anchor ? floatingRef : null,
-    useEvent((panel) => dispatch({ type: ActionTypes.SetPanel, panel }))
+    useEvent((panel) => dispatch({ type: ActionTypes.SetPanel, panel })),
+    setLocalPanelElement
   )
   let ownerDocument = useOwnerDocument(internalPanelRef)
   let mergeRefs = useMergeRefsFn()
@@ -884,7 +897,7 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   let usesOpenClosedState = useOpenClosed()
   let [visible, transitionData] = useTransition(
     transition,
-    state.panel,
+    localPanelElement,
     usesOpenClosedState !== null
       ? (usesOpenClosedState & State.Open) === State.Open
       : state.popoverState === PopoverStates.Open
@@ -1028,7 +1041,10 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
 
           // Ignore sentinel buttons and items inside the panel
           for (let element of combined.slice()) {
-            if (element.dataset.headlessuiFocusGuard === 'true' || state.panel?.contains(element)) {
+            if (
+              element.dataset.headlessuiFocusGuard === 'true' ||
+              localPanelElement?.contains(element)
+            ) {
               let idx = combined.indexOf(element)
               if (idx !== -1) combined.splice(idx, 1)
             }
diff --git a/packages/@headlessui-react/src/components/transition/transition.tsx b/packages/@headlessui-react/src/components/transition/transition.tsx
index 31aaf07..9f0922f 100644
--- a/packages/@headlessui-react/src/components/transition/transition.tsx
+++ b/packages/@headlessui-react/src/components/transition/transition.tsx
@@ -319,12 +319,12 @@ function TransitionChildFn<TTag extends ElementType = typeof DEFAULT_TRANSITION_
 
     ...theirProps
   } = props as typeof props
-  let [containerElement, setContainerElement] = useState<HTMLElement | null>(null)
+  let [localContainerElement, setLocalContainerElement] = useState<HTMLElement | null>(null)
   let container = useRef<HTMLElement | null>(null)
   let requiresRef = shouldForwardRef(props)
 
   let transitionRef = useSyncRefs(
-    ...(requiresRef ? [container, ref, setContainerElement] : ref === null ? [] : [ref])
+    ...(requiresRef ? [container, ref, setLocalContainerElement] : ref === null ? [] : [ref])
   )
   let strategy = theirProps.unmount ?? true ? RenderStrategy.Unmount : RenderStrategy.Hidden
 
@@ -438,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, containerElement, show, { start, end })
+  let [, transitionData] = useTransition(enabled, localContainerElement, show, { start, end })
 
   let ourProps = compact({
     ref: transitionRef,


Test patch:
diff --git a/packages/@headlessui-react/src/components/combobox/combobox.test.tsx b/packages/@headlessui-react/src/components/combobox/combobox.test.tsx
index 00fb495..2eb6660 100644
--- a/packages/@headlessui-react/src/components/combobox/combobox.test.tsx
+++ b/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())
+    })
+  )
+})
diff --git a/packages/@headlessui-react/src/components/disclosure/disclosure.test.tsx b/packages/@headlessui-react/src/components/disclosure/disclosure.test.tsx
index 2c7aeb2..d328449 100644
--- a/packages/@headlessui-react/src/components/disclosure/disclosure.test.tsx
+++ b/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())
+    })
+  )
+})
diff --git a/packages/@headlessui-react/src/components/listbox/listbox.test.tsx b/packages/@headlessui-react/src/components/listbox/listbox.test.tsx
index f603344..d75b01f 100644
--- a/packages/@headlessui-react/src/components/listbox/listbox.test.tsx
+++ b/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())
+    })
+  )
+})
diff --git a/packages/@headlessui-react/src/components/menu/menu.test.tsx b/packages/@headlessui-react/src/components/menu/menu.test.tsx
index 0285959..a57e748 100644
--- a/packages/@headlessui-react/src/components/menu/menu.test.tsx
+++ b/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())
+    })
+  )
+})
diff --git a/packages/@headlessui-react/src/components/popover/popover.test.tsx b/packages/@headlessui-react/src/components/popover/popover.test.tsx
index 85d00d0..86d9d44 100644
--- a/packages/@headlessui-react/src/components/popover/popover.test.tsx
+++ b/packages/@headlessui-react/src/components/popover/popover.test.tsx
@@ -1,4 +1,4 @@
-import { render } from '@testing-library/react'
+import { render, waitFor } from '@testing-library/react'
 import React, { Fragment, act, createElement, useEffect, useRef, useState } from 'react'
 import ReactDOM from 'react-dom'
 import {
@@ -2844,3 +2844,40 @@ describe('Nested popovers', () => {
     })
   )
 })
+
+describe('transitions', () => {
+  it(
+    'should be possible to close the Popover when using the `transition` prop',
+    suppressConsoleLogs(async () => {
+      render(
+        <Popover>
+          <PopoverButton>Toggle</PopoverButton>
+          <PopoverPanel transition>Contents</PopoverPanel>
+        </Popover>
+      )
+
+      // Focus the button
+      await focus(getPopoverButton())
+
+      // Ensure the button is focused
+      assertActiveElement(getPopoverButton())
+
+      // Open the popover
+      await click(getPopoverButton())
+
+      // Ensure the popover is visible
+      assertPopoverPanel({ state: PopoverState.Visible })
+
+      // Close the popover
+      await click(getPopoverButton())
+
+      // Wait for the transition to finish, and the popover to close
+      await waitFor(() => {
+        assertPopoverPanel({ state: PopoverState.InvisibleUnmounted })
+      })
+
+      // Ensure the button got the restored focus
+      assertActiveElement(getPopoverButton())
+    })
+  )
+})