Support hot reload for server components

examples/04_promise/src/components/Counter.tsx

@@ -3,6 +3,8 @@
 
 import { Suspense, useState, use } from 'react';
 
+import { Hello } from './Hello.js';
+
 export const Counter = ({
   delayedMessage,
 }: {
@@ -17,6 +19,7 @@ export const Counter = ({
       <Suspense fallback="Pending...">
         <Message count={count} delayedMessage={delayedMessage} />
       </Suspense>
+      <Hello />
     </div>
   );
 };

examples/04_promise/src/components/Hello.tsx

@@ -0,0 +1,7 @@
+export const Hello = () => {
+  return (
+    <div style={{ border: '3px gray dashed', margin: '1em', padding: '1em' }}>
+      This is a component without {'"use client"'}.
+    </div>
+  );
+};

packages/waku/src/lib/handlers/dev-worker-api.ts

@@ -4,7 +4,7 @@ import type {
 } from 'node:worker_threads';
 
 import type { ResolvedConfig } from '../config.js';
-import type { ModuleImportResult } from '../plugins/vite-plugin-rsc-hmr.js';
+import type { HotUpdatePayload } from '../plugins/vite-plugin-rsc-hmr.js';
 
 export type RenderRequest = {
   input: string;
@@ -37,9 +37,7 @@ export type MessageReq =
     };
 
 export type MessageRes =
-  | { type: 'full-reload' }
-  | { type: 'hot-import'; source: string }
-  | { type: 'module-import'; result: ModuleImportResult }
+  | { type: 'hot-update'; payload: HotUpdatePayload }
   | { id: number; type: 'start'; context: unknown; stream: ReadableStream }
   | { id: number; type: 'err'; err: unknown; statusCode?: number }
   | { id: number; type: 'moduleId'; moduleId: string }
@@ -108,37 +106,13 @@ const getWorker = () => {
   return workerPromise;
 };
 
-export async function registerReloadCallback(
-  fn: (type: 'full-reload') => void,
+export async function registerHotUpdateCallback(
+  fn: (payload: HotUpdatePayload) => void,
 ) {
   const worker = await getWorker();
   const listener = (mesg: MessageRes) => {
-    if (mesg.type === 'full-reload') {
-      fn(mesg.type);
-    }
-  };
-  worker.on('message', listener);
-  return () => worker.off('message', listener);
-}
-
-export async function registerImportCallback(fn: (source: string) => void) {
-  const worker = await getWorker();
-  const listener = (mesg: MessageRes) => {
-    if (mesg.type === 'hot-import') {
-      fn(mesg.source);
-    }
-  };
-  worker.on('message', listener);
-  return () => worker.off('message', listener);
-}
-
-export async function registerModuleCallback(
-  fn: (result: ModuleImportResult) => void,
-) {
-  const worker = await getWorker();
-  const listener = (mesg: MessageRes) => {
-    if (mesg.type === 'module-import') {
-      fn(mesg.result);
+    if (mesg.type === 'hot-update') {
+      fn(mesg.payload);
     }
   };
   worker.on('message', listener);

packages/waku/src/lib/handlers/dev-worker-impl.ts

@@ -20,7 +20,6 @@ import { renderRsc, getSsrConfig } from '../renderers/rsc-renderer.js';
 import { nonjsResolvePlugin } from '../plugins/vite-plugin-nonjs-resolve.js';
 import { rscTransformPlugin } from '../plugins/vite-plugin-rsc-transform.js';
 import { rscEnvPlugin } from '../plugins/vite-plugin-rsc-env.js';
-import { rscReloadPlugin } from '../plugins/vite-plugin-rsc-reload.js';
 import { rscDelegatePlugin } from '../plugins/vite-plugin-rsc-delegate.js';
 import { mergeUserViteConfig } from '../utils/merge-vite-config.js';
 
@@ -107,8 +106,6 @@ const handleGetSsrConfig = async (
 
 const dummyServer = new Server(); // FIXME we hope to avoid this hack
 
-const moduleImports: Set<string> = new Set();
-
 const mergedViteConfig = await mergeUserViteConfig({
   plugins: [
     viteReact(),
@@ -117,21 +114,10 @@ const mergedViteConfig = await mergeUserViteConfig({
     { name: 'rsc-hmr-plugin', enforce: 'post' }, // dummy to match with handler-dev.ts
     nonjsResolvePlugin(),
     rscTransformPlugin({ isBuild: false }),
-    rscReloadPlugin(moduleImports, (type) => {
-      const mesg: MessageRes = { type };
+    rscDelegatePlugin((payload) => {
+      const mesg: MessageRes = { type: 'hot-update', payload };
       parentPort!.postMessage(mesg);
     }),
-    rscDelegatePlugin(
-      moduleImports,
-      (source) => {
-        const mesg: MessageRes = { type: 'hot-import', source };
-        parentPort!.postMessage(mesg);
-      },
-      (result) => {
-        const mesg: MessageRes = { type: 'module-import', result };
-        parentPort!.postMessage(mesg);
-      },
-    ),
   ],
   optimizeDeps: {
     include: ['react-server-dom-webpack/client', 'react-dom'],

packages/waku/src/lib/handlers/handler-dev.ts

@@ -14,19 +14,13 @@ import { renderHtml } from '../renderers/html-renderer.js';
 import { decodeInput, hasStatusCode } from '../renderers/utils.js';
 import {
   initializeWorker,
-  registerReloadCallback,
-  registerImportCallback,
-  registerModuleCallback,
+  registerHotUpdateCallback,
   renderRscWithWorker,
   getSsrConfigWithWorker,
 } from './dev-worker-api.js';
 import { patchReactRefresh } from '../plugins/patch-react-refresh.js';
 import { rscIndexPlugin } from '../plugins/vite-plugin-rsc-index.js';
-import {
-  rscHmrPlugin,
-  hotImport,
-  moduleImport,
-} from '../plugins/vite-plugin-rsc-hmr.js';
+import { rscHmrPlugin, hotUpdate } from '../plugins/vite-plugin-rsc-hmr.js';
 import { rscEnvPlugin } from '../plugins/vite-plugin-rsc-env.js';
 import type { BaseReq, BaseRes, Handler } from './types.js';
 import { mergeUserViteConfig } from '../utils/merge-vite-config.js';
@@ -66,7 +60,6 @@ export function createHandler<
         rscHmrPlugin(),
         { name: 'nonjs-resolve-plugin' }, // dummy to match with dev-worker-impl.ts
         { name: 'rsc-transform-plugin' }, // dummy to match with dev-worker-impl.ts
-        { name: 'rsc-reload-plugin' }, // dummy to match with dev-worker-impl.ts
         { name: 'rsc-delegate-plugin' }, // dummy to match with dev-worker-impl.ts
       ],
       optimizeDeps: {
@@ -89,9 +82,7 @@ export function createHandler<
     });
     const vite = await createViteServer(mergedViteConfig);
     initializeWorker(config);
-    registerReloadCallback((type) => vite.ws.send({ type }));
-    registerImportCallback((source) => hotImport(vite, source));
-    registerModuleCallback((result) => moduleImport(vite, result));
+    registerHotUpdateCallback((payload) => hotUpdate(vite, payload));
     return vite;
   });
 

packages/waku/src/lib/plugins/vite-plugin-rsc-delegate.ts

@@ -2,17 +2,36 @@ import path from 'node:path';
 import type { Plugin, ViteDevServer } from 'vite';
 import * as swc from '@swc/core';
 
-import type { ModuleImportResult } from './vite-plugin-rsc-hmr.js';
+import type { HotUpdatePayload } from './vite-plugin-rsc-hmr.js';
+
+const isClientEntry = (id: string, code: string) => {
+  const ext = path.extname(id);
+  if (['.ts', '.tsx', '.js', '.jsx'].includes(ext)) {
+    const mod = swc.parseSync(code, {
+      syntax: ext === '.ts' || ext === '.tsx' ? 'typescript' : 'ecmascript',
+      tsx: ext === '.tsx',
+    });
+    for (const item of mod.body) {
+      if (
+        item.type === 'ExpressionStatement' &&
+        item.expression.type === 'StringLiteral' &&
+        item.expression.value === 'use client'
+      ) {
+        return true;
+      }
+    }
+  }
+  return false;
+};
 
 // import { CSS_LANGS_RE } from "vite/dist/node/constants.js";
 const CSS_LANGS_RE =
   /\.(css|less|sass|scss|styl|stylus|pcss|postcss|sss)(?:$|\?)/;
 
 export function rscDelegatePlugin(
-  moduleImports: Set<string>,
-  sourceCallback: (source: string) => void,
-  moduleCallback: (result: ModuleImportResult) => void,
+  callback: (payload: HotUpdatePayload) => void,
 ): Plugin {
+  const moduleImports: Set<string> = new Set();
   let mode = 'development';
   let base = '/';
   let server: ViteDevServer;
@@ -25,13 +44,24 @@ export function rscDelegatePlugin(
     configureServer(serverInstance) {
       server = serverInstance;
     },
-    async handleHotUpdate({ file }) {
-      if (moduleImports.has(file)) {
-        // re-inject
-        const transformedResult = await server.transformRequest(file);
-        if (transformedResult) {
-          const { default: source } = await server.ssrLoadModule(file);
-          moduleCallback({ ...transformedResult, source, id: file });
+    async handleHotUpdate(ctx) {
+      if (mode === 'development') {
+        if (moduleImports.has(ctx.file)) {
+          // re-inject
+          const transformedResult = await server.transformRequest(ctx.file);
+          if (transformedResult) {
+            const { default: source } = await server.ssrLoadModule(ctx.file);
+            callback({
+              type: 'custom',
+              event: 'module-import',
+              data: { ...transformedResult, source, id: ctx.file },
+            });
+          }
+        } else if (
+          ctx.modules.length &&
+          !isClientEntry(ctx.file, await ctx.read())
+        ) {
+          callback({ type: 'custom', event: 'rsc-reload' });
         }
       }
     },
@@ -50,7 +80,7 @@ export function rscDelegatePlugin(
             if (item.source.value.startsWith('virtual:')) {
               // HACK this relies on Vite's internal implementation detail.
               const source = base + '@id/__x00__' + item.source.value;
-              sourceCallback(source);
+              callback({ type: 'custom', event: 'hot-import', data: source });
             } else if (CSS_LANGS_RE.test(item.source.value)) {
               const resolvedSource = await server.pluginContainer.resolveId(
                 item.source.value,
@@ -66,11 +96,15 @@ export function rscDelegatePlugin(
                 );
                 if (transformedResult) {
                   moduleImports.add(resolvedSource.id);
-                  moduleCallback({
-                    ...transformedResult,
-                    source,
-                    id: resolvedSource.id,
-                    css: true,
+                  callback({
+                    type: 'custom',
+                    event: 'module-import',
+                    data: {
+                      ...transformedResult,
+                      source,
+                      id: resolvedSource.id,
+                      css: true,
+                    },
                   });
                 }
               }

packages/waku/src/lib/plugins/vite-plugin-rsc-hmr.ts

@@ -5,19 +5,28 @@ import type {
   ViteDevServer,
 } from 'vite';
 
-export type ModuleImportResult = TransformResult & {
+import {
+  joinPath,
+  fileURLToFilePath,
+  decodeFilePathFromAbsolute,
+} from '../utils/path.js';
+
+type ModuleImportResult = TransformResult & {
   id: string;
   // non-transformed result of `TransformResult.code`
   source: string;
   css?: boolean;
 };
 
-const customCode = `
+const injectingHmrCode = `
 import { createHotContext as __vite__createHotContext } from "/@vite/client";
 import.meta.hot = __vite__createHotContext(import.meta.url);
 
 if (import.meta.hot && !globalThis.__WAKU_HMR_CONFIGURED__) {
   globalThis.__WAKU_HMR_CONFIGURED__ = true;
+  import.meta.hot.on('rsc-reload', () => {
+    globalThis.__WAKU_REFETCH_RSC__?.();
+  });
   import.meta.hot.on('hot-import', (data) => import(/* @vite-ignore */ data));
   import.meta.hot.on('module-import', (data) => {
     // remove element with the same 'waku-module-id'
@@ -39,6 +48,9 @@ if (import.meta.hot && !globalThis.__WAKU_HMR_CONFIGURED__) {
 `;
 
 export function rscHmrPlugin(): Plugin {
+  const wakuClientDist = decodeFilePathFromAbsolute(
+    joinPath(fileURLToFilePath(import.meta.url), '../../../client.js'),
+  );
   let viteServer: ViteDevServer;
   return {
     name: 'rsc-hmr-plugin',
@@ -52,17 +64,36 @@ export function rscHmrPlugin(): Plugin {
         {
           tag: 'script',
           attrs: { type: 'module', async: true },
-          children: customCode,
+          children: injectingHmrCode,
           injectTo: 'head',
         },
       ];
     },
+    async transform(code, id) {
+      if (id === wakuClientDist) {
+        // FIXME this is fragile. Can we do it better?
+        const FETCH_RSC_LINE =
+          'export const fetchRSC = (input, searchParamsString, setElements, cache = fetchCache)=>{';
+        return code.replace(
+          FETCH_RSC_LINE,
+          FETCH_RSC_LINE +
+            `
+globalThis.__WAKU_REFETCH_RSC__ = () => {
+  cache.splice(0);
+  const searchParams = new URLSearchParams(searchParamsString);
+  searchParams.delete('waku_router_skip'); // HACK hard coded, FIXME we need event listeners for 'rsc-reload'
+  const data = fetchRSC(input, searchParams.toString(), setElements, cache);
+  setElements((prev) => mergeElements(prev, data));
+};`,
+        );
+      }
+    },
   };
 }
 
 const pendingMap = new WeakMap<ViteDevServer, Set<string>>();
 
-export function hotImport(vite: ViteDevServer, source: string) {
+function hotImport(vite: ViteDevServer, source: string) {
   let sourceSet = pendingMap.get(vite);
   if (!sourceSet) {
     sourceSet = new Set();
@@ -79,10 +110,7 @@ export function hotImport(vite: ViteDevServer, source: string) {
 
 const modulePendingMap = new WeakMap<ViteDevServer, Set<ModuleImportResult>>();
 
-export function moduleImport(
-  viteServer: ViteDevServer,
-  result: ModuleImportResult,
-) {
+function moduleImport(viteServer: ViteDevServer, result: ModuleImportResult) {
   let sourceSet = modulePendingMap.get(viteServer);
   if (!sourceSet) {
     sourceSet = new Set();
@@ -137,3 +165,21 @@ async function generateInitialScripts(
   }
   return scripts;
 }
+
+export type HotUpdatePayload =
+  | { type: 'full-reload' }
+  | { type: 'custom'; event: 'rsc-reload' }
+  | { type: 'custom'; event: 'hot-import'; data: string }
+  | { type: 'custom'; event: 'module-import'; data: ModuleImportResult };
+
+export function hotUpdate(vite: ViteDevServer, payload: HotUpdatePayload) {
+  if (payload.type === 'full-reload') {
+    vite.ws.send(payload);
+  } else if (payload.event === 'rsc-reload') {
+    vite.ws.send(payload);
+  } else if (payload.event === 'hot-import') {
+    hotImport(vite, payload.data);
+  } else if (payload.event === 'module-import') {
+    moduleImport(vite, payload.data);
+  }
+}