fix: Fixes and improvements for HMR and module graph (#597)

## Fix module graph for SSR modules
### 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
There was a bug in the way we parse the vite SSR imports - we weren't matching on cases where there were extra arguments being passed to these functions. As a result, vite wasn't able to construct the module graph correctly. Ultimately, this meant that when some modules were invalidated, their importers should have been invalidated too, but werent (since the link between them in the graph wasn't there).

## Invalidate explicitly for SSR
On changes to files in SSR graph, we now invalidate the relevant module and as well as its importers (recursively), then invalidate the same way for the corresponding virtual modules in the worker module graph, as well as its importers (also recursively).

## Better client module invalidation
We now:
* look at all client updates - rather than only ones that were also worker updates as we did previously; then
* invalid all the modules vite tells us where relevant when a file changes - previously we'd pass control on to vite's own HMR invalidation logic for use client modules and css files, and otherwise short circuit, but this turns out to be necessary

## Improvement: Avoid redundant logic for client/server lookups
We maintain "lookup" modules for modules with `use client` and `use server` directives, in order to find and import them on demand when RSC needs them. We were handling different variations on how these modules might be resolved (with/without `/@id/`, with/without `.js). We now instead avoid the possibility of there being multiple ways for these to be resolved, so that there's only one request id every specified for each - `virtual:use-client-lookup.js` and `virtual:use-server-lookup.js`) => we were able to remove complexity by only needing to resolve these modules 1 way instead of multiple different ways.

Author: justinvdm

pnpm-lock.yaml

@@ -1,6 +1,6 @@
 {
   "name": "rwsdk",
-  "version": "0.1.15",
+  "version": "0.1.15-test.20250714213423",
   "description": "Build fast, server-driven webapps on Cloudflare with SSR, RSC, and realtime",
   "type": "module",
   "bin": {

sdk/package.json

@@ -6,7 +6,7 @@ export const loadModule = memoize(async (id: string) => {
     return await import(/* @vite-ignore */ id);
   } else {
     const { useClientLookup } = await import(
-      "virtual:use-client-lookup" as string
+      "virtual:use-client-lookup.js" as string
     );
 
     const moduleFn = useClientLookup[id];

sdk/src/runtime/imports/client.ts

@@ -2,7 +2,7 @@ import memoize from "lodash/memoize";
 
 export const ssrLoadModule = memoize(async (id: string) => {
   const { useClientLookup } = await import(
-    "virtual:use-client-lookup" as string
+    "virtual:use-client-lookup.js" as string
   );
 
   const moduleFn = useClientLookup[id];

sdk/src/runtime/imports/ssr.ts

@@ -4,7 +4,7 @@ import { ssrWebpackRequire as baseSsrWebpackRequire } from "rwsdk/__ssr_bridge";
 
 export const loadServerModule = memoize(async (id: string) => {
   const { useServerLookup } = await import(
-    "virtual:use-server-lookup" as string
+    "virtual:use-server-lookup.js" as string
   );
 
   const moduleFn = useServerLookup[id];

sdk/src/runtime/imports/worker.ts

@@ -3,7 +3,7 @@ import { createServerReference as baseCreateServerReference } from "react-server
 
 export const loadServerModule = memoize(async (id: string) => {
   const { useServerLookup } = await import(
-    "virtual:use-server-lookup" as string
+    "virtual:use-server-lookup.js" as string
   );
 
   const moduleFn = useServerLookup[id];

sdk/src/runtime/register/ssr.ts

@@ -271,7 +271,7 @@ export const createDirectiveLookupPlugin = async ({
           build.onResolve(
             {
               filter: new RegExp(
-                `^(${escapedVirtualModuleName}|${escapedPrefixedModuleName})$`,
+                `^(${escapedVirtualModuleName}|${escapedPrefixedModuleName})\.js$`,
               ),
             },
             () => {
@@ -281,7 +281,7 @@ export const createDirectiveLookupPlugin = async ({
                   config.virtualModuleName,
                 );
               return {
-                path: config.virtualModuleName,
+                path: `${config.virtualModuleName}.js`,
                 external: true,
               };
             },
@@ -318,15 +318,11 @@ export const createDirectiveLookupPlugin = async ({
     resolveId(source) {
       process.env.VERBOSE && log("Resolving id=%s", source);
 
-      if (
-        source === config.virtualModuleName ||
-        source === `/@id/${config.virtualModuleName}` ||
-        source === `/@id/${config.virtualModuleName}.js`
-      ) {
+      if (source === `${config.virtualModuleName}.js`) {
         log("Resolving %s module", config.virtualModuleName);
         // context(justinvdm, 16 Jun 2025): Include .js extension
         // so it goes through vite processing chain
-        return `${config.virtualModuleName}.js`;
+        return source;
       }
 
       process.env.VERBOSE && log("No resolution for id=%s", source);

sdk/src/vite/createDirectiveLookupPlugin.mts

@@ -38,6 +38,22 @@ export function findSsrImportSpecifiers(
         pattern: `__vite_ssr_dynamic_import__('$SPECIFIER')`,
         list: dynamicImports,
       },
+      {
+        pattern: `__vite_ssr_import__("$SPECIFIER", $$$REST)`,
+        list: imports,
+      },
+      {
+        pattern: `__vite_ssr_import__('$SPECIFIER', $$$REST)`,
+        list: imports,
+      },
+      {
+        pattern: `__vite_ssr_dynamic_import__("$SPECIFIER", $$$REST)`,
+        list: dynamicImports,
+      },
+      {
+        pattern: `__vite_ssr_dynamic_import__('$SPECIFIER', $$$REST)`,
+        list: dynamicImports,
+      },
     ];
 
     for (const { pattern, list } of patterns) {

sdk/src/vite/findSsrSpecifiers.mts

@@ -63,7 +63,7 @@ export const miniflareHMRPlugin = (givenOptions: {
       } = givenOptions;
 
       if (process.env.VERBOSE) {
-        this.environment.logger.info(
+        log(
           `Hot update: (env=${
             this.environment.name
           }) ${ctx.file}\nModule graph:\n\n${dumpFullModuleGraph(
@@ -73,6 +73,21 @@ export const miniflareHMRPlugin = (givenOptions: {
         );
       }
 
+      if (this.environment.name === "ssr") {
+        log("SSR update, invalidating recursively", ctx.file);
+        invalidateModule(ctx.server, "ssr", ctx.file, {
+          invalidateImportersRecursively: true,
+        });
+        invalidateModule(
+          ctx.server,
+          environment,
+          VIRTUAL_SSR_PREFIX +
+            normalizeModulePath(givenOptions.rootDir, ctx.file),
+          { invalidateImportersRecursively: true },
+        );
+        return [];
+      }
+
       if (!["client", environment].includes(this.environment.name)) {
         return [];
       }
@@ -141,37 +156,30 @@ export const miniflareHMRPlugin = (givenOptions: {
         ctx.file === entry ||
         modules.some((module) => hasEntryAsAncestor(module, entry));
 
-      // The worker doesnt need an update
-      // => Short circuit HMR
-      if (!isWorkerUpdate) {
-        return [];
-      }
-
       // 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
       // Why: There may have been changes to css classes referenced, which might css modules to change
       if (this.environment.name === "client") {
-        for (const [_, module] of ctx.server.environments[environment]
-          .moduleGraph.idToModuleMap) {
-          // todo(justinvdm, 13 Dec 2024): We check+update _all_ css files in worker module graph,
-          // but it could just be a subset of css files that are actually affected, depending
-          // on the importers and imports of the changed file. We should be smarter about this.
-          if (module.file && module.file.endsWith(".css")) {
-            const clientModules =
-              ctx.server.environments.client.moduleGraph.getModulesByFile(
-                module.file,
-              );
-
-            for (const clientModule of clientModules ?? []) {
-              invalidateModule(ctx.server, "client", clientModule);
+        if (isWorkerUpdate) {
+          for (const [_, module] of ctx.server.environments[environment]
+            .moduleGraph.idToModuleMap) {
+            // todo(justinvdm, 13 Dec 2024): We check+update _all_ css files in worker module graph,
+            // but it could just be a subset of css files that are actually affected, depending
+            // on the importers and imports of the changed file. We should be smarter about this.
+            if (module.file && module.file.endsWith(".css")) {
+              const clientModules =
+                ctx.server.environments.client.moduleGraph.getModulesByFile(
+                  module.file,
+                );
+
+              for (const clientModule of clientModules ?? []) {
+                invalidateModule(ctx.server, "client", clientModule);
+              }
             }
           }
         }
 
-        // context(justinvdm, 10 Jul 2025): If this isn't a file with a client
-        // directive or a css file, we shouldn't invalidate anything else to
-        // avoid full page reload
-        return hasClientDirective || ctx.file.endsWith(".css") ? undefined : [];
+        return ctx.modules;
       }
 
       // The worker needs an update, and the hot check is for the worker environment

sdk/src/vite/miniflareHMRPlugin.mts

@@ -160,25 +160,35 @@ export const ssrBridgePlugin = ({
             log,
           );
 
-          const allSpecifiers = [...new Set([...imports, ...dynamicImports])];
+          const allSpecifiers = [
+            ...new Set([...imports, ...dynamicImports]),
+          ].map((id) =>
+            id.startsWith("/@id/") ? id.slice("/@id/".length) : id,
+          );
 
           const switchCases = allSpecifiers
             .map(
               (specifier) =>
-                `    case "${specifier}": return import("${VIRTUAL_SSR_PREFIX}${specifier}");`,
+                `    case "${specifier}": import("${VIRTUAL_SSR_PREFIX}${specifier}");`,
             )
             .join("\n");
 
           const transformedCode = `
 await (async function(__vite_ssr_import__, __vite_ssr_dynamic_import__) {${code}})(
-  (id, ...args) => {ssrImport(id); return __vite_ssr_import__('/@id/${VIRTUAL_SSR_PREFIX}' + id, ...args);},
-  (id, ...args) => {ssrImport(id); return __vite_ssr_dynamic_import__('/@id/${VIRTUAL_SSR_PREFIX}' + id, ...args);}
+  __ssrImport.bind(null, false),
+  __ssrImport.bind(null, true)
 );
 
-function ssrImport(id) {
+function __ssrImport(isDynamic, id, ...args) {
+  id = id.startsWith('/@id/') ? id.slice('/@id/'.length) : id;
+
   switch (id) {
 ${switchCases}
   }
+
+  return isDynamic
+    ? __vite_ssr_dynamic_import__("/@id/${VIRTUAL_SSR_PREFIX}" + id, ...args)
+    : __vite_ssr_import__("/@id/${VIRTUAL_SSR_PREFIX}" + id, ...args);
 }
 `;