fix(vite): improve vite-node hmr (#6343)

packages/vite/src/runtime/vite-node.mjs

@@ -20,13 +20,17 @@ export default async (ssrContext) => {
   // Workaround for stub mode
   // https://github.com/nuxt/framework/pull/3983
   process.server = true
+
+  // Invalidate cache for files changed since last rendering
+  const invalidates = await $fetch('/invalidates', {
+    baseURL: viteNodeOptions.baseURL
+  })
+  for (const key of invalidates) {
+    runner.moduleCache.delete(key)
+  }
+
+  // Execute SSR bundle on demand
   render = render || (await runner.executeFile(viteNodeOptions.entryPath)).default
   const result = await render(ssrContext)
-  // reset cache for non-node-modules
-  for (const key of runner.moduleCache.keys()) {
-    if (!key.includes('/node_modules/')) {
-      runner.moduleCache.delete(key)
-    }
-  }
   return result
 }

packages/vite/src/vite-node.ts

@@ -4,7 +4,7 @@ import { ViteNodeServer } from 'vite-node/server'
 import fse from 'fs-extra'
 import { resolve } from 'pathe'
 import { addServerMiddleware } from '@nuxt/kit'
-import type { Plugin as VitePlugin, ViteDevServer } from 'vite'
+import type { ModuleNode, Plugin as VitePlugin, ViteDevServer } from 'vite'
 import { resolve as resolveModule } from 'mlly'
 import { distDir } from './dirs'
 import type { ViteBuildContext } from './vite'
@@ -14,11 +14,28 @@ import { createIsExternal } from './utils/external'
 // TODO: Remove this in favor of registerViteNodeMiddleware
 // after Nitropack or h3 fixed for adding middlewares after setup
 export function viteNodePlugin (ctx: ViteBuildContext): VitePlugin {
+  // Store the invalidates for the next rendering
+  const invalidates = new Set<string>()
   return {
     name: 'nuxt:vite-node-server',
-    enforce: 'pre',
+    enforce: 'post',
     configureServer (server) {
-      server.middlewares.use('/__nuxt_vite_node__', createViteNodeMiddleware(ctx))
+      server.middlewares.use('/__nuxt_vite_node__', createViteNodeMiddleware(ctx, invalidates))
+    },
+    handleHotUpdate ({ file, server }) {
+      function markInvalidate (mod: ModuleNode) {
+        if (invalidates.has(mod.id)) {
+          return
+        }
+        invalidates.add(mod.id)
+        for (const importer of mod.importers) {
+          markInvalidate(importer)
+        }
+      }
+      const mods = server.moduleGraph.getModulesByFile(file) || []
+      for (const mod of mods) {
+        markInvalidate(mod)
+      }
     }
   }
 }
@@ -49,7 +66,7 @@ function getManifest (server: ViteDevServer) {
   }
 }
 
-function createViteNodeMiddleware (ctx: ViteBuildContext) {
+function createViteNodeMiddleware (ctx: ViteBuildContext, invalidates: Set<string> = new Set()) {
   const app = createApp()
 
   app.use('/manifest', defineEventHandler(() => {
@@ -57,6 +74,20 @@ function createViteNodeMiddleware (ctx: ViteBuildContext) {
     return manifest
   }))
 
+  app.use('/invalidates', defineEventHandler(() => {
+    // When a file has been invalidated, we also invalidate the entry module
+    if (invalidates.size) {
+      for (const key of ctx.ssrServer.moduleGraph.fileToModulesMap.keys()) {
+        if (key.startsWith(ctx.nuxt.options.appDir + '/entry')) {
+          invalidates.add(key)
+        }
+      }
+    }
+    const ids = Array.from(invalidates)
+    invalidates.clear()
+    return ids
+  }))
+
   app.use('/module', defineLazyEventHandler(() => {
     const viteServer = ctx.ssrServer
     const node: ViteNodeServer = new ViteNodeServer(viteServer, {