Vite: Fix storysource addon support

code/lib/builder-vite/package.json

@@ -52,6 +52,7 @@
     "@storybook/node-logger": "7.0.0-beta.47",
     "@storybook/preview": "7.0.0-beta.47",
     "@storybook/preview-api": "7.0.0-beta.47",
+    "@storybook/source-loader": "7.0.0-beta.47",
     "@storybook/types": "7.0.0-beta.47",
     "browser-assert": "^1.2.1",
     "es-module-lexer": "^0.9.3",

code/lib/builder-vite/src/plugins/index.ts

@@ -4,3 +4,4 @@ export * from './strip-story-hmr-boundaries';
 export * from './code-generator-plugin';
 export * from './csf-plugin';
 export * from './external-globals-plugin';
+export * from './source-loader-plugin';

code/lib/builder-vite/src/plugins/source-loader-plugin.ts

@@ -0,0 +1,106 @@
+import type { Plugin } from 'vite';
+import sourceLoaderTransform from '@storybook/source-loader';
+import MagicString from 'magic-string';
+import type { Options } from '@storybook/types';
+
+const storyPattern = /\.stories\.[jt]sx?$/;
+const storySourcePattern = /var __STORY__ = "(.*)"/;
+const storySourceReplacement = '--STORY_SOURCE_REPLACEMENT--';
+
+const mockClassLoader = (id: string) => ({
+  // eslint-disable-next-line no-console
+  emitWarning: (message: string) => console.warn(message),
+  resourcePath: id,
+  getOptions: () => ({ injectStoryParameters: true }),
+  extension: `.${id.split('.').pop()}`,
+});
+
+// HACK: Until we can support only node 15+ and use string.prototype.replaceAll
+const replaceAll = (str: string, search: string, replacement: string) => {
+  return str.split(search).join(replacement);
+};
+
+export function sourceLoaderPlugin(options: Options): Plugin | Plugin[] {
+  if (options.configType === 'DEVELOPMENT') {
+    return {
+      name: 'storybook:source-loader-plugin',
+      enforce: 'pre',
+      async transform(src: string, id: string) {
+        if (id.match(storyPattern)) {
+          const code: string = await sourceLoaderTransform.call(mockClassLoader(id), src);
+          const s = new MagicString(src);
+          // Entirely replace with new code
+          s.overwrite(0, src.length, code);
+
+          return {
+            code: s.toString(),
+            map: s.generateMap({ hires: true, source: id }),
+          };
+        }
+        return undefined;
+      },
+    };
+  }
+
+  // In production, we need to be fancier, to avoid vite:define plugin from replacing values inside the `__STORY__` string
+  const storySources = new WeakMap<Options, Map<string, string>>();
+
+  return [
+    {
+      name: 'storybook-vite-source-loader-plugin',
+      enforce: 'pre',
+      buildStart() {
+        storySources.set(options, new Map());
+      },
+      async transform(src: string, id: string) {
+        if (id.match(storyPattern)) {
+          let code: string = await sourceLoaderTransform.call(mockClassLoader(id), src);
+          // eslint-disable-next-line @typescript-eslint/naming-convention
+          const [_, sourceString] = code.match(storySourcePattern) ?? [null, null];
+          if (sourceString) {
+            const map = storySources.get(options);
+            map?.set(id, sourceString);
+
+            // Remove story source so that it is not processed by vite:define plugin
+            code = replaceAll(code, sourceString, storySourceReplacement);
+          }
+
+          const s = new MagicString(src);
+          // Entirely replace with new code
+          s.overwrite(0, src.length, code);
+
+          return {
+            code: s.toString(),
+            map: s.generateMap(),
+          };
+        }
+        return undefined;
+      },
+    },
+    {
+      name: 'storybook-vite-source-loader-plugin-post',
+      enforce: 'post',
+      buildStart() {
+        storySources.set(options, new Map());
+      },
+      async transform(src: string, id: string) {
+        if (id.match(storyPattern)) {
+          const s = new MagicString(src);
+          const map = storySources.get(options);
+          const storySourceStatement = map?.get(id);
+          // Put the previously-extracted source back in
+          if (storySourceStatement) {
+            const newCode = replaceAll(src, storySourceReplacement, storySourceStatement);
+            s.overwrite(0, src.length, newCode);
+          }
+
+          return {
+            code: s.toString(),
+            map: s.generateMap(),
+          };
+        }
+        return undefined;
+      },
+    },
+  ];
+}

code/lib/builder-vite/src/vite-config.ts

@@ -17,6 +17,7 @@ import {
   mdxPlugin,
   stripStoryHMRBoundary,
   externalGlobalsPlugin,
+  sourceLoaderPlugin,
 } from './plugins';
 
 import type { BuilderOptions } from './types';
@@ -77,6 +78,7 @@ export async function pluginConfig(options: Options) {
 
   const plugins = [
     codeGeneratorPlugin(options),
+    sourceLoaderPlugin(options),
     await csfPlugin(options),
     await mdxPlugin(options),
     injectExportOrderPlugin,

code/yarn.lock

@@ -5195,6 +5195,7 @@ __metadata:
     "@storybook/node-logger": 7.0.0-beta.47
     "@storybook/preview": 7.0.0-beta.47
     "@storybook/preview-api": 7.0.0-beta.47
+    "@storybook/source-loader": 7.0.0-beta.47
     "@storybook/types": 7.0.0-beta.47
     "@types/express": ^4.17.13
     "@types/node": ^16.0.0

scripts/tasks/sandbox-parts.ts

@@ -46,10 +46,7 @@ export const essentialsAddons = [
   'viewport',
 ];
 
-export const create: Task['run'] = async (
-  { key, template, sandboxDir },
-  { addon: addons, dryRun, debug, skipTemplateStories }
-) => {
+export const create: Task['run'] = async ({ key, template, sandboxDir }, { dryRun, debug }) => {
   const parentDir = resolve(sandboxDir, '..');
   await ensureDir(parentDir);
 
@@ -68,17 +65,12 @@ export const create: Task['run'] = async (
       debug,
     });
   }
-
-  const cwd = sandboxDir;
-  if (!skipTemplateStories) {
-    for (const addon of addons) {
-      const addonName = `@storybook/addon-${addon}`;
-      await executeCLIStep(steps.add, { argument: addonName, cwd, dryRun, debug });
-    }
-  }
 };
 
-export const install: Task['run'] = async ({ sandboxDir, template }, { link, dryRun, debug }) => {
+export const install: Task['run'] = async (
+  { sandboxDir, template },
+  { link, dryRun, debug, addon: addons, skipTemplateStories }
+) => {
   const cwd = sandboxDir;
   await installYarn2({ cwd, dryRun, debug });
 
@@ -131,6 +123,13 @@ export const install: Task['run'] = async ({ sandboxDir, template }, { link, dry
       break;
     default:
   }
+
+  if (!skipTemplateStories) {
+    for (const addon of addons) {
+      const addonName = `@storybook/addon-${addon}`;
+      await executeCLIStep(steps.add, { argument: addonName, cwd, dryRun, debug });
+    }
+  }
 };
 
 // Ensure that sandboxes can refer to story files defined in `code/`.