fix(ui): store module graph once per environment (#11421)

Co-authored-by: Hiroshi Ogawa <4232207+hi-ogawa@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: hi-ogawa

packages/ui/client/components/FileDetailsModuleGraph.vue

@@ -12,28 +12,16 @@ const props = defineProps<{
 }>()
 
 const hideNodeModules = ref(true)
-const NODE_MODULES_RE = /[/\\]node_modules[/\\]/
 
 const { state: graphData, isLoading } = useAsyncState(
   () =>
     client.rpc.getModuleGraph(props.projectName, props.file.filepath, props.file.viteEnvironment),
   undefined,
 )
 
-const graph = computed(() => {
-  if (!graphData.value) {
-    return { nodes: [], links: [] }
-  }
-  let moduleGraph = graphData.value
-  if (hideNodeModules.value) {
-    moduleGraph = {
-      ...moduleGraph,
-      inlined: moduleGraph.inlined.filter((n) => !NODE_MODULES_RE.test(n)),
-      externalized: moduleGraph.externalized.filter((n) => !NODE_MODULES_RE.test(n)),
-    }
-  }
-  return getModuleGraph(moduleGraph, props.file.filepath)
-})
+const graph = computed(() =>
+  getModuleGraph(graphData.value, props.file.filepath, hideNodeModules.value),
+)
 </script>
 
 <template>

packages/ui/client/composables/client/static.ts

@@ -1,4 +1,5 @@
-import type { ModuleGraphData, RunnerTestFile, SerializedRootConfig } from 'vitest'
+import type { RunnerTestFile, SerializedRootConfig } from 'vitest'
+import type { SharedModuleGraphByProject } from '../../../../vitest/src/types/general'
 import type { VitestClient, VitestClientRpc } from './ws'
 import { decompressSync, strFromU8 } from 'fflate'
 import { parse } from 'flatted'
@@ -8,7 +9,7 @@ import { StateManager } from './state'
 export interface HTMLReportMetadata {
   files: RunnerTestFile[]
   config: SerializedRootConfig
-  moduleGraph: Record<string, Record<string, ModuleGraphData>>
+  moduleGraph: SharedModuleGraphByProject
   unhandledErrors: unknown[]
   testModules: {
     projectName: string
@@ -38,8 +39,13 @@ function deserializeReportMetadata(metadata: HTMLReportMetadata) {
     getConfig: async () => {
       return metadata.config
     },
-    getModuleGraph: async (projectName, id) => {
-      return metadata.moduleGraph[projectName]?.[id]
+    getModuleGraph: async (projectName, id, viteEnvironment) => {
+      // the reporter keys graphs by the environment the file ran in, which is the file's
+      // `viteEnvironment`. It's unset only for files that never ran, which have no graph.
+      const graph = viteEnvironment
+        ? metadata.moduleGraph[projectName]?.[viteEnvironment]
+        : undefined
+      return { modules: graph?.modules ?? {}, roots: graph?.rootsByTestFile[id] ?? [] }
     },
     getUnhandledErrors: async () => {
       return metadata.unhandledErrors

packages/ui/client/composables/module-graph.ts

@@ -9,6 +9,9 @@ export type ModuleNode = GraphNode<ModuleType>
 export type ModuleLink = GraphLink<ModuleType, ModuleNode>
 export type ModuleGraph = Graph<ModuleType, ModuleNode, ModuleLink>
 export type ModuleGraphController = GraphController<ModuleType, ModuleNode, ModuleLink>
+
+const NODE_MODULES_RE = /[/\\]node_modules[/\\]/
+
 function defineExternalModuleNodes(modules: string[]): ModuleNode[] {
   const labels = modules.map((module) => createModuleLabelItem(module))
   const map = calcExternalLabels(labels)
@@ -41,22 +44,47 @@ function defineInlineModuleNode(module: string, isRoot: boolean): ModuleNode {
   })
 }
 
-export function getModuleGraph(data: ModuleGraphData, rootPath: string | undefined): ModuleGraph {
+export function getModuleGraph(
+  data: ModuleGraphData | undefined,
+  rootPath: string | undefined,
+  hideNodeModules = false,
+): ModuleGraph {
   if (!data) {
     return defineGraph({})
   }
 
+  const inlined: string[] = []
+  const externalized: string[] = []
+  const seen = new Set<string>()
+  function visit(id: string) {
+    const module = data!.modules[id]
+    if (seen.has(id) || !module) {
+      return
+    }
+    seen.add(id)
+    if (module.external) {
+      externalized.push(id)
+      return
+    }
+    inlined.push(id)
+    module.imports.forEach(visit)
+  }
+  data.roots.forEach(visit)
+
+  const isVisible = (id: string) => !hideNodeModules || !NODE_MODULES_RE.test(id)
+  const visibleInlined = inlined.filter(isVisible)
+  const visibleExternalized = externalized.filter(isVisible)
   const externalizedNodes = !config.value.experimental?.viteModuleRunner
-    ? defineExternalModuleNodes([...data.inlined, ...data.externalized])
-    : defineExternalModuleNodes(data.externalized)
+    ? defineExternalModuleNodes([...visibleInlined, ...visibleExternalized])
+    : defineExternalModuleNodes(visibleExternalized)
   const inlinedNodes = !config.value.experimental?.viteModuleRunner
     ? []
-    : (data.inlined.map((module) => defineInlineModuleNode(module, module === rootPath)) ?? [])
+    : visibleInlined.map((module) => defineInlineModuleNode(module, module === rootPath))
   const nodes = [...externalizedNodes, ...inlinedNodes]
   const nodeMap = Object.fromEntries(nodes.map((node) => [node.id, node]))
-  const links = Object.entries(data.graph).flatMap(
-    ([module, deps]) =>
-      deps
+  const links = inlined.flatMap(
+    (module) =>
+      data.modules[module].imports
         .map((dep) => {
           const source = nodeMap[module]
           const target = nodeMap[dep]

packages/ui/node/reporter.ts

@@ -14,7 +14,7 @@ import { gzip, constants as zlibConstants } from 'node:zlib'
 import { stringify } from 'flatted'
 import { dirname, relative, resolve } from 'pathe'
 import c from 'tinyrainbow'
-import { getModuleGraph } from '../../vitest/src/utils/graph'
+import { getSharedModuleGraphByProject } from '../../vitest/src/utils/graph'
 import { distClientRoot } from './paths'
 
 const gzipAsync = promisify(gzip)
@@ -120,7 +120,7 @@ async function serializeReportMetadata(
     files: [],
     config: ctx.serializedRootConfig,
     unhandledErrors: [...unhandledErrors],
-    moduleGraph: {},
+    moduleGraph: getSharedModuleGraphByProject(testModules),
     testModules: [],
     sourceCode: {
       codeTable: [],
@@ -147,8 +147,6 @@ async function serializeReportMetadata(
     return index
   }
 
-  const promises: Promise<void>[] = []
-
   for (const testModule of testModules) {
     result.files.push(testModule.task)
 
@@ -167,23 +165,8 @@ async function serializeReportMetadata(
         testModuleCodes[projectName][testModule.relativeModuleId] = getCodeIndex(code)
       } catch {}
     }
-
-    // TODO: https://github.com/vitest-dev/vitest/issues/9763
-    promises.push(
-      (async () => {
-        result.moduleGraph[projectName] ??= {}
-        result.moduleGraph[projectName][testModule.moduleId] = await getModuleGraph(
-          ctx,
-          projectName,
-          testModule.moduleId,
-          testModule.viteEnvironment?.name,
-        )
-      })(),
-    )
   }
 
-  await Promise.all(promises)
-
   return result
 }
 

packages/vitest/src/types/general.ts

@@ -21,10 +21,22 @@ export interface UserConsoleLog {
   size: number
 }
 
+interface ModuleGraphModules {
+  [id: string]: { external: boolean; imports: string[] }
+}
+
 export interface ModuleGraphData {
-  graph: Record<string, string[]>
-  externalized: string[]
-  inlined: string[]
+  modules: ModuleGraphModules
+  roots: string[]
+}
+
+export interface SharedModuleGraphByProject {
+  [projectName: string]: { [environmentName: string]: SharedModuleGraphData }
+}
+
+export interface SharedModuleGraphData {
+  modules: ModuleGraphModules
+  rootsByTestFile: { [testFile: string]: string[] }
 }
 
 export interface ProvidedContext {}

packages/vitest/src/utils/graph.ts

@@ -1,6 +1,12 @@
 import type { DevEnvironment, EnvironmentModuleNode } from 'vite'
 import type { Vitest } from '../node/core'
-import type { ModuleGraphData } from '../types/general'
+import type { TestProject } from '../node/project'
+import type { TestModule } from '../node/reporters/reported-tasks'
+import type {
+  ModuleGraphData,
+  SharedModuleGraphByProject,
+  SharedModuleGraphData,
+} from '../types/general'
 import { getTestFileEnvironment } from './environments'
 
 export async function getModuleGraph(
@@ -9,13 +15,57 @@ export async function getModuleGraph(
   testFilePath: string,
   viteEnvironment?: string,
 ): Promise<ModuleGraphData> {
-  const graph: Record<string, string[]> = {}
-  const externalized = new Set<string>()
-  const inlined = new Set<string>()
-
   const project = ctx.getProjectByName(projectName)
-  const browser = project.config.browser.enabled
+  const environment = getModuleGraphEnvironment(project, testFilePath, viteEnvironment)
+  const collector = createModuleGraphCollector(project, environment)
+  const roots = collector.add(testFilePath)
+  return { modules: collector.data.modules, roots }
+}
 
+export function getSharedModuleGraphByProject(
+  testModules: ReadonlyArray<TestModule>,
+): SharedModuleGraphByProject {
+  const testModulesByProject = new Map<TestProject, TestModule[]>()
+  for (const testModule of testModules) {
+    const projectTestModules = testModulesByProject.get(testModule.project) ?? []
+    projectTestModules.push(testModule)
+    testModulesByProject.set(testModule.project, projectTestModules)
+  }
+
+  const result: SharedModuleGraphByProject = {}
+  for (const [project, projectTestModules] of testModulesByProject) {
+    result[project.name] = getSharedModuleGraphByEnvironment(project, projectTestModules)
+  }
+  return result
+}
+
+function getSharedModuleGraphByEnvironment(
+  project: TestProject,
+  testModules: TestModule[],
+): SharedModuleGraphByProject[string] {
+  const collectors: { [environmentName: string]: ModuleGraphCollector } = {}
+  for (const testModule of testModules) {
+    const environment = getModuleGraphEnvironment(
+      project,
+      testModule.moduleId,
+      testModule.viteEnvironment?.name,
+    )
+    collectors[environment.name] ??= createModuleGraphCollector(project, environment)
+    collectors[environment.name].add(testModule.moduleId)
+  }
+  return Object.fromEntries(
+    Object.entries(collectors).map(([environmentName, collector]) => [
+      environmentName,
+      collector.data,
+    ]),
+  )
+}
+
+function getModuleGraphEnvironment(
+  project: TestProject,
+  testFilePath: string,
+  viteEnvironment?: string,
+): DevEnvironment {
   let environment: DevEnvironment | undefined
 
   if (viteEnvironment) {
@@ -24,15 +74,34 @@ export async function getModuleGraph(
     environment =
       project.config.experimental.viteModuleRunner === false
         ? project.vite.environments.__vitest__
-        : getTestFileEnvironment(project, testFilePath, browser)
+        : getTestFileEnvironment(project, testFilePath, project.config.browser.enabled)
   }
 
   if (!environment) {
     throw new Error(`Cannot find environment for ${testFilePath}`)
   }
+  return environment
+}
+
+interface ModuleGraphCollector {
+  data: SharedModuleGraphData
+  add: (testFilePath: string) => string[]
+}
+
+function createModuleGraphCollector(
+  project: TestProject,
+  environment: DevEnvironment,
+): ModuleGraphCollector {
+  const data: SharedModuleGraphData = { modules: {}, rootsByTestFile: {} }
+  const browser = project.config.browser.enabled
   const seen = new Map<EnvironmentModuleNode, string>()
 
-  function get(mod?: EnvironmentModuleNode) {
+  function addExternal(id: string) {
+    data.modules[id] ??= { external: true, imports: [] }
+    return id
+  }
+
+  function get(mod?: EnvironmentModuleNode): string | undefined {
     if (!mod || !mod.id) {
       return
     }
@@ -50,37 +119,33 @@ export async function getModuleGraph(
     const id = clearId(mod.id)
     seen.set(mod, id)
     if (id.startsWith('__vite-browser-external:')) {
-      const external = id.slice('__vite-browser-external:'.length)
-      externalized.add(external)
-      return external
+      return addExternal(id.slice('__vite-browser-external:'.length))
     }
     const external = project._resolver.wasExternalized(id)
     if (typeof external === 'string') {
-      externalized.add(external)
-      return external
+      return addExternal(external)
     }
     if (browser && mod.file?.includes(project.browser!.vite.config.cacheDir)) {
-      externalized.add(mod.id)
-      return id
+      return addExternal(id)
     }
-    inlined.add(id)
-    const mods = Array.from(mod.importedModules).filter(
-      (i) => i.id && !i.id.includes('/vitest/dist/'),
-    )
-    graph[id] = mods.map((m) => get(m)).filter(Boolean) as string[]
+    data.modules[id] = { external: false, imports: [] }
+    data.modules[id].imports = Array.from(mod.importedModules)
+      .filter((i) => i.id && !i.id.includes('/vitest/dist/'))
+      .map((m) => get(m))
+      .filter((id) => id != null)
     return id
   }
 
-  get(environment.moduleGraph.getModuleById(testFilePath))
-  project.config.setupFiles.forEach((setupFile) => {
-    get(environment.moduleGraph.getModuleById(setupFile))
-  })
-
-  return {
-    graph,
-    externalized: Array.from(externalized),
-    inlined: Array.from(inlined),
+  // graph roots of the test file including setup files
+  function add(testFilePath: string): string[] {
+    const roots = [testFilePath, ...project.config.setupFiles]
+      .map((file) => get(environment.moduleGraph.getModuleById(file)))
+      .filter((id) => id != null)
+    data.rootsByTestFile[testFilePath] = roots
+    return roots
   }
+
+  return { data, add }
 }
 
 function clearId(id?: string | null) {