From 11d7074521585d371b771527f324573feb479120 Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Wed, 25 Mar 2026 17:08:35 +0100 Subject: [PATCH 1/8] Implement module graph functionality and related tests in Vite builder - Added `buildModuleGraph` function to create a module graph from Vite's module nodes. - Introduced `onModuleGraphChange` to allow listeners to react to changes in the module graph. - Created comprehensive tests for module graph behavior in `index.test.ts`. - Updated type definitions to include `ModuleGraph` and `ModuleNode` for better type safety. --- code/builders/builder-vite/src/index.test.ts | 334 +++++++++++++++++++ code/builders/builder-vite/src/index.ts | 121 ++++++- code/core/src/types/modules/core-common.ts | 15 + 3 files changed, 469 insertions(+), 1 deletion(-) create mode 100644 code/builders/builder-vite/src/index.test.ts diff --git a/code/builders/builder-vite/src/index.test.ts b/code/builders/builder-vite/src/index.test.ts new file mode 100644 index 000000000000..2500e7906ed1 --- /dev/null +++ b/code/builders/builder-vite/src/index.test.ts @@ -0,0 +1,334 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { ModuleNode as StorybookModuleNode, Options } from 'storybook/internal/types'; +import type { ViteDevServer } from 'vite'; + +import { bail, buildModuleGraph, onModuleGraphChange, start } from './index'; +import { createViteServer } from './vite-server'; + +vi.mock('./vite-server', () => ({ + createViteServer: vi.fn(), +})); + +type ViteModuleNodeLike = { + file: string | null; + type: StorybookModuleNode['type']; + importers: Set; + importedModules: Set; +}; + +function createViteModuleNode( + file: string | null, + type: StorybookModuleNode['type'] = 'js' +): ViteModuleNodeLike { + return { + file, + type, + importers: new Set(), + importedModules: new Set(), + }; +} + +function createFileToModulesMap(...entries: Array<[string, Set]>) { + return new Map(entries) as ViteDevServer['moduleGraph']['fileToModulesMap']; +} + +function getFirstNode( + file: string, + moduleGraph: ReturnType +): StorybookModuleNode { + const moduleNode = moduleGraph.get(file)?.values().next().value; + if (!moduleNode) { + throw new Error(`Expected module node for ${file}`); + } + return moduleNode; +} + +type WatcherHandler = (...args: unknown[]) => void; +type FakeWatcher = { + on: ReturnType; + off: ReturnType; + emit: (event: string, ...args: unknown[]) => void; + listenerCount: (event: string) => number; +}; + +function createFakeViteServer() { + const watcherListeners = new Map>(); + const watcher = {} as FakeWatcher; + + watcher.on = vi.fn((event: string, handler: WatcherHandler) => { + const handlers = watcherListeners.get(event) ?? new Set(); + handlers.add(handler); + watcherListeners.set(event, handlers); + return watcher; + }); + watcher.off = vi.fn((event: string, handler: WatcherHandler) => { + watcherListeners.get(event)?.delete(handler); + return watcher; + }); + watcher.emit = (event: string, ...args: unknown[]) => { + watcherListeners.get(event)?.forEach((handler) => { + handler(...args); + }); + watcherListeners.get('all')?.forEach((handler) => { + handler(event, ...args); + }); + }; + watcher.listenerCount = (event: string) => watcherListeners.get(event)?.size ?? 0; + + return { + watcher, + moduleGraph: { + fileToModulesMap: createFileToModulesMap(), + }, + middlewares: { + handle: vi.fn(), + }, + transformIndexHtml: vi.fn().mockResolvedValue(''), + waitForRequestsIdle: vi.fn().mockResolvedValue(undefined), + close: vi.fn().mockResolvedValue(undefined), + } as unknown as ViteDevServer & { + watcher: typeof watcher; + }; +} + +function createStartArgs(): Parameters[0] { + return { + startTime: process.hrtime(), + options: {} as Options, + router: { + get: vi.fn(), + use: vi.fn(), + } as unknown as Parameters[0]['router'], + server: {} as Parameters[0]['server'], + channel: {} as Parameters[0]['channel'], + }; +} + +describe('buildModuleGraph', () => { + it('converts vite module nodes into the shared module graph shape', () => { + const entry = createViteModuleNode('/src/entry.ts'); + const component = createViteModuleNode('/src/component.ts'); + const styles = createViteModuleNode('/src/component.css', 'css'); + + entry.importedModules.add(component); + component.importers.add(entry); + component.importedModules.add(styles); + styles.importers.add(component); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap( + ['/src/entry.ts', new Set([entry])], + ['/src/component.ts', new Set([component])], + ['/src/component.css', new Set([styles])] + ) + ); + + const entryNode = getFirstNode('/src/entry.ts', moduleGraph); + const componentNode = getFirstNode('/src/component.ts', moduleGraph); + const styleNode = getFirstNode('/src/component.css', moduleGraph); + + expect(entryNode.file).toBe('/src/entry.ts'); + expect(componentNode.type).toBe('js'); + expect(styleNode.type).toBe('css'); + + expect(entryNode.importedModules).toEqual(new Set([componentNode])); + expect(componentNode.importers).toEqual(new Set([entryNode])); + expect(componentNode.importedModules).toEqual(new Set([styleNode])); + expect(styleNode.importers).toEqual(new Set([componentNode])); + }); + + it('reuses the same converted node identity across relationships', () => { + const shared = createViteModuleNode('/src/shared.ts'); + const importerA = createViteModuleNode('/src/a.ts'); + const importerB = createViteModuleNode('/src/b.ts'); + + importerA.importedModules.add(shared); + importerB.importedModules.add(shared); + shared.importers.add(importerA); + shared.importers.add(importerB); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap( + ['/src/shared.ts', new Set([shared])], + ['/src/a.ts', new Set([importerA])], + ['/src/b.ts', new Set([importerB])] + ) + ); + + const sharedNode = getFirstNode('/src/shared.ts', moduleGraph); + const importerANode = getFirstNode('/src/a.ts', moduleGraph); + const importerBNode = getFirstNode('/src/b.ts', moduleGraph); + + expect(importerANode.importedModules.has(sharedNode)).toBe(true); + expect(importerBNode.importedModules.has(sharedNode)).toBe(true); + expect(sharedNode.importers).toEqual(new Set([importerANode, importerBNode])); + }); + + it('skips related vite module nodes without a file', () => { + const entry = createViteModuleNode('/src/entry.ts'); + const virtualModule = createViteModuleNode(null); + + entry.importedModules.add(virtualModule); + virtualModule.importers.add(entry); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap(['/src/entry.ts', new Set([entry])]) + ); + const entryNode = getFirstNode('/src/entry.ts', moduleGraph); + + expect(moduleGraph.size).toBe(1); + expect(entryNode.importedModules.size).toBe(0); + }); + + it('keeps multiple module identities for the same file', () => { + const clientModule = createViteModuleNode('/src/shared.ts'); + const ssrModule = createViteModuleNode('/src/shared.ts'); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap(['/src/shared.ts', new Set([clientModule, ssrModule])]) + ); + + expect(moduleGraph.get('/src/shared.ts')?.size).toBe(2); + }); +}); + +describe('onModuleGraphChange', () => { + let fakeViteServer: ReturnType; + + beforeEach(() => { + vi.useFakeTimers(); + fakeViteServer = createFakeViteServer(); + vi.mocked(createViteServer).mockResolvedValue(fakeViteServer); + }); + + afterEach(async () => { + await bail(); + vi.useRealTimers(); + vi.clearAllMocks(); + }); + + it('registers callbacks and unsubscribes them', async () => { + const cb = vi.fn(); + const unsubscribe = onModuleGraphChange(cb); + + expect(unsubscribe).toEqual(expect.any(Function)); + + await start(createStartArgs()); + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + + await vi.advanceTimersByTimeAsync(100); + + expect(cb).toHaveBeenCalledTimes(1); + + unsubscribe(); + + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + await vi.advanceTimersByTimeAsync(100); + + expect(cb).toHaveBeenCalledTimes(1); + }); + + it('passes the module graph payload to listeners', async () => { + const entry = createViteModuleNode('/src/Button.tsx'); + fakeViteServer.moduleGraph.fileToModulesMap = createFileToModulesMap([ + '/src/Button.tsx', + new Set([entry]), + ]); + + const cb = vi.fn(); + onModuleGraphChange(cb); + + await start(createStartArgs()); + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + + await vi.advanceTimersByTimeAsync(100); + + const moduleGraph = cb.mock.calls[0]?.[0]; + + expect(cb).toHaveBeenCalledWith(expect.any(Map)); + expect(moduleGraph?.has('/src/Button.tsx')).toBe(true); + }); + + it('triggers change events after the debounce delay', async () => { + const cb = vi.fn(); + onModuleGraphChange(cb); + + await start(createStartArgs()); + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + + await vi.advanceTimersByTimeAsync(50); + expect(cb).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(50); + expect(cb).toHaveBeenCalledTimes(1); + }); + + it('triggers add events after the debounce delay', async () => { + const cb = vi.fn(); + onModuleGraphChange(cb); + + await start(createStartArgs()); + fakeViteServer.watcher.emit('add', '/src/Button.tsx'); + + await vi.advanceTimersByTimeAsync(100); + + expect(cb).toHaveBeenCalledTimes(1); + }); + + it('triggers unlink events after the debounce delay', async () => { + const cb = vi.fn(); + onModuleGraphChange(cb); + + await start(createStartArgs()); + fakeViteServer.watcher.emit('unlink', '/src/Button.tsx'); + + await vi.advanceTimersByTimeAsync(100); + + expect(cb).toHaveBeenCalledTimes(1); + }); + + it('debounces multiple rapid events into a single callback', async () => { + const cb = vi.fn(); + onModuleGraphChange(cb); + + await start(createStartArgs()); + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + fakeViteServer.watcher.emit('add', '/src/Button.tsx'); + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + + await vi.advanceTimersByTimeAsync(100); + + expect(cb).toHaveBeenCalledTimes(1); + }); + + it('notifies multiple listeners', async () => { + const cb1 = vi.fn(); + const cb2 = vi.fn(); + + onModuleGraphChange(cb1); + onModuleGraphChange(cb2); + + await start(createStartArgs()); + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + + await vi.advanceTimersByTimeAsync(100); + + expect(cb1).toHaveBeenCalledTimes(1); + expect(cb2).toHaveBeenCalledTimes(1); + }); + + it('calls watcher.off during bail and clears listeners', async () => { + const cb = vi.fn(); + onModuleGraphChange(cb); + + await start(createStartArgs()); + await bail(); + + fakeViteServer.watcher.emit('change', '/src/Button.tsx'); + await vi.advanceTimersByTimeAsync(100); + + expect(cb).not.toHaveBeenCalled(); + expect(fakeViteServer.watcher.off).toHaveBeenCalledWith('change', expect.any(Function)); + }); +}); diff --git a/code/builders/builder-vite/src/index.ts b/code/builders/builder-vite/src/index.ts index f447e76b0419..9c82b5a9aeb9 100644 --- a/code/builders/builder-vite/src/index.ts +++ b/code/builders/builder-vite/src/index.ts @@ -3,7 +3,7 @@ import { readFile } from 'node:fs/promises'; import { fileURLToPath } from 'node:url'; import { NoStatsForViteDevError } from 'storybook/internal/server-errors'; -import type { Middleware, Options } from 'storybook/internal/types'; +import type { Middleware, ModuleGraph, ModuleNode, Options } from 'storybook/internal/types'; import type { ViteDevServer } from 'vite'; @@ -33,11 +33,113 @@ function iframeHandler(options: Options, server: ViteDevServer): Middleware { } let server: ViteDevServer; +const listeners = new Set<(moduleGraph: ModuleGraph) => void>(); +let debounce: ReturnType | undefined; +let watcherChangeHandler: (() => void) | undefined; + +export function buildModuleGraph( + fileToModulesMap: ViteDevServer['moduleGraph']['fileToModulesMap'] +): ModuleGraph { + const moduleGraph: ModuleGraph = new Map(); + const moduleNodeMap = new WeakMap(); + + const getOrCreateModuleNode = ( + viteModuleNode: { + file: string | null; + type: ModuleNode['type']; + importers: Set; + importedModules: Set; + }, + fallbackFile?: string + ): ModuleNode | undefined => { + const file = viteModuleNode.file ?? fallbackFile; + if (!file) { + return undefined; + } + + const existingNode = moduleNodeMap.get(viteModuleNode); + if (existingNode) { + return existingNode; + } + + const moduleNode: ModuleNode = { + file, + type: viteModuleNode.type, + importers: new Set(), + importedModules: new Set(), + }; + moduleNodeMap.set(viteModuleNode, moduleNode); + + const moduleSet = moduleGraph.get(file) ?? new Set(); + moduleSet.add(moduleNode); + moduleGraph.set(file, moduleSet); + + return moduleNode; + }; + + fileToModulesMap.forEach((viteModuleSet, filePath) => { + viteModuleSet.forEach((viteModuleNode) => { + const moduleNode = getOrCreateModuleNode(viteModuleNode, filePath); + if (moduleNode) { + viteModuleNode.importers.forEach((importer) => { + const importerNode = getOrCreateModuleNode(importer); + if (importerNode) { + moduleNode.importers.add(importerNode); + } + }); + viteModuleNode.importedModules.forEach((importedModule) => { + const importedModuleNode = getOrCreateModuleNode(importedModule); + if (importedModuleNode) { + moduleNode.importedModules.add(importedModuleNode); + } + }); + } + }); + }); + + console.log('MODULE GRAPH:', `(${fileToModulesMap.size} files)`); + console.log( + Object.fromEntries( + Array.from(moduleGraph.entries()) + .sort((a, b) => a[0].localeCompare(b[0])) + .map(([file, [node]]) => [file, Array.from(node.importedModules).map((n) => n.file)]) + .filter(([_, importedModules]) => importedModules.length > 0) + ) + ); + console.log(''); + + return moduleGraph; +} + +function notifyListeners(moduleGraph: ModuleGraph): void { + listeners.forEach((listener) => { + listener(moduleGraph); + }); +} export async function bail(): Promise { + if (watcherChangeHandler) { + server?.watcher.off('change', watcherChangeHandler); + watcherChangeHandler = undefined; + } + + if (debounce) { + clearTimeout(debounce); + debounce = undefined; + } + + listeners.clear(); return server?.close(); } +export const onModuleGraphChange: NonNullable = (cb) => { + listeners.add(cb); + + return () => { + listeners.delete(cb); + }; +}; + export const start: ViteBuilder['start'] = async ({ startTime, options, @@ -49,6 +151,23 @@ export const start: ViteBuilder['start'] = async ({ router.get('/iframe.html', iframeHandler(options as Options, server)); router.use(server.middlewares); + watcherChangeHandler = () => { + clearTimeout(debounce); + debounce = setTimeout(() => { + notifyListeners(buildModuleGraph(server.moduleGraph.fileToModulesMap)); + }, 100); + }; + + server.watcher.on('all', watcherChangeHandler); + + const waitForModuleGraph = setInterval(async () => { + if (server.moduleGraph.fileToModulesMap.size > 0) { + clearInterval(waitForModuleGraph); + await server.waitForRequestsIdle(); + watcherChangeHandler?.(); + } + }, 1000); + return { bail, stats: { diff --git a/code/core/src/types/modules/core-common.ts b/code/core/src/types/modules/core-common.ts index 4b5e046df196..be961b32d656 100644 --- a/code/core/src/types/modules/core-common.ts +++ b/code/core/src/types/modules/core-common.ts @@ -276,6 +276,21 @@ export interface Builder { bail: (e?: Error) => Promise; corePresets?: string[]; overridePresets?: string[]; + onModuleGraphChange?(cb: (moduleGraph: ModuleGraph) => void): () => void; +} + +/** + * Builder-agnostic module graph for dependency tracking. Modeled after Vite's module graph. + * The same file can be imported in multiple ways (e.g. based on query params or import context), + * each representing a unique module identity, hence the value is a Set. + */ +export type ModuleGraph = Map>; + +export interface ModuleNode { + file: string; + type: 'js' | 'css' | 'asset'; + importers: Set; + importedModules: Set; } /** Options for TypeScript usage within Storybook. */ From cb10bacf94f55109118d5501a277f64e69bf838a Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Wed, 25 Mar 2026 17:17:34 +0100 Subject: [PATCH 2/8] Clean up logging --- code/builders/builder-vite/src/index.ts | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/code/builders/builder-vite/src/index.ts b/code/builders/builder-vite/src/index.ts index 9c82b5a9aeb9..b6da94765f94 100644 --- a/code/builders/builder-vite/src/index.ts +++ b/code/builders/builder-vite/src/index.ts @@ -97,17 +97,6 @@ export function buildModuleGraph( }); }); - console.log('MODULE GRAPH:', `(${fileToModulesMap.size} files)`); - console.log( - Object.fromEntries( - Array.from(moduleGraph.entries()) - .sort((a, b) => a[0].localeCompare(b[0])) - .map(([file, [node]]) => [file, Array.from(node.importedModules).map((n) => n.file)]) - .filter(([_, importedModules]) => importedModules.length > 0) - ) - ); - console.log(''); - return moduleGraph; } From c72d208d6a8f3c844a3e84922c8e8d5a77aeefaf Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Wed, 25 Mar 2026 17:24:30 +0100 Subject: [PATCH 3/8] Update event listener management in Vite builder to prevent memory leaks - Changed the event listener removal from 'change' to 'all' in the bail function to ensure proper cleanup. - Updated the test description to reflect the new behavior of removing the all-event watcher during bail. - Added assertions in tests to verify the listener count before and after bail operations. --- code/builders/builder-vite/src/index.test.ts | 10 ++++++++-- code/builders/builder-vite/src/index.ts | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/code/builders/builder-vite/src/index.test.ts b/code/builders/builder-vite/src/index.test.ts index 2500e7906ed1..dfa1e8be8739 100644 --- a/code/builders/builder-vite/src/index.test.ts +++ b/code/builders/builder-vite/src/index.test.ts @@ -318,17 +318,23 @@ describe('onModuleGraphChange', () => { expect(cb2).toHaveBeenCalledTimes(1); }); - it('calls watcher.off during bail and clears listeners', async () => { + it('removes the all-event watcher during bail to avoid leaks across restarts', async () => { const cb = vi.fn(); onModuleGraphChange(cb); await start(createStartArgs()); + expect(fakeViteServer.watcher.listenerCount('all')).toBe(1); + await bail(); + expect(fakeViteServer.watcher.listenerCount('all')).toBe(0); + + await start(createStartArgs()); + expect(fakeViteServer.watcher.listenerCount('all')).toBe(1); fakeViteServer.watcher.emit('change', '/src/Button.tsx'); await vi.advanceTimersByTimeAsync(100); expect(cb).not.toHaveBeenCalled(); - expect(fakeViteServer.watcher.off).toHaveBeenCalledWith('change', expect.any(Function)); + expect(fakeViteServer.watcher.off).toHaveBeenCalledWith('all', expect.any(Function)); }); }); diff --git a/code/builders/builder-vite/src/index.ts b/code/builders/builder-vite/src/index.ts index b6da94765f94..f0ca36626ac8 100644 --- a/code/builders/builder-vite/src/index.ts +++ b/code/builders/builder-vite/src/index.ts @@ -108,7 +108,7 @@ function notifyListeners(moduleGraph: ModuleGraph): void { export async function bail(): Promise { if (watcherChangeHandler) { - server?.watcher.off('change', watcherChangeHandler); + server?.watcher.off('all', watcherChangeHandler); watcherChangeHandler = undefined; } From f1626c82cf0b605674ae0cee87d0e13b2c609db2 Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Wed, 25 Mar 2026 19:31:03 +0100 Subject: [PATCH 4/8] Clear waitForModuleGraph interval on bail --- code/builders/builder-vite/src/index.test.ts | 18 ++++++++++++++++++ code/builders/builder-vite/src/index.ts | 9 ++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/code/builders/builder-vite/src/index.test.ts b/code/builders/builder-vite/src/index.test.ts index dfa1e8be8739..ff649f900571 100644 --- a/code/builders/builder-vite/src/index.test.ts +++ b/code/builders/builder-vite/src/index.test.ts @@ -337,4 +337,22 @@ describe('onModuleGraphChange', () => { expect(cb).not.toHaveBeenCalled(); expect(fakeViteServer.watcher.off).toHaveBeenCalledWith('all', expect.any(Function)); }); + + it('clears the module-graph polling interval during bail', async () => { + await start(createStartArgs()); + + expect(vi.getTimerCount()).toBe(1); + + await bail(); + + fakeViteServer.moduleGraph.fileToModulesMap = createFileToModulesMap([ + '/src/Button.tsx', + new Set([createViteModuleNode('/src/Button.tsx')]), + ]); + + await vi.advanceTimersByTimeAsync(1000); + + expect(vi.getTimerCount()).toBe(0); + expect(fakeViteServer.waitForRequestsIdle).not.toHaveBeenCalled(); + }); }); diff --git a/code/builders/builder-vite/src/index.ts b/code/builders/builder-vite/src/index.ts index f0ca36626ac8..2336ee324945 100644 --- a/code/builders/builder-vite/src/index.ts +++ b/code/builders/builder-vite/src/index.ts @@ -36,6 +36,7 @@ let server: ViteDevServer; const listeners = new Set<(moduleGraph: ModuleGraph) => void>(); let debounce: ReturnType | undefined; let watcherChangeHandler: (() => void) | undefined; +let waitForModuleGraph: ReturnType | undefined; export function buildModuleGraph( fileToModulesMap: ViteDevServer['moduleGraph']['fileToModulesMap'] @@ -112,6 +113,11 @@ export async function bail(): Promise { watcherChangeHandler = undefined; } + if (waitForModuleGraph) { + clearInterval(waitForModuleGraph); + waitForModuleGraph = undefined; + } + if (debounce) { clearTimeout(debounce); debounce = undefined; @@ -149,9 +155,10 @@ export const start: ViteBuilder['start'] = async ({ server.watcher.on('all', watcherChangeHandler); - const waitForModuleGraph = setInterval(async () => { + waitForModuleGraph = setInterval(async () => { if (server.moduleGraph.fileToModulesMap.size > 0) { clearInterval(waitForModuleGraph); + waitForModuleGraph = undefined; await server.waitForRequestsIdle(); watcherChangeHandler?.(); } From 4deb10418a631251efb93ccc17f8a72cc744ce88 Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Wed, 25 Mar 2026 21:54:49 +0100 Subject: [PATCH 5/8] Improve types and add a clarification --- code/builders/builder-vite/src/index.ts | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/code/builders/builder-vite/src/index.ts b/code/builders/builder-vite/src/index.ts index 2336ee324945..79d873cacaf7 100644 --- a/code/builders/builder-vite/src/index.ts +++ b/code/builders/builder-vite/src/index.ts @@ -3,9 +3,15 @@ import { readFile } from 'node:fs/promises'; import { fileURLToPath } from 'node:url'; import { NoStatsForViteDevError } from 'storybook/internal/server-errors'; -import type { Middleware, ModuleGraph, ModuleNode, Options } from 'storybook/internal/types'; +import type { + Builder, + Middleware, + ModuleGraph, + ModuleNode, + Options, +} from 'storybook/internal/types'; -import type { ViteDevServer } from 'vite'; +import type { ViteDevServer, ModuleNode as ViteModuleNode } from 'vite'; import { build as viteBuild } from './build'; import type { ViteBuilder } from './types'; @@ -47,9 +53,9 @@ export function buildModuleGraph( const getOrCreateModuleNode = ( viteModuleNode: { file: string | null; - type: ModuleNode['type']; - importers: Set; - importedModules: Set; + type: ViteModuleNode['type']; + importers: Set; + importedModules: Set; }, fallbackFile?: string ): ModuleNode | undefined => { @@ -127,7 +133,7 @@ export async function bail(): Promise { return server?.close(); } -export const onModuleGraphChange: NonNullable = (cb) => { +export const onModuleGraphChange: NonNullable['onModuleGraphChange']> = (cb) => { listeners.add(cb); return () => { @@ -146,6 +152,7 @@ export const start: ViteBuilder['start'] = async ({ router.get('/iframe.html', iframeHandler(options as Options, server)); router.use(server.middlewares); + // Debounce handler to prevent multiple callback invocations when multiple files are edited watcherChangeHandler = () => { clearTimeout(debounce); debounce = setTimeout(() => { From af5da75aeadf2fad470b0db4b8f9cf0a2532bf21 Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Thu, 26 Mar 2026 09:54:03 +0100 Subject: [PATCH 6/8] Move buildModuleGraph to separate file --- code/builders/builder-vite/src/index.ts | 74 +------------------ .../src/utils/build-module-graph.ts | 66 +++++++++++++++++ 2 files changed, 69 insertions(+), 71 deletions(-) create mode 100644 code/builders/builder-vite/src/utils/build-module-graph.ts diff --git a/code/builders/builder-vite/src/index.ts b/code/builders/builder-vite/src/index.ts index 79d873cacaf7..661201775488 100644 --- a/code/builders/builder-vite/src/index.ts +++ b/code/builders/builder-vite/src/index.ts @@ -3,19 +3,14 @@ import { readFile } from 'node:fs/promises'; import { fileURLToPath } from 'node:url'; import { NoStatsForViteDevError } from 'storybook/internal/server-errors'; -import type { - Builder, - Middleware, - ModuleGraph, - ModuleNode, - Options, -} from 'storybook/internal/types'; +import type { Builder, Middleware, ModuleGraph, Options } from 'storybook/internal/types'; -import type { ViteDevServer, ModuleNode as ViteModuleNode } from 'vite'; +import type { ViteDevServer } from 'vite'; import { build as viteBuild } from './build'; import type { ViteBuilder } from './types'; import { createViteServer } from './vite-server'; +import { buildModuleGraph } from './utils/build-module-graph'; export { withoutVitePlugins } from './utils/without-vite-plugins'; export { hasVitePlugins } from './utils/has-vite-plugins'; @@ -44,69 +39,6 @@ let debounce: ReturnType | undefined; let watcherChangeHandler: (() => void) | undefined; let waitForModuleGraph: ReturnType | undefined; -export function buildModuleGraph( - fileToModulesMap: ViteDevServer['moduleGraph']['fileToModulesMap'] -): ModuleGraph { - const moduleGraph: ModuleGraph = new Map(); - const moduleNodeMap = new WeakMap(); - - const getOrCreateModuleNode = ( - viteModuleNode: { - file: string | null; - type: ViteModuleNode['type']; - importers: Set; - importedModules: Set; - }, - fallbackFile?: string - ): ModuleNode | undefined => { - const file = viteModuleNode.file ?? fallbackFile; - if (!file) { - return undefined; - } - - const existingNode = moduleNodeMap.get(viteModuleNode); - if (existingNode) { - return existingNode; - } - - const moduleNode: ModuleNode = { - file, - type: viteModuleNode.type, - importers: new Set(), - importedModules: new Set(), - }; - moduleNodeMap.set(viteModuleNode, moduleNode); - - const moduleSet = moduleGraph.get(file) ?? new Set(); - moduleSet.add(moduleNode); - moduleGraph.set(file, moduleSet); - - return moduleNode; - }; - - fileToModulesMap.forEach((viteModuleSet, filePath) => { - viteModuleSet.forEach((viteModuleNode) => { - const moduleNode = getOrCreateModuleNode(viteModuleNode, filePath); - if (moduleNode) { - viteModuleNode.importers.forEach((importer) => { - const importerNode = getOrCreateModuleNode(importer); - if (importerNode) { - moduleNode.importers.add(importerNode); - } - }); - viteModuleNode.importedModules.forEach((importedModule) => { - const importedModuleNode = getOrCreateModuleNode(importedModule); - if (importedModuleNode) { - moduleNode.importedModules.add(importedModuleNode); - } - }); - } - }); - }); - - return moduleGraph; -} - function notifyListeners(moduleGraph: ModuleGraph): void { listeners.forEach((listener) => { listener(moduleGraph); diff --git a/code/builders/builder-vite/src/utils/build-module-graph.ts b/code/builders/builder-vite/src/utils/build-module-graph.ts new file mode 100644 index 000000000000..979c2baabc67 --- /dev/null +++ b/code/builders/builder-vite/src/utils/build-module-graph.ts @@ -0,0 +1,66 @@ +import type { ModuleGraph, ModuleNode } from 'storybook/internal/types'; + +import type { ViteDevServer, ModuleNode as ViteModuleNode } from 'vite'; + +export function buildModuleGraph( + fileToModulesMap: ViteDevServer['moduleGraph']['fileToModulesMap'] +): ModuleGraph { + const moduleGraph: ModuleGraph = new Map(); + const moduleNodeMap = new WeakMap(); + + const getOrCreateModuleNode = ( + viteModuleNode: { + file: string | null; + type: ViteModuleNode['type']; + importers: Set; + importedModules: Set; + }, + fallbackFile?: string + ): ModuleNode | undefined => { + const file = viteModuleNode.file ?? fallbackFile; + if (!file) { + return undefined; + } + + const existingNode = moduleNodeMap.get(viteModuleNode); + if (existingNode) { + return existingNode; + } + + const moduleNode: ModuleNode = { + file, + type: viteModuleNode.type, + importers: new Set(), + importedModules: new Set(), + }; + moduleNodeMap.set(viteModuleNode, moduleNode); + + const moduleSet = moduleGraph.get(file) ?? new Set(); + moduleSet.add(moduleNode); + moduleGraph.set(file, moduleSet); + + return moduleNode; + }; + + fileToModulesMap.forEach((viteModuleSet, filePath) => { + viteModuleSet.forEach((viteModuleNode) => { + const moduleNode = getOrCreateModuleNode(viteModuleNode, filePath); + if (moduleNode) { + viteModuleNode.importers.forEach((importer) => { + const importerNode = getOrCreateModuleNode(importer); + if (importerNode) { + moduleNode.importers.add(importerNode); + } + }); + viteModuleNode.importedModules.forEach((importedModule) => { + const importedModuleNode = getOrCreateModuleNode(importedModule); + if (importedModuleNode) { + moduleNode.importedModules.add(importedModuleNode); + } + }); + } + }); + }); + + return moduleGraph; +} From 89f3d9f1ca46dafb043b1a46ad84eb26e23be957 Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Mon, 30 Mar 2026 09:16:18 +0200 Subject: [PATCH 7/8] Move tests to separate file --- code/builders/builder-vite/src/index.test.ts | 101 +------------- .../src/utils/build-module-graph.test.ts | 132 ++++++++++++++++++ 2 files changed, 133 insertions(+), 100 deletions(-) create mode 100644 code/builders/builder-vite/src/utils/build-module-graph.test.ts diff --git a/code/builders/builder-vite/src/index.test.ts b/code/builders/builder-vite/src/index.test.ts index ff649f900571..c9e7aaf56f51 100644 --- a/code/builders/builder-vite/src/index.test.ts +++ b/code/builders/builder-vite/src/index.test.ts @@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import type { ModuleNode as StorybookModuleNode, Options } from 'storybook/internal/types'; import type { ViteDevServer } from 'vite'; -import { bail, buildModuleGraph, onModuleGraphChange, start } from './index'; +import { bail, onModuleGraphChange, start } from './index'; import { createViteServer } from './vite-server'; vi.mock('./vite-server', () => ({ @@ -33,17 +33,6 @@ function createFileToModulesMap(...entries: Array<[string, Set -): StorybookModuleNode { - const moduleNode = moduleGraph.get(file)?.values().next().value; - if (!moduleNode) { - throw new Error(`Expected module node for ${file}`); - } - return moduleNode; -} - type WatcherHandler = (...args: unknown[]) => void; type FakeWatcher = { on: ReturnType; @@ -105,94 +94,6 @@ function createStartArgs(): Parameters[0] { }; } -describe('buildModuleGraph', () => { - it('converts vite module nodes into the shared module graph shape', () => { - const entry = createViteModuleNode('/src/entry.ts'); - const component = createViteModuleNode('/src/component.ts'); - const styles = createViteModuleNode('/src/component.css', 'css'); - - entry.importedModules.add(component); - component.importers.add(entry); - component.importedModules.add(styles); - styles.importers.add(component); - - const moduleGraph = buildModuleGraph( - createFileToModulesMap( - ['/src/entry.ts', new Set([entry])], - ['/src/component.ts', new Set([component])], - ['/src/component.css', new Set([styles])] - ) - ); - - const entryNode = getFirstNode('/src/entry.ts', moduleGraph); - const componentNode = getFirstNode('/src/component.ts', moduleGraph); - const styleNode = getFirstNode('/src/component.css', moduleGraph); - - expect(entryNode.file).toBe('/src/entry.ts'); - expect(componentNode.type).toBe('js'); - expect(styleNode.type).toBe('css'); - - expect(entryNode.importedModules).toEqual(new Set([componentNode])); - expect(componentNode.importers).toEqual(new Set([entryNode])); - expect(componentNode.importedModules).toEqual(new Set([styleNode])); - expect(styleNode.importers).toEqual(new Set([componentNode])); - }); - - it('reuses the same converted node identity across relationships', () => { - const shared = createViteModuleNode('/src/shared.ts'); - const importerA = createViteModuleNode('/src/a.ts'); - const importerB = createViteModuleNode('/src/b.ts'); - - importerA.importedModules.add(shared); - importerB.importedModules.add(shared); - shared.importers.add(importerA); - shared.importers.add(importerB); - - const moduleGraph = buildModuleGraph( - createFileToModulesMap( - ['/src/shared.ts', new Set([shared])], - ['/src/a.ts', new Set([importerA])], - ['/src/b.ts', new Set([importerB])] - ) - ); - - const sharedNode = getFirstNode('/src/shared.ts', moduleGraph); - const importerANode = getFirstNode('/src/a.ts', moduleGraph); - const importerBNode = getFirstNode('/src/b.ts', moduleGraph); - - expect(importerANode.importedModules.has(sharedNode)).toBe(true); - expect(importerBNode.importedModules.has(sharedNode)).toBe(true); - expect(sharedNode.importers).toEqual(new Set([importerANode, importerBNode])); - }); - - it('skips related vite module nodes without a file', () => { - const entry = createViteModuleNode('/src/entry.ts'); - const virtualModule = createViteModuleNode(null); - - entry.importedModules.add(virtualModule); - virtualModule.importers.add(entry); - - const moduleGraph = buildModuleGraph( - createFileToModulesMap(['/src/entry.ts', new Set([entry])]) - ); - const entryNode = getFirstNode('/src/entry.ts', moduleGraph); - - expect(moduleGraph.size).toBe(1); - expect(entryNode.importedModules.size).toBe(0); - }); - - it('keeps multiple module identities for the same file', () => { - const clientModule = createViteModuleNode('/src/shared.ts'); - const ssrModule = createViteModuleNode('/src/shared.ts'); - - const moduleGraph = buildModuleGraph( - createFileToModulesMap(['/src/shared.ts', new Set([clientModule, ssrModule])]) - ); - - expect(moduleGraph.get('/src/shared.ts')?.size).toBe(2); - }); -}); - describe('onModuleGraphChange', () => { let fakeViteServer: ReturnType; diff --git a/code/builders/builder-vite/src/utils/build-module-graph.test.ts b/code/builders/builder-vite/src/utils/build-module-graph.test.ts new file mode 100644 index 000000000000..bed1ba859756 --- /dev/null +++ b/code/builders/builder-vite/src/utils/build-module-graph.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, it, vi } from 'vitest'; + +import type { ModuleNode as StorybookModuleNode } from 'storybook/internal/types'; +import type { ViteDevServer } from 'vite'; + +import { buildModuleGraph } from './build-module-graph'; + +vi.mock('./vite-server', () => ({ + createViteServer: vi.fn(), +})); + +type ViteModuleNodeLike = { + file: string | null; + type: StorybookModuleNode['type']; + importers: Set; + importedModules: Set; +}; + +function createViteModuleNode( + file: string | null, + type: StorybookModuleNode['type'] = 'js' +): ViteModuleNodeLike { + return { + file, + type, + importers: new Set(), + importedModules: new Set(), + }; +} + +function createFileToModulesMap(...entries: Array<[string, Set]>) { + return new Map(entries) as ViteDevServer['moduleGraph']['fileToModulesMap']; +} + +function getFirstNode( + file: string, + moduleGraph: ReturnType +): StorybookModuleNode { + const moduleNode = moduleGraph.get(file)?.values().next().value; + if (!moduleNode) { + throw new Error(`Expected module node for ${file}`); + } + return moduleNode; +} + +describe('buildModuleGraph', () => { + it('converts vite module nodes into the shared module graph shape', () => { + const entry = createViteModuleNode('/src/entry.ts'); + const component = createViteModuleNode('/src/component.ts'); + const styles = createViteModuleNode('/src/component.css', 'css'); + + entry.importedModules.add(component); + component.importers.add(entry); + component.importedModules.add(styles); + styles.importers.add(component); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap( + ['/src/entry.ts', new Set([entry])], + ['/src/component.ts', new Set([component])], + ['/src/component.css', new Set([styles])] + ) + ); + + const entryNode = getFirstNode('/src/entry.ts', moduleGraph); + const componentNode = getFirstNode('/src/component.ts', moduleGraph); + const styleNode = getFirstNode('/src/component.css', moduleGraph); + + expect(entryNode.file).toBe('/src/entry.ts'); + expect(componentNode.type).toBe('js'); + expect(styleNode.type).toBe('css'); + + expect(entryNode.importedModules).toEqual(new Set([componentNode])); + expect(componentNode.importers).toEqual(new Set([entryNode])); + expect(componentNode.importedModules).toEqual(new Set([styleNode])); + expect(styleNode.importers).toEqual(new Set([componentNode])); + }); + + it('reuses the same converted node identity across relationships', () => { + const shared = createViteModuleNode('/src/shared.ts'); + const importerA = createViteModuleNode('/src/a.ts'); + const importerB = createViteModuleNode('/src/b.ts'); + + importerA.importedModules.add(shared); + importerB.importedModules.add(shared); + shared.importers.add(importerA); + shared.importers.add(importerB); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap( + ['/src/shared.ts', new Set([shared])], + ['/src/a.ts', new Set([importerA])], + ['/src/b.ts', new Set([importerB])] + ) + ); + + const sharedNode = getFirstNode('/src/shared.ts', moduleGraph); + const importerANode = getFirstNode('/src/a.ts', moduleGraph); + const importerBNode = getFirstNode('/src/b.ts', moduleGraph); + + expect(importerANode.importedModules.has(sharedNode)).toBe(true); + expect(importerBNode.importedModules.has(sharedNode)).toBe(true); + expect(sharedNode.importers).toEqual(new Set([importerANode, importerBNode])); + }); + + it('skips related vite module nodes without a file', () => { + const entry = createViteModuleNode('/src/entry.ts'); + const virtualModule = createViteModuleNode(null); + + entry.importedModules.add(virtualModule); + virtualModule.importers.add(entry); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap(['/src/entry.ts', new Set([entry])]) + ); + const entryNode = getFirstNode('/src/entry.ts', moduleGraph); + + expect(moduleGraph.size).toBe(1); + expect(entryNode.importedModules.size).toBe(0); + }); + + it('keeps multiple module identities for the same file', () => { + const clientModule = createViteModuleNode('/src/shared.ts'); + const ssrModule = createViteModuleNode('/src/shared.ts'); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap(['/src/shared.ts', new Set([clientModule, ssrModule])]) + ); + + expect(moduleGraph.get('/src/shared.ts')?.size).toBe(2); + }); +}); From 76cbb128f6289d193d4371fb7bb08e47520b70b2 Mon Sep 17 00:00:00 2001 From: Gert Hengeveld Date: Mon, 30 Mar 2026 10:02:21 +0200 Subject: [PATCH 8/8] Preserve edges for related Vite module nodes discovered before their file path is known --- .../src/utils/build-module-graph.test.ts | 21 +++++++++++++++++++ .../src/utils/build-module-graph.ts | 20 ++++++++++++++++-- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/code/builders/builder-vite/src/utils/build-module-graph.test.ts b/code/builders/builder-vite/src/utils/build-module-graph.test.ts index bed1ba859756..4be9409f0114 100644 --- a/code/builders/builder-vite/src/utils/build-module-graph.test.ts +++ b/code/builders/builder-vite/src/utils/build-module-graph.test.ts @@ -119,6 +119,27 @@ describe('buildModuleGraph', () => { expect(entryNode.importedModules.size).toBe(0); }); + it('preserves edges for related vite module nodes discovered before their file path is known', () => { + const entry = createViteModuleNode('/src/entry.ts'); + const component = createViteModuleNode(null); + + entry.importedModules.add(component); + component.importers.add(entry); + + const moduleGraph = buildModuleGraph( + createFileToModulesMap( + ['/src/entry.ts', new Set([entry])], + ['/src/component.ts', new Set([component])] + ) + ); + + const entryNode = getFirstNode('/src/entry.ts', moduleGraph); + const componentNode = getFirstNode('/src/component.ts', moduleGraph); + + expect(entryNode.importedModules).toEqual(new Set([componentNode])); + expect(componentNode.importers).toEqual(new Set([entryNode])); + }); + it('keeps multiple module identities for the same file', () => { const clientModule = createViteModuleNode('/src/shared.ts'); const ssrModule = createViteModuleNode('/src/shared.ts'); diff --git a/code/builders/builder-vite/src/utils/build-module-graph.ts b/code/builders/builder-vite/src/utils/build-module-graph.ts index 979c2baabc67..4d14235bd88e 100644 --- a/code/builders/builder-vite/src/utils/build-module-graph.ts +++ b/code/builders/builder-vite/src/utils/build-module-graph.ts @@ -7,6 +7,18 @@ export function buildModuleGraph( ): ModuleGraph { const moduleGraph: ModuleGraph = new Map(); const moduleNodeMap = new WeakMap(); + const getModuleFileFromMap = (viteModuleNode: { + file: string | null; + type: ViteModuleNode['type']; + importers: Set; + importedModules: Set; + }): string | undefined => { + for (const [filePath, viteModuleSet] of fileToModulesMap.entries()) { + if (viteModuleSet.has(viteModuleNode as ViteModuleNode)) { + return filePath; + } + } + }; const getOrCreateModuleNode = ( viteModuleNode: { @@ -47,13 +59,17 @@ export function buildModuleGraph( const moduleNode = getOrCreateModuleNode(viteModuleNode, filePath); if (moduleNode) { viteModuleNode.importers.forEach((importer) => { - const importerNode = getOrCreateModuleNode(importer); + const importerNode = + getOrCreateModuleNode(importer) ?? + getOrCreateModuleNode(importer, getModuleFileFromMap(importer)); if (importerNode) { moduleNode.importers.add(importerNode); } }); viteModuleNode.importedModules.forEach((importedModule) => { - const importedModuleNode = getOrCreateModuleNode(importedModule); + const importedModuleNode = + getOrCreateModuleNode(importedModule) ?? + getOrCreateModuleNode(importedModule, getModuleFileFromMap(importedModule)); if (importedModuleNode) { moduleNode.importedModules.add(importedModuleNode); }