From dc20fb5ac9f93f5a73b6e28845867cfbc9e0468a Mon Sep 17 00:00:00 2001 From: Vladimir Sheremet Date: Wed, 17 Dec 2025 16:50:06 +0100 Subject: [PATCH 1/2] feat(experimental): add `onModuleRunner` hook to `worker.init` --- packages/vitest/src/runtime/listeners.ts | 20 ++++++++++++++++++++ packages/vitest/src/runtime/worker.ts | 6 +++--- packages/vitest/src/runtime/workers/base.ts | 3 +++ packages/vitest/src/runtime/workers/init.ts | 4 ++++ packages/vitest/src/runtime/workers/types.ts | 3 ++- packages/vitest/src/runtime/workers/vm.ts | 10 ++++------ 6 files changed, 36 insertions(+), 10 deletions(-) create mode 100644 packages/vitest/src/runtime/listeners.ts diff --git a/packages/vitest/src/runtime/listeners.ts b/packages/vitest/src/runtime/listeners.ts new file mode 100644 index 000000000000..71ad5c56cff3 --- /dev/null +++ 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 { + await Promise.all([...cleanupListeners].map(l => l())) +} + +export function onModuleRunner(cb: (runner: ModuleRunner) => unknown): void { + moduleRunnerListeners.add(cb) +} + +export async function emitModuleRunner(moduleRunner: ModuleRunner): Promise { + await Promise.all([...moduleRunnerListeners].map(l => l(moduleRunner))) +} diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts index c17229a32a66..62888bf6cf92 100644 --- 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() -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 { - 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 index e0dc6af90a98..7b5bef7758f6 100644 --- 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, }) + await 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 index c97402e1d1c4..a526fc3fb5bc 100644 --- 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 | undefined let isRunning = false diff --git a/packages/vitest/src/runtime/workers/types.ts b/packages/vitest/src/runtime/workers/types.ts index 82dea0000909..17db1595827d 100644 --- 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 collectTests: (state: WorkerGlobalState, traces: Traces) => Awaitable - + onModuleRunner?: (moduleRunner: ModuleRunner) => Awaitable setup?: (context: WorkerSetupContext) => Promise<() => Promise> } diff --git a/packages/vitest/src/runtime/workers/vm.ts b/packages/vitest/src/runtime/workers/vm.ts index 80af0e4eb99e..60b87bb9fe31 100644 --- 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, }) + await 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, From 6b895120ace53e7c8132d65f18b532422c6ce4b5 Mon Sep 17 00:00:00 2001 From: Vladimir Sheremet Date: Tue, 13 Jan 2026 11:10:07 +0100 Subject: [PATCH 2/2] chore: make emitModuleRunner sync --- packages/vitest/src/runtime/listeners.ts | 4 ++-- packages/vitest/src/runtime/workers/base.ts | 2 +- packages/vitest/src/runtime/workers/vm.ts | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/vitest/src/runtime/listeners.ts b/packages/vitest/src/runtime/listeners.ts index 71ad5c56cff3..dd472d9a45c0 100644 --- a/packages/vitest/src/runtime/listeners.ts +++ b/packages/vitest/src/runtime/listeners.ts @@ -15,6 +15,6 @@ export function onModuleRunner(cb: (runner: ModuleRunner) => unknown): void { moduleRunnerListeners.add(cb) } -export async function emitModuleRunner(moduleRunner: ModuleRunner): Promise { - await Promise.all([...moduleRunnerListeners].map(l => l(moduleRunner))) +export function emitModuleRunner(moduleRunner: ModuleRunner): void { + moduleRunnerListeners.forEach(l => l(moduleRunner)) } diff --git a/packages/vitest/src/runtime/workers/base.ts b/packages/vitest/src/runtime/workers/base.ts index 7b5bef7758f6..7c9e3cc0a27c 100644 --- a/packages/vitest/src/runtime/workers/base.ts +++ b/packages/vitest/src/runtime/workers/base.ts @@ -117,7 +117,7 @@ export async function runBaseTests(method: 'run' | 'collect', state: WorkerGloba traces, }) - await emitModuleRunner(moduleRunner as any) + emitModuleRunner(moduleRunner as any) await run( method, diff --git a/packages/vitest/src/runtime/workers/vm.ts b/packages/vitest/src/runtime/workers/vm.ts index 60b87bb9fe31..97adb6dca221 100644 --- a/packages/vitest/src/runtime/workers/vm.ts +++ b/packages/vitest/src/runtime/workers/vm.ts @@ -99,7 +99,7 @@ export async function runVmTests(method: 'run' | 'collect', state: WorkerGlobalS traces, }) - await emitModuleRunner(moduleRunner as any) + emitModuleRunner(moduleRunner as any) Object.defineProperty(context, VITEST_VM_CONTEXT_SYMBOL, { value: {