Focus not returned to triggering button when closing modal

[eroltibern]

What package within Headless UI are you using?

@headlessui/react

What version of that package are you using?

v1.7.17 - v.2.x.x

What browser are you using?

Chrome

Reproduction URL

Sandbox

Describe your issue

When opening a Dialog by clicking on a button with its content wrapped in a span, focus is not properly restored to the triggering button after the dialog closes. Focus is instead returned to body.

<button onClick={open}>
  <span>Click me</span>
</button>

I believe this is due to the span element being added as the latest focused element in the active element history, and when trying to focus the span the browser falls back to focusing body instead.

A solution could be to find the first focusable element in the history array and trying to focus that? Something like this:

return useEvent(() => {
    return (
      localHistory.current.find(
        (x) => x != null && x.isConnected && x.matches(focusableSelector)
      ) ?? null
    )
  })

Here: https://github.com/tailwindlabs/headlessui/blob/6ac6930116c47cbd72402c33bfe327698c8e2780/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx#L271

What do you think? I can open a PR if I'm not missing something important here.