fix: improve client HMR (#1029)

#1021 was incomplete. This should fix it.

Author: dai-shi

e2e/hot-reload.spec.ts

@@ -112,11 +112,9 @@ test.describe('hot reload', () => {
     await modifyFile('src/pages/index.tsx', 'Modified Page', 'Edited Page');
     await expect(page.getByText('Edited Page')).toBeVisible();
     await page.waitForTimeout(500); // need to wait not to full reload
-    // FIXME The following should pass but not for now.
-    //       It's probably because Vite adds `?t=...` timestamp with HMR.
-    // await expect(page.getByTestId('count')).toHaveText('3');
-    // await page.getByTestId('increment').click();
-    // await expect(page.getByTestId('count')).toHaveText('4');
+    await expect(page.getByTestId('count')).toHaveText('3');
+    await page.getByTestId('increment').click();
+    await expect(page.getByTestId('count')).toHaveText('4');
     await terminate(pid!);
   });
 });

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

@@ -77,12 +77,17 @@ export function rscDelegatePlugin(
     },
     async handleHotUpdate(ctx) {
       if (mode === 'development') {
+        if (ctx.file.endsWith('/pages.gen.ts')) {
+          // auto generated file by fsRouterTypegenPlugin
+          return [];
+        }
         await updateAllStyles(); // FIXME is this too aggressive?
         if (moduleImports.has(ctx.file)) {
           // re-inject
           const transformedResult = await server.transformRequest(ctx.file);
           if (transformedResult) {
             const { default: source } = await server.ssrLoadModule(ctx.file);
+            console.log('[rsc] module import', ctx.file);
             callback({
               type: 'custom',
               event: 'module-import',
@@ -93,6 +98,7 @@ export function rscDelegatePlugin(
           ctx.modules.length &&
           !isClientEntry(ctx.file, await ctx.read())
         ) {
+          console.log('[rsc] hot reload', ctx.file);
           callback({ type: 'custom', event: 'rsc-reload' });
         }
       }

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

@@ -26,6 +26,22 @@ 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('vite:afterUpdate', (data) => {
+    if (data.type === 'update') {
+      for (const update of data.updates) {
+        if (
+          update.type === 'js-update' &&
+          globalThis.__WAKU_CLIENT_MODULE_LOADING__.has(update.path)
+        ) {
+          globalThis.__WAKU_CLIENT_MODULE_LOADING__.set(update.path,
+            globalThis.__WAKU_CLIENT_IMPORT__(update.path + '?t=' + update.timestamp).then((m) => {
+              globalThis.__WAKU_CLIENT_MODULE_CACHE__.set(update.path, m);
+            })
+          );
+        }
+      }
+    }
+  });
   import.meta.hot.on('rsc-reload', () => {
     globalThis.__WAKU_RSC_RELOAD_LISTENERS__?.forEach((l) => l());
   });