Vite: Improve handling of preview annotations

code/.storybook/preview.tsx

@@ -130,9 +130,6 @@ const ThemedSetRoot = () => {
   return null;
 };
 
-// eslint-disable-next-line no-underscore-dangle
-const preview = (window as any).__STORYBOOK_PREVIEW__ as PreviewWeb<ReactRenderer> | undefined;
-const channel = (window as any).__STORYBOOK_ADDONS_CHANNEL__ as Channel | undefined;
 const loaders = [
   /**
    * This loader adds a DocsContext to the story, which is required for the most Blocks to work. A
@@ -147,6 +144,9 @@ const loaders = [
    * The DocsContext will then be added via the decorator below.
    */
   async ({ parameters: { relativeCsfPaths, attached = true } }) => {
+    // eslint-disable-next-line no-underscore-dangle
+    const preview = (window as any).__STORYBOOK_PREVIEW__ as PreviewWeb<ReactRenderer> | undefined;
+    const channel = (window as any).__STORYBOOK_ADDONS_CHANNEL__ as Channel | undefined;
     // __STORYBOOK_PREVIEW__ and __STORYBOOK_ADDONS_CHANNEL__ is set in the PreviewWeb constructor
     // which isn't loaded in portable stories/vitest
     if (!relativeCsfPaths || !preview || !channel) {

code/builders/builder-vite/src/codegen-modern-iframe-script.test.ts

@@ -0,0 +1,181 @@
+import { describe, expect, it } from 'vitest';
+
+import { generateModernIframeScriptCodeFromPreviews } from './codegen-modern-iframe-script';
+
+const projectRoot = 'projectRoot';
+
+describe('generateModernIframeScriptCodeFromPreviews', () => {
+  it('handle one annotation', async () => {
+    const result = await generateModernIframeScriptCodeFromPreviews({
+      previewAnnotations: ['/user/.storybook/preview'],
+      projectRoot,
+      frameworkName: 'frameworkName',
+      isCsf4: false,
+    });
+    expect(result).toMatchInlineSnapshot(`
+      "import { setup } from 'storybook/internal/preview/runtime';
+
+       import 'virtual:/@storybook/builder-vite/setup-addons.js';
+
+       setup();
+
+       import { composeConfigs, PreviewWeb } from 'storybook/internal/preview-api';
+       import { isPreview } from 'storybook/internal/csf';
+       import { importFn } from 'virtual:/@storybook/builder-vite/storybook-stories.js';
+
+       import * as preview_2408 from "/user/.storybook/preview";
+       const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+         const configs = [
+           hmrPreviewAnnotationModules[0] ?? preview_2408
+         ]
+         return composeConfigs(configs);
+       }
+
+       window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
+
+       window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
+
+       if (import.meta.hot) {
+         import.meta.hot.accept('virtual:/@storybook/builder-vite/storybook-stories.js', (newModule) => {
+           // importFn has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
+         });
+
+         import.meta.hot.accept(["/user/.storybook/preview"], (previewAnnotationModules) => {
+           // getProjectAnnotations has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
+         });
+       };"
+    `);
+  });
+
+  it('handle one annotation CSF4', async () => {
+    const result = await generateModernIframeScriptCodeFromPreviews({
+      previewAnnotations: ['/user/.storybook/preview'],
+      projectRoot,
+      frameworkName: 'frameworkName',
+      isCsf4: true,
+    });
+    expect(result).toMatchInlineSnapshot(`
+      "import { setup } from 'storybook/internal/preview/runtime';
+
+       import 'virtual:/@storybook/builder-vite/setup-addons.js';
+
+       setup();
+
+       import { composeConfigs, PreviewWeb } from 'storybook/internal/preview-api';
+       import { isPreview } from 'storybook/internal/csf';
+       import { importFn } from 'virtual:/@storybook/builder-vite/storybook-stories.js';
+
+       import * as preview_2408 from "/user/.storybook/preview";
+       const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+         const preview = hmrPreviewAnnotationModules[0] ?? preview_2408;
+         return preview.default.composed;
+       }
+
+       window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
+
+       window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
+
+       if (import.meta.hot) {
+         import.meta.hot.accept('virtual:/@storybook/builder-vite/storybook-stories.js', (newModule) => {
+           // importFn has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
+         });
+
+         import.meta.hot.accept(["/user/.storybook/preview"], (previewAnnotationModules) => {
+           // getProjectAnnotations has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
+         });
+       };"
+    `);
+  });
+
+  it('handle multiple annotations', async () => {
+    const result = await generateModernIframeScriptCodeFromPreviews({
+      previewAnnotations: ['/user/previewAnnotations1', '/user/.storybook/preview'],
+      projectRoot,
+      frameworkName: 'frameworkName',
+      isCsf4: false,
+    });
+    expect(result).toMatchInlineSnapshot(`
+      "import { setup } from 'storybook/internal/preview/runtime';
+
+       import 'virtual:/@storybook/builder-vite/setup-addons.js';
+
+       setup();
+
+       import { composeConfigs, PreviewWeb } from 'storybook/internal/preview-api';
+       import { isPreview } from 'storybook/internal/csf';
+       import { importFn } from 'virtual:/@storybook/builder-vite/storybook-stories.js';
+
+       import * as previewAnnotations1_2526 from "/user/previewAnnotations1";
+       import * as preview_2408 from "/user/.storybook/preview";
+       const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+         const configs = [
+           hmrPreviewAnnotationModules[0] ?? previewAnnotations1_2526,
+           hmrPreviewAnnotationModules[1] ?? preview_2408
+         ]
+         return composeConfigs(configs);
+       }
+
+       window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
+
+       window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
+
+       if (import.meta.hot) {
+         import.meta.hot.accept('virtual:/@storybook/builder-vite/storybook-stories.js', (newModule) => {
+           // importFn has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
+         });
+
+         import.meta.hot.accept(["/user/previewAnnotations1","/user/.storybook/preview"], (previewAnnotationModules) => {
+           // getProjectAnnotations has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
+         });
+       };"
+    `);
+  });
+
+  it('handle multiple annotations CSF4', async () => {
+    const result = await generateModernIframeScriptCodeFromPreviews({
+      previewAnnotations: ['/user/previewAnnotations1', '/user/.storybook/preview'],
+      projectRoot,
+      frameworkName: 'frameworkName',
+      isCsf4: true,
+    });
+    expect(result).toMatchInlineSnapshot(`
+      "import { setup } from 'storybook/internal/preview/runtime';
+
+       import 'virtual:/@storybook/builder-vite/setup-addons.js';
+
+       setup();
+
+       import { composeConfigs, PreviewWeb } from 'storybook/internal/preview-api';
+       import { isPreview } from 'storybook/internal/csf';
+       import { importFn } from 'virtual:/@storybook/builder-vite/storybook-stories.js';
+
+       import * as preview_2408 from "/user/.storybook/preview";
+       const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+         const preview = hmrPreviewAnnotationModules[0] ?? preview_2408;
+         return preview.default.composed;
+       }
+
+       window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
+
+       window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
+
+       if (import.meta.hot) {
+         import.meta.hot.accept('virtual:/@storybook/builder-vite/storybook-stories.js', (newModule) => {
+           // importFn has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
+         });
+
+         import.meta.hot.accept(["/user/.storybook/preview"], (previewAnnotationModules) => {
+           // getProjectAnnotations has changed so we need to patch the new one in
+           window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
+         });
+       };"
+    `);
+  });
+});

code/builders/builder-vite/src/codegen-modern-iframe-script.ts

@@ -1,6 +1,9 @@
 import { getFrameworkName, loadPreviewOrConfigFile } from 'storybook/internal/common';
+import { isCsfFactoryPreview, readConfig } from 'storybook/internal/csf-tools';
 import type { Options, PreviewAnnotation } from 'storybook/internal/types';
 
+import { genArrayFromRaw, genImport, genSafeVariableName } from 'knitwork';
+import { filename } from 'pathe/utils';
 import { dedent } from 'ts-dedent';
 
 import { processPreviewAnnotation } from './utils/process-preview-annotation';
@@ -11,38 +14,72 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
   const frameworkName = await getFrameworkName(options);
 
   const previewOrConfigFile = loadPreviewOrConfigFile({ configDir });
+  const previewConfig = await readConfig(previewOrConfigFile!);
+  const isCsf4 = isCsfFactoryPreview(previewConfig);
+
   const previewAnnotations = await presets.apply<PreviewAnnotation[]>(
     'previewAnnotations',
     [],
     options
   );
-  const [previewFileUrl, ...previewAnnotationURLs] = [previewOrConfigFile, ...previewAnnotations]
-    .filter(Boolean)
+  return generateModernIframeScriptCodeFromPreviews({
+    previewAnnotations: [...previewAnnotations, previewOrConfigFile],
+    projectRoot,
+    frameworkName,
+    isCsf4,
+  });
+}
+
+export async function generateModernIframeScriptCodeFromPreviews(options: {
+  previewAnnotations: (PreviewAnnotation | undefined)[];
+  projectRoot: string;
+  frameworkName: string;
+  isCsf4: boolean;
+}) {
+  const { projectRoot, frameworkName } = options;
+  const previewAnnotationURLs = options.previewAnnotations
+    .filter((path) => path !== undefined)
     .map((path) => processPreviewAnnotation(path, projectRoot));
 
+  const variables: string[] = [];
+  const imports: string[] = [];
+  for (const previewAnnotation of previewAnnotationURLs) {
+    const variable =
+      genSafeVariableName(filename(previewAnnotation)).replace(/_(45|46|47)/g, '_') +
+      '_' +
+      hash(previewAnnotation);
+    variables.push(variable);
+    imports.push(genImport(previewAnnotation, { name: '*', as: variable }));
+  }
+
+  const previewFileURL = previewAnnotationURLs[previewAnnotationURLs.length - 1];
+  const previewFileVariable = variables[variables.length - 1];
+  const previewFileImport = imports[imports.length - 1];
+
   // This is pulled out to a variable because it is reused in both the initial page load
-  // 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 = dedent`
-  const getProjectAnnotations = async (hmrPreviewAnnotationModules = []) => {
-    const preview = await import('${previewFileUrl}');
-
-    if (isPreview(preview.default)) {
-      return preview.default.composed;
-    }
-
-    const configs = await Promise.all([${previewAnnotationURLs
-      .map(
+  // and the HMR handler.
+  // The `hmrPreviewAnnotationModules` parameter is used to pass the updated modules from HMR.
+  // However, only the changed modules are provided, the rest are null.
+  const getPreviewAnnotationsFunction = options.isCsf4
+    ? dedent`
+  const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+    const preview = hmrPreviewAnnotationModules[0] ?? ${previewFileVariable};
+    return preview.default.composed;
+  }`
+    : dedent`
+  const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+    const configs = ${genArrayFromRaw(
+      variables.map(
         (previewAnnotation, index) =>
-          // Prefer the updated module from an HMR update, otherwise import the original module
-          `hmrPreviewAnnotationModules[${index}] ?? import('${previewAnnotation}')`
-      )
-      .join(',\n')}])
-    return composeConfigs([...configs, preview]);
+          // Prefer the updated module from an HMR update, otherwise the original module
+          `hmrPreviewAnnotationModules[${index}] ?? ${previewAnnotation}`
+      ),
+      '  '
+    )}
+    return composeConfigs(configs);
   }`;
 
-  // eslint-disable-next-line @typescript-eslint/no-shadow
-  const generateHMRHandler = (frameworkName: string): string => {
+  const generateHMRHandler = (): string => {
     // Web components are not compatible with HMR, so disable HMR, reload page instead.
     if (frameworkName === '@storybook/web-components-vite') {
       return dedent`
@@ -58,8 +95,7 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
         window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
       });
 
-      import.meta.hot.accept(${JSON.stringify(previewAnnotationURLs)}, (previewAnnotationModules) => {
-        ${getPreviewAnnotationsFunction}
+      import.meta.hot.accept(${JSON.stringify(options.isCsf4 ? [previewFileURL] : previewAnnotationURLs)}, (previewAnnotationModules) => {
         // getProjectAnnotations has changed so we need to patch the new one in
         window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
       });
@@ -76,6 +112,7 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
    */
   const code = dedent`
   import { setup } from 'storybook/internal/preview/runtime';
+
   import '${SB_VIRTUAL_FILES.VIRTUAL_ADDON_SETUP_FILE}';
 
   setup();
@@ -84,12 +121,17 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
   import { isPreview } from 'storybook/internal/csf';
   import { importFn } from '${SB_VIRTUAL_FILES.VIRTUAL_STORIES_FILE}';
 
+  ${options.isCsf4 ? previewFileImport : imports.join('\n')}
   ${getPreviewAnnotationsFunction}
 
   window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
 
   window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
-
-  ${generateHMRHandler(frameworkName)};`.trim();
+
+  ${generateHMRHandler()};
+  `.trim();
   return code;
 }
+function hash(value: string) {
+  return value.split('').reduce((acc, char) => acc + char.charCodeAt(0), 0);
+}

code/builders/builder-vite/src/optimizeDeps.ts

@@ -125,6 +125,7 @@ const INCLUDE_CANDIDATES = [
   'qs',
   'react-dom',
   'react-dom/client',
+  'react-dom/test-utils',
   'react-fast-compare',
   'react-is',
   'react-textarea-autosize',
@@ -144,6 +145,7 @@ const INCLUDE_CANDIDATES = [
   'refractor/lang/typescript.js',
   'refractor/lang/yaml.js',
   'regenerator-runtime/runtime.js',
+  'semver', // TODO: Remove once https://github.com/npm/node-semver/issues/712 is fixed
   'sb-original/default-loader',
   'sb-original/image-context',
   'slash',