fix: generalize hydration cleanup

packages/react-router/src/ScriptOnce.tsx

@@ -12,9 +12,8 @@ export function ScriptOnce({ children }: { children: string }) {
   return (
     <script
       nonce={router.options.ssr?.nonce}
-      className="$tsr"
       dangerouslySetInnerHTML={{
-        __html: children + ';typeof $_TSR !== "undefined" && $_TSR.c()',
+        __html: children + ';document.currentScript.remove()',
       }}
     />
   )

packages/react-start-client/src/StartClient.tsx

@@ -1,6 +1,6 @@
 import { Await, RouterProvider } from '@tanstack/react-router'
 
-import { hydrateStart } from '@tanstack/start-client-core/client'
+import { hydrateStart } from './hydrateStart'
 
 import type { AnyRouter } from '@tanstack/router-core'
 

packages/react-start-client/src/hydrateStart.ts

@@ -0,0 +1,12 @@
+import { hydrateStart as coreHydrateStart } from '@tanstack/start-client-core/client'
+import type { AnyRouter } from '@tanstack/router-core'
+
+/**
+ * React-specific wrapper for hydrateStart that signals hydration completion
+ */
+export async function hydrateStart(): Promise<AnyRouter> {
+  const router = await coreHydrateStart()
+  // Signal that router hydration is complete so cleanup can happen if stream has ended
+  window.$_TSR?.h()
+  return router
+}

packages/react-start-client/src/index.tsx

@@ -1,2 +1,3 @@
 export { StartClient } from './StartClient'
 export { renderRsc } from './renderRSC'
+export { hydrateStart } from './hydrateStart'

packages/react-start-client/tsconfig.json

@@ -4,9 +4,5 @@
     "jsx": "react-jsx",
     "module": "esnext"
   },
-  "include": [
-    "src",
-    "vite.config.ts",
-    "../start-client-core/src/startEntry.d.ts"
-  ]
+  "include": ["src", "vite.config.ts"]
 }

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

@@ -3,3 +3,4 @@ export { json } from './json'
 export type { JsonResponse } from './json'
 export { hydrate } from './ssr-client'
 export * from './ssr-client'
+export type { TsrSsrGlobal, DehydratedMatch, DehydratedRouter } from './types'

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

@@ -2,41 +2,20 @@ import invariant from 'tiny-invariant'
 import { batch } from '@tanstack/store'
 import { isNotFound } from '../not-found'
 import { createControlledPromise } from '../utils'
-import type { AnyRouteMatch, MakeRouteMatch } from '../Matches'
+import type { GLOBAL_SEROVAL, GLOBAL_TSR } from './constants'
+import type { DehydratedMatch, TsrSsrGlobal } from './types'
+import type { AnyRouteMatch } from '../Matches'
 import type { AnyRouter } from '../router'
-import type { Manifest } from '../manifest'
 import type { RouteContextOptions } from '../route'
 import type { AnySerializationAdapter } from './serializer/transformer'
-import type { GLOBAL_SEROVAL, GLOBAL_TSR } from './constants'
 
 declare global {
   interface Window {
     [GLOBAL_TSR]?: TsrSsrGlobal
     [GLOBAL_SEROVAL]?: any
-    // Vue sets this to true before TSR scripts run to defer cleanup until after hydration
-    $_TSR_DEFER?: boolean
   }
 }
 
-export interface TsrSsrGlobal {
-  router?: DehydratedRouter
-  // clean scripts; shortened since this is sent for each streamed script
-  c: () => void
-  // push script into buffer; shortened since this is sent for each streamed script as soon as the first custom transformer was invoked
-  p: (script: () => void) => void
-  buffer: Array<() => void>
-  // custom transformers, shortened since this is sent for each streamed value that needs a custom transformer
-  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
-  // called by Vue after hydration to perform deferred cleanup
-  cleanup?: () => void
-}
-
 function hydrateMatch(
   match: AnyRouteMatch,
   deyhydratedMatch: DehydratedMatch,
@@ -49,22 +28,6 @@ function hydrateMatch(
   match.updatedAt = deyhydratedMatch.u
   match.error = deyhydratedMatch.e
 }
-export interface DehydratedMatch {
-  i: MakeRouteMatch['id']
-  b?: MakeRouteMatch['__beforeLoadContext']
-  l?: MakeRouteMatch['loaderData']
-  e?: MakeRouteMatch['error']
-  u: MakeRouteMatch['updatedAt']
-  s: MakeRouteMatch['status']
-  ssr?: MakeRouteMatch['ssr']
-}
-
-export interface DehydratedRouter {
-  manifest: Manifest | undefined
-  dehydratedData?: any
-  lastMatchId?: string
-  matches: Array<DehydratedMatch>
-}
 
 export async function hydrate(router: AnyRouter): Promise<any> {
   invariant(
@@ -183,10 +146,6 @@ 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

@@ -6,10 +6,9 @@ import { GLOBAL_TSR } from './constants'
 import { defaultSerovalPlugins } from './serializer/seroval-plugins'
 import { makeSsrSerovalPlugin } from './serializer/transformer'
 import { TSR_SCRIPT_BARRIER_ID } from './transformStreamWithRouter'
+import type { DehydratedMatch, DehydratedRouter } from './types'
 import type { AnySerializationAdapter } from './serializer/transformer'
 import type { AnyRouter } from '../router'
-import type { DehydratedMatch } from './ssr-client'
-import type { DehydratedRouter } from './client'
 import type { AnyRouteMatch } from '../Matches'
 import type { Manifest, RouterManagedTag } from '../manifest'
 
@@ -91,7 +90,7 @@ class ScriptBuffer {
     if (bufferedScripts.length === 0) {
       return undefined
     }
-    bufferedScripts.push(`${GLOBAL_TSR}.c()`)
+    bufferedScripts.push(`document.currentScript.remove()`)
     const joinedScripts = bufferedScripts.join(';')
     return joinedScripts
   }
@@ -143,7 +142,7 @@ export function attachRouterServerSsrUtils({
         if (!script) {
           return ''
         }
-        return `<script${router.options.ssr?.nonce ? ` nonce='${router.options.ssr.nonce}'` : ''} class='$tsr'>${script}</script>`
+        return `<script${router.options.ssr?.nonce ? ` nonce='${router.options.ssr.nonce}'` : ''}>${script}</script>`
       })
     },
     dehydrate: async () => {
@@ -223,7 +222,7 @@ export function attachRouterServerSsrUtils({
         },
         scopeId: SCOPE_ID,
         onDone: () => {
-          scriptBuffer.enqueue(GLOBAL_TSR + '.streamEnd=true')
+          scriptBuffer.enqueue(GLOBAL_TSR + '.e()')
           p.resolve('')
         },
         onError: (err) => p.reject(err),

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

@@ -1,25 +1,15 @@
 self.$_TSR = {
-  c() {
-    // If Vue has set the defer flag, don't remove scripts yet - wait for Vue to call cleanup()
-    if (self.$_TSR_DEFER) {
-      return
-    }
-    document.querySelectorAll('.\\$tsr').forEach((o) => {
-      o.remove()
-    })
-    if (this.hydrated && this.streamEnd) {
-      delete self.$_TSR
-      delete self.$R['tsr']
-    }
+  h() {
+    this.hydrated = true
+    this.c()
+  },
+  e() {
+    this.streamEnded = true
+    this.c()
   },
-  // Called by Vue after hydration is complete to perform deferred cleanup
-  cleanup() {
-    document.querySelectorAll('.\\$tsr').forEach((o) => {
-      o.remove()
-    })
-    if (this.hydrated && this.streamEnd) {
+  c() {
+    if (this.hydrated && this.streamEnded) {
       delete self.$_TSR
-      delete self.$_TSR_DEFER
       delete self.$R['tsr']
     }
   },

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

@@ -0,0 +1,40 @@
+import type { Manifest } from '../manifest'
+import type { MakeRouteMatch } from '../Matches'
+
+export interface DehydratedMatch {
+  i: MakeRouteMatch['id']
+  b?: MakeRouteMatch['__beforeLoadContext']
+  l?: MakeRouteMatch['loaderData']
+  e?: MakeRouteMatch['error']
+  u: MakeRouteMatch['updatedAt']
+  s: MakeRouteMatch['status']
+  ssr?: MakeRouteMatch['ssr']
+}
+
+export interface DehydratedRouter {
+  manifest: Manifest | undefined
+  dehydratedData?: any
+  lastMatchId?: string
+  matches: Array<DehydratedMatch>
+}
+
+export interface TsrSsrGlobal {
+  router?: DehydratedRouter
+  // Signal that router hydration is complete
+  h: () => void
+  // Signal that stream has ended
+  e: () => void
+  // Cleanup all hydration resources and scripts
+  c: () => void
+  // p: Push script into buffer or execute immediately
+  p: (script: () => void) => void
+  buffer: Array<() => void>
+  // custom transformers, shortened since this is sent for each streamed value that needs a custom transformer
+  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
+  streamEnded?: boolean
+}

packages/react-router/tests/hydrate.test.ts → packages/router-core/tests/hydrate.test.ts

@@ -1,13 +1,8 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { createMemoryHistory } from '@tanstack/history'
 import { hydrate } from '@tanstack/router-core/ssr/client'
-import {
-  createMemoryHistory,
-  createRootRoute,
-  createRoute,
-  createRouter,
-  notFound,
-} from '../src'
-import type { TsrSsrGlobal } from '@tanstack/router-core/ssr/client'
+import { BaseRootRoute, BaseRoute, RouterCore, notFound } from '../src'
+import type { TsrSsrGlobal } from '../src/ssr/types'
 import type { AnyRouteMatch } from '../src'
 
 describe('hydrate', () => {
@@ -25,17 +20,17 @@ describe('hydrate', () => {
 
     const history = createMemoryHistory({ initialEntries: ['/'] })
 
-    const rootRoute = createRootRoute({})
+    const rootRoute = new BaseRootRoute({})
 
-    const indexRoute = createRoute({
+    const indexRoute = new BaseRoute({
       getParentRoute: () => rootRoute,
       path: '/',
       component: () => 'Index',
       notFoundComponent: () => 'Not Found',
       head: mockHead,
     })
 
-    const otherRoute = createRoute({
+    const otherRoute = new BaseRoute({
       getParentRoute: () => indexRoute,
       path: '/other',
       component: () => 'Other',
@@ -45,7 +40,7 @@ describe('hydrate', () => {
       indexRoute.addChildren([otherRoute]),
     ])
 
-    mockRouter = createRouter({ routeTree, history, isServer: true })
+    mockRouter = new RouterCore({ routeTree, history, isServer: true })
   })
 
   afterEach(() => {
@@ -100,6 +95,8 @@ describe('hydrate', () => {
         lastMatchId: '/',
         matches: [],
       },
+      h: vi.fn(),
+      e: vi.fn(),
       c: vi.fn(),
       p: vi.fn(),
       buffer: mockBuffer,
@@ -127,6 +124,8 @@ describe('hydrate', () => {
         lastMatchId: '/',
         matches: [],
       },
+      h: vi.fn(),
+      e: vi.fn(),
       c: vi.fn(),
       p: vi.fn(),
       buffer: [],
@@ -148,6 +147,8 @@ describe('hydrate', () => {
         lastMatchId: '/',
         matches: [],
       },
+      h: vi.fn(),
+      e: vi.fn(),
       c: vi.fn(),
       p: vi.fn(),
       buffer: [],
@@ -199,6 +200,8 @@ describe('hydrate', () => {
         lastMatchId: '/',
         matches: dehydratedMatches,
       },
+      h: vi.fn(),
+      e: vi.fn(),
       c: vi.fn(),
       p: vi.fn(),
       buffer: [],
@@ -242,6 +245,8 @@ describe('hydrate', () => {
           },
         ],
       },
+      h: vi.fn(),
+      e: vi.fn(),
       c: vi.fn(),
       p: vi.fn(),
       buffer: [],

packages/solid-router/src/ScriptOnce.tsx

@@ -15,7 +15,7 @@ export function ScriptOnce({
     <script
       nonce={router.options.ssr?.nonce}
       class="$tsr"
-      innerHTML={children + ';typeof $_TSR !== "undefined" && $_TSR?.c()'}
+      innerHTML={children + ';document.currentScript.remove()'}
     />
   )
 }