refactor: use async transition for client navigation

packages/react-server/src/entry/browser.tsx

@@ -1,4 +1,4 @@
-import { createDebug, memoize } from "@hiogawa/utils";
+import { createDebug } from "@hiogawa/utils";
 import React from "react";
 import ReactDOMClient from "react-dom/client";
 import {
@@ -36,7 +36,7 @@ export async function start() {
   const history = createEncodedBrowserHistory();
   const router = new Router(history);
 
-  let $__setLayout: (v: Promise<ServerRouterData>) => void;
+  let $__setLayout: (v: ServerRouterData) => void;
   let $__startActionTransition: React.TransitionStartFunction;
 
   //
@@ -58,7 +58,7 @@ export async function start() {
       fetch(request),
       { callServer },
     );
-    $__startActionTransition(() => $__setLayout(result));
+    $__startActionTransition(async () => $__setLayout(await result));
     return (await result).action?.data;
   };
 
@@ -72,31 +72,25 @@ export async function start() {
       readStreamScript<string>().pipeThrough(new TextEncoderStream()),
       { callServer },
     );
-  const initialLayoutPromise = Promise.resolve(initialLayout);
 
   //
   // browser root
   //
 
   function LayoutHandler(props: React.PropsWithChildren) {
     const [layoutPromise, setLayoutPromise] =
-      React.useState<Promise<ServerRouterData>>(initialLayoutPromise);
-
-    // very shaky trick to merge with current layout
-    $__setLayout = (nextPromise) => {
-      setLayoutPromise(
-        memoize(async (currentPromise: Promise<ServerRouterData>) => {
-          const current = await currentPromise;
-          const next = await nextPromise;
-          return {
-            action: next.action,
-            layout: {
-              ...current.layout,
-              ...next.layout,
-            },
-          } satisfies ServerRouterData;
-        }),
-      );
+      React.useState<ServerRouterData>(initialLayout);
+
+    $__setLayout = (next) => {
+      setLayoutPromise((current) => {
+        return {
+          action: next.action,
+          layout: {
+            ...current.layout,
+            ...next.layout,
+          },
+        } satisfies ServerRouterData;
+      });
     };
 
     const [isPending, startTransition] = React.useTransition();
@@ -135,9 +129,9 @@ export async function start() {
           revalidate: ROUTER_REVALIDATE_KEY in location.state,
         }),
       );
-      startTransition(() => {
+      startTransition(async () => {
         $__setLayout(
-          ReactClient.createFromFetch<ServerRouterData>(fetch(request), {
+          await ReactClient.createFromFetch<ServerRouterData>(fetch(request), {
             callServer,
           }),
         );

packages/react-server/src/entry/server.tsx

@@ -85,15 +85,13 @@ export async function renderHtml(
 
   const [stream1, stream2] = result.stream.tee();
 
-  const layoutPromise = ReactClient.createFromReadableStream<ServerRouterData>(
-    stream1,
-    {
+  const serverData =
+    await ReactClient.createFromReadableStream<ServerRouterData>(stream1, {
       ssrManifest: {
         moduleMap: createModuleMap(),
         moduleLoading: null,
       },
-    },
-  );
+    });
 
   const url = new URL(request.url);
   const history = createMemoryHistory({
@@ -105,7 +103,7 @@ export async function renderHtml(
 
   const reactRootEl = (
     <RouterContext.Provider value={router}>
-      <LayoutStateContext.Provider value={{ data: layoutPromise }}>
+      <LayoutStateContext.Provider value={{ data: serverData }}>
         <RouteManifestContext.Provider value={routeManifest}>
           <RouteAssetLinks />
           <LayoutRoot />

packages/react-server/src/features/router/client.tsx

@@ -10,7 +10,7 @@ import {
 import { LAYOUT_ROOT_NAME, type ServerRouterData } from "./utils";
 
 type LayoutStateContextType = {
-  data: Promise<ServerRouterData>;
+  data: ServerRouterData;
 };
 
 export const LayoutStateContext = React.createContext<LayoutStateContextType>(
@@ -19,8 +19,7 @@ export const LayoutStateContext = React.createContext<LayoutStateContextType>(
 
 export function LayoutContent(props: { name: string }) {
   const ctx = React.useContext(LayoutStateContext);
-  const data = React.use(ctx.data);
-  return data.layout[props.name];
+  return ctx.data.layout[props.name];
 }
 
 export function LayoutRoot() {

packages/react-server/src/features/server-action/client.tsx

@@ -14,9 +14,8 @@ export function ActionRedirectHandler() {
 // TODO: how to trigger nearest error page on action error?
 function ThrowActionError() {
   const ctx = React.useContext(LayoutStateContext);
-  const data = React.use(ctx.data);
-  if (data.action?.error) {
-    throw createError(data.action?.error);
+  if (ctx.data.action?.error) {
+    throw createError(ctx.data.action?.error);
   }
   return null;
 }