feat: attempt at enabling full transition support for the rsc router

.changeset/lucky-tables-itch.md

@@ -0,0 +1,6 @@
+---
+"integration": patch
+"react-router": patch
+---
+
+feat: enable full transition support for the rsc router

integration/rsc/rsc-test.ts

@@ -508,6 +508,11 @@ implementations.forEach((implementation) => {
                       path: "ssr-error",
                       lazy: () => import("./routes/ssr-error/ssr-error"),
                     },
+                    {
+                      id: "action-transition-state",
+                      path: "action-transition-state",
+                      lazy: () => import("./routes/action-transition-state/home"),
+                    }
                   ],
                 },
               ] satisfies RSCRouteConfig;
@@ -1314,6 +1319,49 @@ implementations.forEach((implementation) => {
                 throw new Error("Error from SSR component");
               }
             `,
+
+            "src/routes/action-transition-state/home.tsx": js`
+              import { Suspense } from "react";
+              import { IncrementButton } from "./client";
+              let count = 0;
+
+              export default function ActionTransitionState() {
+                return (
+                  <div>
+                    <form
+                      action={async () => {
+                        "use server";
+                        await new Promise((r) => setTimeout(r, 1000));
+                        count++;
+                      }}
+                    >
+                      <IncrementButton count={count} />
+                    </form>
+                    <Suspense>
+                      <AsyncComponent count={count} />
+                    </Suspense>
+                  </div>
+                );
+              }
+
+              async function AsyncComponent({ count }) {
+                await new Promise((r) => setTimeout(r, 1000));
+                return <div data-testid="async-count">AsyncCount: {count}</div>;
+              }
+            `,
+            "src/routes/action-transition-state/client.tsx": js`
+              "use client";
+              import { useFormStatus } from "react-dom";
+
+              export function IncrementButton({ count }: { count: number }) {
+                const { pending } = useFormStatus();
+                return (
+                  <button data-testid="increment-button" type="submit" disabled={pending}>
+                    IncrementCount: {pending ? count + 1 : count}
+                  </button>
+                );
+              }
+            `,
           },
         });
       });
@@ -1796,6 +1844,44 @@ implementations.forEach((implementation) => {
           const actionResponse = await actionResponsePromise;
           expect(await actionResponse.headerValue("x-test")).toBe("test");
         });
+
+        test("Supports transition state throughout the revalidation lifecycle", async ({
+          page,
+        }) => {
+          test.skip(
+            implementation.name === "parcel",
+            "Uses inline server actions which parcel doesn't support yet",
+          );
+
+          await page.goto(`http://localhost:${port}/action-transition-state`, {
+            waitUntil: "networkidle",
+          });
+
+          const count0Button = page.getByText("IncrementCount: 0");
+          await expect(count0Button).toBeEnabled();
+          await count0Button.click();
+
+          const count1Button = page.getByText("IncrementCount: 1");
+          await expect(count1Button).toBeDisabled();
+
+          expect(await page.getByTestId("async-count").textContent()).toBe(
+            "AsyncCount: 0",
+          );
+
+          await page.waitForFunction(
+            () =>
+              !(
+                document.querySelector(
+                  '[data-testid="increment-button"]',
+                ) as HTMLButtonElement
+              )?.disabled,
+          );
+          await expect(count1Button).toBeEnabled();
+
+          await expect(page.getByTestId("async-count")).toHaveText(
+            "AsyncCount: 1",
+          );
+        });
       });
 
       test.describe("Errors", () => {

packages/react-router/lib/components.tsx

@@ -291,6 +291,143 @@ export interface RouterProviderProps {
   unstable_onError?: unstable_ClientOnErrorFunction;
 }
 
+function shallowDiff(a: any, b: any) {
+  if (a === b) {
+    return false;
+  }
+  let aKeys = Object.keys(a);
+  let bKeys = Object.keys(b);
+  if (aKeys.length !== bKeys.length) {
+    return true;
+  }
+  for (let key of aKeys) {
+    if (a[key] !== b[key]) {
+      return true;
+    }
+  }
+  return false;
+}
+
+export function UNSTABLE_TransitionEnabledRouterProvider({
+  router,
+  flushSync: reactDomFlushSyncImpl,
+  unstable_onError,
+}: RouterProviderProps) {
+  let fetcherData = React.useRef<Map<string, any>>(new Map());
+  let [revalidating, startRevalidation] = React.useTransition();
+  let [state, setState] = React.useState(router.state);
+
+  (router as any).__setPendingRerender = (promise: Promise<() => void>) =>
+    startRevalidation(
+      // @ts-expect-error - need react 19 types for this to be async
+      async () => {
+        const rerender = await promise;
+        startRevalidation(() => {
+          rerender();
+        });
+      },
+    );
+
+  let navigator = React.useMemo((): Navigator => {
+    return {
+      createHref: router.createHref,
+      encodeLocation: router.encodeLocation,
+      go: (n) => router.navigate(n),
+      push: (to, state, opts) =>
+        router.navigate(to, {
+          state,
+          preventScrollReset: opts?.preventScrollReset,
+        }),
+      replace: (to, state, opts) =>
+        router.navigate(to, {
+          replace: true,
+          state,
+          preventScrollReset: opts?.preventScrollReset,
+        }),
+    };
+  }, [router]);
+
+  let basename = router.basename || "/";
+
+  let dataRouterContext = React.useMemo(
+    () => ({
+      router,
+      navigator,
+      static: false,
+      basename,
+      unstable_onError,
+    }),
+    [router, navigator, basename, unstable_onError],
+  );
+
+  React.useLayoutEffect(() => {
+    return router.subscribe(
+      (newState, { deletedFetchers, flushSync, viewTransitionOpts }) => {
+        newState.fetchers.forEach((fetcher, key) => {
+          if (fetcher.data !== undefined) {
+            fetcherData.current.set(key, fetcher.data);
+          }
+        });
+        deletedFetchers.forEach((key) => fetcherData.current.delete(key));
+
+        const diff = shallowDiff(state, newState);
+
+        if (!diff) return;
+
+        if (flushSync) {
+          if (reactDomFlushSyncImpl) {
+            reactDomFlushSyncImpl(() => setState(newState));
+          } else {
+            setState(newState);
+          }
+        } else {
+          React.startTransition(() => {
+            setState(newState);
+          });
+        }
+      },
+    );
+  }, [router, reactDomFlushSyncImpl, state]);
+
+  // The fragment and {null} here are important!  We need them to keep React 18's
+  // useId happy when we are server-rendering since we may have a <script> here
+  // containing the hydrated server-side staticContext (from StaticRouterProvider).
+  // useId relies on the component tree structure to generate deterministic id's
+  // so we need to ensure it remains the same on the client even though
+  // we don't need the <script> tag
+  return (
+    <>
+      <DataRouterContext.Provider value={dataRouterContext}>
+        <DataRouterStateContext.Provider
+          value={{
+            ...state,
+            revalidation: revalidating ? "loading" : state.revalidation,
+          }}
+        >
+          <FetchersContext.Provider value={fetcherData.current}>
+            {/* <ViewTransitionContext.Provider value={vtContext}> */}
+            <Router
+              basename={basename}
+              location={state.location}
+              navigationType={state.historyAction}
+              navigator={navigator}
+            >
+              <MemoizedDataRoutes
+                routes={router.routes}
+                future={router.future}
+                state={state}
+                unstable_onError={unstable_onError}
+              />
+            </Router>
+            {/* </ViewTransitionContext.Provider> */}
+          </FetchersContext.Provider>
+        </DataRouterStateContext.Provider>
+      </DataRouterContext.Provider>
+      {null}
+    </>
+  );
+}
+
 /**
  * Render the UI for the given {@link DataRouter}. This component should
  * typically be at the top of an app's element tree.