fix: more robust stream handling (#5996)

Author: schiller-manuel

packages/react-router/src/ssr/renderRouterToStream.tsx

@@ -61,10 +61,17 @@ export const renderRouterToStream = async ({
             }),
         onError: (error, info) => {
           console.error('Error in renderToPipeableStream:', error, info)
+          // Destroy the passthrough stream on error
+          if (!reactAppPassthrough.destroyed) {
+            reactAppPassthrough.destroy(
+              error instanceof Error ? error : new Error(String(error)),
+            )
+          }
         },
       })
     } catch (e) {
       console.error('Error in renderToPipeableStream:', e)
+      reactAppPassthrough.destroy(e instanceof Error ? e : new Error(String(e)))
     }
 
     const responseStream = transformPipeableStreamWithRouter(

packages/react-router/src/ssr/renderRouterToString.tsx

@@ -28,5 +28,7 @@ export const renderRouterToString = async ({
       status: 500,
       headers: responseHeaders,
     })
+  } finally {
+    router.serverSsr?.cleanup()
   }
 }

packages/router-core/src/ssr/createRequestHandler.ts

@@ -20,40 +20,54 @@ export function createRequestHandler<TRouter extends AnyRouter>({
 }): RequestHandler<TRouter> {
   return async (cb) => {
     const router = createRouter()
+    // Track whether the callback will handle cleanup
+    let cbWillCleanup = false
 
-    attachRouterServerSsrUtils({
-      router,
-      manifest: await getRouterManifest?.(),
-    })
+    try {
+      attachRouterServerSsrUtils({
+        router,
+        manifest: await getRouterManifest?.(),
+      })
 
-    const url = new URL(request.url, 'http://localhost')
-    const origin = getOrigin(request)
-    const href = url.href.replace(url.origin, '')
+      const url = new URL(request.url, 'http://localhost')
+      const origin = getOrigin(request)
+      const href = url.href.replace(url.origin, '')
 
-    // Create a history for the router
-    const history = createMemoryHistory({
-      initialEntries: [href],
-    })
+      // Create a history for the router
+      const history = createMemoryHistory({
+        initialEntries: [href],
+      })
 
-    // Update the router with the history and context
-    router.update({
-      history,
-      origin: router.options.origin ?? origin,
-    })
+      // Update the router with the history and context
+      router.update({
+        history,
+        origin: router.options.origin ?? origin,
+      })
 
-    await router.load()
+      await router.load()
 
-    await router.serverSsr?.dehydrate()
+      await router.serverSsr?.dehydrate()
 
-    const responseHeaders = getRequestHeaders({
-      router,
-    })
+      const responseHeaders = getRequestHeaders({
+        router,
+      })
 
-    return cb({
-      request,
-      router,
-      responseHeaders,
-    } as any)
+      // Mark that the callback will handle cleanup
+      cbWillCleanup = true
+      return cb({
+        request,
+        router,
+        responseHeaders,
+      } as any)
+    } finally {
+      if (!cbWillCleanup) {
+        // Clean up router SSR state if the callback won't handle it
+        // (e.g., if an error occurred before the callback was invoked).
+        // When the callback runs, it handles cleanup (either via transformStreamWithRouter
+        // for streaming, or directly in renderRouterToString for non-streaming).
+        router.serverSsr?.cleanup()
+      }
+    }
   }
 }
 

packages/router-core/src/ssr/ssr-server.ts

@@ -16,6 +16,7 @@ import type { Manifest, RouterManagedTag } from '../manifest'
 declare module '../router' {
   interface ServerSsr {
     setRenderFinished: () => void
+    cleanup: () => void
   }
   interface RouterEvents {
     onInjectedHtml: {
@@ -55,11 +56,17 @@ const INITIAL_SCRIPTS = [
 ]
 
 class ScriptBuffer {
-  constructor(private router: AnyRouter) {}
+  private router: AnyRouter | undefined
   private _queue: Array<string> = [...INITIAL_SCRIPTS]
   private _scriptBarrierLifted = false
+  private _cleanedUp = false
+
+  constructor(router: AnyRouter) {
+    this.router = router
+  }
 
   enqueue(script: string) {
+    if (this._cleanedUp) return
     if (this._scriptBarrierLifted && this._queue.length === 0) {
       queueMicrotask(() => {
         this.injectBufferedScripts()
@@ -69,7 +76,7 @@ class ScriptBuffer {
   }
 
   liftBarrier() {
-    if (this._scriptBarrierLifted) return
+    if (this._scriptBarrierLifted || this._cleanedUp) return
     this._scriptBarrierLifted = true
     if (this._queue.length > 0) {
       queueMicrotask(() => {
@@ -90,11 +97,18 @@ class ScriptBuffer {
   }
 
   injectBufferedScripts() {
+    if (this._cleanedUp) return
     const scriptsToInject = this.takeAll()
-    if (scriptsToInject) {
-      this.router.serverSsr!.injectScript(() => scriptsToInject)
+    if (scriptsToInject && this.router?.serverSsr) {
+      this.router.serverSsr.injectScript(() => scriptsToInject)
     }
   }
+
+  cleanup() {
+    this._cleanedUp = true
+    this._queue = []
+    this.router = undefined
+  }
 }
 
 export function attachRouterServerSsrUtils({
@@ -203,6 +217,8 @@ export function attachRouterServerSsrUtils({
     onRenderFinished: (listener) => listeners.push(listener),
     setRenderFinished: () => {
       listeners.forEach((l) => l())
+      // Clear listeners after calling them to prevent memory leaks
+      listeners.length = 0
       scriptBuffer.liftBarrier()
     },
     takeBufferedScripts() {
@@ -221,6 +237,14 @@ export function attachRouterServerSsrUtils({
     liftScriptBarrier() {
       scriptBuffer.liftBarrier()
     },
+    cleanup() {
+      // Guard against multiple cleanup calls
+      if (!router.serverSsr) return
+      listeners.length = 0
+      scriptBuffer.cleanup()
+      router.serverSsr.injectedHtml = []
+      router.serverSsr = undefined
+    },
   }
 }
 

packages/router-core/src/ssr/transformStreamWithRouter.ts

@@ -35,7 +35,7 @@ type ReadablePassthrough = {
   destroyed: boolean
 }
 
-function createPassthrough(onCancel?: () => void) {
+function createPassthrough(onCancel: () => void) {
   let controller: ReadableStreamDefaultController<any>
   const encoder = new TextEncoder()
   const stream = new ReadableStream({
@@ -44,27 +44,33 @@ function createPassthrough(onCancel?: () => void) {
     },
     cancel() {
       res.destroyed = true
-      onCancel?.()
+      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)
       }
-      controller.close()
       res.destroyed = true
+      controller.close()
     },
     destroy: (error) => {
+      // Don't destroy already destroyed stream
+      if (res.destroyed) return
       res.destroyed = true
       controller.error(error)
     },
@@ -105,14 +111,20 @@ export function transformStreamWithRouter(
 ) {
   let stopListeningToInjectedHtml: (() => void) | undefined = undefined
   let timeoutHandle: NodeJS.Timeout
+  let cleanedUp = false
 
-  const finalPassThrough = createPassthrough(() => {
+  function cleanup() {
+    if (cleanedUp) return
+    cleanedUp = true
     if (stopListeningToInjectedHtml) {
       stopListeningToInjectedHtml()
       stopListeningToInjectedHtml = undefined
     }
     clearTimeout(timeoutHandle)
-  })
+    router.serverSsr?.cleanup()
+  }
+
+  const finalPassThrough = createPassthrough(cleanup)
   const textDecoder = new TextDecoder()
 
   let isAppRendering = true
@@ -149,13 +161,16 @@ export function transformStreamWithRouter(
   )
 
   function handleInjectedHtml() {
+    // Don't process if already cleaned up
+    if (cleanedUp) return
+
     router.serverSsr!.injectedHtml.forEach((promise) => {
       processingCount++
 
       promise
         .then((html) => {
-          // Don't write to destroyed stream
-          if (finalPassThrough.destroyed) {
+          // Don't write to destroyed stream or after cleanup
+          if (cleanedUp || finalPassThrough.destroyed) {
             return
           }
           if (isAppRendering) {
@@ -165,10 +180,7 @@ export function transformStreamWithRouter(
           }
         })
         .catch((err) => {
-          // Only reject if not already settled
-          if (!finalPassThrough.destroyed) {
-            injectedHtmlDonePromise.reject(err)
-          }
+          injectedHtmlDonePromise.reject(err)
         })
         .finally(() => {
           processingCount--
@@ -183,6 +195,11 @@ export function transformStreamWithRouter(
 
   injectedHtmlDonePromise
     .then(() => {
+      // Don't process if already cleaned up or destroyed
+      if (cleanedUp || finalPassThrough.destroyed) {
+        return
+      }
+
       clearTimeout(timeoutHandle)
       const finalHtml =
         leftover + leftoverHtml + getBufferedRouterStream() + pendingClosingTags
@@ -194,19 +211,24 @@ export function transformStreamWithRouter(
       finalPassThrough.end(finalHtml)
     })
     .catch((err) => {
+      // Don't process if already cleaned up
+      if (cleanedUp || finalPassThrough.destroyed) {
+        return
+      }
+
       console.error('Error reading routerStream:', err)
       finalPassThrough.destroy(err)
     })
-    .finally(() => {
-      if (stopListeningToInjectedHtml) {
-        stopListeningToInjectedHtml()
-        stopListeningToInjectedHtml = undefined
-      }
-    })
+    .finally(cleanup)
 
   // Transform the appStream
   readStream(appStream, {
     onData: (chunk) => {
+      // Don't process if already cleaned up
+      if (cleanedUp || finalPassThrough.destroyed) {
+        return
+      }
+
       const text = decodeChunk(chunk.value)
       const chunkString = leftover + text
       const bodyEndMatch = chunkString.match(patternBodyEnd)
@@ -266,8 +288,8 @@ export function transformStreamWithRouter(
       }
     },
     onEnd: () => {
-      // Don't process if stream was already destroyed/cancelled
-      if (finalPassThrough.destroyed) {
+      // Don't process if stream was already destroyed/cancelled or cleaned up
+      if (cleanedUp || finalPassThrough.destroyed) {
         return
       }
 
@@ -288,6 +310,11 @@ export function transformStreamWithRouter(
       }
     },
     onError: (error) => {
+      // Don't process if already cleaned up
+      if (cleanedUp) {
+        return
+      }
+
       console.error('Error reading appStream:', error)
       isAppRendering = false
       router.serverSsr!.setRenderFinished()

packages/solid-router/src/ssr/renderRouterToString.tsx

@@ -40,5 +40,7 @@ export const renderRouterToString = async ({
       status: 500,
       headers: responseHeaders,
     })
+  } finally {
+    router.serverSsr?.cleanup()
   }
 }

packages/start-server-core/src/createStartHandler.ts

@@ -81,6 +81,8 @@ export function createStartHandler<TRegister = Register>(
     requestOpts,
   ) => {
     let router: AnyRouter | null = null as AnyRouter | null
+    // Track whether the callback will handle cleanup
+    let cbWillCleanup = false as boolean
     try {
       const origin = getOrigin(request)
 
@@ -207,6 +209,8 @@ export function createStartHandler<TRegister = Register>(
                   await router.serverSsr!.dehydrate()
 
                   const responseHeaders = getStartResponseHeaders({ router })
+                  // Mark that the callback will handle cleanup
+                  cbWillCleanup = true
                   const response = await cb({
                     request,
                     router,
@@ -312,9 +316,14 @@ export function createStartHandler<TRegister = Register>(
 
       return response
     } finally {
-      if (router) {
-        router = null
+      if (router && !cbWillCleanup) {
+        // Clean up router SSR state if it was set up but won't be cleaned up by the callback
+        // (e.g., in redirect cases or early returns before the callback is invoked).
+        // When the callback runs, it handles cleanup (either via transformStreamWithRouter
+        // for streaming, or directly in renderRouterToString for non-streaming).
+        router.serverSsr?.cleanup()
       }
+      router = null
     }
   }