Fix hydration of components inside <Suspense> (#2663)

* Add repro for suspense

wip

* Refactor to wildcard import

* Targeted fix for react 18 + suspense

* Update changelog

* Update types

* Add styling

* update styling

packages/@headlessui-react/CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fix incorrectly focused `Combobox.Input` component on page load ([#2654](https://github.com/tailwindlabs/headlessui/pull/2654))
 - Ensure `appear` works using the `Transition` component (even when used with SSR) ([#2646](https://github.com/tailwindlabs/headlessui/pull/2646))
 - Improve resetting values when using the `nullable` prop on the `Combobox` component ([#2660](https://github.com/tailwindlabs/headlessui/pull/2660))
+- Fix hydration of components inside `<Suspense>` ([#2663](https://github.com/tailwindlabs/headlessui/pull/2663))
 
 ## [1.7.16] - 2023-07-27
 

packages/@headlessui-react/src/hooks/use-server-handoff-complete.ts

@@ -1,8 +1,41 @@
-import { useState, useEffect } from 'react'
+import * as React from 'react'
 import { env } from '../utils/env'
 
+/**
+ * This is used to determine if we're hydrating in React 18.
+ *
+ * The `useServerHandoffComplete` hook doesn't work with `<Suspense>`
+ * because it assumes all hydration happens at one time during page load.
+ *
+ * Given that the problem only exists in React 18 we can rely
+ * on newer APIs to determine if hydration is happening.
+ */
+function useIsHydratingInReact18(): boolean {
+  let isServer = typeof document === 'undefined'
+
+  // React < 18 doesn't have any way to figure this out afaik
+  if (!('useSyncExternalStore' in React)) {
+    return false
+  }
+
+  // This weird pattern makes sure bundlers don't throw at build time
+  // because `useSyncExternalStore` isn't defined in React < 18
+  const useSyncExternalStore = ((r) => r.useSyncExternalStore)(React)
+
+  // @ts-ignore
+  let result = useSyncExternalStore(
+    () => () => {},
+    () => false,
+    () => (isServer ? false : true)
+  )
+
+  return result
+}
+
+// TODO: We want to get rid of this hook eventually
 export function useServerHandoffComplete() {
-  let [complete, setComplete] = useState(env.isHandoffComplete)
+  let isHydrating = useIsHydratingInReact18()
+  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
@@ -11,13 +44,17 @@ export function useServerHandoffComplete() {
     setComplete(false)
   }
 
-  useEffect(() => {
+  React.useEffect(() => {
     if (complete === true) return
     setComplete(true)
   }, [complete])
 
   // Transition from pending to complete (forcing a re-render when server rendering)
-  useEffect(() => env.handoff(), [])
+  React.useEffect(() => env.handoff(), [])
+
+  if (isHydrating) {
+    return false
+  }
 
   return complete
 }

packages/playground-react/pages/suspense/portal.tsx

@@ -0,0 +1,72 @@
+'use client'
+
+import { Portal } from '@headlessui/react'
+import { Suspense, lazy } from 'react'
+
+function MyComponent({ children }: { children(message: string): JSX.Element }) {
+  return <>{children('test')}</>
+}
+
+let MyComponentLazy = lazy(async () => {
+  await new Promise((resolve) => setTimeout(resolve, 4000))
+
+  return { default: MyComponent }
+})
+
+export default function Index() {
+  return (
+    <div>
+      <h1 className="p-8 text-3xl font-bold">Suspense + Portals</h1>
+
+      <Portal>
+        <div className="absolute top-24 right-48 z-10 flex h-32 w-32 flex-col items-center justify-center rounded border border-black/5 bg-white bg-clip-padding p-px shadow">
+          <div className="w-full rounded-t-sm bg-gray-100 p-1 text-center text-gray-700">
+            Instant
+          </div>
+          <div className="flex w-full flex-1 items-center justify-center text-3xl font-bold text-gray-400">
+            1
+          </div>
+        </div>
+      </Portal>
+      <Portal>
+        <div className="absolute top-24 right-8 z-10 flex h-32 w-32 flex-col items-center justify-center rounded border border-black/5 bg-white bg-clip-padding p-px shadow">
+          <div className="w-full rounded-t-sm bg-gray-100 p-1 text-center text-gray-700">
+            Instant
+          </div>
+          <div className="flex w-full flex-1 items-center justify-center text-3xl font-bold text-gray-400">
+            2
+          </div>
+        </div>
+      </Portal>
+
+      <Suspense fallback={<span>Loading ...</span>}>
+        <MyComponentLazy>
+          {(env) => (
+            <div>
+              <Portal>
+                <div className="absolute top-64 right-48 z-10 flex h-32 w-32 flex-col items-center justify-center rounded border border-black/5 bg-white bg-clip-padding p-px shadow">
+                  <div className="w-full rounded-t-sm bg-gray-100 p-1 text-center text-gray-700">
+                    Suspense
+                  </div>
+                  <div className="flex w-full flex-1 items-center justify-center text-3xl font-bold text-gray-400">
+                    {env} 1
+                  </div>
+                </div>
+              </Portal>
+              <Portal>
+                <div className="absolute top-64 right-8 z-10 flex h-32 w-32 flex-col items-center justify-center rounded border border-black/5 bg-white bg-clip-padding p-px shadow">
+                  <div className="w-full rounded-t-sm bg-gray-100 p-1 text-center text-gray-700">
+                    Suspense
+                  </div>
+                  <div className="flex w-full flex-1 items-center justify-center text-3xl font-bold text-gray-400">
+                    {env} 2
+                  </div>
+                </div>
+              </Portal>
+            </div>
+          )}
+        </MyComponentLazy>
+      </Suspense>
+    </div>
+  )
+}