DEV: loading indicator to mitigate FOUC

e2e/ssr-basic.spec.ts

@@ -63,14 +63,14 @@
 
     test('increase counter', async ({ page }) => {
       await page.goto(`http://localhost:${port}/`);
-
       await expect(page.getByTestId('app-name')).toHaveText('Waku');
-
+      // hydration is delayed 500ms at most in dev.
+      await expect(page.locator('#waku-module-spinner')).toBeHidden();
       await expect(page.getByTestId('count')).toHaveText('0');
       await page.getByTestId('increment').click();
       await page.getByTestId('increment').click();
       await page.getByTestId('increment').click();
       await expect(page.getByTestId('count')).toHaveText('3');
     });
 
     test('no js environment should have first screen', async ({ browser }) => {
@@ -81,7 +81,7 @@
       await page.goto(`http://localhost:${port}/`);
       await expect(page.getByTestId('app-name')).toHaveText('Waku');
       await expect(page.getByTestId('count')).toHaveText('0');
-      await page.getByTestId('increment').click();
+      await page.getByTestId('increment').click({ force: true });
       await expect(page.getByTestId('count')).toHaveText('0');
       await page.close();
       await context.close();

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 './types.js';
+import type { ModuleImportResult } from '../plugins/vite-plugin-rsc-hmr.js';
 
 export type RenderRequest = {
   input: string;

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

@@ -121,13 +121,17 @@ const mergedViteConfig = await mergeUserViteConfig({
       const mesg: MessageRes = { type };
       parentPort!.postMessage(mesg);
     }),
-    rscDelegatePlugin(moduleImports, (resultOrSource) => {
-      const mesg: MessageRes =
-        typeof resultOrSource === 'object'
-          ? { type: 'module-import', result: resultOrSource }
-          : { type: 'hot-import', source: resultOrSource };
-      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

@@ -63,7 +63,7 @@ export function createHandler<
         patchReactRefresh(viteReact()),
         rscEnvPlugin({ config, hydrate: ssr }),
         rscIndexPlugin(config),
-        rscHmrPlugin(),
+        rscHmrPlugin(config),
         { 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

packages/waku/src/lib/handlers/types.ts

@@ -1,5 +1,3 @@
-import type { TransformResult } from 'vite';
-
 export type BaseReq = {
   stream: ReadableStream;
   url: URL;
@@ -18,7 +16,3 @@ export type Handler<Req extends BaseReq, Res extends BaseRes> = (
   res: Res,
   next: (err?: unknown) => void,
 ) => void;
-
-export type ModuleImportResult = TransformResult & {
-  id: string;
-};

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

@@ -1,15 +1,17 @@
 import path from 'node:path';
 import type { Plugin, ViteDevServer } from 'vite';
 import * as swc from '@swc/core';
-import type { ModuleImportResult } from '../handlers/types.js';
+
+import type { ModuleImportResult } from './vite-plugin-rsc-hmr.js';
 
 // 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>,
-  importCallback: (source: string | ModuleImportResult) => void,
+  sourceCallback: (source: string) => void,
+  moduleCallback: (result: ModuleImportResult) => void,
 ): Plugin {
   let mode = 'development';
   let base = '/';
@@ -27,7 +29,9 @@ export function rscDelegatePlugin(
       if (moduleImports.has(file)) {
         // re-inject
         const transformedResult = await server.transformRequest(file);
-        transformedResult && importCallback({ ...transformedResult, id: file });
+        if (transformedResult) {
+          moduleCallback({ ...transformedResult, id: file });
+        }
       }
     },
     async transform(code, id) {
@@ -45,7 +49,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;
-              importCallback(source);
+              sourceCallback(source);
             } else if (CSS_LANGS_RE.test(item.source.value)) {
               const resolvedSource = await server.pluginContainer.resolveId(
                 item.source.value,
@@ -58,9 +62,10 @@ export function rscDelegatePlugin(
                 );
                 if (transformedResult) {
                   moduleImports.add(resolvedSource.id);
-                  importCallback({
+                  moduleCallback({
                     ...transformedResult,
                     id: resolvedSource.id,
+                    css: true,
                   });
                 }
               }

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

@@ -1,5 +1,10 @@
-import type { Plugin, ViteDevServer } from 'vite';
-import type { ModuleImportResult } from '../handlers/types.js';
+import path from 'node:path';
+import type { Plugin, TransformResult, ViteDevServer } from 'vite';
+
+export type ModuleImportResult = TransformResult & {
+  id: string;
+  css?: boolean;
+};
 
 const customCode = `
 import { createHotContext as __vite__createHotContext } from "/@vite/client";
@@ -8,27 +13,35 @@ 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('hot-import', (data) => import(/* @vite-ignore */ data));
-  import.meta.hot.on('module', (data) => {
+  const removeSpinner = () => {
+    const spinner = document.getElementById('waku-module-spinner');
+    spinner?.nextSibling?.remove();
+    spinner?.remove();
+  }
+  setTimeout(removeSpinner, 500);
+  import.meta.hot.on('module-import', (data) => {
     // remove element with the same 'waku-module-id'
-    let script = document.querySelector(
-      'script[waku-module-id="' + data.id + '"]',
-    );
+    let script = document.querySelector('script[waku-module-id="' + data.id + '"]');
     script?.remove();
-
     const code = data.code;
     script = document.createElement('script');
     script.type = 'module';
     script.text = code;
     script.setAttribute('waku-module-id', data.id);
     document.head.appendChild(script);
+    if (data.css) removeSpinner();
   });
 }
 `;
 
-export function rscHmrPlugin(): Plugin {
+export function rscHmrPlugin(opts: { srcDir: string; mainJs: string }): Plugin {
+  let mainJsFile: string;
   return {
     name: 'rsc-hmr-plugin',
     enforce: 'post',
+    configResolved(config) {
+      mainJsFile = path.posix.join(config.root, opts.srcDir, opts.mainJs);
+    },
     transformIndexHtml() {
       return [
         {
@@ -37,8 +50,43 @@ export function rscHmrPlugin(): Plugin {
           children: customCode,
           injectTo: 'head',
         },
+        {
+          tag: 'div',
+          attrs: {
+            id: 'waku-module-spinner',
+            style:
+              'position: fixed; top: 0; left: 0; right: 0; bottom: 0; background: rgba(0, 0, 0, 0.8); display: flex; align-items: center; justify-content: center; font-family: sans-serif; font-size: 2rem; color: white; cursor: wait;',
+          },
+          children: 'Loading...',
+          injectTo: 'head',
+        },
       ];
     },
+    transform(code, id, options) {
+      if (options?.ssr) return;
+      if (id === mainJsFile) {
+        // FIXME this is pretty fragile, should we patch react-dom/client?
+        return code.replace(
+          'hydrateRoot(document.body, rootElement);',
+          `
+{
+  const spinner = document.getElementById('waku-module-spinner');
+  if (spinner) {
+    const observer = new MutationObserver(() => {
+      if (!document.contains(spinner)) {
+        observer.disconnect();
+        hydrateRoot(document.body, rootElement);
+      }
+    });
+    observer.observe(document, { childList: true, subtree: true });
+  } else {
+    hydrateRoot(document.body, rootElement);
+  }
+}
+        `,
+        );
+      }
+    },
   };
 }
 
@@ -71,10 +119,14 @@ export function moduleImport(
     modulePendingMap.set(viteServer, sourceSet);
     viteServer.ws.on('connection', () => {
       for (const result of sourceSet!) {
-        viteServer.ws.send({ type: 'custom', event: 'module', data: result });
+        viteServer.ws.send({
+          type: 'custom',
+          event: 'module-import',
+          data: result,
+        });
       }
     });
   }
   sourceSet.add(result);
-  viteServer.ws.send({ type: 'custom', event: 'module', data: result });
+  viteServer.ws.send({ type: 'custom', event: 'module-import', data: result });
 }