From 35012893e959c1c7ec86b1e25a1343c890b46144 Mon Sep 17 00:00:00 2001
From: Jordan Pittman <jordan@cryptica.me>
Date: Wed, 2 Aug 2023 10:20:22 -0400
Subject: [PATCH] Fix hydration of components inside `<Suspense>` (#2633)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Only use useServerHandoffComplete in React < 18

It’s only useful for the useId hook. It is not compatible with `<Suspense>` because hydration is delayed then.

* Make sure portals first render matches the server by rendering nothing

Since Portals cannot SSR the first render MUST also return `null`. React really needs an `isHydrating` API.

* Lazily resolve root containers

This fixes a problem where clicks were assumed to be outside because of the delayed `<Portal>` render. The second portal render doesn’t cause the dialog to re-render thus the initial ref values were stale.

* Update changelog

---------

Co-authored-by: Robin Malfait <malfait.robin@gmail.com>
---
 packages/@headlessui-react/CHANGELOG.md       |  1 +
 .../src/components/dialog/dialog.tsx          |  6 +-
 .../src/components/focus-trap/focus-trap.tsx  |  5 --
 .../src/components/portal/portal.tsx          |  6 +-
 .../transitions/transition.test.tsx           | 73 -------------------
 .../src/components/transitions/transition.tsx | 13 ----
 .../@headlessui-react/src/hooks/use-id.ts     | 24 +++++-
 .../src/hooks/use-root-containers.tsx         |  9 ++-
 .../src/hooks/use-server-handoff-complete.ts  | 23 ------
 9 files changed, 37 insertions(+), 123 deletions(-)
 delete mode 100644 packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts

diff --git a/packages/@headlessui-react/CHANGELOG.md b/packages/@headlessui-react/CHANGELOG.md
index 5537504bf..63b1e7262 100644
--- a/packages/@headlessui-react/CHANGELOG.md
+++ b/packages/@headlessui-react/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Use correct value when resetting `<Listbox multiple>` and `<Combobox multiple>` ([#2626](https://github.com/tailwindlabs/headlessui/pull/2626))
 - Render `<MainTreeNode />` in `Popover.Group` component only ([#2634](https://github.com/tailwindlabs/headlessui/pull/2634))
+- Fix hydration of components inside `<Suspense>` ([#2633](https://github.com/tailwindlabs/headlessui/pull/2633))
 
 ## [1.7.16] - 2023-07-27
 
diff --git a/packages/@headlessui-react/src/components/dialog/dialog.tsx b/packages/@headlessui-react/src/components/dialog/dialog.tsx
index d1909e0de..0463b15d7 100644
--- a/packages/@headlessui-react/src/components/dialog/dialog.tsx
+++ b/packages/@headlessui-react/src/components/dialog/dialog.tsx
@@ -37,7 +37,6 @@ 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'
-import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { StackProvider, StackMessage } from '../../internal/stack-context'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 import { useOwnerDocument } from '../../hooks/use-owner'
@@ -205,8 +204,7 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
 
   let setTitleId = useEvent((id: string | null) => dispatch({ type: ActionTypes.SetTitleId, id }))
 
-  let ready = useServerHandoffComplete()
-  let enabled = ready ? (__demoMode ? false : dialogState === DialogStates.Open) : false
+  let enabled = __demoMode ? false : dialogState === DialogStates.Open
   let hasNestedDialogs = nestedDialogCount > 1 // 1 is the current dialog
   let hasParentDialog = useContext(DialogContext) !== null
   let [portals, PortalWrapper] = useNestedPortals()
@@ -216,7 +214,7 @@ function DialogFn<TTag extends ElementType = typeof DEFAULT_DIALOG_TAG>(
     MainTreeNode,
   } = useRootContainers({
     portals,
-    defaultContainers: [state.panelRef.current ?? internalDialogRef.current],
+    defaultContainers: [() => state.panelRef.current ?? internalDialogRef.current],
   })
 
   // If there are multiple dialogs, then you can be the root, the leaf or one
diff --git a/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx b/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx
index 88a2362fd..3b816bc0f 100644
--- a/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx
+++ b/packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx
@@ -10,7 +10,6 @@ import React, {
 
 import { Props } from '../../types'
 import { forwardRefWithAs, HasDisplayName, RefProp, render } from '../../utils/render'
-import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
 import { Features as HiddenFeatures, Hidden } from '../../internal/hidden'
 import { focusElement, focusIn, Focus, FocusResult } from '../../utils/focus-management'
@@ -82,10 +81,6 @@ function FocusTrapFn<TTag extends ElementType = typeof DEFAULT_FOCUS_TRAP_TAG>(
   let focusTrapRef = useSyncRefs(container, ref)
   let { initialFocus, containers, features = Features.All, ...theirProps } = props
 
-  if (!useServerHandoffComplete()) {
-    features = Features.None
-  }
-
   let ownerDocument = useOwnerDocument(container)
 
   useRestoreFocus({ ownerDocument }, Boolean(features & Features.RestoreFocus))
diff --git a/packages/@headlessui-react/src/components/portal/portal.tsx b/packages/@headlessui-react/src/components/portal/portal.tsx
index 77d854939..3b6607184 100644
--- a/packages/@headlessui-react/src/components/portal/portal.tsx
+++ b/packages/@headlessui-react/src/components/portal/portal.tsx
@@ -19,7 +19,6 @@ import { Props } from '../../types'
 import { forwardRefWithAs, RefProp, HasDisplayName, render } from '../../utils/render'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { usePortalRoot } from '../../internal/portal-force-root'
-import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { optionalRef, useSyncRefs } from '../../hooks/use-sync-refs'
 import { useOnUnmount } from '../../hooks/use-on-unmount'
 import { useOwnerDocument } from '../../hooks/use-owner'
@@ -91,7 +90,6 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     env.isServer ? null : ownerDocument?.createElement('div') ?? null
   )
   let parent = useContext(PortalParentContext)
-  let ready = useServerHandoffComplete()
 
   useIsoMorphicEffect(() => {
     if (!target || !element) return
@@ -123,7 +121,9 @@ function PortalFn<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     }
   })
 
-  if (!ready) return null
+  let [isFirstRender, setIsFirstRender] = useState(true)
+  useEffect(() => setIsFirstRender(false), [])
+  if (isFirstRender) return null
 
   let ourProps = { ref: portalRef }
 
diff --git a/packages/@headlessui-react/src/components/transitions/transition.test.tsx b/packages/@headlessui-react/src/components/transitions/transition.test.tsx
index 9ee464201..61a3562c0 100644
--- a/packages/@headlessui-react/src/components/transitions/transition.test.tsx
+++ b/packages/@headlessui-react/src/components/transitions/transition.test.tsx
@@ -155,27 +155,6 @@ describe('Setup API', () => {
         </span>
       `)
     })
-
-    it(
-      'should yell at us when we forget to forward the ref when using a render prop',
-      suppressConsoleLogs(() => {
-        expect.assertions(1)
-
-        function Dummy(props: any) {
-          return <span {...props}>Children</span>
-        }
-
-        expect(() => {
-          render(
-            <Transition show={true} as={Fragment}>
-              {() => <Dummy />}
-            </Transition>
-          )
-        }).toThrowErrorMatchingInlineSnapshot(
-          `"Did you forget to passthrough the \`ref\` to the actual DOM node?"`
-        )
-      })
-    )
   })
 
   describe('nested', () => {
@@ -349,58 +328,6 @@ describe('Setup API', () => {
         </div>
       `)
     })
-
-    it(
-      'should yell at us when we forgot to forward the ref on one of the Transition.Child components',
-      suppressConsoleLogs(() => {
-        expect.assertions(1)
-
-        function Dummy(props: any) {
-          return <div {...props} />
-        }
-
-        expect(() => {
-          render(
-            <div className="My Page">
-              <Transition show={true}>
-                <Transition.Child as={Fragment}>{() => <Dummy>Sidebar</Dummy>}</Transition.Child>
-                <Transition.Child as={Fragment}>{() => <Dummy>Content</Dummy>}</Transition.Child>
-              </Transition>
-            </div>
-          )
-        }).toThrowErrorMatchingInlineSnapshot(
-          `"Did you forget to passthrough the \`ref\` to the actual DOM node?"`
-        )
-      })
-    )
-
-    it(
-      'should yell at us when we forgot to forward a ref on the Transition component',
-      suppressConsoleLogs(() => {
-        expect.assertions(1)
-
-        function Dummy(props: any) {
-          return <div {...props} />
-        }
-
-        expect(() => {
-          render(
-            <div className="My Page">
-              <Transition show={true} as={Fragment}>
-                {() => (
-                  <Dummy>
-                    <Transition.Child>{() => <aside>Sidebar</aside>}</Transition.Child>
-                    <Transition.Child>{() => <section>Content</section>}</Transition.Child>
-                  </Dummy>
-                )}
-              </Transition>
-            </div>
-          )
-        }).toThrowErrorMatchingInlineSnapshot(
-          `"Did you forget to passthrough the \`ref\` to the actual DOM node?"`
-        )
-      })
-    )
   })
 
   describe('transition classes', () => {
diff --git a/packages/@headlessui-react/src/components/transitions/transition.tsx b/packages/@headlessui-react/src/components/transitions/transition.tsx
index 416ff2350..e07425ae2 100644
--- a/packages/@headlessui-react/src/components/transitions/transition.tsx
+++ b/packages/@headlessui-react/src/components/transitions/transition.tsx
@@ -27,7 +27,6 @@ import { match } from '../../utils/match'
 import { useIsMounted } from '../../hooks/use-is-mounted'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { useLatestValue } from '../../hooks/use-latest-value'
-import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
 import { useTransition } from '../../hooks/use-transition'
 import { useEvent } from '../../hooks/use-event'
@@ -348,19 +347,10 @@ function TransitionChildFn<TTag extends ElementType = typeof DEFAULT_TRANSITION_
     afterLeave,
   })
 
-  let ready = useServerHandoffComplete()
-
-  useEffect(() => {
-    if (ready && state === TreeStates.Visible && container.current === null) {
-      throw new Error('Did you forget to passthrough the `ref` to the actual DOM node?')
-    }
-  }, [container, state, ready])
-
   // Skipping initial transition
   let skip = initial && !appear
 
   let transitionDirection = (() => {
-    if (!ready) return 'idle'
     if (skip) return 'idle'
     if (prevShow.current === show) return 'idle'
     return show ? 'enter' : 'leave'
@@ -480,9 +470,6 @@ function TransitionRootFn<TTag extends ElementType = typeof DEFAULT_TRANSITION_C
   let internalTransitionRef = useRef<HTMLElement | null>(null)
   let transitionRef = useSyncRefs(internalTransitionRef, ref)
 
-  // The TransitionChild will also call this hook, and we have to make sure that we are ready.
-  useServerHandoffComplete()
-
   let usesOpenClosedState = useOpenClosed()
 
   if (show === undefined && usesOpenClosedState !== null) {
diff --git a/packages/@headlessui-react/src/hooks/use-id.ts b/packages/@headlessui-react/src/hooks/use-id.ts
index 761d6c1c3..c4ddb9fc3 100644
--- a/packages/@headlessui-react/src/hooks/use-id.ts
+++ b/packages/@headlessui-react/src/hooks/use-id.ts
@@ -1,6 +1,5 @@
 import React from 'react'
 import { useIsoMorphicEffect } from './use-iso-morphic-effect'
-import { useServerHandoffComplete } from './use-server-handoff-complete'
 import { env } from '../utils/env'
 
 // We used a "simple" approach first which worked for SSR and rehydration on the client. However we
@@ -23,3 +22,26 @@ export let useId =
 
     return id != null ? '' + id : undefined
   }
+
+// NOTE: Do NOT use this outside of the `useId` hook
+// It is not compatible with `<Suspense>` (which is in React 18 which has its own `useId` hook)
+function useServerHandoffComplete() {
+  let [complete, setComplete] = React.useState(env.isHandoffComplete)
+
+  if (complete && env.isHandoffComplete === false) {
+    // This means we are in a test environment and we need to reset the handoff state
+    // This kinda breaks the rules of React but this is only used for testing purposes
+    // And should theoretically be fine
+    setComplete(false)
+  }
+
+  React.useEffect(() => {
+    if (complete === true) return
+    setComplete(true)
+  }, [complete])
+
+  // Transition from pending to complete (forcing a re-render when server rendering)
+  React.useEffect(() => env.handoff(), [])
+
+  return complete
+}
diff --git a/packages/@headlessui-react/src/hooks/use-root-containers.tsx b/packages/@headlessui-react/src/hooks/use-root-containers.tsx
index e3b8d3430..cb438558b 100644
--- a/packages/@headlessui-react/src/hooks/use-root-containers.tsx
+++ b/packages/@headlessui-react/src/hooks/use-root-containers.tsx
@@ -3,12 +3,15 @@ import { Hidden, Features as HiddenFeatures } from '../internal/hidden'
 import { useEvent } from './use-event'
 import { useOwnerDocument } from './use-owner'
 
+type Container = HTMLElement | null | MutableRefObject<HTMLElement | null>
+type MaybeContainerFn = Container | (() => Container)
+
 export function useRootContainers({
   defaultContainers = [],
   portals,
   mainTreeNodeRef: _mainTreeNodeRef,
 }: {
-  defaultContainers?: (HTMLElement | null | MutableRefObject<HTMLElement | null>)[]
+  defaultContainers?: MaybeContainerFn[]
   portals?: MutableRefObject<HTMLElement[]>
   mainTreeNodeRef?: MutableRefObject<HTMLElement | null>
 } = {}) {
@@ -21,6 +24,10 @@ export function useRootContainers({
 
     // Resolve default containers
     for (let container of defaultContainers) {
+      if (typeof container === 'function') {
+        container = container()
+      }
+
       if (container === null) continue
       if (container instanceof HTMLElement) {
         containers.push(container)
diff --git a/packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts b/packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts
deleted file mode 100644
index b931ea2d6..000000000
--- a/packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts
+++ /dev/null
@@ -1,23 +0,0 @@
-import { useState, useEffect } from 'react'
-import { env } from '../utils/env'
-
-export function useServerHandoffComplete() {
-  let [complete, setComplete] = useState(env.isHandoffComplete)
-
-  if (complete && env.isHandoffComplete === false) {
-    // This means we are in a test environment and we need to reset the handoff state
-    // This kinda breaks the rules of React but this is only used for testing purposes
-    // And should theoretically be fine
-    setComplete(false)
-  }
-
-  useEffect(() => {
-    if (complete === true) return
-    setComplete(true)
-  }, [complete])
-
-  // Transition from pending to complete (forcing a re-render when server rendering)
-  useEffect(() => env.handoff(), [])
-
-  return complete
-}