Build-Mode: Module automocking manifest generation

code/.storybook/preview.tsx

@@ -34,7 +34,8 @@ import * as templatePreview from '../core/template/stories/preview';
 import '../renderers/react/template/components/index';
 import { isChromatic } from './isChromatic';
 
-sb.mock('../core/src/test/stories/ModuleMocking.utils');
+sb.mock('../core/src/test/stories/ModuleMocking.utils', { spy: true });
+sb.mock('../core/src/test/stories/ModuleAutoMocking.utils');
 
 const { document } = global;
 globalThis.CONFIG_TYPE = 'DEVELOPMENT';

code/core/src/core-server/presets/common-preset.ts

@@ -34,6 +34,7 @@ import { initializeSaveStory } from '../utils/save-story/save-story';
 import { parseStaticDir } from '../utils/server-statics';
 import { type OptionsWithRequiredCache, initializeWhatsNew } from '../utils/whats-new';
 import { viteInjectMockerRuntime } from './vitePlugins/vite-inject-mocker/plugin';
+import { viteMockBuildManifestPlugin } from './vitePlugins/vite-mock-build-manifest/plugin';
 
 const interpolate = (string: string, data: Record<string, string> = {}) =>
   Object.entries(data).reduce((acc, [k, v]) => acc.replace(new RegExp(`%${k}%`, 'g'), v), string);
@@ -340,6 +341,7 @@ export const viteFinal = async (
           utilsObjectNames: ['sb'],
         },
       }),
+      ...(previewConfigPath ? [viteMockBuildManifestPlugin({ previewConfigPath })] : []),
     ],
   };
 };

code/core/src/core-server/presets/vitePlugins/vite-inject-mocker/ModuleMockerBuildInterceptor.ts

@@ -0,0 +1,24 @@
+import { MockerRegistry } from '@vitest/mocker';
+import type { MockedModule } from '@vitest/mocker';
+
+// This interceptor will be used in the production build.
+// It relies on a pre-generated manifest.
+export class ModuleMockerBuildInterceptor {
+  private manifestPromise: Promise<any>;
+  private mocks = new MockerRegistry();
+
+  constructor(options) {
+    // Fetch the manifest we created during the build
+    this.manifestPromise = fetch('/mock-manifest.json').then((res) => res.json());
+  }
+
+  async register(module: MockedModule): Promise<void> {
+    this.mocks.add(module);
+
+    // Here, we can override the default import behavior
+    // This is the most complex part. We might need to integrate with MSW
+    // to redirect requests for mocked modules to our pre-built automock assets.
+  }
+
+  // ... other methods
+}

code/core/src/core-server/presets/vitePlugins/vite-inject-mocker/plugin.ts

@@ -2,7 +2,7 @@ import { readFile } from 'node:fs/promises';
 
 import { exactRegex } from '@rolldown/pluginutils';
 import { dedent } from 'ts-dedent';
-import type { ViteDevServer } from 'vite';
+import type { ResolvedConfig, ViteDevServer } from 'vite';
 
 const entryPath = '/vite-inject-mocker-entry.js';
 
@@ -12,88 +12,94 @@ const entryCode = dedent`
 
 const __STORYBOOK_GLOBAL_THIS_ACCESSOR__ = '__vitest_mocker__';
 
-const runtimeCode = () => dedent`
-  import { ModuleMockerServerInterceptor } from "@vitest/mocker/browser";
-  import { registerModuleMocker } from "@vitest/mocker/register";
-  globalThis.__STORYBOOK_MOCKER__ = registerModuleMocker((globalThisAccessor) => new ModuleMockerServerInterceptor(globalThisAccessor));
+const runtimeCode = (command: ResolvedConfig['command']) => {
+  if (command === 'serve') {
+    return dedent`
+      import { ModuleMockerServerInterceptor } from "@vitest/mocker/browser";
+      import { registerModuleMocker } from "@vitest/mocker/register";
+      globalThis.__STORYBOOK_MOCKER__ = registerModuleMocker((globalThisAccessor) => new ModuleMockerServerInterceptor(globalThisAccessor));
 
-  if (import.meta.hot) {
-    import.meta.hot.on('invalidate-mocker', (payload) => {
-      globalThis.${__STORYBOOK_GLOBAL_THIS_ACCESSOR__}.invalidate();
-    });
+      if (import.meta.hot) {
+        import.meta.hot.on('invalidate-mocker', (payload) => {
+          globalThis.${__STORYBOOK_GLOBAL_THIS_ACCESSOR__}.invalidate();
+        });
+      }
+    `;
+  } else {
+    return dedent`
+      // For 'build', we'll use a custom interceptor that works with a pre-built manifest
+      const { ModuleMockerBuildInterceptor } = await import('./ModuleMockerBuildInterceptor'); // We will create this
+      const { registerModuleMocker } = await import('@vitest/mocker/register');
+      globalThis.__STORYBOOK_MOCKER__ = registerModuleMocker(
+        (accessor) => new ModuleMockerBuildInterceptor({ globalThisAccessor: accessor })
+    `;
   }
-`;
+};
 
 let server: ViteDevServer;
 
 export const viteInjectMockerRuntime = (options: {
   previewConfigPath?: string | null;
-}): import('vite').Plugin => ({
-  name: 'vite:inject-mocker-runtime',
-  config() {
-    return {
-      resolve: {
-        // external: ['msw/browser', 'msw/core/http'],
+}): import('vite').Plugin => {
+  let viteConfig: ResolvedConfig;
+
+  return {
+    name: 'vite:inject-mocker-runtime',
+    configResolved(config) {
+      viteConfig = config;
+    },
+    configureServer(server_) {
+      server = server_;
+      if (options.previewConfigPath) {
+        server.watcher.on('change', (file) => {
+          if (file === options.previewConfigPath) {
+            server.ws.send({
+              type: 'custom',
+              event: 'invalidate-mocker',
+            });
+          }
+        });
+      }
+    },
+    resolveId: {
+      filter: {
+        id: [exactRegex(entryPath)],
       },
-    };
-  },
-  configureServer(server_) {
-    server = server_;
-    if (options.previewConfigPath) {
-      server.watcher.on('change', (file) => {
-        if (file === options.previewConfigPath) {
-          server.ws.send({
-            type: 'custom',
-            event: 'invalidate-mocker',
-          });
+      handler(id) {
+        if (exactRegex(id).test(entryPath)) {
+          return id;
         }
-      });
-    }
-  },
-  resolveId: {
-    filter: {
-      id: [exactRegex(entryPath)],
+        return null;
+      },
     },
-    handler(id) {
+    async load(id) {
       if (exactRegex(id).test(entryPath)) {
-        return id;
+        return runtimeCode(viteConfig.command);
+      }
+
+      if (id.includes('@vitest/mocker/dist/register.js')) {
+        const content = await readFile(require.resolve('@vitest/mocker/dist/register.js'), 'utf-8');
+        const result = content
+          .replace(
+            /__VITEST_GLOBAL_THIS_ACCESSOR__/g,
+            JSON.stringify(__STORYBOOK_GLOBAL_THIS_ACCESSOR__)
+          )
+          .replace('__VITEST_MOCKER_ROOT__', JSON.stringify(server.config.root));
+        return result;
       }
       return null;
     },
-  },
-  async load(id) {
-    if (exactRegex(id).test(entryPath)) {
-      return runtimeCode();
-    }
+    transformIndexHtml(html: string) {
+      const headTag = html.match(/<head[^>]*>/);
 
-    if (id.includes('@vitest/mocker/dist/register.js')) {
-      console.log(id);
-      if (!server) {
-        // mocker doesn't work during build
-        return 'export {}';
+      if (headTag) {
+        const headTagIndex = html.indexOf(headTag[0]);
+        const newHtml =
+          html.slice(0, headTagIndex + headTag[0].length) +
+          entryCode +
+          html.slice(headTagIndex + headTag[0].length);
+        return newHtml;
       }
-
-      const content = await readFile(require.resolve('@vitest/mocker/dist/register.js'), 'utf-8');
-      const result = content
-        .replace(
-          /__VITEST_GLOBAL_THIS_ACCESSOR__/g,
-          JSON.stringify(__STORYBOOK_GLOBAL_THIS_ACCESSOR__)
-        )
-        .replace('__VITEST_MOCKER_ROOT__', JSON.stringify(server.config.root));
-      return result;
-    }
-    return null;
-  },
-  transformIndexHtml(html: string) {
-    const headTag = html.match(/<head[^>]*>/);
-
-    if (headTag) {
-      const headTagIndex = html.indexOf(headTag[0]);
-      const newHtml =
-        html.slice(0, headTagIndex + headTag[0].length) +
-        entryCode +
-        html.slice(headTagIndex + headTag[0].length);
-      return newHtml;
-    }
-  },
-});
+    },
+  };
+};

code/core/src/core-server/presets/vitePlugins/vite-mock-build-manifest/plugin.ts

@@ -0,0 +1,169 @@
+import crypto from 'node:crypto';
+import { readFile } from 'node:fs/promises';
+
+import { automockModule, findMockRedirect } from '@vitest/mocker/node';
+import { walk } from 'estree-walker';
+import type { Plugin, ResolvedConfig } from 'vite';
+
+interface MockBuildManifestPluginOptions {
+  /** The absolute path to the preview.tsx file where mocks are defined. */
+  previewConfigPath: string;
+}
+
+const VIRTUAL_AUTOMOCK_PREFIX = 'virtual:automock:';
+
+/**
+ * A Vite plugin that runs only during build. It scans the preview.tsx file for `sb.mock()` calls,
+ * processes them, and creates a manifest for runtime use.
+ */
+export function viteMockBuildManifestPlugin(options: MockBuildManifestPluginOptions): Plugin {
+  let viteConfig: ResolvedConfig;
+
+  // Temporary maps to hold chunk references during the build.
+  const redirectChunkRefs = new Map<string, string>();
+  const automockChunkRefs = new Map<string, string>();
+
+  return {
+    name: 'storybook:mock-build-manifest',
+    apply: 'build',
+
+    buildStart() {
+      redirectChunkRefs.clear();
+      automockChunkRefs.clear();
+    },
+
+    configResolved(config) {
+      viteConfig = config;
+    },
+
+    async transform(code, id) {
+      // We only care about the specific `preview.tsx` file where mocks are defined.
+      if (id !== options.previewConfigPath) {
+        return null;
+      }
+
+      const ast = this.parse(code);
+      const mockCalls: Array<{ path: string; isFactoryMock: boolean }> = [];
+
+      // Use an AST walker to find all `sb.mock()` expressions.
+      walk(ast as any, {
+        enter(node) {
+          if (
+            node.type === 'CallExpression' &&
+            node.callee.type === 'MemberExpression' &&
+            node.callee.object.type === 'Identifier' &&
+            node.callee.object.name === 'sb' &&
+            node.callee.property.type === 'Identifier' &&
+            node.callee.property.name === 'mock'
+          ) {
+            if (node.arguments.length > 0 && node.arguments[0].type === 'Literal') {
+              mockCalls.push({
+                path: node.arguments[0].value as string,
+                isFactoryMock: node.arguments.length > 1,
+              });
+            }
+          }
+        },
+      });
+
+      if (mockCalls.length === 0) {
+        return null;
+      }
+
+      for (const call of mockCalls) {
+        // Factory mocks (e.g., `sb.mock('path', () => ({ ... }))`) are handled at runtime
+        // because the factory function needs to exist in the browser's memory with its scope.
+        // We only need to process path-only mocks here.
+        if (call.isFactoryMock) {
+          continue;
+        }
+
+        const resolved = await this.resolve(call.path, id);
+        if (!resolved) {
+          this.warn(`[vitest-mocker] Could not resolve mock path "${call.path}" in ${id}`);
+          continue;
+        }
+
+        const resolvedId = resolved.id;
+
+        // Check for a corresponding file in a `__mocks__` directory.
+        const redirectPath = findMockRedirect(viteConfig.root, resolvedId, null);
+
+        if (redirectPath) {
+          const chunkRef = this.emitFile({
+            type: 'chunk',
+            id: redirectPath,
+            importer: options.previewConfigPath,
+          });
+
+          // The manifest maps the original module ID to the URL of the newly created asset.
+          redirectChunkRefs.set(resolvedId, chunkRef);
+        } else {
+          // If there's no redirect, it's an automock. We'll generate the mocked version.
+          const virtualId = `${VIRTUAL_AUTOMOCK_PREFIX}${resolvedId}`;
+
+          const chunkRef = this.emitFile({
+            type: 'chunk',
+            id: virtualId, // The entry point is our virtual module ID.
+            importer: options.previewConfigPath,
+          });
+
+          // The manifest maps the original module ID to the URL of the automocked asset.
+          automockChunkRefs.set(resolvedId, chunkRef);
+        }
+      }
+
+      // We return null because we don't need to change the preview.tsx file itself.
+      // All our work is done by calling `this.emitFile`.
+      return null;
+    },
+
+    // The resolveId and load hooks are used to provide the source for our virtual automock modules.
+    resolveId(id) {
+      if (id.startsWith(VIRTUAL_AUTOMOCK_PREFIX)) {
+        return id;
+      }
+      return null;
+    },
+
+    async load(id) {
+      if (id.startsWith(VIRTUAL_AUTOMOCK_PREFIX)) {
+        const originalId = id.slice(VIRTUAL_AUTOMOCK_PREFIX.length);
+        const originalCode = await readFile(originalId, 'utf-8');
+        // Generate the mocked source code. Vite will transform this code after we return it.
+        const mocked = automockModule(originalCode, 'automock', this.parse, {
+          globalThisAccessor: JSON.stringify('__vitest_mocker__'),
+        });
+        return mocked.toString();
+      }
+      return null;
+    },
+
+    generateBundle() {
+      // This hook runs after all chunks have been generated and filenames are finalized.
+      const manifest = {
+        redirects: {} as Record<string, string>,
+        automocks: {} as Record<string, string>,
+      };
+
+      // Resolve the final URLs for our __mocks__ redirects.
+      for (const [originalId, chunkRef] of redirectChunkRefs.entries()) {
+        const finalFileName = this.getFileName(chunkRef);
+        manifest.redirects[originalId] = `/${finalFileName}`;
+      }
+
+      // Resolve the final URLs for our automocked modules.
+      for (const [originalId, chunkRef] of automockChunkRefs.entries()) {
+        const finalFileName = this.getFileName(chunkRef);
+        manifest.automocks[originalId] = `/${finalFileName}`;
+      }
+
+      // Emit the final manifest file.
+      this.emitFile({
+        type: 'asset',
+        fileName: 'mock-manifest.json',
+        source: JSON.stringify(manifest, null, 2),
+      });
+    },
+  };
+}