Fix client-only module HMR in RSC Framework Mode

Author: markdalgleish

integration/vite-css-test.ts

@@ -599,8 +599,7 @@ async function hmrWorkflow({
       // changing non-React modules imported by the route.
       if (
         templateName.includes("rsc") &&
-        (file === "styles.module.css" ||
-          file === "styles-postcss-linked.css" ||
+        (file === "styles-postcss-linked.css" ||
           file === "styles-vanilla-global.css.ts")
       ) {
         continue;

packages/react-router-dev/vite/rsc/plugin.ts

@@ -20,6 +20,7 @@ import {
   isVirtualClientRouteModuleId,
   CLIENT_NON_COMPONENT_EXPORTS,
 } from "./virtual-route-modules";
+import { hmrInvalidateClientOnlyModulesInRsc } from "./plugins/hmr-invalidate-client-only-modules-in-rsc";
 import validatePluginOrder from "../plugins/validate-plugin-order";
 
 export function reactRouterRSCVitePlugin(): Vite.PluginOption[] {
@@ -321,6 +322,7 @@ export function reactRouterRSCVitePlugin(): Vite.PluginOption[] {
         return modules;
       },
     },
+    hmrInvalidateClientOnlyModulesInRsc(),
     validatePluginOrder(),
     rsc({ entries: getRscEntries() }),
   ];

packages/react-router-dev/vite/rsc/plugins/hmr-invalidate-client-only-modules-in-rsc.ts

@@ -0,0 +1,61 @@
+import type * as Vite from "vite";
+
+export function hmrInvalidateClientOnlyModulesInRsc(): Vite.Plugin {
+  return {
+    name: "react-router/rsc/hmr/invalidate-client-only-modules-in-rsc",
+    async hotUpdate(ctx) {
+      // We only want to invalidate ancestors of client-only modules in the RSC
+      // graph, so bail out if we're not in the RSC environment
+      if (this.environment.name !== "rsc") {
+        return;
+      }
+
+      const updatedServerModule = this.environment.moduleGraph.getModuleById(
+        ctx.file,
+      );
+
+      // If this file is in the RSC graph, it's not a client-only module and
+      // changes will already be picked up, so bail out
+      if (updatedServerModule) {
+        return;
+      }
+
+      // Find the corresponding client module for this file so we can walk the
+      // module graph
+      const updatedClientModule =
+        ctx.server.environments.client.moduleGraph.getModuleById(ctx.file);
+
+      // If this file is not in the client graph, it's not a client-only module
+      // and we don't need to invalidate anything, so bail out
+      if (!updatedClientModule) {
+        return;
+      }
+
+      const visited = new Set<Vite.EnvironmentModuleNode>();
+      const walk = (module: Vite.EnvironmentModuleNode) => {
+        if (!module || visited.has(module) || !module.id) {
+          return;
+        }
+
+        visited.add(module);
+
+        // If this module is in the RSC graph, invalidate it and stop walking
+        const serverModule = this.environment.moduleGraph.getModuleById(
+          module.id,
+        );
+        if (serverModule) {
+          this.environment.moduleGraph.invalidateModule(serverModule);
+          return;
+        }
+
+        if (module.importers) {
+          for (const importer of module.importers) {
+            walk(importer);
+          }
+        }
+      };
+
+      walk(updatedClientModule);
+    },
+  };
+}