fix: stable client reference id

packages/vinext/src/server/app-browser-client-loader.ts

@@ -0,0 +1,202 @@
+import { createElement, forwardRef } from "react";
+import { setRequireModule } from "@vitejs/plugin-rsc/core/browser";
+import * as clientReferences from "virtual:vite-rsc/client-references";
+import type { CachedRscResponse } from "../shims/navigation.js";
+
+declare const __vite_rsc_raw_import__: (id: string) => Promise<unknown>;
+
+type TrackedPromise<T> = Promise<T> & {
+  reason?: unknown;
+  status?: "fulfilled" | "pending" | "rejected";
+  value?: T;
+};
+
+type MemoLikeValue = {
+  $$typeof: symbol;
+  displayName?: string;
+  type?: unknown;
+};
+
+const CLIENT_REFERENCE_ROW_PATTERN = /(?:^|\n)[0-9a-fA-F]+:I\[("(?:\\.|[^"\\])*")/g;
+const REACT_FORWARD_REF_TYPE = Symbol.for("react.forward_ref");
+const REACT_MEMO_TYPE = Symbol.for("react.memo");
+const textDecoder = new TextDecoder();
+const clientModuleCache = new Map<string, TrackedPromise<unknown>>();
+const stableMemoExportCache = new Map<string, unknown>();
+const stableModuleExportCache = new Map<string, unknown>();
+let installed = false;
+
+function withTrailingSlash(path: string): string {
+  return path[path.length - 1] === "/" ? path : `${path}/`;
+}
+
+function normalizeClientReferenceId(id: string): string {
+  return id.split("$$cache=")[0];
+}
+
+function trackPromise<T>(promise: Promise<T>): TrackedPromise<T> {
+  const tracked = promise as TrackedPromise<T>;
+  if (tracked.status) {
+    return tracked;
+  }
+
+  tracked.status = "pending";
+  promise.then(
+    (value) => {
+      tracked.status = "fulfilled";
+      tracked.value = value;
+    },
+    (reason) => {
+      tracked.status = "rejected";
+      tracked.reason = reason;
+    },
+  );
+  return tracked;
+}
+
+function isObjectRecord(value: unknown): value is Record<string, unknown> {
+  return !!value && typeof value === "object";
+}
+
+function isForwardRefLikeValue(value: unknown): value is { $$typeof: symbol } {
+  return isObjectRecord(value) && "$$typeof" in value && value.$$typeof === REACT_FORWARD_REF_TYPE;
+}
+
+function isMemoLikeValue(value: unknown): value is MemoLikeValue {
+  return isObjectRecord(value) && "$$typeof" in value && value.$$typeof === REACT_MEMO_TYPE;
+}
+
+function createStableMemoProxy(
+  normalizedId: string,
+  exportName: string,
+  value: MemoLikeValue,
+): unknown {
+  const cacheKey = `${normalizedId}#${exportName}`;
+  const cached = stableMemoExportCache.get(cacheKey);
+  if (cached) {
+    return cached;
+  }
+
+  const target = value as unknown as Parameters<typeof createElement>[0];
+  const supportsRef = isForwardRefLikeValue(value.type);
+
+  const proxy = supportsRef
+    ? forwardRef<unknown, Record<string, unknown>>(function VinextStableMemoProxy(props, ref) {
+        return createElement(target, Object.assign({}, props, { ref }) as never);
+      })
+    : function VinextStableMemoProxy(props: Record<string, unknown>) {
+        return createElement(target, props);
+      };
+
+  const namedProxy = proxy as { displayName?: string };
+  namedProxy.displayName =
+    value.displayName ??
+    (typeof value.type === "function" && value.type.name ? value.type.name : exportName);
+
+  stableMemoExportCache.set(cacheKey, namedProxy);
+  return namedProxy;
+}
+
+function stabilizeClientModuleExports(normalizedId: string, moduleExports: unknown): unknown {
+  const cached = stableModuleExportCache.get(normalizedId);
+  if (cached) {
+    return cached;
+  }
+
+  if (isMemoLikeValue(moduleExports)) {
+    const stable = createStableMemoProxy(normalizedId, "default", moduleExports);
+    stableModuleExportCache.set(normalizedId, stable);
+    return stable;
+  }
+
+  if (
+    !moduleExports ||
+    (typeof moduleExports !== "object" && typeof moduleExports !== "function")
+  ) {
+    stableModuleExportCache.set(normalizedId, moduleExports);
+    return moduleExports;
+  }
+
+  let changed = false;
+  const next = Object.create(Object.getPrototypeOf(moduleExports));
+  const descriptors = Object.getOwnPropertyDescriptors(moduleExports);
+
+  for (const [exportName, descriptor] of Object.entries(descriptors)) {
+    if ("value" in descriptor && isMemoLikeValue(descriptor.value)) {
+      descriptor.value = createStableMemoProxy(normalizedId, exportName, descriptor.value);
+      changed = true;
+    }
+    Object.defineProperty(next, exportName, descriptor);
+  }
+
+  const stableModuleExports = changed ? next : moduleExports;
+  stableModuleExportCache.set(normalizedId, stableModuleExports);
+  return stableModuleExports;
+}
+
+function loadClientReference(normalizedId: string): Promise<unknown> {
+  if (!import.meta.env.__vite_rsc_build__) {
+    return __vite_rsc_raw_import__(
+      withTrailingSlash(import.meta.env.BASE_URL) + normalizedId.slice(1),
+    ).then((moduleExports) => stabilizeClientModuleExports(normalizedId, moduleExports));
+  }
+
+  const importReference = clientReferences.default[normalizedId] as
+    | (() => Promise<unknown>)
+    | undefined;
+  if (!importReference) {
+    throw new Error(`client reference not found '${normalizedId}'`);
+  }
+  return importReference().then((moduleExports) =>
+    stabilizeClientModuleExports(normalizedId, moduleExports),
+  );
+}
+
+function getTrackedClientModule(id: string): TrackedPromise<unknown> {
+  const normalizedId = normalizeClientReferenceId(id);
+  const cached = clientModuleCache.get(normalizedId);
+  if (cached) {
+    return cached;
+  }
+
+  const tracked = trackPromise(Promise.resolve().then(() => loadClientReference(normalizedId)));
+  clientModuleCache.set(normalizedId, tracked);
+  return tracked;
+}
+
+export function installVinextBrowserClientLoader(): void {
+  if (installed) {
+    return;
+  }
+
+  installed = true;
+  setRequireModule({
+    load(id) {
+      return getTrackedClientModule(id);
+    },
+  });
+}
+
+export async function prewarmClientReferencesFromSnapshot(
+  snapshot: CachedRscResponse,
+): Promise<void> {
+  const payload = textDecoder.decode(snapshot.buffer);
+  const pending: Promise<unknown>[] = [];
+
+  CLIENT_REFERENCE_ROW_PATTERN.lastIndex = 0;
+  let match: RegExpExecArray | null;
+  while ((match = CLIENT_REFERENCE_ROW_PATTERN.exec(payload))) {
+    try {
+      const id = JSON.parse(match[1]) as string;
+      pending.push(getTrackedClientModule(id));
+    } catch {
+      // Ignore malformed rows and let the regular RSC decoder surface errors.
+    }
+  }
+
+  if (pending.length === 0) {
+    return;
+  }
+
+  await Promise.allSettled(pending);
+}

packages/vinext/src/server/app-browser-entry.ts

@@ -16,7 +16,7 @@ import {
   createTemporaryReferenceSet,
   encodeReply,
   setServerCallback,
-} from "@vitejs/plugin-rsc/browser";
+} from "@vitejs/plugin-rsc/react/browser";
 import { hydrateRoot } from "react-dom/client";
 import "../client/instrumentation-client.js";
 import { notifyAppRouterTransitionStart } from "../client/instrumentation-client-state.js";
@@ -46,6 +46,10 @@ import {
   createProgressiveRscStream,
   getVinextBrowserGlobal,
 } from "./app-browser-stream.js";
+import {
+  installVinextBrowserClientLoader,
+  prewarmClientReferencesFromSnapshot,
+} from "./app-browser-client-loader.js";
 
 type SearchParamInput = ConstructorParameters<typeof URLSearchParams>[0];
 
@@ -488,11 +492,15 @@ async function readInitialRscStream(): Promise<ReadableStream<Uint8Array>> {
 
   restoreHydrationNavigationContext(window.location.pathname, window.location.search, params);
 
-  if (!rscResponse.body) {
+  const responseSnapshot = await snapshotRscResponse(rscResponse);
+  await prewarmClientReferencesFromSnapshot(responseSnapshot);
+
+  const restoredResponse = restoreRscResponse(responseSnapshot, false);
+  if (!restoredResponse.body) {
     throw new Error("[vinext] Initial RSC response had no body");
   }
 
-  return rscResponse.body;
+  return restoredResponse.body;
 }
 
 function registerServerActionCallback(): void {
@@ -534,8 +542,10 @@ function registerServerActionCallback(): void {
 
     clearClientNavigationCaches();
 
+    const responseSnapshot = await snapshotRscResponse(fetchResponse);
+    await prewarmClientReferencesFromSnapshot(responseSnapshot);
     const result = await createFromFetch<ServerActionResult | ReactNode>(
-      Promise.resolve(fetchResponse),
+      Promise.resolve(restoreRscResponse(responseSnapshot)),
       { temporaryReferences },
     );
 
@@ -573,6 +583,7 @@ function registerServerActionCallback(): void {
 }
 
 async function main(): Promise<void> {
+  installVinextBrowserClientLoader();
   registerServerActionCallback();
 
   const rscStream = await readInitialRscStream();
@@ -642,6 +653,7 @@ async function main(): Promise<void> {
         // wrapping only) — no stale-navigation recheck needed between here and the
         // next await.
         const cachedNavigationSnapshot = createClientNavigationRenderSnapshot(href, cachedParams);
+        await prewarmClientReferencesFromSnapshot(cachedRoute.response);
         const cachedPayload = await createFromFetch<ReactNode>(
           Promise.resolve(restoreRscResponse(cachedRoute.response)),
         );
@@ -726,6 +738,10 @@ async function main(): Promise<void> {
 
       if (navId !== activeNavigationId) return;
 
+      await prewarmClientReferencesFromSnapshot(responseSnapshot);
+
+      if (navId !== activeNavigationId) return;
+
       const rscPayload = await createFromFetch<ReactNode>(
         Promise.resolve(restoreRscResponse(responseSnapshot)),
       );
@@ -801,8 +817,12 @@ async function main(): Promise<void> {
     import.meta.hot.on("rsc:update", async () => {
       try {
         clearClientNavigationCaches();
+        const responseSnapshot = await snapshotRscResponse(
+          await fetch(toRscUrl(window.location.pathname + window.location.search)),
+        );
+        await prewarmClientReferencesFromSnapshot(responseSnapshot);
         const rscPayload = await createFromFetch<ReactNode>(
-          fetch(toRscUrl(window.location.pathname + window.location.search)),
+          Promise.resolve(restoreRscResponse(responseSnapshot)),
         );
         // HMR updates skip renderNavigationPayload — no snapshot activated.
         updateBrowserTree(

tests/e2e/app-router/nextjs-compat/layout-search-params.spec.ts

@@ -0,0 +1,60 @@
+/**
+ * Next.js compat: layout state across search param changes.
+ *
+ * Based on Next.js: test/e2e/app-dir/search-params-react-key/layout-params.test.ts
+ * https://github.com/vercel/next.js/blob/canary/test/e2e/app-dir/search-params-react-key/layout-params.test.ts
+ *
+ * Extends the same expectation to a parent client layout rendered by a server layout:
+ * query-only push/replace should not remount that layout.
+ */
+
+import { expect, test } from "@playwright/test";
+import { waitForAppRouterHydration } from "../../helpers";
+
+const BASE = "http://localhost:4174";
+
+test.describe("Next.js compat: layout state across search param changes", () => {
+  test("router.push() keeps parent client layout mounted on query-only navigation", async ({
+    page,
+  }) => {
+    await page.goto(`${BASE}/nextjs-compat/layout-search-params/demo`);
+    await waitForAppRouterHydration(page);
+
+    await page.click("#layout-increment");
+    await page.click("#layout-increment");
+    await expect(page.locator("#layout-count")).toHaveText("2");
+    await expect(page.locator("#layout-mount-count")).toHaveText("1");
+
+    await page.click("#layout-push");
+
+    await expect(async () => {
+      expect(page.url()).toContain("foo=bar");
+    }).toPass({ timeout: 10_000 });
+
+    await expect(page.locator("#search-params")).toContainText('"foo":"bar"');
+    await expect(page.locator("#layout-count")).toHaveText("2");
+    await expect(page.locator("#layout-mount-count")).toHaveText("1");
+  });
+
+  test("router.replace() keeps parent client layout mounted on query-only navigation", async ({
+    page,
+  }) => {
+    await page.goto(`${BASE}/nextjs-compat/layout-search-params/demo`);
+    await waitForAppRouterHydration(page);
+
+    await page.click("#layout-increment");
+    await page.click("#layout-increment");
+    await expect(page.locator("#layout-count")).toHaveText("2");
+    await expect(page.locator("#layout-mount-count")).toHaveText("1");
+
+    await page.click("#layout-replace");
+
+    await expect(async () => {
+      expect(page.url()).toContain("foo=baz");
+    }).toPass({ timeout: 10_000 });
+
+    await expect(page.locator("#search-params")).toContainText('"foo":"baz"');
+    await expect(page.locator("#layout-count")).toHaveText("2");
+    await expect(page.locator("#layout-mount-count")).toHaveText("1");
+  });
+});

tests/fixtures/app-basic/app/nextjs-compat/layout-search-params/[id]/layout-shell.tsx

@@ -0,0 +1,42 @@
+"use client";
+
+import React, { useEffect, useState } from "react";
+import { useRouter } from "next/navigation";
+
+declare global {
+  interface Window {
+    __vinextLayoutSearchParamsMountCount__?: number;
+  }
+}
+
+function LayoutShellInner({ children }: { children: React.ReactNode }) {
+  const router = useRouter();
+  const [count, setCount] = useState(0);
+  const [mountCount, setMountCount] = useState(0);
+
+  useEffect(() => {
+    window.__vinextLayoutSearchParamsMountCount__ =
+      (window.__vinextLayoutSearchParamsMountCount__ ?? 0) + 1;
+    setMountCount(window.__vinextLayoutSearchParamsMountCount__);
+  }, []);
+
+  return (
+    <div>
+      <h1>Layout Search Params</h1>
+      <div id="layout-count">{count}</div>
+      <div id="layout-mount-count">{mountCount}</div>
+      <button id="layout-increment" onClick={() => setCount((value) => value + 1)}>
+        Increment layout state
+      </button>
+      <button id="layout-push" onClick={() => router.push("?foo=bar")}>
+        Push foo=bar
+      </button>
+      <button id="layout-replace" onClick={() => router.replace("?foo=baz")}>
+        Replace foo=baz
+      </button>
+      {children}
+    </div>
+  );
+}
+
+export const LayoutShell = React.memo(LayoutShellInner);