refactor(react-server): layout invalidation on server

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

@@ -3,13 +3,9 @@ import { createBrowserHistory } from "@tanstack/history";
 import React from "react";
 import reactDomClient from "react-dom/client";
 import { LayoutRoot, LayoutStateContext } from "../features/router/client";
-import {
-  type ServerRouterData,
-  createLayoutContentRequest,
-  getNewLayoutContentKeys,
-} from "../features/router/utils";
+import type { ServerRouterData } from "../features/router/utils";
 import { injectActionId } from "../features/server-action/utils";
-import { wrapRscRequestUrl } from "../features/server-component/utils";
+import { wrapStreamRequestUrl } from "../features/server-component/utils";
 import { initializeWebpackBrowser } from "../features/use-client/browser";
 import { RootErrorBoundary } from "../lib/client/error-boundary";
 import { Router, RouterContext, useRouter } from "../lib/client/router";
@@ -51,11 +47,10 @@ export async function start() {
       tinyassert(args[0] instanceof FormData);
       injectActionId(args[0], id);
     }
-    // TODO: for now, we invalidate only leaf content
-    const pathname = history.location.pathname;
-    const newKeys = getNewLayoutContentKeys(pathname, pathname);
     const request = new Request(
-      wrapRscRequestUrl(history.location.href, newKeys),
+      wrapStreamRequestUrl(history.location.href, {
+        lastPathname: history.location.pathname,
+      }),
       {
         method: "POST",
         body: args[0],
@@ -131,15 +126,16 @@ export async function start() {
       const lastPathname = lastLocation.current.pathname;
       lastLocation.current = location;
 
-      // TODO: (refactor) server can find `newKeys` based on from/to location
-      const pathname = location.pathname;
-      let newKeys = getNewLayoutContentKeys(lastPathname, pathname);
-      if (RSC_HMR_STATE_KEY in location.state) {
-        newKeys = Object.keys(createLayoutContentRequest(pathname));
-      }
-      debug("[navigation]", location, { pathname, lastPathname, newKeys });
-
-      const request = new Request(wrapRscRequestUrl(location.href, newKeys));
+      debug("[navigation]", location, {
+        pathname: location.pathname,
+        lastPathname,
+      });
+      const request = new Request(
+        wrapStreamRequestUrl(location.href, {
+          lastPathname,
+          invalidateAll: RSC_HMR_STATE_KEY in location.state,
+        }),
+      );
       startTransition(() => {
         __setLayout(
           reactServerDomClient.createFromFetch<ServerRouterData>(

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

@@ -1,20 +1,14 @@
-import {
-  createDebug,
-  objectMapKeys,
-  objectMapValues,
-  objectPickBy,
-} from "@hiogawa/utils";
+import { createDebug, objectMapKeys, objectMapValues } from "@hiogawa/utils";
 import type { RenderToReadableStreamOptions } from "react-dom/server";
 import reactServerDomServer from "react-server-dom-webpack/server.edge";
 import {
   type ActionResult,
   type LayoutRequest,
   type ServerRouterData,
-  createLayoutContentRequest,
 } from "../features/router/utils";
 import { type ActionContext } from "../features/server-action/react-server";
 import { ejectActionId } from "../features/server-action/utils";
-import { unwrapRscRequest } from "../features/server-component/utils";
+import { unwrapStreamRequest } from "../features/server-component/utils";
 import { createBundlerConfig } from "../features/use-client/react-server";
 import {
   DEFAULT_ERROR_CONTEXT,
@@ -38,7 +32,6 @@ export interface ReactServerHandlerContext {
 
 export interface ReactServerHandlerStreamResult {
   stream: ReadableStream<Uint8Array>;
-  layoutRequest: LayoutRequest;
   actionResult?: ActionResult;
 }
 
@@ -47,28 +40,17 @@ export type ReactServerHandlerResult =
   | ReactServerHandlerStreamResult;
 
 export const handler: ReactServerHandler = async (ctx) => {
-  // check rsc-only request
-  const rscOnly = unwrapRscRequest(ctx.request);
-
   // action
   let actionResult: ActionResult | undefined;
   if (ctx.request.method === "POST") {
     actionResult = await actionHandler(ctx);
   }
 
-  const request = rscOnly?.request ?? ctx.request;
-  const url = new URL(request.url);
-  let layoutRequest = createLayoutContentRequest(url.pathname);
-
-  if (rscOnly) {
-    layoutRequest = objectPickBy(layoutRequest, (_v, k) =>
-      rscOnly.newKeys.includes(k),
-    );
-  }
-
+  // check stream only request
+  const { request, layoutRequest, isStream } = unwrapStreamRequest(ctx.request);
   const stream = await render({ request, layoutRequest, actionResult });
 
-  if (rscOnly) {
+  if (isStream) {
     return new Response(stream, {
       headers: {
         ...actionResult?.responseHeaders,
@@ -77,7 +59,7 @@ export const handler: ReactServerHandler = async (ctx) => {
     });
   }
 
-  return { stream, layoutRequest, actionResult };
+  return { stream, actionResult };
 };
 
 //

packages/react-server/src/features/server-component/utils.tsx

@@ -1,25 +1,49 @@
+import { objectPickBy } from "@hiogawa/utils";
+import {
+  createLayoutContentRequest,
+  getNewLayoutContentKeys,
+} from "../router/utils";
+
 // TODO: use accept header x-component?
 const RSC_PARAM = "__rsc";
 
-export function wrapRscRequestUrl(url: string, newKeys: string[] = []): string {
+// TODO: allow invalidating each layout layer
+type StreamRequestParam = {
+  lastPathname?: string;
+  invalidateAll?: boolean; // currently used for server component HMR
+};
+
+export function wrapStreamRequestUrl(
+  url: string,
+  param: StreamRequestParam,
+): string {
   const newUrl = new URL(url, window.location.href);
-  newUrl.searchParams.set(RSC_PARAM, JSON.stringify(newKeys));
+  newUrl.searchParams.set(RSC_PARAM, JSON.stringify(param));
   return newUrl.toString();
 }
 
-export function unwrapRscRequest(request: Request) {
+export function unwrapStreamRequest(request: Request) {
   const url = new URL(request.url);
   const rscParam = url.searchParams.get(RSC_PARAM);
+  url.searchParams.delete(RSC_PARAM);
+
+  let layoutRequest = createLayoutContentRequest(url.pathname);
   if (rscParam) {
-    url.searchParams.delete(RSC_PARAM);
-    const newKeys: string[] = JSON.parse(rscParam);
-    return {
-      request: new Request(url, {
-        method: request.method,
-        headers: request.headers,
-      }),
-      newKeys,
-    };
+    const param = JSON.parse(rscParam);
+    if (param.lastPathname && !param.invalidateAll) {
+      const newKeys = getNewLayoutContentKeys(param.lastPathname, url.pathname);
+      layoutRequest = objectPickBy(layoutRequest, (_v, k) =>
+        newKeys.includes(k),
+      );
+    }
   }
-  return;
+
+  return {
+    request: new Request(url, {
+      method: request.method,
+      headers: request.headers,
+    }),
+    layoutRequest,
+    isStream: Boolean(rscParam),
+  };
 }