Address TanStack hydration review follow-ups

Author: justin808

CHANGELOG.md

@@ -46,7 +46,7 @@ After a release, run `/update-changelog` in Claude Code to analyze commits, writ
 
 #### Fixed
 
-- **[Pro]** **TanStack Router hydration no longer bails to a full client re-render**: Chunk-preload sequencing for post-hydration navigation is preserved by awaiting the preload promise in `runPostHydrationLoad` before `router.load()`. `TanStackHydrationApp` previously wrapped `RouterProvider` in a `Suspense` + `RouteChunkPreloadGate` chain, but `serverRender.ts`'s `buildAppElement` emits `AppWrapper > RouterProvider` directly with no Suspense boundary. The extra client-side boundary produced a tree-shape mismatch during hydration, causing React to discard the SSR HTML and re-render from scratch. `RouterProvider` is now rendered directly so the client tree mirrors the server output exactly. [PR 3213](https://github.com/shakacode/react_on_rails/pull/3213) by [Seifeldin7](https://github.com/Seifeldin7).
+- **[Pro]** **TanStack Router hydration no longer bails to a full client re-render**: TanStack Router SSR pages no longer discard server-rendered HTML during hydration because the client tree now renders `RouterProvider` with the same shape as the server output. Post-hydration navigation still waits for matched lazy route chunks before `router.load()`. [PR 3213](https://github.com/shakacode/react_on_rails/pull/3213) by [Seifeldin7](https://github.com/Seifeldin7).
 - **[Pro]** **Allow patched ruby-jwt releases**: React on Rails Pro now requires `jwt >= 3.2.0`, removing the previous `~> 2.7` cap so applications can resolve the patched ruby-jwt release for the empty-key HMAC advisory. [PR 3322](https://github.com/shakacode/react_on_rails/pull/3322) by [ihabadham](https://github.com/ihabadham).
 - **[Pro]** **Pro migration generator rewrites all base-package references and preserves Gemfile pins**: `rails generate react_on_rails:pro` now rewrites Jest/Vitest mock helpers (`jest.mock`, `vi.mock`, `requireActual`/`importActual`, and the rest) and TypeScript `declare module 'react-on-rails'` blocks alongside its existing `import`/`require`/dynamic-import handling, and the Gemfile swap now preserves the user's existing version pin (and other gem options) instead of overwriting them with the running gem's version. `react_on_rails:doctor` is widened to match: it also flags stale side-effect imports (`import 'react-on-rails';`), Jest/Vitest mock helpers, and `declare module` blocks, and the new side-effect-import pattern keeps the doctor a superset of the rewriter so anything the rewriter doesn't reach gets surfaced. Closes [Issue 3104](https://github.com/shakacode/react_on_rails/issues/3104). [PR 3232](https://github.com/shakacode/react_on_rails/pull/3232) by [justin808](https://github.com/justin808).
 - **[Pro]** **Pro migration scans TypeScript 4.7 `.mts` and `.cts` modules**: `react_on_rails:doctor` and the Pro migration rewriter now include `.mts`/`.cts` source files (and their `.d.mts`/`.d.cts` declaration counterparts) when looking for stale `react-on-rails` references, matching the existing `.mjs`/`.cjs` coverage. Fixes [Issue 3250](https://github.com/shakacode/react_on_rails/issues/3250). [PR 3334](https://github.com/shakacode/react_on_rails/pull/3334) by [justin808](https://github.com/justin808).

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

@@ -16,8 +16,15 @@ type ReactWithOptionalInsertionEffect = typeof React & {
   useInsertionEffect?: typeof useEffect;
 };
 
+// Read from the React namespace instead of using a named import because the Pro
+// package supports React >= 16, where useInsertionEffect does not exist as an
+// export. In React 18+, useInsertionEffect is currently excluded from the
+// StrictMode passive/layout effect replay cycle relied on by CSS-in-JS
+// libraries; that makes it the narrowest hook for "real unmount only" cleanup.
+// React versions without this hook also do not have that StrictMode replay
+// behavior, so useEffect is an adequate fallback there.
 const { useInsertionEffect } = React as ReactWithOptionalInsertionEffect;
-const useRealUnmountEffect: typeof useEffect =
+const useSynchronousRealUnmountEffect: typeof useEffect =
   typeof useInsertionEffect === 'function' ? useInsertionEffect : useEffect;
 
 type TanStackRouterHydrationInternals = TanStackRouter & {
@@ -167,7 +174,7 @@ function TanStackHydrationApp({
   const routerRef = useRef<TanStackRouter | null>(null);
   const didTriggerPostHydrationLoadRef = useRef(false);
   const didSetSsrFlagRef = useRef(false);
-  const latestEffectRunIdRef = useRef(0);
+  const latestEffectRunIdRef = useRef(0); // 0 = no post-hydration effect run yet.
   // Set during render-phase SSR init; awaited in runPostHydrationLoad before
   // router.load() so post-hydration navigation waits for matched lazy chunks.
   const routeChunkPreloadPromiseRef = useRef<Promise<void> | null>(null);
@@ -310,12 +317,10 @@ function TanStackHydrationApp({
 
   const router = routerRef.current;
 
-  // Clear the temporary router.ssr flag synchronously on real unmount. In React
-  // 18+, useInsertionEffect is not replayed by the StrictMode passive-effect
-  // remount cycle, so this cleanup does not race the second useEffect setup.
-  // React versions without useInsertionEffect also do not have that StrictMode
-  // passive-effect remount behavior, so useEffect is an adequate fallback.
-  useRealUnmountEffect(() => {
+  // Clear the temporary router.ssr flag synchronously on real unmount. This is
+  // the unmount path; the async finally() below is the normal settled/cancelled
+  // path. didSetSsrFlagRef is the shared latch so exactly one path clears.
+  useSynchronousRealUnmountEffect(() => {
     if (!router || !hasSsrPayload) {
       return undefined;
     }
@@ -372,6 +377,10 @@ function TanStackHydrationApp({
           return;
         }
       }
+      // No final cancellation check is needed for the no-await fast path:
+      // without pending hydration or preload promises, cleanup cannot run between
+      // the checks above and this call. If unmount happens after router.load()
+      // starts, the cleanup's cancelLoad() call handles the in-flight load.
       await router.load();
     };
 
@@ -382,11 +391,14 @@ function TanStackHydrationApp({
         }
       })
       .finally(() => {
-        // Always clear temporary router.ssr set by this module, regardless of
-        // cancellation state, unless a StrictMode effect re-mount already
-        // re-armed a newer post-hydration load run.
-        // The didSetSsrFlagRef guard ensures we only clear values this module
-        // created, preserving user-provided router.ssr from createRouter().
+        // Invariant: temporary router.ssr is cleared by exactly one path:
+        //   1. this finally block after the post-hydration load settles/cancels;
+        //   2. the synchronous real-unmount cleanup above.
+        // didSetSsrFlagRef is the shared latch, preserving user-provided
+        // router.ssr from createRouter(). latestEffectRunIdRef prevents stale
+        // StrictMode passive-effect finally blocks from racing a remount; React
+        // runs passive cleanup/setup back-to-back before queued promise
+        // continuations drain.
         if (latestEffectRunIdRef.current === effectRunId && didSetSsrFlagRef.current) {
           router.ssr = undefined;
           didSetSsrFlagRef.current = false;

packages/react-on-rails-pro/tests/tanstackRouter.test.ts

@@ -1120,6 +1120,8 @@ describe('tanstack-router integration (Pro)', () => {
     };
     const deps = {
       RouterProvider: ({ router: providerRouter }: { router: TanStackRouter }) => {
+        // Intentional render-phase mutation for test isolation: simulate TanStack
+        // Router clearing the private flag before our post-hydration effect.
         providerRouter.ssr = undefined;
         return React.createElement('div');
       },
@@ -1670,6 +1672,7 @@ describe('tanstack-router integration (Pro)', () => {
     expect(router.ssr).toEqual({ manifest: undefined });
 
     await compatAct(async () => {
+      expect(resolveChunks).toHaveLength(1);
       resolveChunks.forEach((resolve) => resolve());
       await new Promise((r) => {
         setTimeout(r, 0);