Fix re-focusing element after close (#1186)

* fix restoreElement logic

The code for React already worked, let's update the Vue code to make it
similar which properly restores focus.

* update changelog

CHANGELOG.md

@@ -34,6 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Improve outside click support ([#1175](https://github.com/tailwindlabs/headlessui/pull/1175))
 - Reset Combobox Input when the value gets reset ([#1181](https://github.com/tailwindlabs/headlessui/pull/1181))
 - Adjust active {item,option} index ([#1184](https://github.com/tailwindlabs/headlessui/pull/1184))
+- Fix re-focusing element after close ([#1186](https://github.com/tailwindlabs/headlessui/pull/1186))
 
 ## [@headlessui/react@v1.5.0] - 2022-02-17
 

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

@@ -24,7 +24,7 @@ import { dom } from '../../utils/dom'
 import { useOpenClosed, State, useOpenClosedProvider } from '../../internal/open-closed'
 import { match } from '../../utils/match'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
-import { sortByDomNode } from '../../utils/focus-management'
+import { FocusableMode, isFocusableElement, sortByDomNode } from '../../utils/focus-management'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 
 enum ListboxStates {
@@ -247,14 +247,15 @@ export let Listbox = defineComponent({
     }
 
     // Handle outside click
-    useOutsideClick(buttonRef, (event, target) => {
-      let active = document.activeElement
-
+    useOutsideClick([buttonRef, optionsRef], (event, target) => {
       if (listboxState.value !== ListboxStates.Open) return
 
-      if (!dom(optionsRef)?.contains(target)) api.closeListbox()
-      if (active !== document.body && active?.contains(target)) return // Keep focus on newly clicked/focused element
-      if (!event.defaultPrevented) dom(buttonRef)?.focus({ preventScroll: true })
+      api.closeListbox()
+
+      if (!isFocusableElement(target, FocusableMode.Loose)) {
+        event.preventDefault()
+        dom(buttonRef)?.focus()
+      }
     })
 
     // @ts-expect-error Types of property 'dataRef' are incompatible.

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

@@ -22,7 +22,7 @@ import { useTreeWalker } from '../../hooks/use-tree-walker'
 import { useOpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { match } from '../../utils/match'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
-import { sortByDomNode } from '../../utils/focus-management'
+import { FocusableMode, isFocusableElement, sortByDomNode } from '../../utils/focus-management'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 
 enum MenuStates {
@@ -201,14 +201,15 @@ export let Menu = defineComponent({
     }
 
     // Handle outside click
-    useOutsideClick(buttonRef, (event, target) => {
-      let active = document.activeElement
-
+    useOutsideClick([buttonRef, itemsRef], (event, target) => {
       if (menuState.value !== MenuStates.Open) return
 
-      if (!dom(itemsRef)?.contains(target)) api.closeMenu()
-      if (active !== document.body && active?.contains(target)) return // Keep focus on newly clicked/focused element
-      if (!event.defaultPrevented) dom(buttonRef)?.focus({ preventScroll: true })
+      api.closeMenu()
+
+      if (!isFocusableElement(target, FocusableMode.Loose)) {
+        event.preventDefault()
+        dom(buttonRef)?.focus()
+      }
     })
 
     // @ts-expect-error Types of property 'dataRef' are incompatible.

packages/@headlessui-vue/src/hooks/use-focus-trap.ts

@@ -24,8 +24,8 @@ export function useFocusTrap(
   function handleFocus() {
     if (!enabled.value) return
     if (containers.value.size !== 1) return
-    let { initialFocus } = options.value
 
+    let { initialFocus } = options.value
     let activeElement = document.activeElement as HTMLElement
 
     if (initialFocus) {
@@ -36,7 +36,10 @@ export function useFocusTrap(
       return // Already focused within Dialog
     }
 
-    restoreElement.value = activeElement
+    if (!restoreElement.value) {
+      // We already have a restore element
+      restoreElement.value = activeElement
+    }
 
     // Try to focus the initialFocus ref
     if (initialFocus) {