fix: cleanup streamed values (#5870)

Author: schiller-manuel

packages/react-router/src/HeadContent.tsx

@@ -155,7 +155,7 @@ export const useTags = () => {
     structuralSharing: true as any,
   })
 
-  const headScripts = useRouterState({
+  const headScripts: Array<RouterManagedTag> = useRouterState({
     select: (state) =>
       (
         state.matches
@@ -173,12 +173,29 @@ export const useTags = () => {
     structuralSharing: true as any,
   })
 
+  let serverHeadScript: RouterManagedTag | undefined = undefined
+
+  if (router.serverSsr) {
+    const bufferedScripts = router.serverSsr.takeBufferedScripts()
+    if (bufferedScripts) {
+      serverHeadScript = {
+        tag: 'script',
+        attrs: {
+          nonce,
+          className: '$tsr',
+        },
+        children: bufferedScripts,
+      }
+    }
+  }
+
   return uniqBy(
     [
       ...meta,
       ...preloadMeta,
       ...links,
       ...styles,
+      ...(serverHeadScript ? [serverHeadScript] : []),
       ...headScripts,
     ] as Array<RouterManagedTag>,
     (d) => {

packages/react-router/src/ScriptOnce.tsx

@@ -2,7 +2,6 @@ import { useRouter } from './useRouter'
 
 /**
  * Server-only helper to emit a script tag exactly once during SSR.
- * Appends an internal marker to signal hydration completion.
  */
 export function ScriptOnce({ children }: { children: string }) {
   const router = useRouter()

packages/router-core/src/router.ts

@@ -754,6 +754,7 @@ export interface ServerSsr {
   isDehydrated: () => boolean
   onRenderFinished: (listener: () => void) => void
   dehydrate: () => Promise<void>
+  takeBufferedScripts: () => string | undefined
 }
 
 export type AnyRouterWithContext<TContext> = RouterCore<

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

@@ -1 +1,2 @@
 export const GLOBAL_TSR = '$_TSR'
+export declare const GLOBAL_SEROVAL: '$R'

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

@@ -6,11 +6,12 @@ import type { AnyRouter } from '../router'
 import type { Manifest } from '../manifest'
 import type { RouteContextOptions } from '../route'
 import type { AnySerializationAdapter } from './serializer/transformer'
-import type { GLOBAL_TSR } from './constants'
+import type { GLOBAL_SEROVAL, GLOBAL_TSR } from './constants'
 
 declare global {
   interface Window {
     [GLOBAL_TSR]?: TsrSsrGlobal
+    [GLOBAL_SEROVAL]?: any
   }
 }
 
@@ -25,6 +26,10 @@ export interface TsrSsrGlobal {
   t?: Map<string, (value: any) => any>
   // this flag indicates whether the transformers were initialized
   initialized?: boolean
+  // router is hydrated and doesnt need the streamed values anymore
+  hydrated?: boolean
+  // stream has ended
+  streamEnd?: boolean
 }
 
 function hydrateMatch(
@@ -165,6 +170,10 @@ export async function hydrate(router: AnyRouter): Promise<any> {
   // Allow the user to handle custom hydration data
   await router.options.hydrate?.(dehydratedData)
 
+  window.$_TSR.hydrated = true
+  // potentially clean up streamed values IF stream has ended already
+  window.$_TSR.c()
+
   // now that all necessary data is hydrated:
   // 1) fully reconstruct the route context
   // 2) execute `head()` and `scripts()` for each match

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

@@ -48,6 +48,54 @@ export function dehydrateMatch(match: AnyRouteMatch): DehydratedMatch {
   return dehydratedMatch
 }
 
+const INITIAL_SCRIPTS = [
+  getCrossReferenceHeader(SCOPE_ID),
+  minifiedTsrBootStrapScript,
+]
+
+class ScriptBuffer {
+  constructor(private router: AnyRouter) {}
+  private _queue: Array<string> = [...INITIAL_SCRIPTS]
+  private _scriptBarrierLifted = false
+
+  enqueue(script: string) {
+    if (this._scriptBarrierLifted && this._queue.length === 0) {
+      queueMicrotask(() => {
+        this.injectBufferedScripts()
+      })
+    }
+    this._queue.push(script)
+  }
+
+  liftBarrier() {
+    if (this._scriptBarrierLifted) return
+    this._scriptBarrierLifted = true
+    if (this._queue.length > 0) {
+      queueMicrotask(() => {
+        this.injectBufferedScripts()
+      })
+    }
+  }
+
+  takeAll() {
+    const bufferedScripts = this._queue
+    this._queue = []
+    if (bufferedScripts.length === 0) {
+      return undefined
+    }
+    bufferedScripts.push(`${GLOBAL_TSR}.c()`)
+    const joinedScripts = bufferedScripts.join(';')
+    return joinedScripts
+  }
+
+  injectBufferedScripts() {
+    const scriptsToInject = this.takeAll()
+    if (scriptsToInject) {
+      this.router.serverSsr!.injectScript(() => scriptsToInject)
+    }
+  }
+}
+
 export function attachRouterServerSsrUtils({
   router,
   manifest,
@@ -58,16 +106,9 @@ export function attachRouterServerSsrUtils({
   router.ssr = {
     manifest,
   }
-  let initialScriptSent = false
-  const getInitialScript = () => {
-    if (initialScriptSent) {
-      return ''
-    }
-    initialScriptSent = true
-    return `${getCrossReferenceHeader(SCOPE_ID)};${minifiedTsrBootStrapScript};`
-  }
   let _dehydrated = false
   const listeners: Array<() => void> = []
+  const scriptBuffer = new ScriptBuffer(router)
 
   router.serverSsr = {
     injectedHtml: [],
@@ -84,7 +125,10 @@ export function attachRouterServerSsrUtils({
     injectScript: (getScript) => {
       return router.serverSsr!.injectHtml(async () => {
         const script = await getScript()
-        return `<script ${router.options.ssr?.nonce ? `nonce='${router.options.ssr.nonce}'` : ''} class='$tsr'>${getInitialScript()}${script};$_TSR.c()</script>`
+        if (!script) {
+          return ''
+        }
+        return `<script${router.options.ssr?.nonce ? `nonce='${router.options.ssr.nonce}' ` : ''} class='$tsr'>${script}</script>`
       })
     },
     dehydrate: async () => {
@@ -104,7 +148,10 @@ export function attachRouterServerSsrUtils({
       if (lastMatchId) {
         dehydratedRouter.lastMatchId = lastMatchId
       }
-      dehydratedRouter.dehydratedData = await router.options.dehydrate?.()
+      const dehydratedData = await router.options.dehydrate?.()
+      if (dehydratedData) {
+        dehydratedRouter.dehydratedData = dehydratedData
+      }
       _dehydrated = true
 
       const p = createControlledPromise<string>()
@@ -115,6 +162,7 @@ export function attachRouterServerSsrUtils({
             | Array<AnySerializationAdapter>
             | undefined
         )?.map((t) => makeSsrSerovalPlugin(t, trackPlugins)) ?? []
+
       crossSerializeStream(dehydratedRouter, {
         refs: new Map(),
         plugins: [...plugins, ...defaultSerovalPlugins],
@@ -123,10 +171,13 @@ export function attachRouterServerSsrUtils({
           if (trackPlugins.didRun) {
             serialized = GLOBAL_TSR + '.p(()=>' + serialized + ')'
           }
-          router.serverSsr!.injectScript(() => serialized)
+          scriptBuffer.enqueue(serialized)
         },
         scopeId: SCOPE_ID,
-        onDone: () => p.resolve(''),
+        onDone: () => {
+          scriptBuffer.enqueue(GLOBAL_TSR + '.streamEnd=true')
+          p.resolve('')
+        },
         onError: (err) => p.reject(err),
       })
       // make sure the stream is kept open until the promise is resolved
@@ -138,6 +189,12 @@ export function attachRouterServerSsrUtils({
     onRenderFinished: (listener) => listeners.push(listener),
     setRenderFinished: () => {
       listeners.forEach((l) => l())
+      scriptBuffer.liftBarrier()
+    },
+    takeBufferedScripts() {
+      const scripts = scriptBuffer.takeAll()
+      scriptBuffer.liftBarrier()
+      return scripts
     },
   }
 }

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

@@ -3,6 +3,10 @@ self.$_TSR = {
     document.querySelectorAll('.\\$tsr').forEach((o) => {
       o.remove()
     })
+    if (this.hydrated && this.streamEnd) {
+      delete self.$_TSR
+      delete self.$R['tsr']
+    }
   },
   p(script) {
     !this.initialized ? this.buffer.push(script) : script()

packages/solid-router/src/HeadContent.tsx

@@ -166,13 +166,30 @@ export const useTags = () => {
       })),
   })
 
+  let serverHeadScript: RouterManagedTag | undefined = undefined
+
+  if (router.serverSsr) {
+    const bufferedScripts = router.serverSsr.takeBufferedScripts()
+    if (bufferedScripts) {
+      serverHeadScript = {
+        tag: 'script',
+        attrs: {
+          nonce,
+          class: '$tsr',
+        },
+        children: bufferedScripts,
+      }
+    }
+  }
+
   return () =>
     uniqBy(
       [
         ...meta(),
         ...preloadMeta(),
         ...links(),
         ...styles(),
+        ...(serverHeadScript ? [serverHeadScript] : []),
         ...headScripts(),
       ] as Array<RouterManagedTag>,
       (d) => {

packages/start-server-core/src/router-manifest.ts

@@ -36,19 +36,24 @@ export async function getStartManifest() {
   })
 
   const manifest = {
-    ...startManifest,
     routes: Object.fromEntries(
       Object.entries(startManifest.routes).map(([k, v]) => {
         const { preloads, assets } = v
         const result = {} as {
           preloads?: Array<string>
           assets?: Array<RouterManagedTag>
         }
-        if (preloads) {
+        let hasData = false
+        if (preloads && preloads.length > 0) {
           result['preloads'] = preloads
+          hasData = true
         }
-        if (assets) {
+        if (assets && assets.length > 0) {
           result['assets'] = assets
+          hasData = true
+        }
+        if (!hasData) {
+          return []
         }
         return [k, result]
       }),