fix: More HMR fixes for robustness (#598)

## Avoid js transforms for css
### Context
To support SSR having its own module resolution behaviour (e.g. not using `react-server` import condition or RSC react runtimes) while still existing in the same runtime env as the rest of the worker (where `react-sever` import condition and RSC react runtimes should apply), we have separate `worker` and `ssr` environments, then have an "ssr bridge" concept: in the `worker` environment, for all modules to be evaluated in SSR context, we fetch code from the `ssr` environment for that module.

We parse the code we fetch from the SSR environment to find all imports vite does (`__vite_ssr_import__`, `__vite_ssr_dynamic_import__`), and add `import`s for these fetches so that vite will build the module graph correctly.

### The fix
We do this fetching and transforming too for css modules too - except in these cases the result of the fetch is an empty string, and should remain that way rather than getting any other code added to it, in order to avoid adding js to the css.

## Opt out of custom HMR logic for non js and css files
If the js files are not js (e.g. js, jsx, mjs, cjs, ts, mts, etc) and not css, we should ignore them and pass control on to other plugins. We were trying to process .sqlite files during HMR, for e.g, which was resulting in HMR of these modules causing miniflare to crash.

## Skip HMR for SSR env when files are not in "use client" import chains
HMR for SSR should only look at "use client" files, or files imported by "use client" files. It was processing for other files (e.g. .sqlite files - while we avoid non js/css files now regardless, it still shows the need for more thorough checking for the use client graph for HMR in SSR env).

## Skip HMR for client env when files are not in "use client" import chains
It is possible for the client environment to contain files that are neither marked with "use client" nor imported by files that are. For example, if there is a style.css imported in worker module graph. For this specific example, to be honest, I'm not entirely sure why for this example case, I noticed this when inspecting the module graph. Regardless, it shows the need for more thorough checking for whether the module is in fact part of "use client" import chains.

For this reason, we short circuit HMR for the client env if we the file is not in a "use client" import chain.

## Avoid recursive invalidating for SSR modules
Turns out we only need to invalidate the SSR modules themselves - not their importers as well - for HMR to work correctly for SSR.

Author: justinvdm

sdk/src/vite/miniflareHMRPlugin.mts

@@ -45,6 +45,40 @@ const hasEntryAsAncestor = (
   return false;
 };
 
+const isInUseClientGraph = ({
+  file,
+  clientFiles,
+  server,
+}: {
+  file: string;
+  clientFiles: Set<string>;
+  server: ViteDevServer;
+}) => {
+  const id = normalizeModulePath(server.config.root, file);
+  if (clientFiles.has(id)) {
+    return true;
+  }
+
+  const modules = server.environments.client.moduleGraph.getModulesByFile(file);
+
+  if (!modules) {
+    return false;
+  }
+
+  for (const m of modules) {
+    for (const importer of m.importers) {
+      if (
+        importer.file &&
+        isInUseClientGraph({ file: importer.file, clientFiles, server })
+      ) {
+        return true;
+      }
+    }
+  }
+
+  return false;
+};
+
 export const miniflareHMRPlugin = (givenOptions: {
   clientFiles: Set<string>;
   serverFiles: Set<string>;
@@ -73,17 +107,28 @@ export const miniflareHMRPlugin = (givenOptions: {
         );
       }
 
+      if (!isJsFile(ctx.file) && !ctx.file.endsWith(".css")) {
+        return;
+      }
+
       if (this.environment.name === "ssr") {
         log("SSR update, invalidating recursively", ctx.file);
-        invalidateModule(ctx.server, "ssr", ctx.file, {
-          invalidateImportersRecursively: true,
+        const isUseClientUpdate = isInUseClientGraph({
+          file: ctx.file,
+          clientFiles,
+          server: ctx.server,
         });
+
+        if (!isUseClientUpdate) {
+          return [];
+        }
+
+        invalidateModule(ctx.server, "ssr", ctx.file);
         invalidateModule(
           ctx.server,
           environment,
           VIRTUAL_SSR_PREFIX +
             normalizeModulePath(givenOptions.rootDir, ctx.file),
-          { invalidateImportersRecursively: true },
         );
         return [];
       }
@@ -152,9 +197,7 @@ export const miniflareHMRPlugin = (givenOptions: {
         ) ?? [],
       );
 
-      const isWorkerUpdate =
-        ctx.file === entry ||
-        modules.some((module) => hasEntryAsAncestor(module, entry));
+      const isWorkerUpdate = Boolean(modules);
 
       // The worker needs an update, but this is the client environment
       // => Notify for HMR update of any css files imported by in worker, that are also in the client module graph
@@ -179,6 +222,16 @@ export const miniflareHMRPlugin = (givenOptions: {
           }
         }
 
+        const isUseClientUpdate = isInUseClientGraph({
+          file: ctx.file,
+          clientFiles,
+          server: ctx.server,
+        });
+
+        if (!isUseClientUpdate && !ctx.file.endsWith(".css")) {
+          return [];
+        }
+
         return ctx.modules;
       }
 

sdk/src/vite/ssrBridgePlugin.mts

@@ -2,6 +2,7 @@ import type { Plugin, ViteDevServer } from "vite";
 import debug from "debug";
 import { SSR_BRIDGE_PATH } from "../lib/constants.mjs";
 import { findSsrImportSpecifiers } from "./findSsrSpecifiers.mjs";
+import { isJsFile } from "./isJsFile.mjs";
 
 const log = debug("rwsdk:vite:ssr-bridge-plugin");
 
@@ -152,6 +153,13 @@ export const ssrBridgePlugin = ({
             log("Fetch module result: id=%s, result=%O", realId, result);
 
           const code = "code" in result ? result.code : undefined;
+
+          if (realId.endsWith(".css")) {
+            process.env.VERBOSE &&
+              log("Not a JS file, returning code: %s", code);
+            return code ?? "";
+          }
+
           log("Fetched SSR module code length: %d", code?.length || 0);
 
           const { imports, dynamicImports } = findSsrImportSpecifiers(
@@ -169,7 +177,7 @@ export const ssrBridgePlugin = ({
           const switchCases = allSpecifiers
             .map(
               (specifier) =>
-                `    case "${specifier}": import("${VIRTUAL_SSR_PREFIX}${specifier}");`,
+                `    case "${specifier}": void import("${VIRTUAL_SSR_PREFIX}${specifier}");`,
             )
             .join("\n");