fix(router-core): preserve retained beforeLoad context while pending

packages/react-router/tests/retained-ancestor-beforeload-context.test.tsx

@@ -0,0 +1,153 @@
+import { cleanup, render, screen, waitFor } from '@testing-library/react'
+import { afterEach, expect, test } from 'vitest'
+import { createControlledPromise } from '@tanstack/router-core'
+import {
+  Outlet,
+  RouterProvider,
+  createMemoryHistory,
+  createRootRoute,
+  createRoute,
+  createRouter,
+} from '../src'
+
+afterEach(cleanup)
+
+// While a child route shows its pending fallback, ancestor routes stay mounted
+// and must keep the context their own `beforeLoad` produced. Losing it is not
+// cosmetic: ancestors that reach into a nested context value - the usual shape
+// for an auth/session context - throw while they are still on screen.
+test('retained ancestor keeps its beforeLoad context while a child route is pending', async () => {
+  const childLoader = createControlledPromise<void>()
+  // Errors thrown while rendering the mounted ancestor. Caught here only so the
+  // failure surfaces as an assertion instead of taking out the route's error
+  // boundary; in an app this is an uncaught render crash.
+  const renderErrors: Array<string> = []
+
+  const rootRoute = createRootRoute({
+    beforeLoad: () => ({ auth: { user: 'ada' } }),
+    component: function RootLayout() {
+      const { auth } = rootRoute.useRouteContext()
+      let user
+      try {
+        user = auth.user
+      } catch (error) {
+        renderErrors.push(String(error))
+      }
+      return (
+        <>
+          <div>{`user:${user ?? 'MISSING'}`}</div>
+          <Outlet />
+        </>
+      )
+    },
+  })
+  const indexRoute = createRoute({
+    getParentRoute: () => rootRoute,
+    path: '/',
+    component: () => <div>Home</div>,
+  })
+  const childRoute = createRoute({
+    getParentRoute: () => rootRoute,
+    path: '/child',
+    pendingMs: 0,
+    pendingComponent: () => <div>Pending</div>,
+    loader: () => childLoader,
+    component: () => <div>Child</div>,
+  })
+  const router = createRouter({
+    routeTree: rootRoute.addChildren([indexRoute, childRoute]),
+    history: createMemoryHistory({ initialEntries: ['/'] }),
+  })
+
+  render(<RouterProvider router={router} />)
+  await screen.findByText('Home')
+  await waitFor(() => expect(router.state.status).toBe('idle'))
+
+  const navigation = router.navigate({ to: '/child' })
+  await screen.findByText('Pending')
+
+  expect(renderErrors).toEqual([])
+  expect(screen.getByText('user:ada')).toBeInTheDocument()
+
+  childLoader.resolve()
+  await navigation
+})
+
+// Contextualization walks the lane serially, so an ancestor whose `beforeLoad`
+// is still in flight parks the walk above every deeper ancestor. The pending
+// fallback is published from that parked state, so the ancestors the walk has
+// not reached yet must already be presentable.
+test('retained ancestor keeps its beforeLoad context while an ancestor above it is pending', async () => {
+  const childLoader = createControlledPromise<void>()
+  const rootBeforeLoad = createControlledPromise<void>()
+  let rootResolved = false
+  let rootBeforeLoadStarted = false
+  const renderErrors: Array<string> = []
+
+  const rootRoute = createRootRoute({
+    beforeLoad: async () => {
+      // Only the navigation blocks; the initial load must settle normally.
+      if (rootResolved) {
+        rootBeforeLoadStarted = true
+        await rootBeforeLoad
+      }
+      rootResolved = true
+      return { session: 'live' }
+    },
+    component: () => <Outlet />,
+  })
+  const dashRoute = createRoute({
+    getParentRoute: () => rootRoute,
+    path: '/dash',
+    beforeLoad: () => ({ auth: { user: 'ada' } }),
+    component: function DashLayout() {
+      const { auth } = dashRoute.useRouteContext()
+      let user
+      try {
+        user = auth.user
+      } catch (error) {
+        renderErrors.push(String(error))
+      }
+      return (
+        <>
+          <div>{`user:${user ?? 'MISSING'}`}</div>
+          <Outlet />
+        </>
+      )
+    },
+  })
+  const overviewRoute = createRoute({
+    getParentRoute: () => dashRoute,
+    path: 'overview',
+    component: () => <div>Overview</div>,
+  })
+  const detailRoute = createRoute({
+    getParentRoute: () => dashRoute,
+    path: 'detail',
+    pendingMs: 0,
+    pendingComponent: () => <div>Pending</div>,
+    loader: () => childLoader,
+    component: () => <div>Detail</div>,
+  })
+  const router = createRouter({
+    routeTree: rootRoute.addChildren([
+      dashRoute.addChildren([overviewRoute, detailRoute]),
+    ]),
+    history: createMemoryHistory({ initialEntries: ['/dash/overview'] }),
+  })
+
+  render(<RouterProvider router={router} />)
+  await screen.findByText('Overview')
+  await waitFor(() => expect(router.state.status).toBe('idle'))
+
+  const navigation = router.navigate({ to: '/dash/detail' })
+  await screen.findByText('Pending')
+
+  expect(rootBeforeLoadStarted).toBe(true)
+  expect(renderErrors).toEqual([])
+  expect(screen.getByText('user:ada')).toBeInTheDocument()
+
+  rootBeforeLoad.resolve()
+  childLoader.resolve()
+  await navigation
+})

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

@@ -362,6 +362,12 @@ async function contextualize(
   const [location, matches] = lane
   const signal = options[0 /* controller */].signal
   const preload = !!options[4 /* preload */]
+  for (let index = 0; index < retainedEnd; index++) {
+    const retained = options[3 /* base */][index]?.context
+    if (retained) {
+      matches[index]!.context = retained
+    }
+  }
   for (let index = options[7 /* resolvedPrefix */] ?? 0; index < end; index++) {
     const match = matches[index]!
     const route = getRoute(router, match)
@@ -397,8 +403,13 @@ async function contextualize(
         ...parentContext,
         ...routeContext,
       }
-      match.context = context
+      const retained =
+        index < retainedEnd && route.options.beforeLoad
+          ? options[3 /* base */][index]?.context
+          : undefined
+      match.context = retained ?? context
     } catch (cause) {
+      match.context = context
       releaseFlight(router, match)
       return [index, normalizeLaneError(route, cause, options)]
     }
@@ -408,6 +419,7 @@ async function contextualize(
     }
     const validationError = match.paramsError ?? match.searchError
     if (validationError !== undefined) {
+      match.context = context
       releaseFlight(router, match)
       return [index, normalizeLaneError(route, validationError, options)]
     }
@@ -448,6 +460,7 @@ async function contextualize(
       }
       const outcome = normalize(result, false, route.id)
       if (outcome[0 /* kind */] !== SUCCESS) {
+        match.context = context
         releaseFlight(router, match)
         return [index, outcome]
       }
@@ -456,6 +469,7 @@ async function contextualize(
         ...result,
       }
     } catch (cause) {
+      match.context = context
       releaseFlight(router, match)
       return [index, normalizeLaneError(route, cause, options)]
     } finally {