Fix crash when using as={Fragment} on MenuButton, ListboxButton, DisclosureButton or Button components (#3478)

This PR fixes an issue where a maximum update depth exceeded error was
thrown when using `as={Fragment}` on button related components.

The issue here is that the `ref` on a element would re-fire every render
_if_ the a function was used _and_ the function is a new function (aka
not a stable function).

This resulted in the `ref` being called with the DOM element, then
`null`, then the DOM element, then `null`, and so on.

To solve this, we have to make sure that the `ref` is always a stable
reference.

Fixes: #3476
Fixes: #3439

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 crash when using `as={Fragment}` on `MenuButton`, `ListboxButton`, `DisclosureButton` or `Button` components ([#3478](https://github.com/tailwindlabs/headlessui/pull/3478))
 
 ## [2.1.7] - 2024-09-11
 

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

@@ -1,5 +1,5 @@
 import { render, screen } from '@testing-library/react'
-import React from 'react'
+import React, { Fragment } from 'react'
 import { Button } from './button'
 
 describe('Rendering', () => {
@@ -35,5 +35,15 @@ describe('Rendering', () => {
 
       expect(screen.getByRole('button')).toHaveAttribute('data-autofocus')
     })
+
+    it('should be possible to render a Button using as={Fragment}', async () => {
+      render(
+        <Button as={Fragment}>
+          <button>Toggle</button>
+        </Button>
+      )
+
+      expect(screen.getByRole('button')).toHaveAttribute('type')
+    })
   })
 })

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

@@ -10,6 +10,7 @@ import {
   forwardRefWithAs,
   mergeProps,
   render,
+  useMergeRefsFn,
   type HasDisplayName,
   type RefProp,
 } from '../../utils/render'
@@ -41,6 +42,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   ref: Ref<HTMLElement>
 ) {
   let providedDisabled = useDisabled()
+  let mergeRefs = useMergeRefsFn()
   let { disabled = providedDisabled || false, autoFocus = false, ...theirProps } = props
 
   let { isFocusVisible: focus, focusProps } = useFocusRing({ autoFocus })
@@ -64,6 +66,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   }, [disabled, hover, focus, active, autoFocus])
 
   return render({
+    mergeRefs,
     ourProps,
     theirProps,
     slot,

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

@@ -1211,6 +1211,31 @@ describe('Rendering', () => {
         expect(getComboboxButton()).not.toHaveAttribute('type')
       })
     })
+
+    it(
+      'should be possible to render a ComboboxButton using as={Fragment}',
+      suppressConsoleLogs(async () => {
+        render(
+          <Combobox>
+            <ComboboxInput />
+            <ComboboxButton as={Fragment}>
+              <button>Toggle</button>
+            </ComboboxButton>
+            <ComboboxOptions>
+              <ComboboxOption value="a">Option A</ComboboxOption>
+              <ComboboxOption value="b">Option B</ComboboxOption>
+              <ComboboxOption value="c">Option C</ComboboxOption>
+            </ComboboxOptions>
+          </Combobox>
+        )
+
+        assertComboboxButton({ state: ComboboxState.InvisibleUnmounted })
+
+        await click(getComboboxButton())
+
+        assertComboboxButton({ state: ComboboxState.Visible })
+      })
+    )
   })
 
   describe('Combobox.Options', () => {

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

@@ -70,6 +70,7 @@ import {
   forwardRefWithAs,
   mergeProps,
   render,
+  useMergeRefsFn,
   type HasDisplayName,
   type PropsForFeatures,
   type RefProp,
@@ -1495,6 +1496,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   let data = useData('Combobox.Button')
   let actions = useActions('Combobox.Button')
   let buttonRef = useSyncRefs(ref, actions.setButtonElement)
+  let mergeRefs = useMergeRefsFn()
 
   let internalId = useId()
   let {
@@ -1616,6 +1618,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   )
 
   return render({
+    mergeRefs,
     ourProps,
     theirProps,
     slot,

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

@@ -1,5 +1,5 @@
 import { render, waitFor } from '@testing-library/react'
-import React, { Suspense, createElement, useEffect, useRef } from 'react'
+import React, { Fragment, Suspense, createElement, useEffect, useRef } from 'react'
 import {
   DisclosureState,
   assertActiveElement,
@@ -439,6 +439,28 @@ describe('Rendering', () => {
         expect(getDisclosureButton()).not.toHaveAttribute('type')
       })
     })
+
+    it(
+      'should be possible to render a DisclosureButton using as={Fragment}',
+      suppressConsoleLogs(async () => {
+        render(
+          <Disclosure>
+            <DisclosureButton as={Fragment}>
+              <button>Toggle</button>
+            </DisclosureButton>
+            <DisclosurePanel>Contents</DisclosurePanel>
+          </Disclosure>
+        )
+
+        assertDisclosureButton({ state: DisclosureState.InvisibleUnmounted })
+        assertDisclosurePanel({ state: DisclosureState.InvisibleUnmounted })
+
+        await click(getDisclosureButton())
+
+        assertDisclosureButton({ state: DisclosureState.Visible })
+        assertDisclosurePanel({ state: DisclosureState.Visible })
+      })
+    )
   })
 
   describe('Disclosure.Panel', () => {

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

@@ -1,5 +1,5 @@
 import { render, waitFor } from '@testing-library/react'
-import React, { createElement, useEffect, useState } from 'react'
+import React, { Fragment, createElement, useEffect, useState } from 'react'
 import {
   ListboxMode,
   ListboxState,
@@ -760,6 +760,32 @@ describe('Rendering', () => {
         expect(getListboxButton()).not.toHaveAttribute('type')
       })
     })
+
+    it(
+      'should be possible to render a ListboxButton using as={Fragment}',
+      suppressConsoleLogs(async () => {
+        render(
+          <Listbox>
+            <ListboxButton as={Fragment}>
+              <button>Toggle</button>
+            </ListboxButton>
+            <ListboxOptions>
+              <ListboxOption value="a">Option A</ListboxOption>
+              <ListboxOption value="b">Option B</ListboxOption>
+              <ListboxOption value="c">Option C</ListboxOption>
+            </ListboxOptions>
+          </Listbox>
+        )
+
+        assertListboxButton({ state: ListboxState.InvisibleUnmounted })
+        assertListbox({ state: ListboxState.InvisibleUnmounted })
+
+        await click(getListboxButton())
+
+        assertListboxButton({ state: ListboxState.Visible })
+        assertListbox({ state: ListboxState.Visible })
+      })
+    )
   })
 
   describe('Listbox.Options', () => {

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

@@ -75,6 +75,7 @@ import {
   forwardRefWithAs,
   mergeProps,
   render,
+  useMergeRefsFn,
   type HasDisplayName,
   type PropsForFeatures,
   type RefProp,
@@ -785,6 +786,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     autoFocus = false,
     ...theirProps
   } = props
+  let mergeRefs = useMergeRefsFn()
   let buttonRef = useSyncRefs(ref, useFloatingReference(), actions.setButtonElement)
   let getFloatingReferenceProps = useFloatingReferenceProps()
 
@@ -880,6 +882,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   )
 
   return render({
+    mergeRefs,
     ourProps,
     theirProps,
     slot,

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

@@ -1,5 +1,5 @@
 import { render, waitFor } from '@testing-library/react'
-import React, { createElement, useEffect } from 'react'
+import React, { Fragment, createElement, useEffect } from 'react'
 import {
   MenuState,
   assertActiveElement,
@@ -306,6 +306,32 @@ describe('Rendering', () => {
         expect(getMenuButton()).not.toHaveAttribute('type')
       })
     })
+
+    it(
+      'should be possible to render a MenuButton using as={Fragment}',
+      suppressConsoleLogs(async () => {
+        render(
+          <Menu>
+            <MenuButton as={Fragment}>
+              <button>Toggle</button>
+            </MenuButton>
+            <MenuItems>
+              <MenuItem as="a">Item A</MenuItem>
+              <MenuItem as="a">Item B</MenuItem>
+              <MenuItem as="a">Item C</MenuItem>
+            </MenuItems>
+          </Menu>
+        )
+
+        assertMenuButton({ state: MenuState.InvisibleUnmounted })
+        assertMenu({ state: MenuState.InvisibleUnmounted })
+
+        await click(getMenuButton())
+
+        assertMenuButton({ state: MenuState.Visible })
+        assertMenu({ state: MenuState.Visible })
+      })
+    )
   })
 
   describe('Menu.Items', () => {

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

@@ -68,6 +68,7 @@ import {
   forwardRefWithAs,
   mergeProps,
   render,
+  useMergeRefsFn,
   type HasDisplayName,
   type RefProp,
 } from '../../utils/render'
@@ -483,6 +484,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   } = props
   let [state, dispatch] = useMenuContext('Menu.Button')
   let getFloatingReferenceProps = useFloatingReferenceProps()
+  let mergeRefs = useMergeRefsFn()
   let buttonRef = useSyncRefs(
     ref,
     useFloatingReference(),
@@ -570,6 +572,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
   )
 
   return render({
+    mergeRefs,
     ourProps,
     theirProps,
     slot,