Fix components not properly closing when using the transition prop

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 properly closing 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,20 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
   }
 
   let [floatingRef, style] = useFloatingPanel(anchor)
+  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,
@@ -450,12 +451,14 @@ function PanelFn<TTag extends ElementType = typeof DEFAULT_PANEL_TAG>(
   let [state, dispatch] = useDisclosureContext('Disclosure.Panel')
   let { close } = useDisclosureAPIContext('Disclosure.Panel')
   let mergeRefs = useMergeRefsFn()
+  let [localPanelElement, setLocalPanelElement] = useState<HTMLElement | null>(null)
 
   let panelRef = useSyncRefs(
     ref,
     useEvent((element) => {
       startTransition(() => dispatch({ type: ActionTypes.SetPanelElement, element }))
-    })
+    }),
+    setLocalPanelElement
   )
 
   useEffect(() => {
@@ -468,7 +471,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,
@@ -931,6 +932,7 @@ function OptionsFn<TTag extends ElementType = typeof DEFAULT_OPTIONS_TAG>(
     ...theirProps
   } = props
   let anchor = useResolvedAnchor(rawAnchor)
+  let [localOptionsElement, setLocalOptionsElement] = useState<HTMLElement | null>(null)
 
   // Always enable `portal` functionality, when `anchor` is enabled
   if (anchor) {
@@ -945,7 +947,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 +1025,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())
+    })
+  )
+})

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,13 @@ function ItemsFn<TTag extends ElementType = typeof DEFAULT_ITEMS_TAG>(
   let [state, dispatch] = useMenuContext('Menu.Items')
   let [floatingRef, style] = useFloatingPanel(anchor)
   let getFloatingPanelProps = useFloatingPanelProps()
+  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 +639,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