From 79a55fbe96edadbbefb12788e77b08c947c744ce Mon Sep 17 00:00:00 2001
From: Valentin Palkovic <valentin@chromatic.com>
Date: Thu, 23 May 2024 14:04:22 +0200
Subject: [PATCH] Vite: Fix HMR issue for Storybook preview files

---
 .../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..1d7f44a587f9 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,12 @@ 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) =>
+          `hmrPreviewAnnotationModules.at(${index}) ?? import('${previewAnnotation}')`
+      )
       .join(',\n')}])
     return composeConfigs(configs);
   }`;
@@ -45,10 +48,11 @@ 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) => {
+      console.log({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();
   };