fix: hydration improvements

TanStack · Jul 31, 2023 · 4bd5c2b · 4bd5c2b
1 parent 26041c5
commit 4bd5c2b
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 20 deletions.
diff --git a/packages/react-start/src/client.tsx b/packages/react-start/src/client.tsx
@@ -1,5 +1,5 @@
 import * as React from 'react'
-import { AnyRouter, RouterProvider } from '@tanstack/router'
+import { AnyRouter, RouterProvider, useRouter } from '@tanstack/router'
 // @ts-ignore
 import cprc from '@gisatcz/cross-package-react-context'
 
@@ -8,10 +8,14 @@ export function StartClient(props: { router: AnyRouter }) {
 }
 
 export function DehydrateRouter() {
-  const dehydrated = React.useContext(
+  const router = useRouter()
+
+  const dehydratedCtx = React.useContext(
     cprc.getContext('TanStackRouterHydrationContext', {}),
   )
 
+  const dehydrated = router.dehydratedData || dehydratedCtx
+
   return (
     <script
       id="__TSR_DEHYDRATED__"

diff --git a/packages/router/src/react.tsx b/packages/router/src/react.tsx
@@ -368,10 +368,6 @@ export function RouterProvider<
   )
 }
 
-export function useHydrationContext() {
-  return cprc.getContext('TanStackRouterHydrationContext', {}).payload
-}
-
 export function useRouter(): RegisteredRouter {
   const value = React.useContext(routerContext)
   warning(value, 'useRouter must be used inside a <Router> component!')

diff --git a/packages/router/src/router.ts b/packages/router/src/router.ts
@@ -274,6 +274,7 @@ export class Router<
   state: RouterState<TRoutesInfo>
   startedLoadingAt = Date.now()
   resolveNavigation: () => void = () => {}
+  dehydratedData?: TDehydrated
 
   constructor(options: RouterConstructorOptions<TRouteTree, TDehydrated>) {
     this.options = {
@@ -582,7 +583,6 @@ export class Router<
         componentTypes.some((d) => route.options[d]?.preload)
       )
 
-
       const routeMatch: RouteMatch = {
         id: matchId,
         routeId: route.id,
@@ -1104,13 +1104,7 @@ export class Router<
 
   dehydrate = (): DehydratedRouter => {
     return {
-      state: {
-        ...pick(this.state, ['location', 'status', 'lastUpdated']),
-        // matches: this.state.matches.map((m) => ({
-        //   id: m.id,
-        //   promiseKeys: Object.keys(m.__promisesByKey),
-        // })),
-      },
+      state: pick(this.state, ['location', 'status', 'lastUpdated']),
     }
   }
 
@@ -1127,7 +1121,7 @@ export class Router<
     )
 
     const ctx = _ctx
-
+    this.dehydratedData = ctx.payload as any
     this.options.hydrate?.(ctx.payload as any)
 
     this.__store.setState((s) => {
@@ -1141,11 +1135,6 @@ export class Router<
 
     await this.load()
 
-    // this.state.matches.forEach((m) => {
-    //   m.__promiseKeys =
-    //     ctx.router.state.matches.find((d) => d.id === m.id)?.promiseKeys ?? []
-    // })
-
     return
   }