fix(core): preserve lazy compilation HMR state

e2e/cases/lazy-compilation/add-initial-chunk/rsbuild.config.ts

@@ -3,9 +3,14 @@ import { pluginReact } from '@rsbuild/plugin-react';
 
 export default defineConfig({
   plugins: [pluginReact()],
+  dev: {
+    lazyCompilation: {
+      entries: true,
+      imports: true,
+    },
+  },
   tools: {
     rspack: {
-      lazyCompilation: true,
       output: {
         asyncChunks: false,
       },

e2e/cases/lazy-compilation/dynamic-import/index.test.ts

@@ -4,6 +4,10 @@ const BUILD_FOO = 'building src/foo.js';
 
 test('should lazy compile dynamic imported modules', async ({ page, devOnly }) => {
   const rsbuild = await devOnly();
+  await page.addInitScript(() => {
+    const key = 'lazy-compilation-document-loads';
+    sessionStorage.setItem(key, String(Number(sessionStorage.getItem(key) ?? 0) + 1));
+  });
 
   // initial build
   await rsbuild.expectBuildEnd();
@@ -16,4 +20,7 @@ test('should lazy compile dynamic imported modules', async ({ page, devOnly }) =
   await rsbuild.expectBuildEnd();
   const value = await page.evaluate(() => window.foo);
   expect(value).toBe(42);
+  expect(
+    await page.evaluate(() => Number(sessionStorage.getItem('lazy-compilation-document-loads'))),
+  ).toBe(1);
 });

e2e/cases/lazy-compilation/hmr-lifecycle/index.test.ts

@@ -0,0 +1,145 @@
+import { join } from 'node:path';
+import { expect, expectPoll, gotoPage, test } from '@e2e/helper';
+
+const BUILD_PAGE1 = 'building test-temp-src/page1/index.jsx';
+
+declare global {
+  interface Window {
+    __rsbuildHmrWebSockets: {
+      generation: number;
+      sockets: WebSocket[];
+    };
+  }
+}
+
+const createLazyEntryConfig = (tempSrc: string) => ({
+  dev: { lazyCompilation: true },
+  source: {
+    entry: {
+      page1: join(tempSrc, 'page1/index.jsx'),
+      page2: join(tempSrc, 'page2/index.js'),
+    },
+  },
+});
+
+test('preserves source and CSS HMR state after activating a lazy entry', async ({
+  page,
+  devOnly,
+  editFile,
+  copySrcDir,
+}) => {
+  const tempSrc = await copySrcDir();
+  const rsbuild = await devOnly({ config: createLazyEntryConfig(tempSrc) });
+
+  rsbuild.clearLogs();
+  await gotoPage(page, rsbuild, 'page1');
+  await rsbuild.expectLog(BUILD_PAGE1, { posix: true });
+  await rsbuild.expectBuildEnd();
+  await expect(page.locator('#test')).toHaveText('Lazy source');
+
+  const documentId = await page.evaluate<string>('window.__lazyHmrDocumentId');
+  await page.locator('#increment').click();
+  await expect(page.locator('#count')).toHaveText('1');
+
+  await editFile(join(tempSrc, 'page1/App.jsx'), (code) =>
+    code.replace('Lazy source', 'Updated source'),
+  );
+  await expect(page.locator('#test')).toHaveText('Updated source');
+  await expect(page.locator('#count')).toHaveText('1');
+  expect(await page.evaluate<string>('window.__lazyHmrDocumentId')).toBe(documentId);
+
+  await editFile(join(tempSrc, 'page1/App.css'), () => '#test { color: rgb(0, 0, 255); }');
+  await expect(page.locator('#test')).toHaveCSS('color', 'rgb(0, 0, 255)');
+  await expect(page.locator('#count')).toHaveText('1');
+  expect(await page.evaluate<string>('window.__lazyHmrDocumentId')).toBe(documentId);
+});
+
+test('reconnects a loaded lazy entry without reloading when the hash matches', async ({
+  page,
+  devOnly,
+  editFile,
+  copySrcDir,
+}) => {
+  const tempSrc = await copySrcDir();
+  await page.addInitScript(() => {
+    const NativeWebSocket = window.WebSocket;
+    const state = { generation: 0, sockets: [] as WebSocket[] };
+
+    window.__rsbuildHmrWebSockets = state;
+    window.WebSocket = new Proxy(NativeWebSocket, {
+      construct(target, args) {
+        const socket = Reflect.construct(target, args) as WebSocket;
+        state.generation += 1;
+        state.sockets.push(socket);
+        return socket;
+      },
+    });
+  });
+
+  const rsbuild = await devOnly({ config: createLazyEntryConfig(tempSrc) });
+  rsbuild.clearLogs();
+  await gotoPage(page, rsbuild, 'page1');
+  await rsbuild.expectLog(BUILD_PAGE1, { posix: true });
+  await rsbuild.expectBuildEnd();
+  await expect(page.locator('#test')).toHaveText('Lazy source');
+  const documentId = await page.evaluate<string>('window.__lazyHmrDocumentId');
+
+  const initialGeneration = await page.evaluate(() => {
+    window.__rsbuildHmrWebSockets.sockets.at(-1)?.close();
+    return window.__rsbuildHmrWebSockets.generation;
+  });
+
+  await expectPoll(() =>
+    page.evaluate(() => window.__rsbuildHmrWebSockets.generation),
+  ).toBeGreaterThan(initialGeneration);
+  await expectPoll(() =>
+    page.evaluate(
+      () => window.__rsbuildHmrWebSockets.sockets.at(-1)?.readyState === window.WebSocket.OPEN,
+    ),
+  ).toBe(true);
+
+  expect(await page.evaluate<string>('window.__lazyHmrDocumentId')).toBe(documentId);
+
+  await editFile(join(tempSrc, 'page1/App.jsx'), (code) =>
+    code.replace('Lazy source', 'Updated source'),
+  );
+  await expect(page.locator('#test')).toHaveText('Updated source');
+  expect(await page.evaluate<string>('window.__lazyHmrDocumentId')).toBe(documentId);
+
+  await rsbuild.close();
+});
+
+test('reloads after reconnect when the loaded lazy entry hash is stale', async ({
+  page,
+  devOnly,
+  editFile,
+  copySrcDir,
+}) => {
+  const tempSrc = await copySrcDir();
+  const config = createLazyEntryConfig(tempSrc);
+
+  const firstServer = await devOnly({ config });
+  firstServer.clearLogs();
+  await gotoPage(page, firstServer, 'page1');
+  await firstServer.expectLog(BUILD_PAGE1, { posix: true });
+  await firstServer.expectBuildEnd();
+  await expect(page.locator('#test')).toHaveText('Lazy source');
+  const documentId = await page.evaluate<string>('window.__lazyHmrDocumentId');
+
+  await firstServer.close();
+  await editFile(join(tempSrc, 'page1/App.jsx'), (code) =>
+    code.replace('Lazy source', 'Updated while disconnected'),
+  );
+  const secondServer = await devOnly({
+    config: {
+      ...config,
+      server: { port: firstServer.port },
+    },
+  });
+
+  await secondServer.expectLog(BUILD_PAGE1, { posix: true });
+  await expect(page.locator('#test')).toHaveText('Updated while disconnected');
+  expect(await page.evaluate<string>('window.__lazyHmrDocumentId')).not.toBe(documentId);
+
+  await secondServer.close();
+});

e2e/cases/lazy-compilation/hmr-lifecycle/rsbuild.config.ts

@@ -0,0 +1,6 @@
+import { defineConfig } from '@rsbuild/core';
+import { pluginReact } from '@rsbuild/plugin-react';
+
+export default defineConfig({
+  plugins: [pluginReact()],
+});

e2e/cases/lazy-compilation/hmr-lifecycle/src/page1/App.css

@@ -0,0 +1,3 @@
+#test {
+  color: rgb(255, 0, 0);
+}

e2e/cases/lazy-compilation/hmr-lifecycle/src/page1/App.jsx

@@ -0,0 +1,18 @@
+import { useState } from 'react';
+import './App.css';
+
+const App = () => {
+  const [count, setCount] = useState(0);
+
+  return (
+    <>
+      <div id="test">Lazy source</div>
+      <button id="increment" type="button" onClick={() => setCount((value) => value + 1)}>
+        Increment
+      </button>
+      <div id="count">{count}</div>
+    </>
+  );
+};
+
+export default App;

e2e/cases/lazy-compilation/hmr-lifecycle/src/page1/index.jsx

@@ -0,0 +1,10 @@
+import React from 'react';
+import { createRoot } from 'react-dom/client';
+import App from './App';
+
+window.__lazyHmrDocumentId ??= crypto.randomUUID();
+
+const container = document.getElementById('root');
+if (container) {
+  createRoot(container).render(React.createElement(App));
+}

e2e/cases/lazy-compilation/hmr-lifecycle/src/page2/index.js

@@ -0,0 +1 @@
+document.getElementById('root').textContent = 'Page 2';

packages/core/src/client/hmr.ts

@@ -1,12 +1,14 @@
 import type {
   ClientMessage,
   ClientMessageError,
+  HmrUpdateCause,
   ServerMessage,
   ServerMessageErrors,
   ServerMessageFullReload,
   ServerMessageResolvedClientError,
 } from '../server/socketServer';
 import type { LogLevel, NormalizedClientConfig, WebSocketUrlResolver } from '../types';
+import { createHmrHashState, reduceHmrHashState } from './hmrHashState';
 import { logger } from './log';
 
 let createOverlay: undefined | ((title: string, content: string) => void);
@@ -18,6 +20,17 @@ declare const RSPACK_INTERCEPT_MODULE_EXECUTION: ((options: {
   module: { hot: Rspack.Hot };
 }) => void)[];
 
+type LazyCompilationApplyOptions = Rspack.ApplyOptions & {
+  preserveDisposedModuleFactories: boolean;
+};
+
+// TODO(rspack#14772): Use `Rspack.ApplyOptions` directly after the minimum Rspack
+// version exposes `preserveDisposedModuleFactories`.
+const lazyCompilationApplyOptions: LazyCompilationApplyOptions = {
+  ignoreUnaccepted: true,
+  preserveDisposedModuleFactories: true,
+};
+
 const getErrorField = (error: unknown, field: keyof Error): string | undefined => {
   if (error instanceof Error) {
     const value = error[field];
@@ -98,7 +111,7 @@ export function init(
   const customListenersMap: CustomListenersMap = new Map();
 
   // Hash of the last successful build
-  let lastHash: string | undefined;
+  let hashState = createHmrHashState();
   let hasBuildErrors = false;
   const base = serverBase.endsWith('/') ? serverBase : `${serverBase}/`;
 
@@ -164,6 +177,10 @@ export function init(
   function handleErrors({ text, html }: ServerMessageErrors['data']) {
     clearBuildErrors();
     hasBuildErrors = true;
+    hashState = reduceHmrHashState(hashState, {
+      type: 'reset',
+      appliedHash: BUILD_HASH,
+    }).state;
 
     // Also log them to the console.
     for (const error of text) {
@@ -207,17 +224,42 @@ export function init(
     );
   }
 
+  const failClosedReload = () => {
+    hashState = reduceHmrHashState(hashState, {
+      type: 'reset',
+      appliedHash: BUILD_HASH,
+    }).state;
+    fullReload();
+  };
+
+  const rememberHash = (hash: string, cause: HmrUpdateCause) => {
+    const transition = reduceHmrHashState(hashState, {
+      type: 'hash',
+      hash,
+      cause,
+      appliedHash: BUILD_HASH,
+    });
+    hashState = transition.state;
+    if (transition.reload) {
+      failClosedReload();
+    }
+  };
+
+  const advancePendingHashModes = () => {
+    hashState = reduceHmrHashState(hashState, { type: 'applied', hash: BUILD_HASH }).state;
+  };
+
   // BUILD_HASH is replaced with import.meta.rspackHash when the client is prebuilt,
   // then resolved to the current compilation hash.
-  const shouldUpdate = () => lastHash !== BUILD_HASH;
+  const shouldUpdate = () => hashState.lastHash !== BUILD_HASH;
 
   const handleApplyUpdates = (err: unknown, updatedModules: (string | number)[] | null) => {
     const forcedReload = err || !updatedModules;
     if (forcedReload) {
       if (err) {
         logger.error('[rsbuild] HMR update failed, performing full reload:', err);
       }
-      fullReload();
+      failClosedReload();
       return;
     }
 
@@ -227,6 +269,7 @@ export function init(
 
   // Attempt to update code on the fly, fall back to a hard reload.
   function tryApplyUpdates() {
+    advancePendingHashModes();
     // detect is there a newer version of this code available
     if (!shouldUpdate()) {
       return;
@@ -238,9 +281,18 @@ export function init(
         return;
       }
 
-      // https://rspack.rs/api/runtime-api/module-variables#importmetawebpackhot
-      import.meta.webpackHot.check(true).then(
+      const applyOptions =
+        hashState.pending[0]?.cause === 'lazy' ? lazyCompilationApplyOptions : true;
+      let update: Promise<(string | number)[] | null>;
+      try {
+        update = import.meta.webpackHot.check(applyOptions);
+      } catch (err) {
+        handleApplyUpdates(err, null);
+        return;
+      }
+      update.then(
         (updatedModules) => {
+          advancePendingHashModes();
           handleApplyUpdates(null, updatedModules);
         },
         (err: unknown) => {
@@ -252,7 +304,7 @@ export function init(
 
     // HotModuleReplacementPlugin is not registered in Rspack configuration
     // fallback to reload page
-    fullReload();
+    failClosedReload();
   }
 
   let socket: WebSocket | null = null;
@@ -290,8 +342,8 @@ export function init(
 
     switch (message.type) {
       case 'hash':
-        // Update the last compilation hash
-        lastHash = message.data;
+      case 'lazy-compilation-hash':
+        rememberHash(message.data, message.type === 'hash' ? 'normal' : 'lazy');
 
         if (clearOverlay && shouldUpdate()) {
           clearOverlay();
@@ -333,6 +385,7 @@ export function init(
   }
 
   function onClose() {
+    hashState = reduceHmrHashState(hashState, { type: 'disconnect' }).state;
     if (reconnectCount >= config.reconnect) {
       if (config.reconnect > 0) {
         logger.warn('[rsbuild] WebSocket connection failed after maximum retry attempts.');