fix(hmr): preserve environment snapshot during server restart

packages/vite/src/node/server/__tests__/hmr.spec.ts

@@ -0,0 +1,55 @@
+import path from 'node:path'
+import { expect, onTestFinished, test } from 'vitest'
+import { handleHMRUpdate } from '../hmr'
+import { type ServerOptions, createServer } from '../index'
+
+async function testRestartDuringHotUpdate(serverOptions: ServerOptions = {}) {
+  let hookCalls = 0
+  const server = await createServer({
+    configFile: false,
+    root: import.meta.dirname,
+    logLevel: 'silent',
+    server: {
+      middlewareMode: true,
+      ws: false,
+      ...serverOptions,
+    },
+    plugins: [
+      {
+        name: 'restart-during-hot-update',
+        async hotUpdate() {
+          if (hookCalls++ === 0) {
+            await server.restart()
+            throw new Error('hot update interrupted by restart')
+          }
+        },
+      },
+    ],
+  })
+  onTestFinished(async () => {
+    await server.close()
+  })
+
+  await expect(
+    handleHMRUpdate(
+      'update',
+      path.join(import.meta.dirname, 'fixture.js'),
+      server,
+    ),
+  ).resolves.toBeUndefined()
+}
+
+test('cancels HMR when the server restarts during a hot update', async () => {
+  await testRestartDuringHotUpdate()
+})
+
+test('does not schedule stale HMR with a custom environment handler', async () => {
+  let hotUpdateEnvironmentsCalls = 0
+  await testRestartDuringHotUpdate({
+    async hotUpdateEnvironments(server, hmr) {
+      hotUpdateEnvironmentsCalls++
+      await Promise.all(Object.values(server.environments).map(hmr))
+    },
+  })
+  expect(hotUpdateEnvironmentsCalls).toBe(0)
+})

packages/vite/src/node/server/hmr.ts

@@ -416,7 +416,11 @@ export async function handleHMRUpdate(
   const { config } = server
   const mixedModuleGraph = ignoreDeprecationWarnings(() => server.moduleGraph)
 
-  const environments = Object.values(server.environments)
+  const environmentSnapshot = server.environments
+  const environments = Object.values(environmentSnapshot)
+  // A plugin hook may restart the server, replacing the environments and
+  // invalidating this HMR transaction.
+  const isStale = () => server.environments !== environmentSnapshot
   const shortFile = getShortName(file, config.root)
 
   const isConfig = file === config.configFile
@@ -509,14 +513,13 @@ export async function handleHMRUpdate(
   const clientHotUpdateOptions = hotMap.get(clientEnvironment)!.options
   const ssrHotUpdateOptions = hotMap.get(ssrEnvironment)?.options
   try {
-    for (const plugin of getSortedHotUpdatePlugins(
-      server.environments.client,
-    )) {
+    for (const plugin of getSortedHotUpdatePlugins(clientEnvironment)) {
       if (plugin.hotUpdate) {
         const filteredModules = await getHookHandler(plugin.hotUpdate).call(
           clientContext,
           clientHotUpdateOptions,
         )
+        if (isStale()) return
         if (filteredModules) {
           clientHotUpdateOptions.modules = filteredModules
           // Invalidate the hmrContext to force compat modules to be updated
@@ -552,6 +555,7 @@ export async function handleHMRUpdate(
         const filteredModules = await getHookHandler(
           plugin.handleHotUpdate!,
         ).call(contextForHandleHotUpdate, mixedHmrContext)
+        if (isStale()) return
         if (filteredModules) {
           mixedHmrContext.modules = filteredModules
           clientHotUpdateOptions.modules =
@@ -590,7 +594,8 @@ export async function handleHMRUpdate(
       }
     }
   } catch (error) {
-    hotMap.get(server.environments.client)!.error = error
+    if (isStale()) return
+    hotMap.get(clientEnvironment)!.error = error
   }
 
   for (const environment of environments) {
@@ -604,17 +609,20 @@ export async function handleHMRUpdate(
             context,
             hot.options,
           )
+          if (isStale()) return
           if (filteredModules) {
             hot.options.modules = filteredModules
           }
         }
       }
     } catch (error) {
+      if (isStale()) return
       hot.error = error
     }
   }
 
   async function hmr(environment: DevEnvironment) {
+    if (isStale()) return
     try {
       const { options, error } = hotMap.get(environment)!
       if (error) {
@@ -654,6 +662,8 @@ export async function handleHMRUpdate(
     }
   }
 
+  if (isStale()) return
+
   const hotUpdateEnvironments =
     server.config.server.hotUpdateEnvironments ??
     ((server, hmr) => {