Improve dialog and server handoff

CHANGELOG.md

@@ -11,6 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Introduce Open/Closed state, to simplify component communication ([#466](https://github.com/tailwindlabs/headlessui/pull/466))
 
+### Fixes
+
+- Improve SSR for `Dialog` ([#477](https://github.com/tailwindlabs/headlessui/pull/477))
+- Delay focus trap initialization ([#477](https://github.com/tailwindlabs/headlessui/pull/477))
+
 ## [Unreleased - Vue]
 
 ### Added 

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

@@ -33,6 +33,7 @@ import { contains } from '../../internal/dom-containers'
 import { Description, useDescriptions } from '../description/description'
 import { useWindowEvent } from '../../hooks/use-window-event'
 import { useOpenClosed, State } from '../../internal/open-closed'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 
 enum DialogStates {
   Open,
@@ -235,7 +236,8 @@ let DialogRoot = forwardRefWithAs(function Dialog<
     return () => observer.disconnect()
   }, [dialogState, internalDialogRef, close])
 
-  let enabled = dialogState === DialogStates.Open
+  let ready = useServerHandoffComplete()
+  let enabled = ready && dialogState === DialogStates.Open
 
   useFocusTrap(containers, enabled, { initialFocus })
   useInertOthers(internalDialogRef, enabled)

packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx

@@ -9,6 +9,7 @@ import {
 import { Props } from '../../types'
 import { render } from '../../utils/render'
 import { useFocusTrap } from '../../hooks/use-focus-trap'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 
 let DEFAULT_FOCUS_TRAP_TAG = 'div' as const
 
@@ -18,7 +19,8 @@ export function FocusTrap<TTag extends ElementType = typeof DEFAULT_FOCUS_TRAP_T
   let containers = useRef<Set<HTMLElement>>(new Set())
   let { initialFocus, ...passthroughProps } = props
 
-  useFocusTrap(containers, true, { initialFocus })
+  let ready = useServerHandoffComplete()
+  useFocusTrap(containers, ready, { initialFocus })
 
   let propsWeControl = {
     ref(element: HTMLElement | null) {

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

@@ -16,6 +16,7 @@ import { render } from '../../utils/render'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { useElementStack, StackProvider } from '../../internal/stack-context'
 import { usePortalRoot } from '../../internal/portal-force-root'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 
 function usePortalTarget(): HTMLElement | null {
   let forceInRoot = usePortalRoot()
@@ -57,6 +58,8 @@ export function Portal<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     typeof window === 'undefined' ? null : document.createElement('div')
   )
 
+  let ready = useServerHandoffComplete()
+
   useElementStack(element)
 
   useIsoMorphicEffect(() => {
@@ -77,16 +80,14 @@ export function Portal<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
     }
   }, [target, element])
 
+  if (!ready) return null
+
   return (
     <StackProvider>
       {!target || !element
         ? null
         : createPortal(
-            render({
-              props: passthroughProps,
-              defaultTag: DEFAULT_PORTAL_TAG,
-              name: 'Portal',
-            }),
+            render({ props: passthroughProps, defaultTag: DEFAULT_PORTAL_TAG, name: 'Portal' }),
             element
           )}
     </StackProvider>

packages/@headlessui-react/src/components/transitions/transition.tsx

@@ -23,6 +23,7 @@ import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { Features, PropsForFeatures, render, RenderStrategy } from '../../utils/render'
 import { Reason, transition } from './utils/transition'
 import { OpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
+import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 
 type ID = ReturnType<typeof useId>
 
@@ -260,11 +261,13 @@ function TransitionChild<TTag extends ElementType = typeof DEFAULT_TRANSITION_CH
 
   let events = useEvents({ beforeEnter, afterEnter, beforeLeave, afterLeave })
 
+  let ready = useServerHandoffComplete()
+
   useEffect(() => {
-    if (state === TreeStates.Visible && container.current === null) {
+    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])
+  }, [container, state, ready])
 
   // Skipping initial transition
   let skip = initial && !appear

packages/@headlessui-react/src/hooks/use-focus-trap.ts

@@ -2,10 +2,10 @@ import {
   useRef,
   // Types
   MutableRefObject,
+  useEffect,
 } from 'react'
 
 import { Keys } from '../components/keyboard'
-import { useIsoMorphicEffect } from './use-iso-morphic-effect'
 import { focusElement, focusIn, Focus, FocusResult } from '../utils/focus-management'
 import { contains } from '../internal/dom-containers'
 import { useWindowEvent } from './use-window-event'
@@ -22,7 +22,7 @@ export function useFocusTrap(
   let mounted = useRef(false)
 
   // Handle initial focus
-  useIsoMorphicEffect(() => {
+  useEffect(() => {
     if (!enabled) return
     if (containers.current.size !== 1) return
 

packages/@headlessui-react/src/hooks/use-id.ts

@@ -1,28 +1,25 @@
-import { useState, useEffect } from 'react'
+import { useState } from 'react'
 import { useIsoMorphicEffect } from './use-iso-morphic-effect'
+import { useServerHandoffComplete } from './use-server-handoff-complete'
 
 // We used a "simple" approach first which worked for SSR and rehydration on the client. However we
 // didn't take care of the Suspense case. To fix this we used the approach the @reach-ui/auto-id
 // uses.
 //
 // Credits: https://github.com/reach/reach-ui/blob/develop/packages/auto-id/src/index.tsx
 
-let state = { serverHandoffComplete: false }
 let id = 0
 function generateId() {
   return ++id
 }
 
 export function useId() {
-  let [id, setId] = useState(state.serverHandoffComplete ? generateId : null)
+  let ready = useServerHandoffComplete()
+  let [id, setId] = useState(ready ? generateId : null)
 
   useIsoMorphicEffect(() => {
     if (id === null) setId(generateId())
   }, [id])
 
-  useEffect(() => {
-    if (state.serverHandoffComplete === false) state.serverHandoffComplete = true
-  }, [])
-
   return id != null ? '' + id : undefined
 }

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

@@ -0,0 +1,19 @@
+import { useState, useEffect } from 'react'
+
+let state = { serverHandoffComplete: false }
+
+export function useServerHandoffComplete() {
+  let [serverHandoffComplete, setServerHandoffComplete] = useState(state.serverHandoffComplete)
+
+  useEffect(() => {
+    if (serverHandoffComplete === true) return
+
+    setServerHandoffComplete(true)
+  }, [serverHandoffComplete])
+
+  useEffect(() => {
+    if (state.serverHandoffComplete === false) state.serverHandoffComplete = true
+  }, [])
+
+  return serverHandoffComplete
+}