fix(react-server): noramlize client reference with ?t= query (hmr timestamp)

packages/react-server/examples/basic/e2e/basic.test.ts

@@ -256,26 +256,28 @@ test("rsc + client + rsc hmr @dev", async ({ page }) => {
 
   // edit client
   await editFile("./src/components/counter.tsx", (s) =>
-    s.replace("test-hmr-div", "test-hmr-edit-div"),
+    s.replace("test-hmr-div", "test-hmr-edit1-div"),
   );
-  await page.getByText("test-hmr-edit-div").click();
+  await page.getByText("test-hmr-edit1-div").click();
   await page.getByText("Count: 1").click();
 
-  // edit server again re-mounts client
+  // edit server again
   await editFile("./src/routes/test/page.tsx", (s) =>
     s.replace("Server (EDIT 1) Time", "Server (EDIT 2) Time"),
   );
   await page.getByText("Server (EDIT 2) Time").click();
-  await page.getByText("Count: 0").click();
-
-  // edit client again should work
-  await page.getByRole("button", { name: "+" }).click();
   await page.getByText("Count: 1").click();
+
+  // edit client again
   await editFile("./src/components/counter.tsx", (s) =>
-    s.replace("test-hmr-edit-div", "test-hmr-edit-edit-div"),
+    s.replace("test-hmr-edit1-div", "test-hmr-edit2-div"),
   );
-  await page.getByText("test-hmr-edit-edit-div").click();
+  await page.getByText("test-hmr-edit2-div").click();
   await page.getByText("Count: 1").click();
+
+  // check no hydration error after reload
+  await page.reload();
+  await waitForHydration(page);
 });
 
 test("module invalidation @dev", async ({ page }) => {
@@ -599,12 +601,9 @@ test("client module used at boundary and non-boundary hmr @dev", async ({
   );
   await page.getByText("Client2Context [okok]").click();
 
-  // TODO: still dual package after HMR due to
-  // `<id>?t=...` imported by client component
-  // `<id>` imported by server component (as client reference)
   await page.reload();
   await waitForHydration(page);
-  await page.getByText("Client2Context [not-ok]").click();
+  await page.getByText("Client2Context [okok]").click();
 });
 
 test("RouteProps.request", async ({ page }) => {

packages/react-server/src/features/use-client/plugin.ts

@@ -1,7 +1,12 @@
 import fs from "node:fs";
 import nodePath from "node:path";
 import { createDebug, memoize, tinyassert } from "@hiogawa/utils";
-import { type Plugin, type PluginOption, parseAstAsync } from "vite";
+import {
+  type Plugin,
+  type PluginOption,
+  type ViteDevServer,
+  parseAstAsync,
+} from "vite";
 import type { ReactServerManager } from "../../plugin";
 import { USE_CLIENT_RE, getExportNames } from "../../plugin/ast-utils";
 import { hashString } from "../../plugin/utils";
@@ -26,6 +31,11 @@ export function vitePluginServerUseClient({
   manager: ReactServerManager;
   runtimePath: string;
 }): PluginOption {
+  // TODO:
+  // eventually we should try entirely virtual module approach for client reference (not only node_modules)
+  // so that we can delegate precise resolution (e.g. `?v=` deps optimization hash, `?t=` hmr timestamp)
+  // to actual client (browser, ssr) environment instead of faking out things on RSC module graph
+
   // intercept Vite's node resolve to virtualize "use client" in node_modules
   const pluginUseClientNodeModules: Plugin = {
     name: "server-virtual-use-client-node-modules",
@@ -73,7 +83,7 @@ export function vitePluginServerUseClient({
         meta.exportNames = exportNames;
         // we need to transform to client reference directly
         // otherwise `soruce` will be resolved infinitely by recursion
-        id = noramlizeClientReferenceId(id);
+        id = wrapId(id);
         const result = generateClientReferenceCode(
           id,
           exportNames,
@@ -105,7 +115,8 @@ export function vitePluginServerUseClient({
       // normalize client reference during dev
       // to align with Vite's import analysis
       if (!manager.buildType) {
-        id = noramlizeClientReferenceId(id);
+        tinyassert(manager.parentServer);
+        id = await noramlizeClientReferenceId(id, manager.parentServer);
       } else {
         // obfuscate reference
         id = hashString(id);
@@ -165,21 +176,32 @@ function generateClientReferenceCode(
 
 // Apply same noramlizaion as Vite's dev import analysis
 // to avoid dual package with "/xyz" and "/@fs/xyz" for example.
-// For now this tries to cover simple cases
 // https://github.com/vitejs/vite/blob/0c0aeaeb3f12d2cdc3c47557da209416c8d48fb7/packages/vite/src/node/plugins/importAnalysis.ts#L327-L399
-export function noramlizeClientReferenceId(id: string) {
-  const root = process.cwd(); // TODO: pass vite root config
+export async function noramlizeClientReferenceId(
+  id: string,
+  parentServer: ViteDevServer,
+) {
+  const root = parentServer.config.root;
   if (id.startsWith(root)) {
     id = id.slice(root.length);
   } else if (nodePath.isAbsolute(id)) {
     id = "/@fs" + id;
   } else {
-    // aka wrapId
-    id = id.startsWith(`/@id`) ? id : `/@id/${id.replace("\0", "__x00__")}`;
+    id = wrapId(id);
+  }
+  // this is needed only for browser, so we'll strip it off
+  // during ssr client reference import
+  const mod = await parentServer.moduleGraph.getModuleByUrl(id);
+  if (mod && mod.lastHMRTimestamp > 0) {
+    id += `?t=${mod.lastHMRTimestamp}`;
   }
   return id;
 }
 
+function wrapId(id: string) {
+  return id.startsWith(`/@id`) ? id : `/@id/${id.replace("\0", "__x00__")}`;
+}
+
 const VIRTUAL_PREFIX = "virtual:use-client-node-module/";
 
 export function vitePluginClientUseClient({

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

@@ -24,6 +24,8 @@ const ssrWebpackRequire: WebpackRequire = memoize(ssrImport, {
 async function ssrImport(id: string) {
   debug("[__webpack_require__]", { id });
   if (import.meta.env.DEV) {
+    // strip off `?t=` added for browser by noramlizeClientReferenceId
+    id = id.split("?t=")[0]!;
     // transformed to "ssrLoadModule" during dev
     return import(/* @vite-ignore */ id);
   } else {

packages/react-server/src/plugin/index.ts

@@ -53,6 +53,8 @@ const RUNTIME_REACT_SERVER_PATH = fileURLToPath(
 export type { ReactServerManager };
 
 class ReactServerManager {
+  parentServer?: ViteDevServer;
+
   buildType?: "rsc" | "client" | "ssr";
 
   // expose "use client" node modules to client via virtual modules
@@ -249,6 +251,7 @@ export function vitePluginReactServer(options?: {
     },
     async configureServer(server) {
       parentServer = server;
+      manager.parentServer = server;
     },
     async buildStart(_options) {
       if (parentEnv.command === "serve") {