Fix TanStack Router SSR hydration mismatch in Pro async path

packages/react-on-rails-pro/src/tanstack-router/clientHydrate.ts

@@ -1,101 +1,56 @@
-import { createElement, useEffect, useRef, type ComponentType, type ReactElement } from 'react';
+import { createElement, useEffect, useRef, type ReactElement } from 'react';
 import type {
   DehydratedRouterState,
   TanStackHistory,
   TanStackRouter,
   TanStackRouterOptions,
-  TanStackSsrRouterState,
 } from './types.ts';
 import type { RailsContext } from 'react-on-rails/types';
 
-/* eslint-disable import/prefer-default-export */
+/* eslint-disable import/prefer-default-export, no-underscore-dangle */
 
 /**
  * Client-side hydration for a TanStack Router app.
  *
+ * Uses RouterProvider directly (not RouterClient) to match the server-rendered
+ * component tree. RouterClient wraps RouterProvider in <Await> which always
+ * suspends on first render via defer(), causing a hydration mismatch because
+ * the server renders RouterProvider directly without an <Await> wrapper.
+ *
  * Flow:
  * 1. Create router with browser history
- * 2. Hydrate from dehydrated state provided by serverRenderTanStackAppAsync
- * 3. Set router.ssr to skip auto-load on mount (Transitioner behavior, legacy path only)
- * 4. After hydration, trigger router.load() to enable client-side navigation
- * 5. Return a React component that renders RouterProvider
+ * 2. If SSR payload exists, synchronously inject route matches to match server output
+ * 3. Set router.ssr to prevent Transitioner auto-load during hydration
+ * 4. Hydrate with dehydrated router state if available
+ * 5. After mount, trigger router.load() and clear SSR flag once settled
  */
 
 interface TanStackHydrationAppProps {
   options: TanStackRouterOptions;
   incomingProps: Record<string, unknown>;
   railsContext: RailsContext & { serverSide: false };
   RouterProvider: React.ComponentType<{ router: TanStackRouter }>;
-  RouterClient?: ComponentType<{ router: TanStackRouter }>;
   createBrowserHistory: () => TanStackHistory;
 }
 
-interface TanStackSsrGlobal {
-  router?: TanStackSsrRouterState;
-  buffer: Array<() => void>;
-  initialized?: boolean;
-  hydrated?: boolean;
-  streamEnded?: boolean;
-  // These short keys are TanStack Router's upstream $_TSR bootstrap contract:
-  // h = hydrated, e = stream ended, c = cleanup, p = run-or-buffer script.
-  h: () => void;
-  e: () => void;
-  c: () => void;
-  p: (script: () => void) => void;
-}
-
-function setTanStackSsrGlobal(ssrRouter: TanStackSsrRouterState): void {
-  const globalState = window as unknown as Record<string, unknown>;
-  const existing = globalState.$_TSR as TanStackSsrGlobal | undefined;
-  const tsrGlobal: TanStackSsrGlobal = existing ?? {
-    buffer: [],
-    h() {
-      this.hydrated = true;
-    },
-    e() {
-      this.streamEnded = true;
-    },
-    c() {},
-    p(script: () => void) {
-      if (!this.initialized) {
-        this.buffer.push(script);
-        return;
-      }
-
-      script();
-    },
-  };
-
-  tsrGlobal.router = ssrRouter;
-  globalState.$_TSR = tsrGlobal;
-}
-
 function TanStackHydrationApp({
   options,
   incomingProps,
   // eslint-disable-next-line @typescript-eslint/no-unused-vars -- Required by TanStackHydrationAppProps interface.
   railsContext: _railsContext,
   RouterProvider,
-  RouterClient,
   createBrowserHistory,
 }: TanStackHydrationAppProps): ReactElement {
-  // eslint-disable-next-line no-underscore-dangle -- Internal hydration payload key injected by server-side render.
   const dehydratedState = incomingProps.__tanstackRouterDehydratedState as
     | DehydratedRouterState
     | null
     | undefined;
-  const ssrRouter = dehydratedState?.ssrRouter;
   const hasSsrPayload = dehydratedState != null;
-  const hasSsrRouter = ssrRouter !== undefined;
   const hasDehydratedRouter =
     dehydratedState?.dehydratedRouter !== undefined && dehydratedState.dehydratedRouter !== null;
 
   const routerRef = useRef<TanStackRouter | null>(null);
-  // This only deduplicates the manual load within one mounted hydration instance.
-  // A full remount creates a fresh router and may legitimately issue another load.
   const didTriggerPostHydrationLoadRef = useRef(false);
-  const didInitializeSsrGlobalRef = useRef(false);
-  const didSetLegacySsrFlagRef = useRef(false);
 
   if (routerRef.current === null) {
     const router = options.createRouter();
@@ -104,54 +59,70 @@ function TanStackHydrationApp({
     const browserHistory = createBrowserHistory();
     router.update({ history: browserHistory });
 
-    // Hydrate router with dehydrated state from server.
-    if (hasSsrRouter) {
-      // RouterClient will hydrate the route matches from window.$_TSR.
-    } else if (hasDehydratedRouter && typeof router.hydrate === 'function') {
-      router.hydrate(dehydratedState.dehydratedRouter);
-    } else if (hasDehydratedRouter) {
-      throw new Error(
-        'react-on-rails-pro/tanstack-router: Cannot hydrate SSR payload. ' +
-          'router.hydrate() is required but not available on your TanStack Router version. ' +
-          'Ensure @tanstack/react-router >=1.139.0 <2.0.0 is installed and provides router.hydrate().',
-      );
-    }
+    // Only apply SSR hydration when a server payload exists.
+    // Client-only renders (prerender: false) must not set router.ssr or
+    // inject matches — the Transitioner handles initial loading for those.
+    if (hasSsrPayload) {
+      // Validate internal APIs before using them.
+      if (typeof router.matchRoutes !== 'function' || !router.__store?.setState) {
+        throw new Error(
+          'react-on-rails-pro/tanstack-router: router.matchRoutes() and router.__store are required ' +
+            'but not available. Ensure @tanstack/react-router >=1.139.0 is installed.',
+        );
+      }
 
-    // Legacy hydration path only: signal SSR mode so the Transitioner skips its
-    // initial router.load() call, preventing a hydration mismatch.  The object
-    // shape matches TanStack Router's internal $_TSR hydration contract (the
-    // Transitioner only checks truthiness).  The new ssrRouter/RouterClient path
-    // does not need this — RouterClient sets router.ssr internally via its own
-    // hydrate() function.
-    if (!hasSsrRouter && hasSsrPayload && !router.ssr) {
+      // Synchronously inject route matches to match server-rendered output.
+      // The server fully loads routes (via router.load()) before rendering, so
+      // all matches are resolved with status 'success'. We replicate this on the
+      // client so the initial render produces the same component tree as the
+      // server HTML. This uses the same __store.setState pattern as TanStack
+      // Router's own hydrate() in @tanstack/router-core/ssr/ssr-client.js.
+      //
+      // IMPORTANT: matchRoutes() returns matches with status 'pending' for routes
+      // that have loaders, beforeLoad, lazyFn, or preloadable components. Since the
+      // server already loaded everything, we must override to 'success' — otherwise
+      // MatchInner throws the loadPromise causing Suspense to suspend and a
+      // hydration mismatch.
+      const matches = router.matchRoutes(router.state.location);
+      for (const match of matches) {
+        const m = match as Record<string, unknown>;
+        if (m.status === 'pending') {
+          m.status = 'success';
+        }
+      }
+      router.__store.setState((s: Record<string, unknown>) => ({
+        ...s,
+        status: 'idle',
+        resolvedLocation: (s as { location: unknown }).location,
+        matches,
+      }));
+
+      // Set SSR flag so the Transitioner skips its initial router.load() call,
+      // preventing a state update during hydration that would cause a mismatch.
+      // The shape matches TanStack Router's internal $_TSR hydration contract
+      // (the Transitioner only checks truthiness).
       router.ssr = { manifest: undefined };
-      didSetLegacySsrFlagRef.current = true;
+
+      // Hydrate router with dehydrated state from server if available.
+      // Note: router.hydrate() is not a built-in TanStack Router method — it
+      // would only exist if the user's createRouter() attaches one. The guard
+      // is defensive; in practice this is currently a no-op since the standard
+      // router doesn't have a hydrate method (router.options.hydrate is the
+      // user callback, invoked differently by ssr-client.js).
+      if (hasDehydratedRouter && typeof router.hydrate === 'function') {
+        router.hydrate(dehydratedState.dehydratedRouter);
+      }
     }
 
     routerRef.current = router;
   }
 
   const router = routerRef.current;
 
-  if (hasSsrRouter && typeof window !== 'undefined' && !didInitializeSsrGlobalRef.current) {
-    setTanStackSsrGlobal(ssrRouter);
-    didInitializeSsrGlobalRef.current = true;
-  }
-
   // After mount, trigger router.load() to enable client-side navigation.
   // The SSR flag prevented auto-loading, so we do it manually here.
   useEffect(() => {
-    if (!router) {
-      return undefined;
-    }
-
-    if (hasSsrRouter) {
-      return undefined;
-    }
-
-    // Only SSR hydration needs a manual load call.
-    // For client-only renders, Transitioner handles initial loading.
-    if (!hasSsrPayload) {
+    if (!router || !hasSsrPayload) {
       return undefined;
     }
 
@@ -161,8 +132,6 @@ function TanStackHydrationApp({
     didTriggerPostHydrationLoadRef.current = true;
 
     let cancelled = false;
-    // `cancelled` only suppresses logging for a discarded mount. The in-flight load still
-    // completes unless the router exposes a best-effort cancelLoad() hook.
     router
       .load()
       .catch((err: unknown) => {
@@ -171,14 +140,10 @@ function TanStackHydrationApp({
         }
       })
       .finally(() => {
-        // Legacy hydration only: clear the temporary SSR hint after the first
-        // client load has completed so it cannot influence later navigations.
-        // Only clear when this module set it, so pre-existing router.ssr state
-        // from user code or upstream router internals is preserved.
-        if (didSetLegacySsrFlagRef.current) {
-          router.ssr = undefined;
-          didSetLegacySsrFlagRef.current = false;
-        }
+        // Clear the SSR flag after the load settles (success or failure) so it
+        // doesn't permanently block the Transitioner from calling router.load()
+        // on subsequent navigations.
+        router.ssr = undefined;
       });
 
     return () => {
@@ -188,15 +153,15 @@ function TanStackHydrationApp({
         cancellableRouter.cancelLoad();
       }
     };
-  }, [hasSsrPayload, hasSsrRouter, router]);
-
-  const RouterRoot =
-    hasSsrRouter && typeof window !== 'undefined' && RouterClient ? RouterClient : RouterProvider;
+  }, [hasSsrPayload, router]);
 
-  let app: ReactElement = createElement(RouterRoot, { router });
+  // Always use RouterProvider directly — matching the server-rendered tree.
+  // RouterClient is NOT used because it wraps RouterProvider in <Await> which
+  // introduces a Suspense boundary that doesn't exist in the server HTML,
+  // causing React hydration mismatch errors.
+  let app: ReactElement = createElement(RouterProvider, { router });
   if (options.AppWrapper) {
     const wrapperProps = { ...incomingProps } as Record<string, unknown>;
-    // eslint-disable-next-line no-underscore-dangle -- Internal hydration payload key should not reach user AppWrapper props.
     delete wrapperProps.__tanstackRouterDehydratedState;
     app = createElement(options.AppWrapper, wrapperProps, app);
   }
@@ -209,15 +174,16 @@ export function clientHydrateTanStackApp(
   props: Record<string, unknown>,
   railsContext: RailsContext & { serverSide: false },
   RouterProvider: React.ComponentType<{ router: TanStackRouter }>,
-  RouterClient: ComponentType<{ router: TanStackRouter }> | undefined,
+  // RouterClient is accepted for backward compatibility but intentionally unused.
+  // See TanStackHydrationApp JSDoc for why RouterProvider is used directly.
+  _RouterClient: unknown,
   createBrowserHistory: () => TanStackHistory,
 ): ReactElement {
   return createElement(TanStackHydrationApp, {
     options,
     incomingProps: props,
     railsContext,
     RouterProvider,
-    RouterClient,
     createBrowserHistory,
   });
 }

packages/react-on-rails-pro/src/tanstack-router/index.ts

@@ -46,9 +46,10 @@ interface TanStackRouterDeps {
    */
   RouterProvider: React.ComponentType<{ router: TanStackRouter }>;
   /**
-   * Optional RouterClient component from @tanstack/react-router/ssr/client.
-   * When provided, it enables TanStack Router's SSR client hydration path for
-   * router versions that do not expose router.dehydrate()/router.hydrate().
+   * @deprecated No longer used for hydration. RouterProvider is always used
+   * directly to match the server-rendered tree. RouterClient caused hydration
+   * mismatches because it wraps RouterProvider in <Await> which suspends.
+   * Kept for backward compatibility only.
    */
   RouterClient?: React.ComponentType<{ router: TanStackRouter }>;
   /**

packages/react-on-rails-pro/src/tanstack-router/types.ts

@@ -8,6 +8,10 @@ import type { ComponentType, ReactNode } from 'react';
 export interface TanStackRouter {
   update: (opts: { history: TanStackHistory }) => void;
   load: () => Promise<void>;
+  matchRoutes: (location: unknown) => unknown[];
+  __store: {
+    setState: (updater: (s: Record<string, unknown>) => Record<string, unknown>) => void;
+  };
   state: {
     status: string;
     location: {
@@ -23,10 +27,9 @@ export interface TanStackRouter {
   dehydrate?: () => unknown;
   hydrate?: (data: unknown) => void;
   // TanStack Router's Transitioner checks this field (truthiness only) to skip
-  // auto-loading on mount.  The canonical shape is { manifest?: unknown }, set
-  // internally by TanStack's hydrate() during RouterClient hydration.  We set
-  // it on the legacy (non-RouterClient) hydration path to prevent a duplicate
-  // initial load that causes hydration mismatch.
+  // auto-loading on mount.  The canonical shape is { manifest?: unknown }.
+  // Set during client hydration to prevent a duplicate initial load that
+  // causes hydration mismatch.
   ssr?: { manifest?: unknown };
 }