Ensure moving focus within a Portal component, does not close the Popover component (#2492)

* abstract resolving root containers to hook

This way we can reuse it in other components when needed.

* allow registering a `Portal` component to a parent

This allows us to find all the `Portal` components that are nested in a
given component without manually adding refs to every `Portal` component
itself.

This will come in handy in the `Popover` component where we will allow
focus in the child `Portal` components otherwise a focus outside of the
`Popover` will close the it. In other components we often crawl the DOM
directly using `[data-headlessui-portal]` data attributes, however this
will fetch _all_ the `Portal` components, not the ones that started in
the current component.

* allow focus in portalled containers

The `Popover` component will close by default if focus is moved outside
of it. However, if you use a `Portal` comopnent inside the
`Popover.Panel` then from a DOM perspective you are moving the focus
outside of the `Popover.Panel`. This prevents the closing, and allows
the focus into the `Portal`.

It currently only allows for `Portal` components that originated from
the `Popover` component. This means that if you open a `Dialog`
component from within the `Popover` component, the `Dialog` already
renders a `Portal` but since this is part of the `Dialog` and not the
`Popover` it will close the `Popover` when focus is moved to the
`Dialog` component.

* ensure `useNestedPortals` register/unregister with the parent

This ensures that if you have a structure like this:

```jsx
<Dialog> {/* Renders a portal internally */}
   <Popover>
      <Portal> {/* First level */}
         <Popover.Panel>
            <Menu>
               <Portal> {/* Second level */}
                  <Menu.Items>
                  {/* ... */}
                  </Menu.Items>
               </Portal>
            </Menu>
         </Popover.Panel>
      </Portal>
   </Popover>
</Dialog>
```

That working with the `Menu` doesn't close the `Popover` or the `Dialog`.

* cleanup `useRootContainers` hook

This will allow you to pass in portal elements as well. + cleanup of
the resolving of all DOM nodes.

* handle nested portals in `Dialog` component

* expose `contains` function from `useRootContainers`

Shorthand to check if any of the root containers contains the given
element.

* add tests to verify that actions in `Portal` components won't close the `Popover`

* update changelog

* re-order use-outside-click logic

To make it similar between React & Vue

* inject the `PortalWrapper` context in the correct spot

* ensure to forward the incoming `attrs`

Author: RobinMalfait

packages/@headlessui-react/CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Stop `<Transition appear>` from overwriting classes on re-render ([#2457](https://github.com/tailwindlabs/headlessui/pull/2457))
 - Improve control over `Menu` and `Listbox` options while searching ([#2471](https://github.com/tailwindlabs/headlessui/pull/2471))
 - Consider clicks inside iframes to be "outside" ([#2485](https://github.com/tailwindlabs/headlessui/pull/2485))
+- Ensure moving focus within a `Portal` component, does not close the `Popover` component ([#2492](https://github.com/tailwindlabs/headlessui/pull/2492))
 
 ### Changed
 

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

@@ -33,7 +33,7 @@ import { Keys } from '../keyboard'
 import { isDisabledReactIssue7711 } from '../../utils/bugs'
 import { useId } from '../../hooks/use-id'
 import { FocusTrap } from '../../components/focus-trap/focus-trap'
-import { Portal } from '../../components/portal/portal'
+import { Portal, useNestedPortals } from '../../components/portal/portal'
 import { ForcePortalRoot } from '../../internal/portal-force-root'
 import { ComponentDescription, Description, useDescriptions } from '../description/description'
 import { useOpenClosed, State } from '../../internal/open-closed'
@@ -42,10 +42,10 @@ import { StackProvider, StackMessage } from '../../internal/stack-context'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 import { useOwnerDocument } from '../../hooks/use-owner'
 import { useEventListener } from '../../hooks/use-event-listener'
-import { Hidden, Features as HiddenFeatures } from '../../internal/hidden'
 import { useEvent } from '../../hooks/use-event'
 import { useDocumentOverflowLockedEffect } from '../../hooks/document-overflow/use-document-overflow'
 import { useInert } from '../../hooks/use-inert'
+import { useRootContainers } from '../../hooks/use-root-containers'
 
 enum DialogStates {
   Open,
@@ -158,9 +158,6 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
   let internalDialogRef = useRef<HTMLDivElement | null>(null)
   let dialogRef = useSyncRefs(internalDialogRef, ref)
 
-  // Reference to a node in the "main" tree, not in the portalled Dialog tree.
-  let mainTreeNode = useRef<HTMLDivElement | null>(null)
-
   let ownerDocument = useOwnerDocument(internalDialogRef)
 
   // Validations
@@ -212,6 +209,15 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
   let enabled = ready ? (__demoMode ? false : dialogState === DialogStates.Open) : false
   let hasNestedDialogs = nestedDialogCount > 1 // 1 is the current dialog
   let hasParentDialog = useContext(DialogContext) !== null
+  let [portals, PortalWrapper] = useNestedPortals()
+  let {
+    resolveContainers: resolveRootContainers,
+    mainTreeNodeRef,
+    MainTreeNode,
+  } = useRootContainers({
+    portals,
+    defaultContainers: [state.panelRef.current ?? internalDialogRef.current],
+  })
 
   // If there are multiple dialogs, then you can be the root, the leaf or one
   // in between. We only care abou whether you are the top most one or not.
@@ -238,9 +244,9 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
       if (root.id === 'headlessui-portal-root') return false
 
       // Find the root of the main tree node
-      return root.contains(mainTreeNode.current) && root instanceof HTMLElement
+      return root.contains(mainTreeNodeRef.current) && root instanceof HTMLElement
     }) ?? null) as HTMLElement | null
-  }, [mainTreeNode])
+  }, [mainTreeNodeRef])
   useInert(resolveRootOfMainTreeNode, inertOthersEnabled)
 
   // This would mark the parent dialogs as inert
@@ -250,34 +256,18 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
   })()
   let resolveRootOfParentDialog = useCallback(() => {
     return (Array.from(ownerDocument?.querySelectorAll('[data-headlessui-portal]') ?? []).find(
-      (root) => root.contains(mainTreeNode.current) && root instanceof HTMLElement
+      (root) => root.contains(mainTreeNodeRef.current) && root instanceof HTMLElement
     ) ?? null) as HTMLElement | null
-  }, [mainTreeNode])
+  }, [mainTreeNodeRef])
   useInert(resolveRootOfParentDialog, inertParentDialogs)
 
-  let resolveRootContainers = useEvent(() => {
-    // Third party roots
-    let rootContainers = Array.from(
-      ownerDocument?.querySelectorAll('html > *, body > *, [data-headlessui-portal]') ?? []
-    ).filter((container) => {
-      if (container === document.body) return false // Skip `<body>`
-      if (container === document.head) return false // Skip `<head>`
-      if (!(container instanceof HTMLElement)) return false // Skip non-HTMLElements
-      if (container.contains(mainTreeNode.current)) return false // Skip if it is the main app
-      if (state.panelRef.current && container.contains(state.panelRef.current)) return false
-      return true // Keep
-    })
-
-    return [...rootContainers, state.panelRef.current ?? internalDialogRef.current] as HTMLElement[]
-  })
-
   // Close Dialog on outside click
   let outsideClickEnabled = (() => {
     if (!enabled) return false
     if (hasNestedDialogs) return false
     return true
   })()
-  useOutsideClick(() => resolveRootContainers(), close, outsideClickEnabled)
+  useOutsideClick(resolveRootContainers, close, outsideClickEnabled)
 
   // Handle `Escape` to close
   let escapeToCloseEnabled = (() => {
@@ -375,23 +365,25 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
                         : FocusTrap.features.None
                     }
                   >
-                    {render({
-                      ourProps,
-                      theirProps,
-                      slot,
-                      defaultTag: DEFAULT_DIALOG_TAG,
-                      features: DialogRenderFeatures,
-                      visible: dialogState === DialogStates.Open,
-                      name: 'Dialog',
-                    })}
+                    <PortalWrapper>
+                      {render({
+                        ourProps,
+                        theirProps,
+                        slot,
+                        defaultTag: DEFAULT_DIALOG_TAG,
+                        features: DialogRenderFeatures,
+                        visible: dialogState === DialogStates.Open,
+                        name: 'Dialog',
+                      })}
+                    </PortalWrapper>
                   </FocusTrap>
                 </DescriptionProvider>
               </ForcePortalRoot>
             </Portal.Group>
           </DialogContext.Provider>
         </Portal>
       </ForcePortalRoot>
-      <Hidden features={HiddenFeatures.Hidden} ref={mainTreeNode} />
+      <MainTreeNode />
     </StackProvider>
   )
 }

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

@@ -2623,6 +2623,80 @@ describe('Mouse interactions', () => {
       assertActiveElement(getPopoverButton())
     })
   )
+
+  it(
+    'should not close the Popover if the focus is moved outside of the Popover but still in the same React tree using Portals',
+    suppressConsoleLogs(async () => {
+      let clickFn = jest.fn()
+      render(
+        <Popover>
+          <Popover.Button>Toggle</Popover.Button>
+          <Popover.Panel>
+            <Portal>
+              <button onClick={clickFn}>foo</button>
+            </Portal>
+          </Popover.Panel>
+        </Popover>
+      )
+
+      // Open the popover
+      await click(getPopoverButton())
+
+      // Verify it is open
+      assertPopoverPanel({ state: PopoverState.Visible })
+
+      // Click the button outside the Popover (DOM) but inside (Portal / React tree)
+      await click(getByText('foo'))
+
+      // Verify it is still open
+      assertPopoverPanel({ state: PopoverState.Visible })
+
+      // Verify the button was clicked
+      expect(clickFn).toHaveBeenCalled()
+    })
+  )
+
+  it(
+    'should not close the Popover if the focus is moved outside of the Popover but still in the same React tree using nested Portals',
+    suppressConsoleLogs(async () => {
+      let clickFn = jest.fn()
+      render(
+        <Popover>
+          <Popover.Button>Toggle</Popover.Button>
+          <Popover.Panel>
+            Level 0
+            <Portal>
+              Level 1
+              <Portal>
+                Level 2
+                <Portal>
+                  Level 3
+                  <Portal>
+                    Level 4<button onClick={clickFn}>foo</button>
+                  </Portal>
+                </Portal>
+              </Portal>
+            </Portal>
+          </Popover.Panel>
+        </Popover>
+      )
+
+      // Open the popover
+      await click(getPopoverButton())
+
+      // Verify it is open
+      assertPopoverPanel({ state: PopoverState.Visible })
+
+      // Click the button outside the Popover (DOM) but inside (Portal / React tree)
+      await click(getByText('foo'))
+
+      // Verify it is still open
+      assertPopoverPanel({ state: PopoverState.Visible })
+
+      // Verify the button was clicked
+      expect(clickFn).toHaveBeenCalled()
+    })
+  )
 })
 
 describe('Nested popovers', () => {

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

@@ -54,6 +54,8 @@ import { useTabDirection, Direction as TabDirection } from '../../hooks/use-tab-
 import { microTask } from '../../utils/micro-task'
 import { useLatestValue } from '../../hooks/use-latest-value'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
+import { useRootContainers } from '../../hooks/use-root-containers'
+import { useNestedPortals } from '../../components/portal/portal'
 
 type MouseEvent<T> = Parameters<MouseEventHandler<T>>[0]
 
@@ -309,18 +311,26 @@ function PopoverFn<TTag extends ElementType = typeof DEFAULT_POPOVER_TAG>(
 
   useEffect(() => registerPopover?.(registerBag), [registerPopover, registerBag])
 
+  let [portals, PortalWrapper] = useNestedPortals()
+  let root = useRootContainers({
+    portals,
+    defaultContainers: [button, panel],
+  })
+
   // Handle focus out
   useEventListener(
     ownerDocument?.defaultView,
     'focus',
     (event) => {
+      if (event.target === window) return
+      if (!(event.target instanceof HTMLElement)) return
       if (popoverState !== PopoverStates.Open) return
       if (isFocusWithinPopoverGroup()) return
       if (!button) return
       if (!panel) return
-      if (event.target === window) return
-      if (beforePanelSentinel.current?.contains?.(event.target as HTMLElement)) return
-      if (afterPanelSentinel.current?.contains?.(event.target as HTMLElement)) return
+      if (root.contains(event.target)) return
+      if (beforePanelSentinel.current?.contains?.(event.target)) return
+      if (afterPanelSentinel.current?.contains?.(event.target)) return
 
       dispatch({ type: ActionTypes.ClosePopover })
     },
@@ -329,7 +339,7 @@ function PopoverFn<TTag extends ElementType = typeof DEFAULT_POPOVER_TAG>(
 
   // Handle outside click
   useOutsideClick(
-    [button, panel],
+    root.resolveContainers,
     (event, target) => {
       dispatch({ type: ActionTypes.ClosePopover })
 
@@ -385,13 +395,16 @@ function PopoverFn<TTag extends ElementType = typeof DEFAULT_POPOVER_TAG>(
               [PopoverStates.Closed]: State.Closed,
             })}
           >
-            {render({
-              ourProps,
-              theirProps,
-              slot,
-              defaultTag: DEFAULT_POPOVER_TAG,
-              name: 'Popover',
-            })}
+            <PortalWrapper>
+              {render({
+                ourProps,
+                theirProps,
+                slot,
+                defaultTag: DEFAULT_POPOVER_TAG,
+                name: 'Popover',
+              })}
+              <root.MainTreeNode />
+            </PortalWrapper>
           </OpenClosedProvider>
         </PopoverAPIContext.Provider>
       </PopoverContext.Provider>

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

@@ -10,6 +10,8 @@ import React, {
   ElementType,
   MutableRefObject,
   Ref,
+  useMemo,
+  ContextType,
 } from 'react'
 import { createPortal } from 'react-dom'
 
@@ -22,6 +24,7 @@ import { optionalRef, useSyncRefs } from '../../hooks/use-sync-refs'
 import { useOnUnmount } from '../../hooks/use-on-unmount'
 import { useOwnerDocument } from '../../hooks/use-owner'
 import { env } from '../../utils/env'
+import { useEvent } from '../../hooks/use-event'
 
 function usePortalTarget(ref: MutableRefObject<HTMLElement | null>): HTMLElement | null {
   let forceInRoot = usePortalRoot()
@@ -87,7 +90,7 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
   let [element] = useState<HTMLDivElement | null>(() =>
     env.isServer ? null : ownerDocument?.createElement('div') ?? null
   )
-
+  let parent = useContext(PortalParentContext)
   let ready = useServerHandoffComplete()
 
   useIsoMorphicEffect(() => {
@@ -101,6 +104,13 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     }
   }, [target, element])
 
+  useIsoMorphicEffect(() => {
+    if (!element) return
+    if (!parent) return
+
+    return parent.register(element)
+  }, [parent, element])
+
   useOnUnmount(() => {
     if (!target || !element) return
 
@@ -164,6 +174,45 @@ function GroupFn<TTag extends ElementType = typeof DEFAULT_GROUP_TAG>(
 
 // ---
 
+let PortalParentContext = createContext<{
+  register: (portal: HTMLElement) => () => void
+  unregister: (portal: HTMLElement) => void
+  portals: MutableRefObject<HTMLElement[]>
+} | null>(null)
+
+export function useNestedPortals() {
+  let parent = useContext(PortalParentContext)
+  let portals = useRef<HTMLElement[]>([])
+
+  let register = useEvent((portal: HTMLElement) => {
+    portals.current.push(portal)
+    if (parent) parent.register(portal)
+    return () => unregister(portal)
+  })
+
+  let unregister = useEvent((portal: HTMLElement) => {
+    let idx = portals.current.indexOf(portal)
+    if (idx !== -1) portals.current.splice(idx, 1)
+    if (parent) parent.unregister(portal)
+  })
+
+  let api = useMemo<ContextType<typeof PortalParentContext>>(
+    () => ({ register, unregister, portals }),
+    [register, unregister, portals]
+  )
+
+  return [
+    portals,
+    useMemo(() => {
+      return function PortalWrapper({ children }: { children: React.ReactNode }) {
+        return <PortalParentContext.Provider value={api}>{children}</PortalParentContext.Provider>
+      }
+    }, [api]),
+  ] as const
+}
+
+// ---
+
 interface ComponentPortal extends HasDisplayName {
   <TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     props: PortalProps<TTag> & RefProp<typeof PortalFn>