@tanstack/router-core 1.142.6 → 1.142.8

package/src/ssr/transformStreamWithRouter.ts

@@ -1,6 +1,6 @@
 import { ReadableStream } from 'node:stream/web'
 import { Readable } from 'node:stream'
-import { createControlledPromise } from '../utils'
+import { TSR_SCRIPT_BARRIER_ID } from './constants'
 import type { AnyRouter } from '../router'
 
 export function transformReadableStreamWithRouter(
@@ -19,316 +19,394 @@ export function transformPipeableStreamWithRouter(
   )
 }
 
-export const TSR_SCRIPT_BARRIER_ID = '$tsr-stream-barrier'
-
-// regex pattern for matching closing body and html tags
-const patternBodyEnd = /(<\/body>)/
-const patternHtmlEnd = /(<\/html>)/
-// regex pattern for matching closing tags
-const patternClosingTag = /(<\/[a-zA-Z][\w:.-]*?>)/g
-
-type ReadablePassthrough = {
-  stream: ReadableStream
-  write: (chunk: unknown) => void
-  end: (chunk?: string) => void
-  destroy: (error: unknown) => void
-  destroyed: boolean
-}
-
-function createPassthrough(onCancel: () => void) {
-  let controller: ReadableStreamDefaultController<any>
-  const encoder = new TextEncoder()
-  const stream = new ReadableStream({
-    start(c) {
-      controller = c
-    },
-    cancel() {
-      res.destroyed = true
-      onCancel()
-    },
-  })
-
-  const res: ReadablePassthrough = {
-    stream,
-    write: (chunk) => {
-      // Don't write to destroyed stream
-      if (res.destroyed) return
-      if (typeof chunk === 'string') {
-        controller.enqueue(encoder.encode(chunk))
-      } else {
-        controller.enqueue(chunk)
-      }
-    },
-    end: (chunk) => {
-      // Don't end already destroyed stream
-      if (res.destroyed) return
-      if (chunk) {
-        res.write(chunk)
+// Use string constants for simple indexOf matching
+const BODY_END_TAG = '</body>'
+const HTML_END_TAG = '</html>'
+
+// Minimum length of a valid closing tag: </a> = 4 characters
+const MIN_CLOSING_TAG_LENGTH = 4
+
+// Default timeout values (in milliseconds)
+const DEFAULT_SERIALIZATION_TIMEOUT_MS = 60000
+const DEFAULT_LIFETIME_TIMEOUT_MS = 60000
+
+// Module-level encoder (stateless, safe to reuse)
+const textEncoder = new TextEncoder()
+
+/**
+ * Finds the position just after the last valid HTML closing tag in the string.
+ *
+ * Valid closing tags match the pattern: </[a-zA-Z][\w:.-]*>
+ * Examples: </div>, </my-component>, </slot:name.nested>
+ *
+ * @returns Position after the last closing tag, or -1 if none found
+ */
+function findLastClosingTagEnd(str: string): number {
+  const len = str.length
+  if (len < MIN_CLOSING_TAG_LENGTH) return -1
+
+  let i = len - 1
+
+  while (i >= MIN_CLOSING_TAG_LENGTH - 1) {
+    // Look for > (charCode 62)
+    if (str.charCodeAt(i) === 62) {
+      // Look backwards for valid tag name characters
+      let j = i - 1
+
+      // Skip through valid tag name characters
+      while (j >= 1) {
+        const code = str.charCodeAt(j)
+        // Check if it's a valid tag name char: [a-zA-Z0-9_:.-]
+        if (
+          (code >= 97 && code <= 122) || // a-z
+          (code >= 65 && code <= 90) || // A-Z
+          (code >= 48 && code <= 57) || // 0-9
+          code === 95 || // _
+          code === 58 || // :
+          code === 46 || // .
+          code === 45 // -
+        ) {
+          j--
+        } else {
+          break
+        }
       }
-      res.destroyed = true
-      controller.close()
-    },
-    destroy: (error) => {
-      // Don't destroy already destroyed stream
-      if (res.destroyed) return
-      res.destroyed = true
-      controller.error(error)
-    },
-    destroyed: false,
-  }
 
-  return res
-}
-
-async function readStream(
-  stream: ReadableStream,
-  opts: {
-    onData?: (chunk: ReadableStreamReadValueResult<any>) => void
-    onEnd?: () => void
-    onError?: (error: unknown) => void
-  },
-) {
-  const reader = stream.getReader()
-  try {
-    let chunk
-    while (!(chunk = await reader.read()).done) {
-      opts.onData?.(chunk)
+      // Check if the first char after </ is a valid start char (letter only)
+      const tagNameStart = j + 1
+      if (tagNameStart < i) {
+        const startCode = str.charCodeAt(tagNameStart)
+        // Tag name must start with a letter (a-z or A-Z)
+        if (
+          (startCode >= 97 && startCode <= 122) ||
+          (startCode >= 65 && startCode <= 90)
+        ) {
+          // Check for </ (charCodes: < = 60, / = 47)
+          if (
+            j >= 1 &&
+            str.charCodeAt(j) === 47 &&
+            str.charCodeAt(j - 1) === 60
+          ) {
+            return i + 1 // Return position after the closing >
+          }
+        }
+      }
     }
-    opts.onEnd?.()
-  } catch (error) {
-    opts.onError?.(error)
-  } finally {
-    reader.releaseLock()
+    i--
   }
+  return -1
 }
 
 export function transformStreamWithRouter(
   router: AnyRouter,
   appStream: ReadableStream,
   opts?: {
+    /** Timeout for serialization to complete after app render finishes (default: 60000ms) */
     timeoutMs?: number
+    /** Maximum lifetime of the stream transform (default: 60000ms). Safety net for cleanup. */
+    lifetimeMs?: number
   },
 ) {
-  let stopListeningToInjectedHtml: (() => void) | undefined = undefined
-  let timeoutHandle: NodeJS.Timeout
+  let stopListeningToInjectedHtml: (() => void) | undefined
+  let stopListeningToSerializationFinished: (() => void) | undefined
+  let serializationTimeoutHandle: ReturnType<typeof setTimeout> | undefined
+  let lifetimeTimeoutHandle: ReturnType<typeof setTimeout> | undefined
   let cleanedUp = false
 
+  let controller: ReadableStreamDefaultController<any>
+  let isStreamClosed = false
+
+  // Check upfront if serialization already finished synchronously
+  // This is the fast path for routes with no deferred data
+  const serializationAlreadyFinished =
+    router.serverSsr?.isSerializationFinished() ?? false
+
+  /**
+   * Cleanup function with guards against multiple calls.
+   * Unsubscribes listeners, clears timeouts, frees buffers, and cleans up router SSR state.
+   */
   function cleanup() {
+    // Guard against multiple cleanup calls - set flag first to prevent re-entry
     if (cleanedUp) return
     cleanedUp = true
-    if (stopListeningToInjectedHtml) {
-      stopListeningToInjectedHtml()
-      stopListeningToInjectedHtml = undefined
+
+    // Unsubscribe listeners first (wrap in try-catch for safety)
+    try {
+      stopListeningToInjectedHtml?.()
+      stopListeningToSerializationFinished?.()
+    } catch (e) {
+      // Ignore errors during unsubscription
+    }
+    stopListeningToInjectedHtml = undefined
+    stopListeningToSerializationFinished = undefined
+
+    // Clear all timeouts
+    if (serializationTimeoutHandle !== undefined) {
+      clearTimeout(serializationTimeoutHandle)
+      serializationTimeoutHandle = undefined
+    }
+    if (lifetimeTimeoutHandle !== undefined) {
+      clearTimeout(lifetimeTimeoutHandle)
+      lifetimeTimeoutHandle = undefined
     }
-    clearTimeout(timeoutHandle)
+
+    // Clear buffers to free memory
+    pendingRouterHtmlParts = []
+    leftover = ''
+    pendingClosingTags = ''
+
+    // Clean up router SSR state (has its own guard)
     router.serverSsr?.cleanup()
   }
 
-  const finalPassThrough = createPassthrough(cleanup)
   const textDecoder = new TextDecoder()
 
-  let isAppRendering = true
-  let routerStreamBuffer = ''
-  let pendingClosingTags = ''
-  let streamBarrierLifted = false
-  let leftover = ''
-  let leftoverHtml = ''
-
-  function getBufferedRouterStream() {
-    const html = routerStreamBuffer
-    routerStreamBuffer = ''
-    return html
+  function safeEnqueue(chunk: string | Uint8Array) {
+    if (isStreamClosed) return
+    if (typeof chunk === 'string') {
+      controller.enqueue(textEncoder.encode(chunk))
+    } else {
+      controller.enqueue(chunk)
+    }
   }
 
-  function decodeChunk(chunk: unknown): string {
-    if (chunk instanceof Uint8Array) {
-      return textDecoder.decode(chunk, { stream: true })
+  function safeClose() {
+    if (isStreamClosed) return
+    isStreamClosed = true
+    try {
+      controller.close()
+    } catch {
+      // Stream may already be errored or closed by consumer - safe to ignore
     }
-    return String(chunk)
   }
 
-  const injectedHtmlDonePromise = createControlledPromise<void>()
-
-  let processingCount = 0
+  function safeError(error: unknown) {
+    if (isStreamClosed) return
+    isStreamClosed = true
+    try {
+      controller.error(error)
+    } catch {
+      // Stream may already be errored or closed by consumer - safe to ignore
+    }
+  }
 
-  // Process any already-injected HTML
-  handleInjectedHtml()
+  const stream = new ReadableStream({
+    start(c) {
+      controller = c
+    },
+    cancel() {
+      isStreamClosed = true
+      cleanup()
+    },
+  })
 
-  // Listen for any new injected HTML
-  stopListeningToInjectedHtml = router.subscribe(
-    'onInjectedHtml',
-    handleInjectedHtml,
-  )
+  let isAppRendering = true
+  let streamBarrierLifted = false
+  let leftover = ''
+  let pendingClosingTags = ''
+  let serializationFinished = serializationAlreadyFinished
 
-  function handleInjectedHtml() {
-    // Don't process if already cleaned up
-    if (cleanedUp) return
+  let pendingRouterHtmlParts: Array<string> = []
 
-    router.serverSsr!.injectedHtml.forEach((promise) => {
-      processingCount++
+  // Take any HTML that was buffered before we started listening
+  const bufferedHtml = router.serverSsr?.takeBufferedHtml()
+  if (bufferedHtml) {
+    pendingRouterHtmlParts.push(bufferedHtml)
+  }
 
-      promise
-        .then((html) => {
-          // Don't write to destroyed stream or after cleanup
-          if (cleanedUp || finalPassThrough.destroyed) {
-            return
-          }
-          if (isAppRendering) {
-            routerStreamBuffer += html
-          } else {
-            finalPassThrough.write(html)
-          }
-        })
-        .catch((err) => {
-          injectedHtmlDonePromise.reject(err)
-        })
-        .finally(() => {
-          processingCount--
-
-          if (!isAppRendering && processingCount === 0) {
-            injectedHtmlDonePromise.resolve()
-          }
-        })
-    })
-    router.serverSsr!.injectedHtml = []
+  function flushPendingRouterHtml() {
+    if (pendingRouterHtmlParts.length > 0) {
+      safeEnqueue(pendingRouterHtmlParts.join(''))
+      pendingRouterHtmlParts = []
+    }
   }
 
-  injectedHtmlDonePromise
-    .then(() => {
-      // Don't process if already cleaned up or destroyed
-      if (cleanedUp || finalPassThrough.destroyed) {
-        return
-      }
+  /**
+   * Attempts to finish the stream if all conditions are met.
+   */
+  function tryFinish() {
+    // Can only finish when app is done rendering and serialization is complete
+    if (isAppRendering || !serializationFinished) return
+    if (cleanedUp || isStreamClosed) return
+
+    // Clear serialization timeout since we're finishing
+    if (serializationTimeoutHandle !== undefined) {
+      clearTimeout(serializationTimeoutHandle)
+      serializationTimeoutHandle = undefined
+    }
 
-      clearTimeout(timeoutHandle)
-      const finalHtml =
-        leftover + leftoverHtml + getBufferedRouterStream() + pendingClosingTags
+    // Flush any remaining bytes in the TextDecoder
+    const decoderRemainder = textDecoder.decode()
 
-      leftover = ''
-      leftoverHtml = ''
-      pendingClosingTags = ''
+    if (leftover) safeEnqueue(leftover)
+    if (decoderRemainder) safeEnqueue(decoderRemainder)
+    flushPendingRouterHtml()
+    if (pendingClosingTags) safeEnqueue(pendingClosingTags)
 
-      finalPassThrough.end(finalHtml)
-    })
-    .catch((err) => {
-      // Don't process if already cleaned up
-      if (cleanedUp || finalPassThrough.destroyed) {
-        return
-      }
+    safeClose()
+    cleanup()
+  }
 
-      console.error('Error reading routerStream:', err)
-      finalPassThrough.destroy(err)
+  // Set up lifetime timeout as a safety net
+  // This ensures cleanup happens even if the stream is never consumed or gets stuck
+  const lifetimeMs = opts?.lifetimeMs ?? DEFAULT_LIFETIME_TIMEOUT_MS
+  lifetimeTimeoutHandle = setTimeout(() => {
+    if (!cleanedUp && !isStreamClosed) {
+      console.warn(
+        `SSR stream transform exceeded maximum lifetime (${lifetimeMs}ms), forcing cleanup`,
+      )
+      safeError(new Error('Stream lifetime exceeded'))
+      cleanup()
+    }
+  }, lifetimeMs)
+
+  // Only set up listeners if serialization hasn't already finished
+  // This avoids unnecessary subscriptions for the common case of no deferred data
+  if (!serializationAlreadyFinished) {
+    // Listen for injected HTML (for deferred data that resolves later)
+    stopListeningToInjectedHtml = router.subscribe('onInjectedHtml', () => {
+      if (cleanedUp || isStreamClosed) return
+
+      // Retrieve buffered HTML
+      const html = router.serverSsr?.takeBufferedHtml()
+      if (!html) return
+
+      if (isAppRendering) {
+        // Buffer for insertion at next valid position
+        pendingRouterHtmlParts.push(html)
+      } else {
+        // App is done rendering, write directly to output
+        safeEnqueue(html)
+      }
     })
-    .finally(cleanup)
 
-  // Transform the appStream
-  readStream(appStream, {
-    onData: (chunk) => {
-      // Don't process if already cleaned up
-      if (cleanedUp || finalPassThrough.destroyed) {
-        return
-      }
+    // Listen for serialization finished
+    stopListeningToSerializationFinished = router.subscribe(
+      'onSerializationFinished',
+      () => {
+        serializationFinished = true
+        tryFinish()
+      },
+    )
+  }
 
-      const text = decodeChunk(chunk.value)
-      const chunkString = leftover + text
-      const bodyEndMatch = chunkString.match(patternBodyEnd)
-      const htmlEndMatch = chunkString.match(patternHtmlEnd)
-
-      if (!streamBarrierLifted) {
-        const streamBarrierIdIncluded = chunkString.includes(
-          TSR_SCRIPT_BARRIER_ID,
-        )
-        if (streamBarrierIdIncluded) {
-          streamBarrierLifted = true
-          router.serverSsr!.liftScriptBarrier()
+  // Transform the appStream
+  ;(async () => {
+    const reader = appStream.getReader()
+    try {
+      while (true) {
+        const { done, value } = await reader.read()
+        if (done) break
+
+        // Don't process if already cleaned up
+        if (cleanedUp || isStreamClosed) return
+
+        const text =
+          value instanceof Uint8Array
+            ? textDecoder.decode(value, { stream: true })
+            : String(value)
+        const chunkString = leftover + text
+
+        // Check for stream barrier (script placeholder) - use indexOf for efficiency
+        if (!streamBarrierLifted) {
+          if (chunkString.includes(TSR_SCRIPT_BARRIER_ID)) {
+            streamBarrierLifted = true
+            router.serverSsr?.liftScriptBarrier()
+          }
         }
-      }
 
-      // If either the body end or html end is in the chunk,
-      // We need to get all of our data in asap
-      if (
-        bodyEndMatch &&
-        htmlEndMatch &&
-        bodyEndMatch.index! < htmlEndMatch.index!
-      ) {
-        const bodyEndIndex = bodyEndMatch.index!
-        pendingClosingTags = chunkString.slice(bodyEndIndex)
-
-        finalPassThrough.write(
-          chunkString.slice(0, bodyEndIndex) +
-            getBufferedRouterStream() +
-            leftoverHtml,
-        )
-
-        leftover = ''
-        leftoverHtml = ''
-        return
-      }
+        // Check for body/html end tags
+        const bodyEndIndex = chunkString.indexOf(BODY_END_TAG)
+        const htmlEndIndex = chunkString.indexOf(HTML_END_TAG)
 
-      let result: RegExpExecArray | null
-      let lastIndex = 0
-      // Reset regex lastIndex since it's global and stateful across exec() calls
-      patternClosingTag.lastIndex = 0
-      while ((result = patternClosingTag.exec(chunkString)) !== null) {
-        lastIndex = result.index + result[0].length
-      }
+        // If we have both </body> and </html> in proper order,
+        // insert router HTML before </body> and hold the closing tags
+        if (
+          bodyEndIndex !== -1 &&
+          htmlEndIndex !== -1 &&
+          bodyEndIndex < htmlEndIndex
+        ) {
+          pendingClosingTags = chunkString.slice(bodyEndIndex)
 
-      if (lastIndex > 0) {
-        const processed =
-          chunkString.slice(0, lastIndex) +
-          getBufferedRouterStream() +
-          leftoverHtml
+          safeEnqueue(chunkString.slice(0, bodyEndIndex))
+          flushPendingRouterHtml()
 
-        finalPassThrough.write(processed)
-        leftover = chunkString.slice(lastIndex)
-        leftoverHtml = ''
-      } else {
-        leftover = chunkString
-        leftoverHtml += getBufferedRouterStream()
-      }
-    },
-    onEnd: () => {
-      // Don't process if stream was already destroyed/cancelled or cleaned up
-      if (cleanedUp || finalPassThrough.destroyed) {
-        return
+          leftover = ''
+          continue
+        }
+
+        // Handling partial closing tags split across chunks:
+        //
+        // Since `chunkString = leftover + text`, any incomplete tag fragment from the
+        // previous chunk is prepended to the current chunk, allowing split tags like
+        // "</di" + "v>" to be re-detected as a complete "</div>" in the combined string.
+        //
+        // - If a closing tag IS found (lastClosingTagEnd > 0): We enqueue content up to
+        //   the end of that tag, flush router HTML, and store the remainder in `leftover`.
+        //   This remainder may contain a partial tag (e.g., "</sp") which will be
+        //   prepended to the next chunk for re-detection.
+        //
+        // - If NO closing tag is found: The entire chunk is buffered in `leftover` and
+        //   will be prepended to the next chunk. This ensures partial tags are never
+        //   lost and will be detected once the rest of the tag arrives.
+        //
+        // This approach guarantees correct injection points even when closing tags span
+        // chunk boundaries.
+        const lastClosingTagEnd = findLastClosingTagEnd(chunkString)
+
+        if (lastClosingTagEnd > 0) {
+          // Found a closing tag - insert router HTML after it
+          safeEnqueue(chunkString.slice(0, lastClosingTagEnd))
+          flushPendingRouterHtml()
+
+          leftover = chunkString.slice(lastClosingTagEnd)
+        } else {
+          // No closing tag found - buffer the entire chunk
+          leftover = chunkString
+          // Any pending router HTML will be inserted when we find a valid position
+        }
       }
 
+      // Stream ended
+      if (cleanedUp || isStreamClosed) return
+
       // Mark the app as done rendering
       isAppRendering = false
-      router.serverSsr!.setRenderFinished()
+      router.serverSsr?.setRenderFinished()
 
-      // If there are no pending promises, resolve the injectedHtmlDonePromise
-      if (processingCount === 0) {
-        injectedHtmlDonePromise.resolve()
+      // Try to finish if serialization is already done
+      if (serializationFinished) {
+        tryFinish()
       } else {
-        const timeoutMs = opts?.timeoutMs ?? 60000
-        timeoutHandle = setTimeout(() => {
-          injectedHtmlDonePromise.reject(
-            new Error('Injected HTML timeout after app render finished'),
-          )
+        // Set a timeout for serialization to complete
+        const timeoutMs = opts?.timeoutMs ?? DEFAULT_SERIALIZATION_TIMEOUT_MS
+        serializationTimeoutHandle = setTimeout(() => {
+          if (!cleanedUp && !isStreamClosed) {
+            console.error('Serialization timeout after app render finished')
+            safeError(
+              new Error('Serialization timeout after app render finished'),
+            )
+            cleanup()
+          }
         }, timeoutMs)
       }
-    },
-    onError: (error) => {
-      // Don't process if already cleaned up
-      if (cleanedUp) {
-        return
-      }
-
+    } catch (error) {
+      if (cleanedUp) return
       console.error('Error reading appStream:', error)
       isAppRendering = false
-      router.serverSsr!.setRenderFinished()
-      // Clear timeout to prevent it from firing after error
-      clearTimeout(timeoutHandle)
-      // Clear string buffers to prevent memory leaks
-      leftover = ''
-      leftoverHtml = ''
-      routerStreamBuffer = ''
-      pendingClosingTags = ''
-      finalPassThrough.destroy(error)
-      injectedHtmlDonePromise.reject(error)
-    },
+      router.serverSsr?.setRenderFinished()
+      safeError(error)
+      cleanup()
+    } finally {
+      reader.releaseLock()
+    }
+  })().catch((error) => {
+    // Handle any errors that occur outside the try block (e.g., getReader() failure)
+    if (cleanedUp) return
+    console.error('Error in stream transform:', error)
+    safeError(error)
+    cleanup()
   })
 
-  return finalPassThrough.stream
+  return stream
 }