From 281cf97e6b0d988fa2fe5910c810288d46935481 Mon Sep 17 00:00:00 2001
From: patak <matias.capeletto@gmail.com>
Date: Mon, 11 Mar 2024 11:12:14 +0100
Subject: [PATCH] fix: call updateModules for each environmnet

---
 packages/vite/src/node/server/hmr.ts | 140 ++++++++++++++-------------
 1 file changed, 71 insertions(+), 69 deletions(-)

diff --git a/packages/vite/src/node/server/hmr.ts b/packages/vite/src/node/server/hmr.ts
index 74ca97b1e1c864..d2d92068a85d32 100644
--- a/packages/vite/src/node/server/hmr.ts
+++ b/packages/vite/src/node/server/hmr.ts
@@ -228,82 +228,84 @@ export async function handleHMRUpdate(
   // For now, we only call updateModules for the browser. Later on it should
   // also be called for each runtime.
 
-  let mods = server.browserEnvironment.moduleGraph.getModulesByFile(file)
-  if (!mods) {
-    // For now, given that the HMR SSR expects it, try to get the modules from the
-    // server graph if the browser graph doesn't have it
-    mods = server.serverEnvironment.moduleGraph.getModulesByFile(file)
-  }
-
-  // check if any plugin wants to perform custom HMR handling
-  const timestamp = Date.now()
-  const hotContext: HotUpdateContext = {
-    file,
-    timestamp,
-    modules: mods ? [...mods] : [],
-    read: () => readModifiedFile(file),
-    server,
-    // later on hotUpdate will be called for each runtime with a new hotContext
-    environment: 'browser',
-  }
+  server.environments.forEach(async (environment) => {
+    const mods = environment.moduleGraph.getModulesByFile(file)
+
+    // check if any plugin wants to perform custom HMR handling
+    const timestamp = Date.now()
+    const hotContext: HotUpdateContext = {
+      file,
+      timestamp,
+      modules: mods ? [...mods] : [],
+      read: () => readModifiedFile(file),
+      server,
+      // later on hotUpdate will be called for each runtime with a new hotContext
+      environment: 'browser',
+    }
 
-  let hmrContext
+    let hmrContext
 
-  for (const plugin of getSortedHotUpdatePlugins(config)) {
-    if (plugin.hotUpdate) {
-      const filteredModules = await getHookHandler(plugin.hotUpdate)(hotContext)
-      if (filteredModules) {
-        hotContext.modules = filteredModules
-        // Invalidate the hmrContext to force compat modules to be updated
-        hmrContext = undefined
-      }
-    } else {
-      // later on, we'll need: if (runtime === 'browser')
-      // Backward compatibility with mixed client and ssr moduleGraph
-      hmrContext ??= {
-        ...hotContext,
-        modules: hotContext.modules.map((mod) =>
-          server.moduleGraph.getBackwardCompatibleModuleNode(mod),
-        ),
-      } as HmrContext
-      const filteredModules = await getHookHandler(plugin.handleHotUpdate!)(
-        hmrContext,
-      )
-      if (filteredModules) {
-        hmrContext.modules = filteredModules
-        hotContext.modules = filteredModules
-          .map((mod) =>
-            mod.id
-              ? server.browserEnvironment.moduleGraph.getModuleById(mod.id) ??
-                server.serverEnvironment.moduleGraph.getModuleById(mod.id)
-              : undefined,
-          )
-          .filter(Boolean) as EnvironmentModuleNode[]
+    for (const plugin of getSortedHotUpdatePlugins(config)) {
+      if (plugin.hotUpdate) {
+        const filteredModules = await getHookHandler(plugin.hotUpdate)(
+          hotContext,
+        )
+        if (filteredModules) {
+          hotContext.modules = filteredModules
+          // Invalidate the hmrContext to force compat modules to be updated
+          hmrContext = undefined
+        }
+      } else if (environment.id === 'browser') {
+        // later on, we'll need: if (runtime === 'browser')
+        // Backward compatibility with mixed client and ssr moduleGraph
+        hmrContext ??= {
+          ...hotContext,
+          modules: hotContext.modules.map((mod) =>
+            server.moduleGraph.getBackwardCompatibleModuleNode(mod),
+          ),
+        } as HmrContext
+        const filteredModules = await getHookHandler(plugin.handleHotUpdate!)(
+          hmrContext,
+        )
+        if (filteredModules) {
+          hmrContext.modules = filteredModules
+          hotContext.modules = filteredModules
+            .map((mod) =>
+              mod.id
+                ? server.browserEnvironment.moduleGraph.getModuleById(mod.id) ??
+                  server.serverEnvironment.moduleGraph.getModuleById(mod.id)
+                : undefined,
+            )
+            .filter(Boolean) as EnvironmentModuleNode[]
+        }
       }
     }
-  }
 
-  if (!hotContext.modules.length) {
-    // html file cannot be hot updated
-    if (file.endsWith('.html')) {
-      config.logger.info(colors.green(`page reload `) + colors.dim(shortFile), {
-        clear: true,
-        timestamp: true,
-      })
-      hot.send({
-        type: 'full-reload',
-        path: config.server.middlewareMode
-          ? '*'
-          : '/' + normalizePath(path.relative(config.root, file)),
-      })
-    } else {
-      // loaded but not in the module graph, probably not js
-      debugHmr?.(`[no modules matched] ${colors.dim(shortFile)}`)
+    if (!hotContext.modules.length) {
+      // html file cannot be hot updated
+      if (file.endsWith('.html')) {
+        config.logger.info(
+          colors.green(`page reload `) + colors.dim(shortFile),
+          {
+            clear: true,
+            timestamp: true,
+          },
+        )
+        hot.send({
+          type: 'full-reload',
+          path: config.server.middlewareMode
+            ? '*'
+            : '/' + normalizePath(path.relative(config.root, file)),
+        })
+      } else {
+        // loaded but not in the module graph, probably not js
+        debugHmr?.(`[no modules matched] ${colors.dim(shortFile)}`)
+      }
+      return
     }
-    return
-  }
 
-  updateModules(shortFile, hotContext.modules, timestamp, server)
+    updateModules(shortFile, hotContext.modules, timestamp, server)
+  })
 }
 
 type HasDeadEnd = boolean