Improve dialog and SSR (#477)

* delay initialization of Dialog

We were using a useLayoutEffect, now let's use a useEffect instead. It
still moves focus to the correct element, but that process is now a bit
delayed. This means that users will less-likely be urged to "hack"
around the issue by using fake focusable elements which will result in
worse accessibility.

* add hook to deal with server handoff

This will allow us to delay certain features. For example we can delay
the focus trapping until it is fully hydrated. We can also delay
rendering the Portal to ensure hydration works correctly.

* use server handoff complete hook

* update changelog

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
+}