@remix-run/ui 0.7.0 → 0.9.0

package/src/runtime/frame.ts

@@ -10,7 +10,7 @@ import { createRangeRoot, createRoot } from './vdom.ts'
 import { diffNodes } from './diff-dom.ts'
 import { createStyleManager, type StyleManager } from '../style/index.ts'
 import { findFlushMarker, type FlushKind } from './stream-protocol.ts'
-import { getDocumentModulePreloader } from './module-preloader.ts'
+import { getDocumentModulePreloader, type ProcessClientEntryPreloads } from './module-preloader.ts'
 import { unwrapFrameResolution } from './frame-resolution.ts'
 import {
   disposeClientEntryBoundary,
@@ -18,6 +18,8 @@ import {
   setClientEntryBoundaryOwner,
   type ClientEntryIdentity,
 } from './client-entry-boundary.ts'
+import { getDocumentImportMapManager } from './import-map-manager.ts'
+import { reloadDocument } from './document-reload.ts'
 
 type FrameRoot = [Comment, Comment] | Element | Document | DocumentFragment
 
@@ -104,6 +106,12 @@ type FrameReloadResult = {
   redirectedTo?: string
 }
 
+type FrameReloadTransition = {
+  signal: AbortSignal
+  committed: Promise<void>
+  finished: Promise<FrameReloadResult>
+}
+
 type FrameTemplateListener = (fragment: DocumentFragment) => void
 
 const bufferedFrameTemplates = new Map<string, DocumentFragment[]>()
@@ -170,8 +178,15 @@ export type FrameRuntime = {
   moduleLoads: Map<string, Promise<ElementFunction | undefined>>
   frameInstances: WeakMap<Comment, Frame>
   namedFrames: Map<string, FrameHandle>
-  serverFrameReload: { signal: AbortSignal } | undefined
-  reloadForNavigation?: (options?: FrameReloadOptions) => Promise<FrameReloadResult>
+  processClientEntryPreloads?: ProcessClientEntryPreloads
+  serverFrameReload:
+    | {
+        signal: AbortSignal
+        reconciliationTracker?: ReconciliationTracker
+        blockingFrameTracker?: ReconciliationTracker
+      }
+    | undefined
+  reloadForNavigation?: (options?: FrameReloadOptions) => FrameReloadTransition
 }
 
 export function isFrameRuntime(value: unknown): value is FrameRuntime {
@@ -188,7 +203,7 @@ export function isFrameRuntime(value: unknown): value is FrameRuntime {
 export function reloadFrameForNavigation(
   frame: FrameHandle,
   options?: FrameReloadOptions,
-): Promise<FrameReloadResult> {
+): FrameReloadTransition {
   let runtime = frame.$runtime
   invariant(isFrameRuntime(runtime), 'Expected a frame runtime')
   let reload = runtime.reloadForNavigation
@@ -210,11 +225,14 @@ export type FrameContext = {
   moduleLoads: Map<string, Promise<ElementFunction | undefined>>
   frameInstances: WeakMap<Comment, Frame>
   namedFrames: Map<string, FrameHandle>
+  processClientEntryPreloads?: ProcessClientEntryPreloads
   lifecycleSignal: AbortSignal
   regionTailRef?: ChildNode | null
   regionParent?: ParentNode | null
   signal?: AbortSignal
-  isActiveModulePreload?: (node: Node) => boolean
+  shouldPreserveHeadNode?: (node: Node) => boolean
+  reconciliationTracker?: ReconciliationTracker
+  blockingFrameTracker?: ReconciliationTracker
 }
 
 type FrameInit = {
@@ -233,6 +251,7 @@ type FrameInit = {
   moduleLoads: Map<string, Promise<ElementFunction | undefined>>
   frameInstances: WeakMap<Comment, Frame>
   namedFrames: Map<string, FrameHandle>
+  processClientEntryPreloads?: ProcessClientEntryPreloads
 }
 
 export type Frame = {
@@ -259,8 +278,11 @@ export type Frame = {
 }
 
 type RenderOptions = {
+  documentHref?: string
   flushKind?: FlushKind
-  initialHydrationTracker?: InitialHydrationTracker
+  reconciliationTracker?: ReconciliationTracker
+  blockingFrameTracker?: ReconciliationTracker
+  onCommit?: () => void
   signal?: AbortSignal
   contentStatus?: 'pending' | 'resolved'
   data?: RmxData
@@ -281,6 +303,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
   let reloadKind: 'direct' | 'ancestor' | undefined
   let styleManager = init.styleManager ?? createStyleManager()
   let modulePreloader = getDocumentModulePreloader(container.doc)
+  let importMapManager = getDocumentImportMapManager(container.doc)
   let currentMarker = init.marker
   let displayedContentStatus: 'pending' | 'resolved' = init.marker?.status ?? 'resolved'
   let pendingTemplateMarkerId: string | undefined
@@ -291,16 +314,42 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
   let disposed = false
   let lifecycleController = new AbortController()
 
+  async function consumeClientEntryResources(
+    source: ParentNode,
+    documentHref?: string,
+  ): Promise<boolean> {
+    let importMapStatus = importMapManager.consumeImportMaps(source)
+    if (importMapStatus !== 'ready') {
+      lifecycleController.abort()
+      if (importMapStatus === 'conflict') reloadDocument(container.doc, documentHref)
+      return false
+    }
+    await modulePreloader.consumePreloadLinks(source, init.processClientEntryPreloads)
+    return true
+  }
+
+  let initialClientEntryResources: Promise<boolean> | undefined
+  function shouldPreserveManagedHeadNode(node: Node): boolean {
+    return (
+      importMapManager.shouldPreserveHeadNode(node) ||
+      (modulePreloader.hasActivePreloads() && modulePreloader.isActivePreload(node))
+    )
+  }
+
   if (isDocumentNode(container.root)) {
     modulePreloader.adoptInitialPreloadLinks(container.root)
   } else {
-    modulePreloader.consumePreloadLinks(container.root)
+    initialClientEntryResources = consumeClientEntryResources(container.root)
   }
 
   // Merge any rmx-data found in the current document once at startup.
   mergeRmxDataFromDocument(init.data, container.doc)
 
-  let runtime = createFrameRuntime({ ...init, styleManager, reloadForNavigation: reload })
+  let runtime = createFrameRuntime({
+    ...init,
+    styleManager,
+    reloadForNavigation: startReloadTransition,
+  })
 
   let frame = createFrameHandle({
     src: init.src,
@@ -331,6 +380,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
     moduleLoads: init.moduleLoads,
     frameInstances: init.frameInstances,
     namedFrames: init.namedFrames,
+    processClientEntryPreloads: init.processClientEntryPreloads,
     lifecycleSignal: lifecycleController.signal,
     regionTailRef: container.regionTailRef,
     regionParent: container.regionParent,
@@ -385,7 +435,24 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
       }
 
       if (isRenderAborted(options.signal)) return
-      contentRoot.render(content)
+      let previousServerFrameReload = runtime.serverFrameReload
+      if (options.signal) {
+        runtime.serverFrameReload = {
+          signal: options.signal,
+          reconciliationTracker: options.reconciliationTracker,
+          blockingFrameTracker: options.blockingFrameTracker,
+        }
+      }
+
+      try {
+        contentRoot.render(content)
+        await new Promise<void>((resolve) => context.scheduler.enqueueCommitPhase([resolve]))
+        options.onCommit?.()
+      } finally {
+        runtime.serverFrameReload = previousServerFrameReload
+      }
+
+      if (isRenderAborted(options.signal)) return
       displayedContentStatus = options.contentStatus ?? 'resolved'
       return
     }
@@ -416,10 +483,16 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
 
     if (isFullDocumentReload && htmlContent !== undefined) {
       let parsed = new DOMParser().parseFromString(htmlContent, 'text/html')
-      modulePreloader.consumePreloadLinks(parsed)
+      if (!(await consumeClientEntryResources(parsed, options.documentHref))) return
+      if (isRenderAborted(options.signal)) return
       let responseData = options.data
       mergeRmxDataFromDocument(responseData, parsed)
-      let responseContext = { ...context, data: responseData }
+      let responseContext = {
+        ...context,
+        data: responseData,
+        reconciliationTracker: options.reconciliationTracker,
+        blockingFrameTracker: options.blockingFrameTracker,
+      }
       context.styleManager.adoptServerStyles(
         collectFrameServerStyleTags(createElementContainer(parsed)),
       )
@@ -431,9 +504,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
         regionParent: container.doc.documentElement,
         regionTailRef: null,
         signal: options.signal,
-        isActiveModulePreload: modulePreloader.hasActivePreloads()
-          ? modulePreloader.isActivePreload
-          : undefined,
+        shouldPreserveHeadNode: shouldPreserveManagedHeadNode,
       })
       diffNodes([container.doc.body], [parsed.body], {
         ...responseContext,
@@ -447,10 +518,12 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
       scheduleHydrationInContainer(
         bodyContainer,
         responseContext,
-        options.initialHydrationTracker,
+        options.reconciliationTracker,
         options.signal,
       )
-      await createSubFrames(bodyContainer.childNodes, responseContext, options)
+      let subFramesReady = createSubFrames(bodyContainer.childNodes, responseContext, options)
+      options.onCommit?.()
+      await subFramesReady
       if (isRenderAborted(options.signal)) return
       displayedContentStatus = options.contentStatus ?? 'resolved'
       return
@@ -458,14 +531,20 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
 
     let fragment =
       htmlContent !== undefined ? createFragmentFromString(container.doc, htmlContent) : content
-    modulePreloader.consumePreloadLinks(fragment)
+    if (!(await consumeClientEntryResources(fragment, options.documentHref))) return
+    if (isRenderAborted(options.signal)) return
     context.styleManager.adoptServerStyles(
       collectFrameServerStyleTags(createElementContainer(fragment)),
     )
     removeEmptyHeads(fragment)
     let responseData = options.data
     mergeRmxDataFromFragment(responseData, fragment)
-    let responseContext = { ...context, data: responseData }
+    let responseContext = {
+      ...context,
+      data: responseData,
+      reconciliationTracker: options.reconciliationTracker,
+      blockingFrameTracker: options.blockingFrameTracker,
+    }
 
     let nextContainer = createContainer(fragment)
 
@@ -478,14 +557,15 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
       signal: options.signal,
     })
 
-    if (isRenderAborted(options.signal)) return
     scheduleHydrationInContainer(
       container,
       responseContext,
-      options.initialHydrationTracker,
+      options.reconciliationTracker,
       options.signal,
     )
-    await createSubFrames(container.childNodes, responseContext, options)
+    let subFramesReady = createSubFrames(container.childNodes, responseContext, options)
+    options.onCommit?.()
+    await subFramesReady
     if (isRenderAborted(options.signal)) return
     displayedContentStatus = options.contentStatus ?? 'resolved'
   }
@@ -532,11 +612,13 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
   }
 
   async function hydrateInitial(): Promise<void> {
-    let initialHydrationTracker = createInitialHydrationTracker()
+    let reconciliationTracker = createReconciliationTracker()
 
+    if ((await initialClientEntryResources) === false) return
+    if (disposed || context.lifecycleSignal.aborted) return
     context.styleManager.adoptServerStyles(collectFrameServerStyleTags(container))
     let subFramesReady = createSubFrames(container.childNodes, context)
-    scheduleHydrationInContainer(container, context, initialHydrationTracker)
+    scheduleHydrationInContainer(container, context, reconciliationTracker)
 
     try {
       await subFramesReady
@@ -544,11 +626,11 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
       if (disposed || context.lifecycleSignal.aborted) return
 
       if (currentMarker?.status === 'pending') {
-        await watchPendingFrameTemplate(currentMarker, initialHydrationTracker)
+        await watchPendingFrameTemplate(currentMarker, reconciliationTracker)
       }
 
-      initialHydrationTracker.finalize()
-      await initialHydrationTracker.ready()
+      reconciliationTracker.finalize()
+      await reconciliationTracker.ready()
     } finally {
       clearRmxData(context.data)
     }
@@ -615,7 +697,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
     if (marker.status === 'pending') {
       await watchPendingFrameTemplate(
         marker,
-        options?.initialHydrationTracker,
+        options?.reconciliationTracker,
         options?.signal,
         isInheritedReload
           ? () => {
@@ -659,8 +741,25 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
   }
 
   async function reload(options?: FrameReloadOptions): Promise<FrameReloadResult> {
+    let transition = startReloadTransition(options)
+    void transition.committed.catch(() => {})
+    return await transition.finished
+  }
+
+  function startReloadTransition(options?: FrameReloadOptions): FrameReloadTransition {
     let controller = startReload(options?.signal)
-    return await resolveAndRenderReload(controller, options)
+    let committed = Promise.withResolvers<void>()
+    let commitStarted = false
+    let finished = resolveAndRenderReload(controller, options, (ready) => {
+      if (commitStarted) return
+      commitStarted = true
+      void ready.then(committed.resolve, committed.reject)
+    })
+
+    // Settle committed when a reload is aborted or fails before rendering any content.
+    void finished.then(() => committed.resolve(), committed.reject)
+
+    return { signal: controller.signal, committed: committed.promise, finished }
   }
 
   function startReload(signal?: AbortSignal): AbortController {
@@ -722,6 +821,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
   async function resolveAndRenderReload(
     controller: AbortController,
     options?: FrameReloadOptions,
+    resolveCommit?: (ready: Promise<void>) => void,
   ): Promise<FrameReloadResult> {
     try {
       let resolution = await init.resolveFrame(frame.src, {
@@ -736,7 +836,23 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
       if (reloadController !== controller || controller.signal.aborted) {
         return { signal: controller.signal }
       }
-      await render(content, { signal: controller.signal })
+      let reconciliationTracker = createReconciliationTracker()
+      let blockingFrameTracker = createReconciliationTracker()
+      let commitStarted = false
+      await render(content, {
+        documentHref: isDocumentNode(container.root) ? (redirectedTo ?? frame.src) : undefined,
+        signal: controller.signal,
+        reconciliationTracker,
+        blockingFrameTracker,
+        onCommit() {
+          if (commitStarted) return
+          commitStarted = true
+          blockingFrameTracker.finalize()
+          resolveCommit?.(blockingFrameTracker.ready())
+        },
+      })
+      reconciliationTracker.finalize()
+      await reconciliationTracker.ready()
       return {
         signal: controller.signal,
         redirectedTo:
@@ -825,7 +941,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
 
   async function watchPendingFrameTemplate(
     marker: FrameMarkerData,
-    initialHydrationTracker?: InitialHydrationTracker,
+    reconciliationTracker?: ReconciliationTracker,
     signal?: AbortSignal,
     onResolved?: () => void,
   ): Promise<void> {
@@ -838,7 +954,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
     let early = consumeFrameTemplate(marker.id) ?? getEarlyFrameContent(marker.id)
     if (early) {
       clearPendingFrameTemplateWatch()
-      await render(early, { initialHydrationTracker, signal, contentStatus: 'resolved' })
+      await render(early, { reconciliationTracker, signal, contentStatus: 'resolved' })
       if (!disposed && !context.lifecycleSignal.aborted && !signal?.aborted) onResolved?.()
       return
     }
@@ -872,7 +988,7 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
     let buffered = consumeFrameTemplate(marker.id)
     if (buffered) {
       clearPendingFrameTemplateWatch()
-      await render(buffered, { initialHydrationTracker, signal, contentStatus: 'resolved' })
+      await render(buffered, { reconciliationTracker, signal, contentStatus: 'resolved' })
       if (!disposed && !context.lifecycleSignal.aborted && !signal?.aborted) onResolved?.()
     }
   }
@@ -890,7 +1006,8 @@ export function createFrameRuntime(init: {
   moduleLoads: Map<string, Promise<ElementFunction | undefined>>
   frameInstances: WeakMap<Comment, Frame>
   namedFrames: Map<string, FrameHandle>
-  reloadForNavigation?: (options?: FrameReloadOptions) => Promise<FrameReloadResult>
+  processClientEntryPreloads?: ProcessClientEntryPreloads
+  reloadForNavigation?: (options?: FrameReloadOptions) => FrameReloadTransition
 }): FrameRuntime {
   return {
     [FRAME_RUNTIME]: true,
@@ -905,47 +1022,66 @@ export function createFrameRuntime(init: {
     moduleLoads: init.moduleLoads,
     frameInstances: init.frameInstances,
     namedFrames: init.namedFrames,
+    processClientEntryPreloads: init.processClientEntryPreloads,
     serverFrameReload: undefined,
     reloadForNavigation: init.reloadForNavigation,
   }
 }
 
-type InitialHydrationTracker = {
+export type ReconciliationTracker = {
   track: () => () => void
+  waitFor: (task: Promise<void>) => void
   finalize: () => void
   ready: () => Promise<void>
 }
 
-function createInitialHydrationTracker(): InitialHydrationTracker {
+function createReconciliationTracker(): ReconciliationTracker {
   let pending = 0
   let finalized = false
+  let failed = false
+  let failure: unknown
 
   let resolveReady: (() => void) | undefined
-  let readyPromise = new Promise<void>((resolve) => {
+  let rejectReady: ((error: unknown) => void) | undefined
+  let readyPromise = new Promise<void>((resolve, reject) => {
     resolveReady = resolve
+    rejectReady = reject
   })
 
-  function maybeResolve() {
-    if (finalized && pending === 0) {
-      resolveReady?.()
-      resolveReady = undefined
+  function maybeSettle() {
+    if (!finalized || pending !== 0) return
+    if (failed) rejectReady?.(failure)
+    else resolveReady?.()
+    resolveReady = undefined
+    rejectReady = undefined
+  }
+
+  function track(): () => void {
+    pending++
+    let completed = false
+    return () => {
+      if (completed) return
+      completed = true
+      pending--
+      maybeSettle()
     }
   }
 
   return {
-    track() {
-      pending++
-      let completed = false
-      return () => {
-        if (completed) return
-        completed = true
-        pending--
-        maybeResolve()
-      }
+    track,
+    waitFor(task) {
+      let complete = track()
+      void task.then(complete, (error) => {
+        if (!failed) {
+          failed = true
+          failure = error
+        }
+        complete()
+      })
     },
     finalize() {
       finalized = true
-      maybeResolve()
+      maybeSettle()
     },
     ready() {
       return readyPromise
@@ -1006,7 +1142,7 @@ function collectOwnedServerStyleTags(nodes: Node[], styles: HTMLStyleElement[]):
       continue
     }
 
-    if (node instanceof HTMLStyleElement && node.matches('style[data-rmx]')) {
+    if (node instanceof HTMLStyleElement && node.matches('style[data-rmx-style]')) {
       styles.push(node)
       continue
     }
@@ -1049,7 +1185,7 @@ function copyOwnRmxEntries<T>(target: Record<string, T>, source: Record<string,
 function scheduleHydrationInContainer(
   container: FrameContainer,
   context: FrameContext,
-  initialHydrationTracker?: InitialHydrationTracker,
+  reconciliationTracker?: ReconciliationTracker,
   signal?: AbortSignal,
 ): void {
   let hydrationMarkers = findHydrationMarkers(container)
@@ -1061,7 +1197,7 @@ function scheduleHydrationInContainer(
   for (let marker of hydrationMarkers) {
     let entry = hydrationData[marker.id]
     if (!entry) continue
-    scheduleHydrationMarker(marker, entry, context, initialHydrationTracker, signal)
+    scheduleHydrationMarker(marker, entry, context, reconciliationTracker, signal)
   }
 }
 
@@ -1069,12 +1205,12 @@ function scheduleHydrationMarker(
   marker: HydrationMarker,
   entry: HydrationData,
   context: FrameContext,
-  initialHydrationTracker?: InitialHydrationTracker,
+  reconciliationTracker?: ReconciliationTracker,
   signal?: AbortSignal,
 ): void {
   if (signal?.aborted || context.lifecycleSignal.aborted) return
 
-  let done = initialHydrationTracker?.track()
+  let done = reconciliationTracker?.track()
   let key = `${entry.moduleUrl}#${entry.exportName}`
   let identity: ClientEntryIdentity = {
     moduleUrl: entry.moduleUrl,
@@ -1216,30 +1352,41 @@ function hydrateRegion(
 
   context.pendingClientEntries.delete(start)
 
-  // The same marker can be discovered by overlapping hydration passes
-  // (for example, document root + nested frame root). Reuse the existing
-  // virtual root instead of redefining the marker property.
-  let owner = getClientEntryBoundaryOwner(start)
-  if (owner) {
+  // During a server frame reload, expose the reload signal and reconciliation
+  // tracker to client frames created while this entry renders so blocking
+  // frames keep the reload pending until their content arrives.
+  let renderEntry = (root: Pick<VirtualRoot, 'render'>) => {
     if (!signal) {
-      owner.root.render(vElement)
+      root.render(vElement)
       return
     }
 
     let frameRuntime = context.frame.$runtime
     invariant(
       isFrameRuntime(frameRuntime),
-      'Expected frame runtime while rendering a preserved client entry',
+      'Expected frame runtime while rendering a client entry during a reload',
     )
 
     let previousServerFrameReload = frameRuntime.serverFrameReload
 
-    frameRuntime.serverFrameReload = { signal }
+    frameRuntime.serverFrameReload = {
+      signal,
+      reconciliationTracker: context.reconciliationTracker,
+      blockingFrameTracker: context.blockingFrameTracker,
+    }
     try {
-      owner.root.render(vElement)
+      root.render(vElement)
     } finally {
       frameRuntime.serverFrameReload = previousServerFrameReload
     }
+  }
+
+  // The same marker can be discovered by overlapping hydration passes
+  // (for example, document root + nested frame root). Reuse the existing
+  // virtual root instead of redefining the marker property.
+  let owner = getClientEntryBoundaryOwner(start)
+  if (owner) {
+    renderEntry(owner.root)
     return
   }
 
@@ -1254,7 +1401,7 @@ function hydrateRegion(
   })
 
   setClientEntryBoundaryOwner(start, identity, root)
-  root.render(vElement)
+  renderEntry(root)
 }
 
 async function createSubFrames(
@@ -1300,8 +1447,12 @@ async function createSubFrames(
             moduleLoads: context.moduleLoads,
             frameInstances: context.frameInstances,
             namedFrames: context.namedFrames,
+            processClientEntryPreloads: context.processClientEntryPreloads,
           })
           context.frameInstances.set(node, subFrame)
+          if (frameMarker.status === 'resolved') {
+            tasks.push(subFrame.ready())
+          }
         }
       }