From 4ed69f640ce1240931624bab616d8f1319608c17 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Wed, 3 Apr 2024 15:06:41 +0200
Subject: [PATCH] Keep focus inside of the `<ComboboxInput />` component
 (#3073)

* bail the refocus if focus is already on the correct element

* use `mousedown` instead of `click` event

The `mousedown` event happens before the `focus` event. When we
`e.preventDefault()` in this listener, the `focus` event will not fire.

This also means that the focus is not lost on the actual `input`
component which in turn means that we can maintain the selection /
cursor position inside the `input`.

We still use the `refocusInput()` as a fallback in case something else
goes wrong.

* add comments to describe _why_ we use `mousedown`

* ensure we handle mouse buttons correctly

* ensure we handle `Enter` and `Space` explicitly

Now that we use the `mousedown` event instead of the `click` event, we
have to make sure that we handle the `enter` and `space` keys
explicitly.

This used to be covered by the `click` event, but not for the `mousedown` event.

* ensure we focus the first element when using `ArrowDown` on the `ComboboxButton`

We go to the last one on `ArrownUp`, but we forgot to do this on
`ArrowDown`.

* fix tiny typo

Not related to this PR, but noticed it and fixed it anyway.

* update changelog

* ensure we reset the `isTyping` flag

While we are typing, the flag can remain true. But once we stop typing,
the `nextFrame` handler will kick in and set it to `false` again.

It currently behaves as a debounce-like function such that the
`nextFrame` callbacks are cancelled once a new event is fired.

* ensure unique callbacks in the `_disposables` array

This allows us to keep re-adding dispose functions and only register the
callbacks once.

Ideally we can use a `Set`, but we also want to remove a single callback
if the callback is disposed on its own instead of disposing the whole
group. For this we do require an `idx` which is not available in a
`Set` unless you are looping over all disposable functions.

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

Co-authored-by: Jonathan Reinink <jonathan@reinink.ca>

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

Co-authored-by: Jonathan Reinink <jonathan@reinink.ca>

* update comments

* abstract confusing logic to a `useFrameDebounce()` hook

* use correct path import

* add some breathing room

---------

Co-authored-by: Jonathan Reinink <jonathan@reinink.ca>
---
 packages/@headlessui-react/CHANGELOG.md       |  1 +
 .../src/components/combobox/combobox.test.tsx |  2 +-
 .../src/components/combobox/combobox.tsx      | 71 +++++++++++++++----
 .../@headlessui-react/src/components/mouse.ts |  4 ++
 .../src/hooks/use-frame-debounce.ts           | 18 +++++
 .../src/hooks/use-refocusable-input.ts        |  4 ++
 .../src/utils/disposables.ts                  |  4 ++
 7 files changed, 91 insertions(+), 13 deletions(-)
 create mode 100644 packages/@headlessui-react/src/components/mouse.ts
 create mode 100644 packages/@headlessui-react/src/hooks/use-frame-debounce.ts

diff --git a/packages/@headlessui-react/CHANGELOG.md b/packages/@headlessui-react/CHANGELOG.md
index bd9dbba6e0..b592bc329b 100644
--- a/packages/@headlessui-react/CHANGELOG.md
+++ b/packages/@headlessui-react/CHANGELOG.md
@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Prevent unnecessary execution of the `displayValue` callback in the `ComboboxInput` component ([#3048](https://github.com/tailwindlabs/headlessui/pull/3048))
 - Expose missing `data-disabled` and `data-focus` attributes on the `TabsPanel`, `MenuButton`, `PopoverButton` and `DisclosureButton` components ([#3061](https://github.com/tailwindlabs/headlessui/pull/3061))
 - Fix cursor position when re-focusing the `ComboboxInput` component ([#3065](https://github.com/tailwindlabs/headlessui/pull/3065))
+- Keep focus inside of the `<ComboboxInput />` component ([#3073](https://github.com/tailwindlabs/headlessui/pull/3073))
 
 ### Changed
 
diff --git a/packages/@headlessui-react/src/components/combobox/combobox.test.tsx b/packages/@headlessui-react/src/components/combobox/combobox.test.tsx
index 0e824d8562..15a64472af 100644
--- a/packages/@headlessui-react/src/components/combobox/combobox.test.tsx
+++ b/packages/@headlessui-react/src/components/combobox/combobox.test.tsx
@@ -5194,7 +5194,7 @@ describe.each([{ virtual: true }, { virtual: false }])('Mouse interactions %s',
             options={[
               { value: 'alice', children: 'Alice', disabled: false },
               { value: 'bob', children: 'Bob', disabled: true },
-              { value: 'charile', children: 'Charlie', disabled: false },
+              { value: 'charlie', children: 'Charlie', disabled: false },
             ]}
           />
         )
diff --git a/packages/@headlessui-react/src/components/combobox/combobox.tsx b/packages/@headlessui-react/src/components/combobox/combobox.tsx
index 061cdbe58a..d31d66ed16 100644
--- a/packages/@headlessui-react/src/components/combobox/combobox.tsx
+++ b/packages/@headlessui-react/src/components/combobox/combobox.tsx
@@ -27,6 +27,7 @@ import { useControllable } from '../../hooks/use-controllable'
 import { useDisposables } from '../../hooks/use-disposables'
 import { useElementSize } from '../../hooks/use-element-size'
 import { useEvent } from '../../hooks/use-event'
+import { useFrameDebounce } from '../../hooks/use-frame-debounce'
 import { useId } from '../../hooks/use-id'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { useLatestValue } from '../../hooks/use-latest-value'
@@ -69,6 +70,7 @@ import {
 import { useDescribedBy } from '../description/description'
 import { Keys } from '../keyboard'
 import { Label, useLabelledBy, useLabels, type _internal_ComponentLabel } from '../label/label'
+import { MouseButton } from '../mouse'
 
 enum ComboboxState {
   Open,
@@ -1077,8 +1079,13 @@ function InputFn<
     })
   })
 
+  let debounce = useFrameDebounce()
   let handleKeyDown = useEvent((event: ReactKeyboardEvent<HTMLInputElement>) => {
     isTyping.current = true
+    debounce(() => {
+      isTyping.current = false
+    })
+
     switch (event.key) {
       // Ref: https://www.w3.org/WAI/ARIA/apg/patterns/menu/#keyboard-interaction-12
 
@@ -1388,11 +1395,26 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     switch (event.key) {
       // Ref: https://www.w3.org/WAI/ARIA/apg/patterns/menu/#keyboard-interaction-12
 
+      case Keys.Space:
+      case Keys.Enter:
+        event.preventDefault()
+        event.stopPropagation()
+        if (data.comboboxState === ComboboxState.Closed) {
+          actions.openCombobox()
+        }
+
+        return d.nextFrame(() => refocusInput())
+
       case Keys.ArrowDown:
         event.preventDefault()
         event.stopPropagation()
         if (data.comboboxState === ComboboxState.Closed) {
           actions.openCombobox()
+          d.nextFrame(() => {
+            if (!data.value) {
+              actions.goToOption(Focus.First)
+            }
+          })
         }
 
         return d.nextFrame(() => refocusInput())
@@ -1424,16 +1446,28 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
     }
   })
 
-  let handleClick = useEvent((event: ReactMouseEvent<HTMLButtonElement>) => {
-    if (isDisabledReactIssue7711(event.currentTarget)) return event.preventDefault()
-    if (data.comboboxState === ComboboxState.Open) {
-      actions.closeCombobox()
-    } else {
-      event.preventDefault()
-      actions.openCombobox()
+  let handleMouseDown = useEvent((event: ReactMouseEvent<HTMLButtonElement>) => {
+    // We use the `mousedown` event here since it fires before the focus event,
+    // allowing us to cancel the event before focus is moved from the
+    // `ComboboxInput` to the `ComboboxButton`. This keeps the input focused,
+    // preserving the cursor position and any text selection.
+    event.preventDefault()
+
+    if (isDisabledReactIssue7711(event.currentTarget)) return
+
+    // Since we're using the `mousedown` event instead of a `click` event here
+    // to preserve the focus of the `ComboboxInput`, we need to also check
+    // that the `left` mouse button was clicked.
+    if (event.button === MouseButton.Left) {
+      if (data.comboboxState === ComboboxState.Open) {
+        actions.closeCombobox()
+      } else {
+        actions.openCombobox()
+      }
     }
 
-    d.nextFrame(() => refocusInput())
+    // Ensure we focus the input
+    refocusInput()
   })
 
   let labelledBy = useLabelledBy([id])
@@ -1464,7 +1498,7 @@ function ButtonFn<TTag extends ElementType = typeof DEFAULT_BUTTON_TAG>(
       'aria-labelledby': labelledBy,
       disabled: disabled || undefined,
       autoFocus,
-      onClick: handleClick,
+      onMouseDown: handleMouseDown,
       onKeyDown: handleKeyDown,
     },
     focusProps,
@@ -1689,8 +1723,21 @@ function OptionFn<
     /* We also want to trigger this when the position of the active item changes so that we can re-trigger the scrollIntoView */ data.activeOptionIndex,
   ])
 
-  let handleClick = useEvent((event: { preventDefault: Function }) => {
-    if (disabled || data.virtual?.disabled(value)) return event.preventDefault()
+  let handleMouseDown = useEvent((event: ReactMouseEvent<HTMLButtonElement>) => {
+    // We use the `mousedown` event here since it fires before the focus event,
+    // allowing us to cancel the event before focus is moved from the
+    // `ComboboxInput` to the `ComboboxOption`. This keeps the input focused,
+    // preserving the cursor position and any text selection.
+    event.preventDefault()
+
+    // Since we're using the `mousedown` event instead of a `click` event here
+    // to preserve the focus of the `ComboboxInput`, we need to also check
+    // that the `left` mouse button was clicked.
+    if (event.button !== MouseButton.Left) {
+      return
+    }
+
+    if (disabled || data.virtual?.disabled(value)) return
     select()
 
     // We want to make sure that we don't accidentally trigger the virtual keyboard.
@@ -1758,7 +1805,7 @@ function OptionFn<
     // both single and multi-select.
     'aria-selected': selected,
     disabled: undefined, // Never forward the `disabled` prop
-    onClick: handleClick,
+    onMouseDown: handleMouseDown,
     onFocus: handleFocus,
     onPointerEnter: handleEnter,
     onMouseEnter: handleEnter,
diff --git a/packages/@headlessui-react/src/components/mouse.ts b/packages/@headlessui-react/src/components/mouse.ts
new file mode 100644
index 0000000000..8f56ab7ba5
--- /dev/null
+++ b/packages/@headlessui-react/src/components/mouse.ts
@@ -0,0 +1,4 @@
+export enum MouseButton {
+  Left = 0,
+  Right = 2,
+}
diff --git a/packages/@headlessui-react/src/hooks/use-frame-debounce.ts b/packages/@headlessui-react/src/hooks/use-frame-debounce.ts
new file mode 100644
index 0000000000..fa79640fec
--- /dev/null
+++ b/packages/@headlessui-react/src/hooks/use-frame-debounce.ts
@@ -0,0 +1,18 @@
+import { useDisposables } from './use-disposables'
+import { useEvent } from './use-event'
+
+/**
+ * Schedule some task in the next frame.
+ *
+ * - If you call the returned function multiple times, only the last task will
+ *   be executed.
+ * - If the component is unmounted, the task will be cancelled.
+ */
+export function useFrameDebounce() {
+  let d = useDisposables()
+
+  return useEvent((cb: () => void) => {
+    d.dispose()
+    d.nextFrame(() => cb())
+  })
+}
diff --git a/packages/@headlessui-react/src/hooks/use-refocusable-input.ts b/packages/@headlessui-react/src/hooks/use-refocusable-input.ts
index a13b2438ca..2ed48dfbbc 100644
--- a/packages/@headlessui-react/src/hooks/use-refocusable-input.ts
+++ b/packages/@headlessui-react/src/hooks/use-refocusable-input.ts
@@ -29,6 +29,10 @@ export function useRefocusableInput(ref: MutableRefObject<HTMLInputElement | nul
 
   return useEvent(() => {
     let input = ref.current
+
+    // If the input is already focused, we don't need to do anything
+    if (document.activeElement === input) return
+
     if (!(input instanceof HTMLInputElement)) return
     if (!input.isConnected) return
 
diff --git a/packages/@headlessui-react/src/utils/disposables.ts b/packages/@headlessui-react/src/utils/disposables.ts
index 5e14bb9a20..55ca27b6b5 100644
--- a/packages/@headlessui-react/src/utils/disposables.ts
+++ b/packages/@headlessui-react/src/utils/disposables.ts
@@ -59,6 +59,10 @@ export function disposables() {
     },
 
     add(cb: () => void) {
+      if (_disposables.includes(cb)) {
+        return
+      }
+
       _disposables.push(cb)
       return () => {
         let idx = _disposables.indexOf(cb)