feat(experimental): add onModuleRunner hook to worker.init (#9286)

sheremet-va · web-flow · commit e977f3deb695 · 2026-01-13T15:31:21.000+01:00
diff --git a/packages/vitest/src/runtime/listeners.ts b/packages/vitest/src/runtime/listeners.ts
@@ -0,0 +1,20 @@
+import type { ModuleRunner } from 'vite/module-runner'
+
+const cleanupListeners = new Set<() => unknown>()
+const moduleRunnerListeners = new Set<(runner: ModuleRunner) => unknown>()
+
+export function onCleanup(cb: () => unknown): void {
+  cleanupListeners.add(cb)
+}
+
+export async function cleanup(): Promise<void> {
+  await Promise.all([...cleanupListeners].map(l => l()))
+}
+
+export function onModuleRunner(cb: (runner: ModuleRunner) => unknown): void {
+  moduleRunnerListeners.add(cb)
+}
+
+export function emitModuleRunner(moduleRunner: ModuleRunner): void {
+  moduleRunnerListeners.forEach(l => l(moduleRunner))
+}
diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts
@@ -3,11 +3,11 @@ import type { Traces } from '../utils/traces'
 import type { VitestWorker } from './workers/types'
 import { createStackString, parseStacktrace } from '@vitest/utils/source-map'
 import { setupInspect } from './inspector'
+import * as listeners from './listeners'
 import { VitestEvaluatedModules } from './moduleRunner/evaluatedModules'
 import { onCancel, rpcDone } from './rpc'
 
 const resolvingModules = new Set<string>()
-const globalListeners = new Set<() => unknown>()
 
 async function execute(method: 'run' | 'collect', ctx: ContextRPC, worker: VitestWorker, traces: Traces) {
   const prepareStart = performance.now()
@@ -40,7 +40,7 @@ async function execute(method: 'run' | 'collect', ctx: ContextRPC, worker: Vites
       },
       rpc,
       onCancel,
-      onCleanup: listener => globalListeners.add(listener),
+      onCleanup: listeners.onCleanup,
       providedContext: ctx.providedContext,
       onFilterStackTrace(stack) {
         return createStackString(parseStacktrace(stack))
@@ -73,7 +73,7 @@ export function collect(ctx: ContextRPC, worker: VitestWorker, traces: Traces):
 }
 
 export async function teardown(): Promise<void> {
-  await Promise.all([...globalListeners].map(l => l()))
+  await listeners.cleanup()
 }
 
 const env = process.env
diff --git a/packages/vitest/src/runtime/workers/base.ts b/packages/vitest/src/runtime/workers/base.ts
@@ -7,6 +7,7 @@ import { runInThisContext } from 'node:vm'
 import * as spyModule from '@vitest/spy'
 import { setupChaiConfig } from '../../integrations/chai/config'
 import { loadEnvironment } from '../../integrations/env/loader'
+import { emitModuleRunner } from '../listeners'
 import { VitestEvaluatedModules } from '../moduleRunner/evaluatedModules'
 import { createNodeImportMeta } from '../moduleRunner/moduleRunner'
 import { startVitestModuleRunner } from '../moduleRunner/startModuleRunner'
@@ -116,6 +117,8 @@ export async function runBaseTests(method: 'run' | 'collect', state: WorkerGloba
     traces,
   })
 
+  emitModuleRunner(moduleRunner as any)
+
   await run(
     method,
     ctx.files,
diff --git a/packages/vitest/src/runtime/workers/init.ts b/packages/vitest/src/runtime/workers/init.ts
@@ -4,6 +4,7 @@ import type { WorkerSetupContext } from '../../types/worker'
 import type { VitestWorker } from './types'
 import { serializeError } from '@vitest/utils/error'
 import { Traces } from '../../utils/traces'
+import * as listeners from '../listeners'
 import { createRuntimeRpc } from '../rpc'
 import * as entrypoint from '../worker'
 
@@ -20,6 +21,9 @@ let traces!: Traces
 /** @experimental */
 export function init(worker: Options): void {
   worker.on(onMessage)
+  if (worker.onModuleRunner) {
+    listeners.onModuleRunner(worker.onModuleRunner)
+  }
 
   let runPromise: Promise<unknown> | undefined
   let isRunning = false
diff --git a/packages/vitest/src/runtime/workers/types.ts b/packages/vitest/src/runtime/workers/types.ts
@@ -1,5 +1,6 @@
 import type { Awaitable } from '@vitest/utils'
 import type { BirpcOptions } from 'birpc'
+import type { ModuleRunner } from 'vite/module-runner'
 import type { RuntimeRPC } from '../../types/rpc'
 import type { WorkerGlobalState, WorkerSetupContext } from '../../types/worker'
 import type { Traces } from '../../utils/traces'
@@ -12,6 +13,6 @@ type WorkerRpcOptions = Pick<
 export interface VitestWorker extends WorkerRpcOptions {
   runTests: (state: WorkerGlobalState, traces: Traces) => Awaitable<unknown>
   collectTests: (state: WorkerGlobalState, traces: Traces) => Awaitable<unknown>
-
+  onModuleRunner?: (moduleRunner: ModuleRunner) => Awaitable<unknown>
   setup?: (context: WorkerSetupContext) => Promise<() => Promise<unknown>>
 }
diff --git a/packages/vitest/src/runtime/workers/vm.ts b/packages/vitest/src/runtime/workers/vm.ts
@@ -8,6 +8,7 @@ import { loadEnvironment } from '../../integrations/env/loader'
 import { distDir } from '../../paths'
 import { createCustomConsole } from '../console'
 import { ExternalModulesExecutor } from '../external-executor'
+import { emitModuleRunner } from '../listeners'
 import { getDefaultRequestStubs } from '../moduleRunner/moduleEvaluator'
 import { createNodeImportMeta } from '../moduleRunner/moduleRunner'
 import { startVitestModuleRunner, VITEST_VM_CONTEXT_SYMBOL } from '../moduleRunner/startModuleRunner'
@@ -98,6 +99,8 @@ export async function runVmTests(method: 'run' | 'collect', state: WorkerGlobalS
     traces,
   })
 
+  emitModuleRunner(moduleRunner as any)
+
   Object.defineProperty(context, VITEST_VM_CONTEXT_SYMBOL, {
     value: {
       context,
@@ -124,16 +127,11 @@ export async function runVmTests(method: 'run' | 'collect', state: WorkerGlobalS
   const { run } = (await moduleRunner.import(
     entryFile,
   )) as typeof import('../runVmTests')
-  const fileSpecs = ctx.files.map(f =>
-    typeof f === 'string'
-      ? { filepath: f, testLocations: undefined }
-      : f,
-  )
 
   try {
     await run(
       method,
-      fileSpecs,
+      ctx.files,
       ctx.config,
       moduleRunner,
       traces,
