Fix frame stream reload batching (#11353)

Author: ryanflorence

packages/ui/src/runtime/frame.ts

@@ -8,6 +8,7 @@ import type { Scheduler, VirtualRoot } from './vdom.ts'
 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'
 
 type FrameRoot = [Comment, Comment] | Element | Document | DocumentFragment
 
@@ -64,10 +65,6 @@ function stripDoctypeMarkup(html: string): string {
   return html.replace(DOCTYPE_PATTERN, '')
 }
 
-function hasRenderableHtml(html: string): boolean {
-  return stripDoctypeMarkup(html).trim() !== ''
-}
-
 function syncElementAttributes(target: Element, source: Element) {
   for (let attribute of Array.from(target.attributes)) {
     if (!source.hasAttribute(attribute.name)) {
@@ -142,6 +139,7 @@ export type Frame = {
 }
 
 type RenderOptions = {
+  flushKind?: FlushKind
   initialHydrationTracker?: InitialHydrationTracker
   signal?: AbortSignal
 }
@@ -215,9 +213,9 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
     if (options?.signal?.aborted) return
 
     if (content instanceof ReadableStream) {
-      await renderFrameStream(content, container.doc, async (html) => {
+      await renderFrameStream(content, container.doc, async (html, flushKind) => {
         if (options?.signal?.aborted) return
-        await render(html, options)
+        await render(html, { ...options, flushKind })
       })
       return
     }
@@ -241,12 +239,24 @@ export function createFrame(root: FrameRoot, init: FrameInit): Frame {
       contentRoot = undefined
     }
 
+    if (typeof content === 'string') {
+      let flushed = await consumeFlushBatches(content, async (html, flushKind) => {
+        await render(html, { ...options, flushKind })
+      })
+      if (flushed.applied) {
+        if (flushed.remainder !== '') {
+          await render(flushed.remainder, { ...options, flushKind: 'fragment' })
+        }
+        return
+      }
+    }
+
     let htmlContent = typeof content === 'string' ? stripDoctypeMarkup(content) : undefined
 
     let isFullDocumentReload =
       container.root instanceof Document &&
       htmlContent !== undefined &&
-      isFullDocumentHtml(htmlContent)
+      options?.flushKind === 'document'
 
     if (isFullDocumentReload && htmlContent !== undefined) {
       let parsed = new DOMParser().parseFromString(htmlContent, 'text/html')
@@ -982,13 +992,12 @@ function extractTemplatesFromBuffer(
 async function renderFrameStream(
   stream: ReadableStream<Uint8Array>,
   doc: Document,
-  applyHtml: (html: string) => Promise<void>,
+  applyHtml: (html: string, flushKind: FlushKind) => Promise<void>,
 ): Promise<void> {
   let reader = stream.getReader()
   let decoder = new TextDecoder()
   let buffer = ''
   let html = ''
-  let appliedLength = 0
   let appliedOnce = false
 
   try {
@@ -1002,19 +1011,9 @@ async function renderFrameStream(
 
       if (parsed.html !== '') {
         html += parsed.html
-        // A doctype or whitespace prelude can arrive in its own chunk. Wait
-        // until there is actual frame content before applying the stream.
-        if (!hasRenderableHtml(html)) {
-          continue
-        }
-
-        let htmlMarkers = collectHtmlMarkerSummary(html)
-        if (!hasBalancedMarkerSummary(htmlMarkers)) {
-          continue
-        }
-        await applyHtml(html)
-        appliedLength = html.length
-        appliedOnce = true
+        let flushed = await consumeFlushBatches(html, applyHtml)
+        appliedOnce = flushed.applied || appliedOnce
+        html = flushed.remainder
       }
     }
 
@@ -1027,23 +1026,40 @@ async function renderFrameStream(
       buffer = ''
     }
 
-    let hasHtmlToApply = hasRenderableHtml(html)
-
-    if (hasHtmlToApply && html.length > appliedLength) {
-      await applyHtml(html)
+    if (html !== '') {
+      await applyHtml(html, 'fragment')
       appliedOnce = true
     }
 
     // A frame stream can legitimately resolve to empty content. Ensure the
     // existing frame region is cleared instead of treated as a no-op.
-    if (!hasHtmlToApply && !appliedOnce) {
-      await applyHtml('')
+    if (html === '' && !appliedOnce) {
+      await applyHtml('', 'fragment')
     }
   } finally {
     reader.releaseLock()
   }
 }
 
+async function consumeFlushBatches(
+  html: string,
+  applyHtml: (html: string, flushKind: FlushKind) => Promise<void>,
+): Promise<{ applied: boolean; remainder: string }> {
+  let applied = false
+  let cursor = 0
+  let marker = findFlushMarker(html, cursor)
+
+  while (marker) {
+    let batch = html.slice(cursor, marker.index)
+    await applyHtml(batch, marker.kind)
+    applied = true
+    cursor = marker.endIndex
+    marker = findFlushMarker(html, cursor)
+  }
+
+  return { applied, remainder: html.slice(cursor) }
+}
+
 type FrameContainer = {
   doc: Document
   root: ParentNode
@@ -1108,11 +1124,6 @@ function isRemixNodeFrameContent(content: InternalFrameContent): content is Remi
   )
 }
 
-function isFullDocumentHtml(content: string): boolean {
-  let trimmed = content.trimStart()
-  return /^<!doctype html\b/i.test(trimmed) || /^<html[\s>]/i.test(trimmed)
-}
-
 type HydrationMarker = {
   id: string
   start: Comment
@@ -1206,17 +1217,3 @@ function findEndMarker(
   throw new Error('End marker not found')
 }
 
-function collectHtmlMarkerSummary(html: string): Record<string, number> {
-  return {
-    frameStarts: html.match(/<!--\s*rmx:f:/g)?.length ?? 0,
-    frameEnds: html.match(/<!--\s*\/rmx:f\s*-->/g)?.length ?? 0,
-    hydrationStarts: html.match(/<!--\s*rmx:h:/g)?.length ?? 0,
-    hydrationEnds: html.match(/<!--\s*\/rmx:h\s*-->/g)?.length ?? 0,
-  }
-}
-
-function hasBalancedMarkerSummary(summary: Record<string, number>): boolean {
-  return (
-    summary.frameStarts === summary.frameEnds && summary.hydrationStarts === summary.hydrationEnds
-  )
-}

packages/ui/src/runtime/stream-protocol.ts

@@ -0,0 +1,27 @@
+export type FlushKind = 'document' | 'fragment'
+
+const FLUSH_MARKER_PATTERN = /<!--\s*rmx:flush\s+(document|fragment)\s*-->/g
+
+export function appendFlushMarker(html: string, kind: FlushKind): string {
+  return `${html}<!-- rmx:flush ${kind} -->`
+}
+
+export function stripFlushMarkers(html: string): string {
+  FLUSH_MARKER_PATTERN.lastIndex = 0
+  return html.replace(FLUSH_MARKER_PATTERN, '')
+}
+
+export function findFlushMarker(
+  html: string,
+  startIndex: number,
+): { index: number; endIndex: number; kind: FlushKind } | undefined {
+  FLUSH_MARKER_PATTERN.lastIndex = startIndex
+  let match = FLUSH_MARKER_PATTERN.exec(html)
+  if (!match) return undefined
+
+  return {
+    index: match.index,
+    endIndex: FLUSH_MARKER_PATTERN.lastIndex,
+    kind: match[1] as FlushKind,
+  }
+}

packages/ui/src/server/stream.ts

@@ -3,6 +3,7 @@ import type { ElementType, ElementProps, RemixElement } from '../runtime/jsx.ts'
 import { Fragment, createComponent, createFrameHandle, Frame } from '../runtime/component.ts'
 import { isEntry, type EntryComponent } from '../runtime/client-entries.ts'
 import { normalizeSvgAttribute } from '../runtime/svg-attributes.ts'
+import { appendFlushMarker, type FlushKind, stripFlushMarkers } from '../runtime/stream-protocol.ts'
 
 interface VNode {
   type: ElementType
@@ -90,6 +91,7 @@ interface RenderContext {
   unresolvedHydrationData: Map<string, UnresolvedHydrationData>
   frameData: Map<string, FrameData>
   blockingFrameTails: ReadableStream<Uint8Array>[]
+  flushKind: FlushKind
   serverIdScope: string
   serverIdCounter: number
 }
@@ -230,6 +232,7 @@ export function renderToStream(
     unresolvedHydrationData: new Map(),
     frameData: new Map(),
     blockingFrameTails: [],
+    flushKind: 'fragment',
     serverIdScope: crypto.randomUUID().slice(0, 8),
     serverIdCounter: 0,
   }
@@ -243,7 +246,7 @@ export function renderToStream(
         validateClientEntriesForHydration(context)
         let html = serializeSegment(root)
         let finalHtml = finalizeHtml(html, context)
-        let bytes = encoder.encode(finalHtml)
+        let bytes = encoder.encode(appendFlushMarker(finalHtml, context.flushKind))
         controller.enqueue(bytes)
 
         // If we have any tails from blocking frame streams, stream them now.
@@ -350,13 +353,13 @@ async function splitFirstChunk(stream: ReadableStream<Uint8Array>): Promise<Reso
     },
   })
 
-  return { html: stripDoctypeMarkup(decoder.decode(first)), tail }
+  return { html: stripFlushMarkers(stripDoctypeMarkup(decoder.decode(first))), tail }
 }
 
 async function resolveFrameHtml(
   input: string | ReadableStream<Uint8Array>,
 ): Promise<ResolvedFrameHtml> {
-  if (typeof input === 'string') return { html: stripDoctypeMarkup(input) }
+  if (typeof input === 'string') return { html: stripFlushMarkers(stripDoctypeMarkup(input)) }
 
   return await splitFirstChunk(input)
 }
@@ -399,6 +402,7 @@ function buildSegment(node: RemixNode, context: RenderContext, frameState: SsrFr
       let tag = type
 
       if (tag === 'html') {
+        context.flushKind = 'document'
         return buildElementSegment(tag, props, context, frameState)
       }
 
@@ -812,7 +816,11 @@ function buildComponentSegment(
   let [renderedNode] = handle.render(props)
   let childContext = { ...context, parentVNode: vnode }
 
-  return buildSegment(renderedNode, childContext, frameState)
+  let rendered = buildSegment(renderedNode, childContext, frameState)
+  if (childContext.flushKind === 'document') {
+    context.flushKind = 'document'
+  }
+  return rendered
 }
 
 function createHydrationPropsReplacer(context: RenderContext, frameState: SsrFrameState) {
@@ -1094,7 +1102,7 @@ function transformAttributeName(name: string, isSvg: boolean): string {
 }
 
 function finalizeHtml(html: string, context: RenderContext): string {
-  let hasHtmlRoot = html.trimStart().toLowerCase().startsWith('<html')
+  let hasHtmlRoot = context.flushKind === 'document'
 
   let styles = collectStyleTags(context)
   if (styles) {
@@ -1351,11 +1359,13 @@ async function drain(stream: ReadableStream<Uint8Array>): Promise<string> {
  * @returns Rendered HTML.
  */
 export async function renderToString(node: RemixNode): Promise<string> {
-  return drain(
-    renderToStream(node, {
-      onError(error) {
-        throw error
-      },
-    }),
+  return stripFlushMarkers(
+    await drain(
+      renderToStream(node, {
+        onError(error) {
+          throw error
+        },
+      }),
+    ),
   )
 }