feat: implement transformReadableStreamWithRouter

packages/react-router-server/src/server/index.tsx

@@ -1,2 +1,5 @@
 export { StartServer } from './StartServer'
-export { transformStreamWithRouter } from './transformStreamWithRouter'
+export {
+  transformStreamWithRouter,
+  transformReadableStreamWithRouter,
+} from './transformStreamWithRouter'

packages/react-router-server/src/server/transformStreamWithRouter.tsx

@@ -2,83 +2,120 @@ import { Transform } from 'stream'
 import type { AnyRouter } from '@tanstack/react-router'
 
 export function transformStreamWithRouter(router: AnyRouter) {
-  return transformStreamHtmlCallback(async () => {
+  const callbacks = transformHtmlCallbacks(injectorFromRouter(router))
+  return new Transform({
+    transform(chunk, _encoding, callback) {
+      callbacks
+        .transform(chunk, this.push.bind(this))
+        .then(() => callback())
+        .catch((err) => callback(err))
+    },
+    flush(callback) {
+      callbacks
+        .flush(this.push.bind(this))
+        .then(() => callback())
+        .catch((err) => callback(err))
+    },
+  })
+}
+
+export function transformReadableStreamWithRouter(router: AnyRouter) {
+  const callbacks = transformHtmlCallbacks(injectorFromRouter(router))
+  return new TransformStream<string>({
+    transform(chunk, controller) {
+      return callbacks.transform(chunk, (chunkToPush) => {
+        controller.enqueue(chunkToPush)
+        return true
+      })
+    },
+    flush(controller) {
+      return callbacks.flush((chunkToPush) => {
+        controller.enqueue(chunkToPush)
+        return true
+      })
+    },
+  })
+}
+
+function injectorFromRouter(router: AnyRouter) {
+  return async () => {
     const injectorPromises = router.injectedHtml.map((d) =>
       typeof d === 'function' ? d() : d,
     )
     const injectors = await Promise.all(injectorPromises)
     router.injectedHtml = []
     return injectors.join('')
-  })
+  }
 }
-function transformStreamHtmlCallback(injector: () => Promise<string>) {
+
+// regex pattern for matching closing body and html tags
+const patternBody = /(<\/body>)/
+const patternHtml = /(<\/html>)/
+
+// regex pattern for matching closing tags
+const pattern = /(<\/[a-zA-Z][\w:.-]*?>)/g
+
+const textDecoder = new TextDecoder()
+
+function transformHtmlCallbacks(injector: () => Promise<string>) {
   let leftover = ''
   // If a closing tag is split across chunks, store the HTML to add after it
   // This expects that all the HTML that's added is closed properly
   let leftoverHtml = ''
 
-  return new Transform({
-    transform(chunk, encoding, callback) {
-      const chunkString = leftover + chunk.toString()
-
-      // regex pattern for matching closing body and html tags
-      const patternBody = /(<\/body>)/
-      const patternHtml = /(<\/html>)/
+  return {
+    async transform(chunk: any, push: (chunkToPush: string) => boolean) {
+      const chunkString = leftover + textDecoder.decode(chunk)
 
       const bodyMatch = chunkString.match(patternBody)
       const htmlMatch = chunkString.match(patternHtml)
 
-      injector()
-        .then((html) => {
-          // If a </body></html> sequence was found
-          if (bodyMatch && htmlMatch && bodyMatch.index! < htmlMatch.index!) {
-            const bodyIndex = bodyMatch.index! + bodyMatch[0].length
-            const htmlIndex = htmlMatch.index! + htmlMatch[0].length
-
-            // Add the arbitrary HTML before the closing body tag
-            const processed =
-              chunkString.slice(0, bodyIndex) +
-              html +
-              chunkString.slice(bodyIndex, htmlIndex) +
-              chunkString.slice(htmlIndex)
+      try {
+        const html = await injector()
+        // If a </body></html> sequence was found
+        if (bodyMatch && htmlMatch && bodyMatch.index! < htmlMatch.index!) {
+          const bodyIndex = bodyMatch.index! + bodyMatch[0].length
+          const htmlIndex = htmlMatch.index! + htmlMatch[0].length
 
-            this.push(processed)
-            leftover = ''
-          } else {
-            // For all other closing tags, add the arbitrary HTML after them
-            const pattern = /(<\/[a-zA-Z][\w:.-]*?>)/g
-            let result
-            let lastIndex = 0
+          // Add the arbitrary HTML before the closing body tag
+          const processed =
+            chunkString.slice(0, bodyIndex) +
+            html +
+            chunkString.slice(bodyIndex, htmlIndex) +
+            chunkString.slice(htmlIndex)
 
-            while ((result = pattern.exec(chunkString)) !== null) {
-              lastIndex = result.index + result[0].length
-            }
+          push(processed)
+          leftover = ''
+        } else {
+          // For all other closing tags, add the arbitrary HTML after them
+          let result
+          let lastIndex = 0
 
-            // If a closing tag was found, add the arbitrary HTML and send it through
-            if (lastIndex > 0) {
-              const processed =
-                chunkString.slice(0, lastIndex) + html + leftoverHtml
-              this.push(processed)
-              leftover = chunkString.slice(lastIndex)
-            } else {
-              // If no closing tag was found, store the chunk to process with the next one
-              leftover = chunkString
-              leftoverHtml += html
-            }
+          while ((result = pattern.exec(chunkString)) !== null) {
+            lastIndex = result.index + result[0].length
           }
 
-          callback()
-        })
-        .catch((err) => {
-          console.error(err)
-          callback(err)
-        })
+          // If a closing tag was found, add the arbitrary HTML and send it through
+          if (lastIndex > 0) {
+            const processed =
+              chunkString.slice(0, lastIndex) + html + leftoverHtml
+            push(processed)
+            leftover = chunkString.slice(lastIndex)
+          } else {
+            // If no closing tag was found, store the chunk to process with the next one
+            leftover = chunkString
+            leftoverHtml += html
+          }
+        }
+      } catch (err) {
+        console.error(err)
+        throw err
+      }
     },
-    flush(callback) {
+    async flush(push: (chunkToPush: string) => boolean) {
       if (leftover) {
-        this.push(leftover)
+        push(leftover)
       }
-      callback()
     },
-  })
+  }
 }