wip: full bundle dev env

Author: sapphi-red

packages/vite/src/client/client.ts

@@ -1,3 +1,4 @@
+/// <reference types="rolldown/experimental/runtime-types" />
 import type { ErrorPayload, HotPayload } from 'types/hmrPayload'
 import type { ViteHotContext } from 'types/hot'
 import { HMRClient, HMRContext } from '../shared/hmr'
@@ -483,67 +484,26 @@ export function injectQuery(url: string, queryToInject: string): string {
 export { ErrorOverlay }
 
 if (isFullBundleMode) {
-  class DevRuntime {
-    modules: Record<string, { exports: any }> = {}
-
-    static getInstance() {
-      // @ts-expect-error __rolldown_runtime__
-      let instance = globalThis.__rolldown_runtime__
-      if (!instance) {
-        instance = new DevRuntime()
-        // @ts-expect-error __rolldown_runtime__
-        globalThis.__rolldown_runtime__ = instance
-      }
-      return instance
-    }
-
-    createModuleHotContext(moduleId: string) {
+  class ViteDevRuntime extends DevRuntime {
+    override createModuleHotContext(moduleId: string) {
       const ctx = createHotContext(moduleId)
       // @ts-expect-error TODO: support CSS
       ctx._internal = {
         updateStyle,
         removeStyle,
       }
+      // @ts-expect-error TODO: support this function (used by plugin-react)
+      ctx.getExports = async () =>
+        // @ts-expect-error __rolldown_runtime__ / ctx.ownerPath
+        __rolldown_runtime__.loadExports(ctx.ownerPath)
       return ctx
     }
 
-    applyUpdates(_boundaries: string[]) {
-      //
-    }
-
-    registerModule(
-      id: string,
-      module: { exports: Record<string, () => unknown> },
-    ) {
-      this.modules[id] = module
+    override applyUpdates(_boundaries: string[]): void {
+      // TODO: how should this be handled?
+      // noop, handled in the HMR client
     }
-
-    loadExports(id: string) {
-      const module = this.modules[id]
-      if (module) {
-        return module.exports
-      } else {
-        console.warn(`Module ${id} not found`)
-        return {}
-      }
-    }
-
-    // __esmMin
-    // @ts-expect-error need to add typing
-    createEsmInitializer = (fn, res) => () => (fn && (res = fn((fn = 0))), res)
-    // __commonJSMin
-    // @ts-expect-error need to add typing
-    createCjsInitializer = (cb, mod) => () => (
-      mod || cb((mod = { exports: {} }).exports, mod), mod.exports
-    )
-    // @ts-expect-error it is exits
-    __toESM = __toESM
-    // @ts-expect-error it is exits
-    __toCommonJS = __toCommonJS
-    // @ts-expect-error it is exits
-    __export = __export
   }
 
-  // @ts-expect-error __rolldown_runtime__
-  globalThis.__rolldown_runtime__ ||= new DevRuntime()
+  ;(globalThis as any).__rolldown_runtime__ ??= new ViteDevRuntime()
 }

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

@@ -236,7 +236,7 @@ export class DevEnvironment extends BaseEnvironment {
     }
   }
 
-  private invalidateModule(m: {
+  protected invalidateModule(m: {
     path: string
     message?: string
     firstInvalidatedBy: string

packages/vite/src/node/server/environments/fullBundleEnvironment.ts

@@ -1,3 +1,4 @@
+import path from 'node:path'
 import type { RolldownBuild, RolldownOptions } from 'rolldown'
 import type { Update } from 'types/hmrPayload'
 import colors from 'picocolors'
@@ -11,7 +12,7 @@ import { getHmrImplementation } from '../../plugins/clientInjections'
 import { DevEnvironment, type DevEnvironmentContext } from '../environment'
 import type { ResolvedConfig } from '../../config'
 import type { ViteDevServer } from '../../server'
-import { arraify, createDebugger } from '../../utils'
+import { arraify, createDebugger, normalizePath } from '../../utils'
 import { prepareError } from '../middlewares/error'
 
 const debug = createDebugger('vite:full-bundle-mode')
@@ -57,7 +58,6 @@ export class FullBundleDevEnvironment extends DevEnvironment {
   async onFileChange(
     _type: 'create' | 'update' | 'delete',
     file: string,
-    server: ViteDevServer,
   ): Promise<void> {
     if (this.state.type === 'initial') {
       return
@@ -67,15 +67,21 @@ export class FullBundleDevEnvironment extends DevEnvironment {
       debug?.(
         `BUNDLING: file update detected ${file}, retriggering bundle generation`,
       )
-      this.state.abortController.abort()
       this.triggerGenerateBundle(this.state)
       return
     }
     if (this.state.type === 'bundle-error') {
-      debug?.(
-        `BUNDLE-ERROR: file update detected ${file}, retriggering bundle generation`,
-      )
-      this.triggerGenerateBundle(this.state)
+      const files = await this.state.bundle.watchFiles
+      if (files.includes(file)) {
+        debug?.(
+          `BUNDLE-ERROR: file update detected ${file}, retriggering bundle generation`,
+        )
+        this.triggerGenerateBundle(this.state)
+      } else {
+        debug?.(
+          `BUNDLE-ERROR: file update detected ${file}, but ignored as it is not a dependency`,
+        )
+      }
       return
     }
 
@@ -95,35 +101,111 @@ export class FullBundleDevEnvironment extends DevEnvironment {
         type: 'generating-hmr-patch',
         options: this.state.options,
         bundle: this.state.bundle,
+        patched: this.state.patched,
       }
 
       let hmrOutput: HmrOutput
       try {
         // NOTE: only single outputOptions is supported here
-        hmrOutput = (await this.state.bundle.generateHmrPatch([file]))!
+        hmrOutput = await this.state.bundle.generateHmrPatch([file])
       } catch (e) {
         // TODO: support multiple errors
-        server.ws.send({ type: 'error', err: prepareError(e.errors[0]) })
+        this.hot.send({ type: 'error', err: prepareError(e.errors[0]) })
 
         this.state = {
           type: 'bundled',
           options: this.state.options,
           bundle: this.state.bundle,
+          patched: this.state.patched,
         }
         return
       }
 
-      debug?.(`handle hmr output for ${file}`, {
-        ...hmrOutput,
-        code: typeof hmrOutput.code === 'string' ? '[code]' : hmrOutput.code,
-      })
-
       this.handleHmrOutput(file, hmrOutput, this.state)
       return
     }
     this.state satisfies never // exhaustive check
   }
 
+  protected override invalidateModule(m: {
+    path: string
+    message?: string
+    firstInvalidatedBy: string
+  }): void {
+    ;(async () => {
+      if (
+        this.state.type === 'initial' ||
+        this.state.type === 'bundling' ||
+        this.state.type === 'bundle-error'
+      ) {
+        debug?.(
+          `${this.state.type.toUpperCase()}: invalidate received, but ignored`,
+        )
+        return
+      }
+      this.state.type satisfies 'bundled' | 'generating-hmr-patch' // exhaustive check
+
+      debug?.(
+        `${this.state.type.toUpperCase()}: invalidate received, re-triggering HMR`,
+      )
+
+      // TODO: should this be a separate state?
+      this.state = {
+        type: 'generating-hmr-patch',
+        options: this.state.options,
+        bundle: this.state.bundle,
+        patched: this.state.patched,
+      }
+
+      let hmrOutput: HmrOutput
+      try {
+        // NOTE: only single outputOptions is supported here
+        hmrOutput = await this.state.bundle.hmrInvalidate(
+          normalizePath(path.join(this.config.root, m.path)),
+          m.firstInvalidatedBy,
+        )
+      } catch (e) {
+        // TODO: support multiple errors
+        this.hot.send({ type: 'error', err: prepareError(e.errors[0]) })
+
+        this.state = {
+          type: 'bundled',
+          options: this.state.options,
+          bundle: this.state.bundle,
+          patched: this.state.patched,
+        }
+        return
+      }
+
+      if (hmrOutput.isSelfAccepting) {
+        this.logger.info(
+          colors.yellow(`hmr invalidate `) +
+            colors.dim(m.path) +
+            (m.message ? ` ${m.message}` : ''),
+          { timestamp: true },
+        )
+      }
+
+      // TODO: need to check if this is enough
+      this.handleHmrOutput(m.path, hmrOutput, this.state)
+    })()
+  }
+
+  triggerBundleRegenerationIfStale(): boolean {
+    if (
+      (this.state.type === 'bundled' ||
+        this.state.type === 'generating-hmr-patch') &&
+      this.state.patched
+    ) {
+      this.triggerGenerateBundle(this.state)
+      debug?.(
+        `${this.state.type.toUpperCase()}: access to stale bundle, triggered bundle re-generation`,
+      )
+      return true
+    }
+    return false
+  }
+
   override async close(): Promise<void> {
     await Promise.all([
       super.close(),
@@ -161,6 +243,10 @@ export class FullBundleDevEnvironment extends DevEnvironment {
     options,
     bundle,
   }: BundleStateCommonProperties) {
+    if (this.state.type === 'bundling') {
+      this.state.abortController.abort()
+    }
+
     const controller = new AbortController()
     const promise = this.generateBundle(
       options.output,
@@ -211,6 +297,7 @@ export class FullBundleDevEnvironment extends DevEnvironment {
         type: 'bundled',
         bundle: this.state.bundle,
         options: this.state.options,
+        patched: false,
       }
       debug?.('BUNDLED: bundle generated')
 
@@ -234,7 +321,7 @@ export class FullBundleDevEnvironment extends DevEnvironment {
     }
   }
 
-  private async handleHmrOutput(
+  private handleHmrOutput(
     file: string,
     hmrOutput: HmrOutput,
     { options, bundle }: BundleStateCommonProperties,
@@ -255,7 +342,16 @@ export class FullBundleDevEnvironment extends DevEnvironment {
       return
     }
 
-    if (hmrOutput.code) {
+    // TODO: handle `No corresponding module found for changed file path`
+    if (
+      hmrOutput.code &&
+      hmrOutput.code !== '__rolldown_runtime__.applyUpdates([]);'
+    ) {
+      debug?.(`handle hmr output for ${file}`, {
+        ...hmrOutput,
+        code: typeof hmrOutput.code === 'string' ? '[code]' : hmrOutput.code,
+      })
+
       this.memoryFiles.set(hmrOutput.filename, hmrOutput.code)
       if (hmrOutput.sourcemapFilename && hmrOutput.sourcemap) {
         this.memoryFiles.set(hmrOutput.sourcemapFilename, hmrOutput.sourcemap)
@@ -279,7 +375,17 @@ export class FullBundleDevEnvironment extends DevEnvironment {
           colors.dim([...new Set(updates.map((u) => u.path))].join(', ')),
         { clear: !hmrOutput.firstInvalidatedBy, timestamp: true },
       )
+
+      this.state = {
+        type: 'bundled',
+        options,
+        bundle,
+        patched: true,
+      }
+      return
     }
+
+    debug?.(`ignored file change for ${file}`)
   }
 }
 
@@ -296,12 +402,26 @@ type BundleStateBundling = {
   promise: Promise<void>
   abortController: AbortController
 } & BundleStateCommonProperties
-type BundleStateBundled = { type: 'bundled' } & BundleStateCommonProperties
+type BundleStateBundled = {
+  type: 'bundled'
+  /**
+   * Whether a hmr patch was generated.
+   *
+   * In other words, whether the bundle is stale.
+   */
+  patched: boolean
+} & BundleStateCommonProperties
 type BundleStateBundleError = {
   type: 'bundle-error'
 } & BundleStateCommonProperties
 type BundleStateGeneratingHmrPatch = {
   type: 'generating-hmr-patch'
+  /**
+   * Whether a hmr patch was generated.
+   *
+   * In other words, whether the bundle is stale.
+   */
+  patched: boolean
 } & BundleStateCommonProperties
 
 type BundleStateCommonProperties = {

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

@@ -427,7 +427,7 @@ export async function handleHMRUpdate(
   if (config.experimental.fullBundleMode) {
     // TODO: support handleHotUpdate / hotUpdate
     const environment = server.environments.client as FullBundleDevEnvironment
-    environment.onFileChange(type, file, server)
+    environment.onFileChange(type, file)
     return
   }
 

packages/vite/src/node/server/middlewares/indexHtml.ts

@@ -442,9 +442,9 @@ export function indexHtmlMiddleware(
   server: ViteDevServer | PreviewServer,
 ): Connect.NextHandleFunction {
   const isDev = isDevServer(server)
-  const memoryFiles =
+  const fullBundleEnv =
     isDev && server.environments.client instanceof FullBundleDevEnvironment
-      ? server.environments.client.memoryFiles
+      ? server.environments.client
       : undefined
 
   // Keep the named function. The name is visible in debug logs via `DEBUG=connect:dispatcher ...`
@@ -456,13 +456,18 @@ export function indexHtmlMiddleware(
     const url = req.url && cleanUrl(req.url)
     // htmlFallbackMiddleware appends '.html' to URLs
     if (url?.endsWith('.html') && req.headers['sec-fetch-dest'] !== 'script') {
-      if (memoryFiles) {
+      if (fullBundleEnv) {
         const cleanedUrl = cleanUrl(url).slice(1) // remove first /
-        let content = memoryFiles.get(cleanedUrl)
-        if (!content && memoryFiles.size !== 0) {
+        let content = fullBundleEnv.memoryFiles.get(cleanedUrl)
+        if (!content && fullBundleEnv.memoryFiles.size !== 0) {
           return next()
         }
-        content ??= await generateFallbackHtml(server as ViteDevServer)
+        if (
+          fullBundleEnv.triggerBundleRegenerationIfStale() ||
+          content === undefined
+        ) {
+          content = await generateFallbackHtml(server as ViteDevServer)
+        }
 
         const html =
           typeof content === 'string' ? content : Buffer.from(content.buffer)