From 30c7fa473c24275cf2c9b2cf4481a6b08d8b3c3a Mon Sep 17 00:00:00 2001
From: Valentin Palkovic <valentin@chromatic.com>
Date: Thu, 23 May 2024 16:32:36 +0200
Subject: [PATCH] Merge pull request #27256 from
 storybookjs/valentin/fix-hmr-issue-in-preview

Vite: Fix HMR issue for Storybook preview files
(cherry picked from commit 6c04751b8214f7ffb1297310472122293960bacb)
---
 .../builder-vite/src/codegen-modern-iframe-script.ts | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/code/builders/builder-vite/src/codegen-modern-iframe-script.ts b/code/builders/builder-vite/src/codegen-modern-iframe-script.ts
index 39109c148df4..2d5da031e73d 100644
--- a/code/builders/builder-vite/src/codegen-modern-iframe-script.ts
+++ b/code/builders/builder-vite/src/codegen-modern-iframe-script.ts
@@ -21,9 +21,13 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
   // and the HMR handler.  We don't use the hot.accept callback params because only the changed
   // modules are provided, the rest are null.  We can just re-import everything again in that case.
   const getPreviewAnnotationsFunction = `
-  const getProjectAnnotations = async () => {
+  const getProjectAnnotations = async (hmrPreviewAnnotationModules = []) => {
     const configs = await Promise.all([${previewAnnotationURLs
-      .map((previewAnnotation) => `import('${previewAnnotation}')`)
+      .map(
+        (previewAnnotation, index) =>
+          // Prefer the updated module from an HMR update, otherwise import the original module
+          `hmrPreviewAnnotationModules.at(${index}) ?? import('${previewAnnotation}')`
+      )
       .join(',\n')}])
     return composeConfigs(configs);
   }`;
@@ -45,10 +49,10 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
       window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
       });
 
-    import.meta.hot.accept(${JSON.stringify(previewAnnotationURLs)}, () => {
+    import.meta.hot.accept(${JSON.stringify(previewAnnotationURLs)}, (previewAnnotationModules) => {
       ${getPreviewAnnotationsFunction}
       // getProjectAnnotations has changed so we need to patch the new one in
-      window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations });
+      window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
     });
   }`.trim();
   };