fix: suspense client component (#78)

This fixes an issue about adding scripts to the HTML document in
development mode for HMR when the initial page don't use a client
component, but a client component is rendered in a `<Suspense>`
boundary. #77

Adds a test to verify this use case from now on.

Author: lazarv

packages/react-server/lib/plugins/react-server-runtime.mjs

@@ -35,9 +35,12 @@ export default function viteReactServerRuntime() {
           window.$RefreshSig$ = () => (type) => type;
           window.__vite_plugin_react_preamble_installed__ = true;
           console.log("Hot Module Replacement installed.");
-          if (typeof __react_server_hydrate__ !== "undefined") {
-            import(/* @vite-ignore */ "${reactServerDir}/client/entry.client.jsx");
-          }`;
+          self.__react_server_hydrate_init__ = () => {
+            if (typeof __react_server_hydrate__ !== "undefined") {
+              import(/* @vite-ignore */ "${reactServerDir}/client/entry.client.jsx");
+            }
+          };
+          self.__react_server_hydrate_init__();`;
       } else if (id.endsWith("/@__webpack_require__")) {
         return `
           const moduleCache = new Map();

packages/react-server/server/render-dom.mjs

@@ -258,12 +258,16 @@ export const createRenderer = ({
                                 importMap
                               )}</script>`
                             : ""
-                        }${bootstrapModules
-                          .map(
-                            (mod) =>
-                              `<script type="module" src="${mod}" async></script>`
-                          )
-                          .join("")}`
+                        }${
+                          hmr
+                            ? "<script>self.__react_server_hydrate_init__?.();</script>"
+                            : bootstrapModules
+                                .map(
+                                  (mod) =>
+                                    `<script type="module" src="${mod}" async></script>`
+                                )
+                                .join("")
+                        }`
                       );
                       yield script;
                       hydrated = true;

packages/react-server/server/render-rsc.jsx

@@ -449,11 +449,7 @@ export async function render(Component) {
                   `const moduleCache = new Map();
                     self.__webpack_require__ = function (id) {
                       if (!moduleCache.has(id)) {
-                        ${
-                          config.base
-                            ? `const modulePromise = import(("${`/${config.base}/`.replace(/\/+/g, "/")}" + id).replace(/\\/+/g, "/"));`
-                            : `const modulePromise = import(id);`
-                        }
+                        const modulePromise = import(("${`/${config.base ?? ""}/`.replace(/\/+/g, "/")}" + id).replace(/\\/+/g, "/"));
                         modulePromise.then(
                           (module) => {
                             modulePromise.value = module;

test/__test__/basic.spec.mjs

@@ -6,6 +6,7 @@ import {
   serverLogs,
   waitForChange,
   waitForConsole,
+  waitForHydration,
 } from "playground/utils";
 import { expect, test } from "vitest";
 
@@ -239,3 +240,32 @@ test("use cache dynamic", async () => {
   await page.goto(hostname + "?id=1");
   expect(await page.textContent("body")).toBe(time);
 });
+
+test("suspense client", async () => {
+  await server("fixtures/suspense-client.jsx");
+  await page.goto(hostname);
+  await waitForHydration();
+
+  if (process.env.NODE_ENV === "production") {
+    const scripts = await page.$$("script");
+    expect(scripts.length).toBe(2);
+    expect(await scripts[0].getAttribute("src")).toContain("/client/index");
+    expect(await scripts[1].getAttribute("src")).toBe(null);
+  } else {
+    const button = await page.getByRole("button");
+    expect(await button.isVisible()).toBe(true);
+    await button.click();
+    expect(logs).toContain("use client");
+    await waitForChange(
+      () => {},
+      () => page.$$("script")
+    );
+    const scripts = await page.$$("script");
+    // this is flaky and needs a stable solution
+    expect(scripts.length).toBeGreaterThanOrEqual(4);
+    expect(await scripts[0].getAttribute("src")).toBe("/@vite/client");
+    expect(await scripts[1].getAttribute("src")).toBe("/@hmr");
+    expect(await scripts[2].getAttribute("src")).toBe("/@__webpack_require__");
+    expect(await scripts[3].getAttribute("src")).toBe(null);
+  }
+});

test/fixtures/client-component.jsx

@@ -0,0 +1,7 @@
+"use client";
+
+export default function ClientComponent() {
+  return (
+    <button onClick={() => console.log("use client")}>Client Component</button>
+  );
+}

test/fixtures/suspense-client.jsx

@@ -0,0 +1,17 @@
+import { Suspense } from "react";
+import ClientComponent from "./client-component.jsx";
+
+async function AsyncComponent() {
+  await new Promise((resolve) => setTimeout(resolve, 100));
+  return <ClientComponent />;
+}
+
+export default function App() {
+  return (
+    <div>
+      <Suspense fallback={<div>Loading...</div>}>
+        <AsyncComponent />
+      </Suspense>
+    </div>
+  );
+}