fix: inline Suspense boundaries that complete before the stream is read

.changeset/ssr-inline-completed-boundaries.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/redact': patch
+---
+
+Write streamed SSR output when the destination starts reading, as React does. Suspense boundaries that complete before then are inlined, so a reader that waits for `allReady` gets a complete document that renders without JavaScript. Only boundaries still pending at that point stream a hidden segment and reveal script.

packages/redact/src/server/stream.ts

@@ -7,7 +7,7 @@ import {
   currentSSRFrame,
   type SSRFrame,
 } from './dispatcher'
-import { walk, type SuspendedBoundary, type BrowserBailoutCallback } from './walk'
+import { walk, boundaryEdges, type HtmlPart, type SuspendedBoundary, type BrowserBailoutCallback } from './walk'
 import { browserError, isBrowserToken, isBrowserError } from '../core/browser'
 import { BOUNDARY_REVEAL_RUNTIME, revealScript } from './bootstrap-script'
 import { escapeAttr, escapeScript } from './escape'
@@ -47,6 +47,11 @@ export interface OrchestratorState {
   cancelled: Promise<void>
   onShellReady?: () => void
   frame?: SSRFrame
+  // Output is written once the destination starts reading. Until then,
+  // completed boundaries stay inline, as React does before its first flush.
+  flowing?: boolean
+  flush?: () => void
+  outlined?: boolean
 }
 
 function createState(): OrchestratorState {
@@ -85,9 +90,9 @@ export async function streamHtml(
     // is ~100 stream-controller round-trips. Batching collapses those into
     // one encode and one enqueue per shell — measured ~2-4% of total SSR
     // time on CPU profiles.
-    const shellChunks: string[] = []
+    const shellChunks: HtmlPart[] = []
     let head = 0
-    const bufferedEmit: Emit = (chunk) => {
+    const bufferedEmit = (chunk: HtmlPart) => {
       shellChunks.push(chunk)
     }
 
@@ -137,34 +142,42 @@ export async function streamHtml(
       if (suspended) await Promise.race([suspended, state.cancelled])
     }
 
-    const hints = state.frame?.resources?.drain()
-    if (hints) shellChunks.splice(head, 0, hints)
-    if (state.frame) state.frame.shellFlushed = true
-
-    // 2. Inject runtime + bootstrap scripts (once, after shell). Skip the
-    // reveal/event-replay runtime when nothing needs it — no suspensions to
-    // reveal and no bootstrap scripts to guard against early user input. That
-    // keeps fully-static SSR responses byte-equivalent to a plain walk and
-    // matches React's behavior where `renderToReadableStream` of a static
-    // tree emits only the markup.
-    const hasBootstrap =
-      bootstrapScriptContentToArray(options.bootstrapScriptContent).length > 0 ||
-      (options.bootstrapScripts?.length ?? 0) > 0 ||
-      (options.bootstrapModules?.length ?? 0) > 0
-    if (boundaries.length > 0 || hasBootstrap) {
-      shellChunks.push(`<script${nonce ? ` nonce="${escapeAttr(nonce)}"` : ''}>${BOUNDARY_REVEAL_RUNTIME}</script>`)
-      for (const content of bootstrapScriptContentToArray(options.bootstrapScriptContent)) {
-        shellChunks.push(inlineBootstrapTag(content, nonce))
-      }
-      for (const s of options.bootstrapScripts ?? []) {
-        shellChunks.push(bootstrapTag(s, 'script', nonce))
-      }
-      for (const m of options.bootstrapModules ?? []) {
-        shellChunks.push(bootstrapTag(m, 'module', nonce))
+    // 2. Write the shell when the destination starts reading. Boundaries that
+    // completed by then are inlined, and their resources join the shell's.
+    let written = false
+    state.flush = () => {
+      if (written) return
+      written = true
+      const hints = state.frame?.resources?.drain()
+      if (hints) shellChunks.splice(head, 0, hints)
+      if (state.frame) state.frame.shellFlushed = true
+      let html = partsHtml(shellChunks, state)
+
+      // Inject runtime + bootstrap scripts (once, after shell). Skip the
+      // reveal/event-replay runtime when nothing needs it — no pending
+      // boundaries to reveal and no bootstrap scripts to guard against early
+      // user input. That keeps fully-static SSR responses byte-equivalent to
+      // a plain walk and matches React's behavior where
+      // `renderToReadableStream` of a static tree emits only the markup.
+      const hasBootstrap =
+        bootstrapScriptContentToArray(options.bootstrapScriptContent).length > 0 ||
+        (options.bootstrapScripts?.length ?? 0) > 0 ||
+        (options.bootstrapModules?.length ?? 0) > 0
+      if (state.outlined || hasBootstrap) {
+        html += `<script${nonce ? ` nonce="${escapeAttr(nonce)}"` : ''}>${BOUNDARY_REVEAL_RUNTIME}</script>`
+        for (const content of bootstrapScriptContentToArray(options.bootstrapScriptContent)) {
+          html += inlineBootstrapTag(content, nonce)
+        }
+        for (const s of options.bootstrapScripts ?? []) {
+          html += bootstrapTag(s, 'script', nonce)
+        }
+        for (const m of options.bootstrapModules ?? []) {
+          html += bootstrapTag(m, 'module', nonce)
+        }
       }
+      if (html) emit(normalizeDocumentShell(html))
     }
-
-    if (shellChunks.length) emit(normalizeDocumentShell(shellChunks.join('')))
+    if (state.flowing) state.flush()
     state.onShellReady?.()
 
     // 3. Stream suspended boundaries as they resolve
@@ -199,18 +212,22 @@ function streamBoundary(
     const defer = (error: unknown) => {
       if (isBrowserError(error)) options.onBrowserBailout?.(error, { componentStack: (error as any).componentStack || b.componentStack })
       else options.onError?.(error)
-      emit(`<script${nonce ? ` nonce="${escapeAttr(nonce)}"` : ''}>$RB(${b.id}${isBrowserError(error) ? '' : ',1'})</script>`)
+      if (b.written) emit(`<script${nonce ? ` nonce="${escapeAttr(nonce)}"` : ''}>$RB(${b.id}${isBrowserError(error) ? '' : ',1'})</script>`)
+      else b.marker = isBrowserError(error) ? '$!' : '$E'
     }
     if (state.aborted) { defer(state.reason); return }
 
     // Re-render the boundary's children into a string, restoring the
     // provider stack from when the boundary first suspended.
-    const parts: string[] = []
+    // A boundary not yet written keeps its content for an inline write; its
+    // resources stay with the unwritten shell or segment that will contain it.
+    const written = b.written
+    const parts: HtmlPart[] = []
     const sub: SuspendedBoundary[] = []
     const resourceScope: ResourceScope = {}
     let styles: string[][] | undefined
     const previous = beginSSR(options.identifierPrefix, state.frame)
-    state.frame?.resources?.takeStylesheets()
+    if (written) state.frame?.resources?.takeStylesheets()
     const restore = applyContextSnapshot(b.contextSnapshot)
     try {
       walk(b.children, {
@@ -225,18 +242,22 @@ function streamBoundary(
       state.frame?.resources?.mergeScope(resourceScope)
     } catch (err) {
       if (isThenable(err)) {
-        streamBoundary({ ...b, thenable: err }, emit, options, state)
+        b.thenable = err
+        streamBoundary(b, emit, options, state)
       } else defer(err)
       return
     } finally {
       restore()
       endSSR(previous)
-      const hints = state.frame?.resources?.drain()
-      styles = state.frame?.resources?.takeStylesheets()
-      if (hints) emit(hints)
+      if (written) {
+        const hints = state.frame?.resources?.drain()
+        styles = state.frame?.resources?.takeStylesheets()
+        if (hints) emit(hints)
+      }
     }
 
-    emit(`<div hidden id="S:${b.id}">${parts.join('')}</div>${revealScript(b.id, nonce, styles)}`)
+    if (written) emit(`<div hidden id="S:${b.id}">${partsHtml(parts, state)}</div>${revealScript(b.id, nonce, styles)}`)
+    else b.content = parts
 
     // Recurse: any nested suspensions inside the now-revealed content
     for (const s of sub) streamBoundary(s, emit, options, state)
@@ -245,6 +266,24 @@ function streamBoundary(
   task.then(() => state.pending.delete(task), () => state.pending.delete(task))
 }
 
+// Serialize buffered output. Boundaries that completed or failed before this
+// write are inlined; pending ones keep their fallback and are marked written,
+// so their later completion streams a hidden segment and reveal script.
+function partsHtml(parts: HtmlPart[], state: OrchestratorState): string {
+  let html = ''
+  for (const part of parts) {
+    if (typeof part === 'string') {
+      html += part
+      continue
+    }
+    const marker = part.content ? '$' : part.marker ?? '$?'
+    if (marker === '$?') part.written = state.outlined = true
+    const [open, close] = boundaryEdges(part.id, marker)
+    html += open + partsHtml(part.content ?? part.fallback, state) + close
+  }
+  return html
+}
+
 async function drain(state: OrchestratorState): Promise<void> {
   while (state.pending.size > 0) {
     await Promise.race(state.pending)
@@ -350,6 +389,11 @@ export function renderToReadableStream(
   })
   state.onShellReady = () => queueMicrotask(() => shellResolve(Object.assign(stream, { allReady })))
 
+  const close = (controller: ReadableStreamDefaultController<Uint8Array>) => {
+    try {
+      controller.close()
+    } catch {}
+  }
   const stream = new ReadableStream<Uint8Array>({
     start(controller) {
       const emit = (chunk: string) => {
@@ -370,9 +414,7 @@ export function renderToReadableStream(
       streamHtml(children, emit, options, state).then(
         () => {
           detachAbortListener()
-          try {
-            controller.close()
-          } catch {}
+          if (state.flowing) close(controller)
           allReadyResolve()
         },
         (err) => {
@@ -386,11 +428,16 @@ export function renderToReadableStream(
       )
 
     },
+    pull(controller) {
+      state.flowing = true
+      state.flush?.()
+      if (state.closed && !state.errored) close(controller)
+    },
     cancel(reason) {
       state.abort(reason)
       detachAbortListener()
     },
-  })
+  }, { highWaterMark: 0 })
 
   return shell
 }
@@ -416,7 +463,6 @@ export function renderToPipeableStream(
 ): PipeableHandle {
   const state = createState()
 
-  const buffers: string[] = []
   let dest: PipeableWritable | null = null
   let shellReady = false
   let shellFailed = false
@@ -429,16 +475,9 @@ export function renderToPipeableStream(
     else target.end()
   }
 
-  const flushTo = (w: PipeableWritable) => {
-    if (!buffers.length) return
-    for (const b of buffers) w.write(b)
-    buffers.length = 0
-  }
-
+  // Nothing is emitted before pipe() starts the flow.
   const emit: Emit = (chunk) => {
-    if (finished) return
-    if (dest) dest.write(chunk)
-    else buffers.push(chunk)
+    if (!failed) dest!.write(chunk)
   }
 
   state.onShellReady = () => queueMicrotask(() => {
@@ -470,7 +509,8 @@ export function renderToPipeableStream(
   return {
     pipe<T extends PipeableWritable>(target: T): T {
       dest = target
-      if (!failed) flushTo(target)
+      state.flowing = true
+      if (!failed) state.flush?.()
       if (finished) finishDestination(target)
       return target
     },

packages/redact/src/server/walk.ts

@@ -47,17 +47,25 @@ export interface ServerErrorInfo {
 
 export type BrowserBailoutCallback = (error: Error, errorInfo: ServerErrorInfo) => void
 
+/** Streamed output keeps pending boundaries addressable until it is written. */
+export type HtmlPart = string | SuspendedBoundary
+
 export interface SuspendedBoundary {
   id: number
-  fallbackHTML: string
+  fallback: HtmlPart[]
   children: ReactNode
   thenable: Promise<any>
   contextSnapshot: ContextSnapshot
   componentStack: string
+  // Set by the stream: completed content and failure markers are inlined
+  // when the boundary is written; `written` boundaries need a later reveal.
+  content?: HtmlPart[]
+  marker?: string
+  written?: boolean
 }
 
 export interface WalkOptions {
-  emit: (chunk: string) => void
+  emit: (chunk: HtmlPart) => void
   onHead?: () => void
   onSuspend?: ((boundary: SuspendedBoundary) => void) | undefined
   onBrowserBailout?: BrowserBailoutCallback | undefined
@@ -455,10 +463,10 @@ function walkSuspense(
 
   // Try to render the children synchronously. If a thenable is thrown,
   // record the boundary and emit the fallback.
-  const childParts: string[] = []
+  const childParts: HtmlPart[] = []
   const nested: SuspendedBoundary[] = []
   const thrownInfo: { componentStack?: string | undefined } = {}
-  const childEmit = (s: string) => childParts.push(s)
+  const childEmit = (s: HtmlPart) => childParts.push(s)
   const resourceScope: ResourceScope | undefined = currentSSRFrame().shellFlushed ? {} : undefined
   try {
     walkNode(props.children, {
@@ -476,21 +484,23 @@ function walkSuspense(
       if (isBrowserError(thenable)) opts.onBrowserBailout?.(thenable, { componentStack: (thenable as any).componentStack || opts.componentStack || '' })
       else opts.onError?.(thenable)
     }
-    const fallbackParts: string[] = []
+    const fallbackParts: HtmlPart[] = []
     walkNode(props.fallback, {
       ...opts,
       emit: (s) => fallbackParts.push(s),
       onSuspend: opts.onSuspend ? (boundary) => nested.push(boundary) : undefined,
       nextBoundaryId: opts.nextBoundaryId,
     })
-    if (!suspended) {
-      opts.emit(`<!--${isBrowserError(thenable) ? '$!' : '$E'}${id}-->${fallbackParts.join('')}<!--/$-->`)
-    } else emitBoundary(opts, id, fallbackParts.join(''))
-    for (const boundary of nested) opts.onSuspend?.(boundary)
-    if (suspended) opts.onSuspend?.({
-      id, fallbackHTML: fallbackParts.join(''), children: props.children, thenable, contextSnapshot,
+    // A streamed suspension is emitted as itself, so the stream can still
+    // inline its content if it completes before these bytes are written.
+    const boundary: SuspendedBoundary | undefined = suspended && opts.onSuspend ? {
+      id, fallback: fallbackParts, children: props.children, thenable, contextSnapshot,
       componentStack: thrownInfo.componentStack || opts.componentStack || '',
-    })
+    } : undefined
+    if (boundary) opts.emit(boundary)
+    else emitBoundary(opts, id, fallbackParts, suspended ? '$?' : isBrowserError(thenable) ? '$!' : '$E')
+    for (const inner of nested) opts.onSuspend?.(inner)
+    if (boundary) opts.onSuspend!(boundary)
     return
   }
 
@@ -503,17 +513,25 @@ function walkSuspense(
   // range to adopt on resolve — it creates fresh DOM next to the SSR content,
   // producing visible duplicates (e.g. double navbar logos). Markers let the
   // client treat this as a resolved boundary and hydrate in-place.
-  opts.emit(`<!--$${id}-->`)
-  opts.emit(childParts.join(''))
-  opts.emit(`<!--/$-->`)
+  emitBoundary(opts, id, childParts, '$')
   for (const boundary of nested) opts.onSuspend?.(boundary)
 }
 
-function emitBoundary(opts: WalkOptions, id: number, fallbackHTML: string): void {
-  // Visible div wrapper so the fallback UI shows; B: id lets $RC locate it on
-  // reveal. The leading/trailing comments let hydration detect a pending
-  // boundary and register a reveal callback.
-  opts.emit(`<!--$?${id}--><div id="B:${id}">`)
-  opts.emit(fallbackHTML)
-  opts.emit(`</div><!--/$-->`)
+function emitBoundary(opts: WalkOptions, id: number, parts: HtmlPart[], marker: string): void {
+  const [open, close] = boundaryEdges(id, marker)
+  opts.emit(open)
+  for (const part of parts) opts.emit(part)
+  opts.emit(close)
+}
+
+/**
+ * Boundary markers: `$` completed, `$E`/`$!` client-rendered, `$?` pending.
+ * A pending fallback gets a visible div wrapper so the fallback UI shows;
+ * B: id lets $RC locate it on reveal. The leading/trailing comments let
+ * hydration detect a pending boundary and register a reveal callback.
+ */
+export function boundaryEdges(id: number, marker: string): [string, string] {
+  return marker === '$?'
+    ? [`<!--$?${id}--><div id="B:${id}">`, '</div><!--/$-->']
+    : [`<!--${marker}${id}-->`, '<!--/$-->']
 }

scripts/size-check.mjs

@@ -63,11 +63,17 @@ const STORE_HYDRATION_SIZE_ALLOWANCE = {
   'redact/dom-client (hydration=stub)': [3, 0],
 }
 
+// Measured cost of inlining boundaries that complete before their bytes are
+// written, against main 37c051c, [source, built] gzip bytes above the cap.
+const INLINE_BOUNDARY_SIZE_ALLOWANCE = {
+  'redact/server': [188, 191],
+}
+
 let failed = false
 for (const tree of ['src', 'dist']) {
   for (const row of await measureSizes(tree)) {
     const baseBudget = BUDGETS[row.name]
-    const budget = baseBudget === undefined ? undefined : baseBudget + (STORE_HYDRATION_SIZE_ALLOWANCE[row.name]?.[tree === 'src' ? 0 : 1] ?? 0) + (ROOT_TRANSITION_SIZE_ALLOWANCE[row.name]?.[tree === 'src' ? 0 : 1] ?? 0) + (QUEUE_SIZE_ALLOWANCE[row.name]?.[tree === 'src' ? 0 : 1] ?? 0)
+    const budget = baseBudget === undefined ? undefined : baseBudget + (INLINE_BOUNDARY_SIZE_ALLOWANCE[row.name]?.[tree === 'src' ? 0 : 1] ?? 0) + (STORE_HYDRATION_SIZE_ALLOWANCE[row.name]?.[tree === 'src' ? 0 : 1] ?? 0) + (ROOT_TRANSITION_SIZE_ALLOWANCE[row.name]?.[tree === 'src' ? 0 : 1] ?? 0) + (QUEUE_SIZE_ALLOWANCE[row.name]?.[tree === 'src' ? 0 : 1] ?? 0)
     const ok = budget !== undefined && row.gz <= budget
     failed ||= !ok
     console.log(`${ok ? 'PASS' : 'FAIL'} ${tree} ${row.name}: ${row.gz} B gzip, budget ${budget ?? 'missing'}`)