fix: FOUC

e2e/ssr-basic.spec.ts

@@ -64,8 +64,6 @@ for (const { build, command } of commands) {
     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();
@@ -81,7 +79,7 @@ for (const { build, command } of commands) {
       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({ force: true });
+      await page.getByTestId('increment').click();
       await expect(page.getByTestId('count')).toHaveText('0');
       await page.close();
       await context.close();

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

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

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

@@ -30,7 +30,8 @@ export function rscDelegatePlugin(
         // re-inject
         const transformedResult = await server.transformRequest(file);
         if (transformedResult) {
-          moduleCallback({ ...transformedResult, id: file });
+          const { default: source } = await server.ssrLoadModule(file);
+          moduleCallback({ ...transformedResult, source, id: file });
         }
       }
     },
@@ -57,13 +58,17 @@ export function rscDelegatePlugin(
                 { ssr: true },
               );
               if (resolvedSource?.id) {
+                const { default: source } = await server.ssrLoadModule(
+                  resolvedSource.id,
+                );
                 const transformedResult = await server.transformRequest(
                   resolvedSource.id,
                 );
                 if (transformedResult) {
                   moduleImports.add(resolvedSource.id);
                   moduleCallback({
                     ...transformedResult,
+                    source,
                     id: resolvedSource.id,
                     css: true,
                   });

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

@@ -1,8 +1,14 @@
-import path from 'node:path';
-import type { Plugin, TransformResult, ViteDevServer } from 'vite';
+import type {
+  HtmlTagDescriptor,
+  Plugin,
+  TransformResult,
+  ViteDevServer,
+} from 'vite';
 
 export type ModuleImportResult = TransformResult & {
   id: string;
+  // non-transformed result of `TransformResult.code`
+  source: string;
   css?: boolean;
 };
 
@@ -13,80 +19,44 @@ 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));
-  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 style = document.querySelector('style[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();
+    // avoid HMR flash by first applying the new and removing the old styles 
+    if (style) {
+      queueMicrotask(style.remove);
+    }
   });
 }
 `;
 
-export function rscHmrPlugin(opts: { srcDir: string; mainJs: string }): Plugin {
-  let mainJsFile: string;
+export function rscHmrPlugin(): Plugin {
+  let viteServer: ViteDevServer;
   return {
     name: 'rsc-hmr-plugin',
     enforce: 'post',
-    configResolved(config) {
-      mainJsFile = path.posix.join(config.root, opts.srcDir, opts.mainJs);
+    configureServer(server) {
+      viteServer = server;
     },
-    transformIndexHtml() {
+    async transformIndexHtml() {
       return [
+        ...(await generateInitialScripts(viteServer)),
         {
           tag: 'script',
           attrs: { type: 'module', async: true },
           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);
-  }
-}
-        `,
-        );
-      }
-    },
   };
 }
 
@@ -117,16 +87,28 @@ export function moduleImport(
   if (!sourceSet) {
     sourceSet = new Set();
     modulePendingMap.set(viteServer, sourceSet);
-    viteServer.ws.on('connection', () => {
-      for (const result of sourceSet!) {
-        viteServer.ws.send({
-          type: 'custom',
-          event: 'module-import',
-          data: result,
-        });
-      }
-    });
   }
   sourceSet.add(result);
   viteServer.ws.send({ type: 'custom', event: 'module-import', data: result });
 }
+
+async function generateInitialScripts(
+  viteServer: ViteDevServer,
+): Promise<HtmlTagDescriptor[]> {
+  const sourceSet = modulePendingMap.get(viteServer);
+
+  if (!sourceSet) {
+    return [];
+  }
+
+  const scripts: HtmlTagDescriptor[] = [];
+  for (const result of sourceSet) {
+    scripts.push({
+      tag: 'style',
+      attrs: { type: 'text/css', 'waku-module-id': result.id },
+      children: result.source,
+      injectTo: 'head-prepend',
+    });
+  }
+  return scripts;
+}

packages/website/package.json

@@ -17,7 +17,7 @@
     "next-mdx-remote": "^4.4.1",
     "react": "18.3.0-canary-b30030471-20240117",
     "react-dom": "18.3.0-canary-b30030471-20240117",
-    "waku": "https://pkg.csb.dev/dai-shi/waku/commit/e5791fa5/waku"
+    "waku": "https://pkg.csb.dev/dai-shi/waku/commit/8a4b89af/waku"
   },
   "devDependencies": {
     "@types/react": "^18.2.48",