fix button focus management

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

@@ -743,6 +743,9 @@ describe('Keyboard interactions', () => {
         // Verify it is closed
         assertListboxButton({ state: ListboxState.Closed })
         assertListbox({ state: ListboxState.Closed })
+
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
       })
     )
 
@@ -801,6 +804,9 @@ describe('Keyboard interactions', () => {
         expect(handleChange).toHaveBeenCalledTimes(1)
         expect(handleChange).toHaveBeenCalledWith('a')
 
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
+
         // Open listbox again
         await click(getListboxButton())
 
@@ -1086,6 +1092,9 @@ describe('Keyboard interactions', () => {
         expect(handleChange).toHaveBeenCalledTimes(1)
         expect(handleChange).toHaveBeenCalledWith('a')
 
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
+
         // Open listbox again
         await click(getListboxButton())
 
@@ -1131,6 +1140,9 @@ describe('Keyboard interactions', () => {
         // Verify it is closed
         assertListboxButton({ state: ListboxState.Closed })
         assertListbox({ state: ListboxState.Closed })
+
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
       })
     )
   })
@@ -2652,6 +2664,9 @@ describe('Mouse interactions', () => {
 
       // Should be closed now
       assertListbox({ state: ListboxState.Closed })
+
+      // Verify the button is focused again
+      assertActiveElement(getListboxButton())
     })
   )
 
@@ -2723,6 +2738,9 @@ describe('Mouse interactions', () => {
 
       // Should be closed now
       assertListbox({ state: ListboxState.Closed })
+
+      // Verify the button is focused again
+      assertActiveElement(getListboxButton())
     })
   )
 
@@ -2979,6 +2997,9 @@ describe('Mouse interactions', () => {
       expect(handleChange).toHaveBeenCalledTimes(1)
       expect(handleChange).toHaveBeenCalledWith('bob')
 
+      // Verify the button is focused again
+      assertActiveElement(getListboxButton())
+
       // Open listbox again
       await click(getListboxButton())
 

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

@@ -222,8 +222,8 @@ export function Listbox<
 
       if (!optionsRef.current?.contains(event.target as HTMLElement)) {
         dispatch({ type: ActionTypes.CloseListbox })
-        if (!event.defaultPrevented) buttonRef.current?.focus()
       }
+      if (!event.defaultPrevented) d.nextFrame(() => buttonRef.current?.focus())
     }
 
     window.addEventListener('click', handler)
@@ -309,6 +309,7 @@ const Button = forwardRefWithAs(function Button<
     (event: MouseEvent) => {
       if (state.listboxState === ListboxStates.Open) {
         dispatch({ type: ActionTypes.CloseListbox })
+        d.nextFrame(() => state.buttonRef.current?.focus())
       } else {
         event.preventDefault()
         dispatch({ type: ActionTypes.OpenListbox })
@@ -441,7 +442,7 @@ const Options = forwardRefWithAs(function Options<
             const { dataRef } = state.options[state.activeOptionIndex]
             state.propsRef.current.onChange(dataRef.current.value)
           }
-          d.nextFrame(() => state.buttonRef.current?.focus())
+          disposables().nextFrame(() => state.buttonRef.current?.focus())
           break
 
         case Keys.ArrowDown:
@@ -535,7 +536,6 @@ function Option<TTag extends React.ElementType = typeof DEFAULT_OPTION_TAG, TTyp
 ) {
   const { disabled = false, value, className, ...passthroughProps } = props
   const [state, dispatch] = useListboxContext([Listbox.name, Option.name].join('.'))
-  const d = useDisposables()
   const id = `headlessui-listbox-option-${useId()}`
   const active =
     state.activeOptionIndex !== null ? state.options[state.activeOptionIndex].id === id : false
@@ -581,9 +581,9 @@ function Option<TTag extends React.ElementType = typeof DEFAULT_OPTION_TAG, TTyp
       if (disabled) return event.preventDefault()
       select()
       dispatch({ type: ActionTypes.CloseListbox })
-      d.nextFrame(() => state.buttonRef.current?.focus())
+      disposables().nextFrame(() => state.buttonRef.current?.focus())
     },
-    [d, dispatch, state.buttonRef, disabled, select]
+    [dispatch, state.buttonRef, disabled, select]
   )
 
   const handleFocus = React.useCallback(() => {

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

@@ -573,6 +573,9 @@ describe('Keyboard interactions', () => {
         // Verify it is closed
         assertMenuButton({ state: MenuState.Closed })
         assertMenu({ state: MenuState.Closed })
+
+        // Verify the button is focused again
+        assertActiveElement(getMenuButton())
       })
     )
 
@@ -616,6 +619,9 @@ describe('Keyboard interactions', () => {
         assertMenuButton({ state: MenuState.Closed })
         assertMenu({ state: MenuState.Closed })
 
+        // Verify the button is focused again
+        assertActiveElement(getMenuButton())
+
         // Verify the "click" went through on the `a` tag
         expect(clickHandler).toHaveBeenCalled()
       })
@@ -668,6 +674,9 @@ describe('Keyboard interactions', () => {
       // Verify the button got "clicked"
       expect(clickHandler).toHaveBeenCalledTimes(1)
 
+      // Verify the button is focused again
+      assertActiveElement(getMenuButton())
+
       // Click the menu button again
       await click(getMenuButton())
 
@@ -679,6 +688,9 @@ describe('Keyboard interactions', () => {
 
       // Verify the button got "clicked"
       expect(clickHandler).toHaveBeenCalledTimes(2)
+
+      // Verify the button is focused again
+      assertActiveElement(getMenuButton())
     })
   )
 
@@ -888,6 +900,9 @@ describe('Keyboard interactions', () => {
         // Verify it is closed
         assertMenuButton({ state: MenuState.Closed })
         assertMenu({ state: MenuState.Closed })
+
+        // Verify the button is focused again
+        assertActiveElement(getMenuButton())
       })
     )
 
@@ -933,6 +948,9 @@ describe('Keyboard interactions', () => {
 
         // Verify the "click" went through on the `a` tag
         expect(clickHandler).toHaveBeenCalled()
+
+        // Verify the button is focused again
+        assertActiveElement(getMenuButton())
       })
     )
   })
@@ -972,6 +990,9 @@ describe('Keyboard interactions', () => {
         // Verify it is closed
         assertMenuButton({ state: MenuState.Closed })
         assertMenu({ state: MenuState.Closed })
+
+        // Verify the button is focused again
+        assertActiveElement(getMenuButton())
       })
     )
   })
@@ -2324,6 +2345,9 @@ describe('Mouse interactions', () => {
 
       // Should be closed now
       assertMenu({ state: MenuState.Closed })
+
+      // Verify the button is focused again
+      assertActiveElement(getMenuButton())
     })
   )
 
@@ -2350,6 +2374,9 @@ describe('Mouse interactions', () => {
 
       // Should be closed now
       assertMenu({ state: MenuState.Closed })
+
+      // Verify the button is focused again
+      assertActiveElement(getMenuButton())
     })
   )
 

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

@@ -4,6 +4,7 @@ import * as React from 'react'
 import { Props } from '../../types'
 import { match } from '../../utils/match'
 import { forwardRefWithAs, render } from '../../utils/render'
+import { disposables } from '../../utils/disposables'
 import { useDisposables } from '../../hooks/use-disposables'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
@@ -213,8 +214,8 @@ export function Menu<TTag extends React.ElementType = typeof DEFAULT_MENU_TAG>(
 
       if (!itemsRef.current?.contains(event.target as HTMLElement)) {
         dispatch({ type: ActionTypes.CloseMenu })
-        if (!event.defaultPrevented) buttonRef.current?.focus()
       }
+      if (!event.defaultPrevented) d.nextFrame(() => buttonRef.current?.focus())
     }
 
     window.addEventListener('click', handler)
@@ -294,6 +295,7 @@ const Button = forwardRefWithAs(function Button<
     (event: MouseEvent) => {
       if (state.menuState === MenuStates.Open) {
         dispatch({ type: ActionTypes.CloseMenu })
+        d.nextFrame(() => state.buttonRef.current?.focus())
       } else {
         event.preventDefault()
         dispatch({ type: ActionTypes.OpenMenu })
@@ -357,7 +359,6 @@ const Items = forwardRefWithAs(function Items<
   const itemsRef = useSyncRefs(state.itemsRef, ref)
 
   const id = `headlessui-menu-items-${useId()}`
-  const d = useDisposables()
   const searchDisposables = useDisposables()
 
   const handleKeyDown = React.useCallback(
@@ -380,8 +381,8 @@ const Items = forwardRefWithAs(function Items<
           if (state.activeItemIndex !== null) {
             const { id } = state.items[state.activeItemIndex]
             document.getElementById(id)?.click()
-            d.nextFrame(() => state.buttonRef.current?.focus())
           }
+          disposables().nextFrame(() => state.buttonRef.current?.focus())
           break
 
         case Keys.ArrowDown:
@@ -405,7 +406,7 @@ const Items = forwardRefWithAs(function Items<
         case Keys.Escape:
           event.preventDefault()
           dispatch({ type: ActionTypes.CloseMenu })
-          d.nextFrame(() => state.buttonRef.current?.focus())
+          disposables().nextFrame(() => state.buttonRef.current?.focus())
           break
 
         case Keys.Tab:
@@ -419,7 +420,7 @@ const Items = forwardRefWithAs(function Items<
           break
       }
     },
-    [d, dispatch, searchDisposables, state]
+    [dispatch, searchDisposables, state]
   )
 
   const propsBag = React.useMemo(() => ({ open: state.menuState === MenuStates.Open }), [state])

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

@@ -801,6 +801,9 @@ describe('Keyboard interactions', () => {
         // Verify it is closed
         assertListboxButton({ state: ListboxState.Closed })
         assertListbox({ state: ListboxState.Closed })
+
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
       })
     )
 
@@ -855,6 +858,9 @@ describe('Keyboard interactions', () => {
         expect(handleChange).toHaveBeenCalledTimes(1)
         expect(handleChange).toHaveBeenCalledWith('a')
 
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
+
         // Open listbox again
         await click(getListboxButton())
 
@@ -1154,6 +1160,9 @@ describe('Keyboard interactions', () => {
         expect(handleChange).toHaveBeenCalledTimes(1)
         expect(handleChange).toHaveBeenCalledWith('a')
 
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
+
         // Open listbox again
         await click(getListboxButton())
 
@@ -1202,6 +1211,9 @@ describe('Keyboard interactions', () => {
         // Verify it is closed
         assertListboxButton({ state: ListboxState.Closed })
         assertListbox({ state: ListboxState.Closed })
+
+        // Verify the button is focused again
+        assertActiveElement(getListboxButton())
       })
     )
   })
@@ -2846,6 +2858,9 @@ describe('Mouse interactions', () => {
 
       // Should be closed now
       assertListbox({ state: ListboxState.Closed })
+
+      // Verify the button is focused again
+      assertActiveElement(getListboxButton())
     })
   )
 
@@ -2923,6 +2938,9 @@ describe('Mouse interactions', () => {
 
       // Should be closed now
       assertListbox({ state: ListboxState.Closed })
+
+      // Verify the button is focused again
+      assertActiveElement(getListboxButton())
     })
   )
 
@@ -3197,6 +3215,9 @@ describe('Mouse interactions', () => {
       expect(handleChange).toHaveBeenCalledTimes(1)
       expect(handleChange).toHaveBeenCalledWith('bob')
 
+      // Verify the button is focused again
+      assertActiveElement(getListboxButton())
+
       // Open listbox again
       await click(getListboxButton())
 

packages/@headlessui-vue/src/components/listbox/listbox.ts

@@ -198,8 +198,8 @@ export const Listbox = defineComponent({
 
         if (!optionsRef.value?.contains(event.target as HTMLElement)) {
           api.closeListbox()
-          if (!event.defaultPrevented) buttonRef.value?.focus()
         }
+        if (!event.defaultPrevented) nextTick(() => buttonRef.value?.focus())
       }
 
       window.addEventListener('click', handler)
@@ -318,6 +318,7 @@ export const ListboxButton = defineComponent({
     function handlePointerUp(event: MouseEvent) {
       if (api.listboxState.value === ListboxStates.Open) {
         api.closeListbox()
+        nextTick(() => api.buttonRef.value?.focus())
       } else {
         event.preventDefault()
         api.openListbox()
@@ -407,8 +408,8 @@ export const ListboxOptions = defineComponent({
           if (api.activeOptionIndex.value !== null) {
             const { dataRef } = api.options.value[api.activeOptionIndex.value]
             api.select(dataRef.value)
-            nextTick(() => api.buttonRef.value?.focus())
           }
+          nextTick(() => api.buttonRef.value?.focus())
           break
 
         case Keys.ArrowDown: