feat(coverage): v8 to track node:child_process and node:worker_threads contexts

Author: AriPerkkio

packages/coverage-v8/package.json

@@ -32,6 +32,10 @@
       "types": "./dist/browser.d.ts",
       "default": "./dist/browser.js"
     },
+    "./intercept-new-run-context": {
+      "types": "./dist/intercept-new-run-context.d.ts",
+      "default": "./dist/intercept-new-run-context.js"
+    },
     "./*": "./*"
   },
   "main": "./dist/index.js",
@@ -57,13 +61,15 @@
     "@bcoe/v8-coverage": "^1.0.2",
     "@vitest/utils": "workspace:*",
     "ast-v8-to-istanbul": "^1.0.0",
+    "get-port-please": "catalog:",
     "istanbul-lib-coverage": "catalog:",
     "istanbul-lib-report": "catalog:",
     "istanbul-reports": "catalog:",
     "magicast": "catalog:",
     "obug": "catalog:",
     "std-env": "catalog:",
-    "tinyrainbow": "catalog:"
+    "tinyrainbow": "catalog:",
+    "ws": "catalog:"
   },
   "devDependencies": {
     "@types/istanbul-lib-coverage": "catalog:",

packages/coverage-v8/rollup.config.js

@@ -10,9 +10,10 @@ const require = createRequire(import.meta.url)
 const pkg = require('./package.json')
 
 const entries = {
-  index: 'src/index.ts',
-  browser: 'src/browser.ts',
-  provider: 'src/provider.ts',
+  'index': 'src/index.ts',
+  'browser': 'src/browser.ts',
+  'provider': 'src/provider.ts',
+  'intercept-new-run-context': 'src/intercept-new-run-context.ts',
 }
 
 const external = [

packages/coverage-v8/src/index.ts

@@ -3,21 +3,41 @@ import type { CoverageProviderModule } from 'vitest/node'
 import type { ScriptCoverageWithOffset, V8CoverageProvider } from './provider'
 import inspector from 'node:inspector/promises'
 import { fileURLToPath } from 'node:url'
+import { getPort } from 'get-port-please'
 import { normalize } from 'pathe'
 import { provider } from 'std-env'
+import { WebSocketServer } from 'ws'
 import { loadProvider } from './load-provider'
 
-const session = new inspector.Session()
+let session: inspector.Session | null = null
 let enabled = false
 
-const mod: CoverageProviderModule = {
-  async startCoverage({ isolate }) {
+const mod: CoverageProviderModule & { wss: WebSocketServer | undefined; extendedContextCoverage: Profiler.ScriptCoverage[] } = {
+  wss: undefined,
+  extendedContextCoverage: [],
+
+  async startCoverage({ isolate, trackProcessAndWorker }) {
     if (isolate === false && enabled) {
       return
     }
 
     enabled = true
 
+    if (trackProcessAndWorker) {
+      const port = await getPort()
+      this.wss = new WebSocketServer({ port })
+
+      this.wss.on('connection', socket => socket.on('message', (raw) => {
+        const result: ScriptCoverageWithOffset[] = JSON.parse(raw.toString())
+        this.extendedContextCoverage.push(...(result || []))
+      }))
+
+      process.env.NODE_OPTIONS ||= ''
+      process.env.NODE_OPTIONS += ' --import @vitest/coverage-v8/intercept-new-run-context'
+      process.env.VITEST_WS_PORT = `${port}`
+    }
+
+    session ||= new inspector.Session()
     session.connect()
     await session.post('Profiler.enable')
     await session.post('Profiler.startPreciseCoverage', { callCount: true, detailed: true })
@@ -28,11 +48,16 @@ const mod: CoverageProviderModule = {
       return { result: [] }
     }
 
+    if (!session) {
+      throw new Error('V8 provider missing inspector session.')
+    }
+
+    this.wss?.clients.forEach(client => client.send('take-coverage'))
     const coverage = await session.post('Profiler.takePreciseCoverage')
     const result: ScriptCoverageWithOffset[] = []
 
     // Reduce amount of data sent over rpc by doing some early result filtering
-    for (const entry of coverage.result) {
+    for (const entry of [...coverage.result, ...this.extendedContextCoverage.splice(0)]) {
       if (filterResult(entry)) {
         result.push({
           ...entry,
@@ -49,9 +74,14 @@ const mod: CoverageProviderModule = {
       return
     }
 
+    if (!session) {
+      throw new Error('V8 provider missing inspector session.')
+    }
+
     await session.post('Profiler.stopPreciseCoverage')
     await session.post('Profiler.disable')
     session.disconnect()
+    this.wss?.close()
   },
 
   async getProvider(): Promise<V8CoverageProvider> {

packages/coverage-v8/src/intercept-new-run-context.ts

@@ -0,0 +1,51 @@
+import type { ScriptCoverageWithOffset } from './provider'
+import { WebSocket } from 'ws'
+import provider from './index'
+
+// eslint-disable-next-line antfu/no-top-level-await -- This should be blocking module loading
+await initialize().catch((error) => {
+  console.error('[vitest-coverage] Error initializing process/thread intercepting:', error)
+  throw error
+})
+
+async function initialize() {
+  let reportedCoverage = false
+
+  const ws = new WebSocket(`ws://localhost:${Number(process.env.VITEST_WS_PORT)}`)
+
+  // @ts-expect-error -- untyped
+  ws.on('open', () => ws._socket?.unref?.())
+
+  await provider.startCoverage?.({
+    isolate: true,
+
+    // Environment options that were set by parent should inherit, no need to add more ws servers
+    trackProcessAndWorker: false,
+  })
+
+  onMessage(message => message === 'take-coverage' && takeCoverage())
+  process.on('beforeExit', takeCoverage)
+
+  async function takeCoverage() {
+    if (reportedCoverage) {
+      return
+    }
+
+    reportedCoverage = true
+
+    const coverage = await provider.takeCoverage?.({
+    // Start offset should be 0 as these run outside of Vite
+      moduleExecutionInfo: undefined,
+    }) as { result: ScriptCoverageWithOffset[] }
+
+    ws.send(JSON.stringify(coverage.result.map(entry => ({ ...entry, isExtendedContext: true }))))
+
+    await provider.stopCoverage?.({ isolate: true })
+
+    ws.close()
+  }
+
+  async function onMessage(callback: (message: unknown) => void) {
+    ws.on('message', raw => callback(raw.toString()))
+  }
+}

packages/coverage-v8/src/provider.ts

@@ -20,6 +20,9 @@ import { version } from '../package.json' with { type: 'json' }
 
 export interface ScriptCoverageWithOffset extends Profiler.ScriptCoverage {
   startOffset: number
+
+  /** Whether script ran outside Vite, e.g. in sub-processes or worker threads */
+  isExtendedContext?: boolean
 }
 
 interface RawCoverage { result: ScriptCoverageWithOffset[] }
@@ -331,8 +334,9 @@ export class V8CoverageProvider extends BaseCoverageProvider implements Coverage
 
   private async getSources(
     url: string,
-    onTransform: (filepath: string) => Promise<Vite.TransformResult | undefined | null>,
+    onTransform: (filepath: string, isExtendedContext?: ScriptCoverageWithOffset['isExtendedContext']) => Promise<Vite.TransformResult | undefined | null>,
     functions: Profiler.FunctionCoverage[] = [],
+    isExtendedContext: ScriptCoverageWithOffset['isExtendedContext'] = false,
   ): Promise<{
     code: string
     map?: Vite.Rollup.SourceMap
@@ -342,7 +346,7 @@ export class V8CoverageProvider extends BaseCoverageProvider implements Coverage
       ? url.slice(8)
       : removeStartsWith(url, FILE_PROTOCOL)
     // TODO: do we still need to "catch" here? why would it fail?
-    const transformResult = await onTransform(filepath).catch(() => null)
+    const transformResult = await onTransform(filepath, isExtendedContext).catch(() => null)
 
     const map = transformResult?.map as Vite.Rollup.SourceMap | undefined
     const code = transformResult?.code
@@ -385,8 +389,8 @@ export class V8CoverageProvider extends BaseCoverageProvider implements Coverage
       throw new Error(`Cannot access browser module graph because it was torn down.`)
     }
 
-    const onTransform = async (filepath: string) => {
-      const result = await this.transformFile(filepath, project, environment)
+    const onTransform = async (filepath: string, isExtendedContext: ScriptCoverageWithOffset['isExtendedContext'] = false) => {
+      const result = await this.transformFile(filepath, project, environment, !isExtendedContext)
       if (result && environment === '__browser__' && project.browser) {
         return { ...result, code: `${result.code}// <inline-source-map>` }
       }
@@ -423,7 +427,7 @@ export class V8CoverageProvider extends BaseCoverageProvider implements Coverage
       }
 
       await Promise.all(
-        chunk.map(async ({ url, functions, startOffset }) => {
+        chunk.map(async ({ url, functions, startOffset, isExtendedContext }) => {
           let timeout: ReturnType<typeof setTimeout> | undefined
           let start: number | undefined
 
@@ -436,6 +440,7 @@ export class V8CoverageProvider extends BaseCoverageProvider implements Coverage
             url,
             onTransform,
             functions,
+            isExtendedContext,
           )
 
           coverageMap.merge(await this.remapCoverage(

packages/vitest/src/defaults.ts

@@ -53,6 +53,7 @@ export const coverageConfigDefaults: Required<Pick<CoverageOptions, FieldsWithDe
     branches: [50, 80],
     lines: [50, 80],
   },
+  trackProcessAndWorker: false,
 }
 
 export const fakeTimersDefaults: NonNullable<UserConfig['fakeTimers']> = {

packages/vitest/src/integrations/coverage.ts

@@ -10,7 +10,10 @@ export async function startCoverageInsideWorker(
   const coverageModule = await resolveCoverageProviderModule(options, loader)
 
   if (coverageModule) {
-    return coverageModule.startCoverage?.(runtimeOptions)
+    return coverageModule.startCoverage?.({
+      ...runtimeOptions,
+      trackProcessAndWorker: options?.trackProcessAndWorker ?? false,
+    })
   }
 
   return null

packages/vitest/src/node/cli/cli-config.ts

@@ -326,6 +326,9 @@ export const cliOptionsConfig: VitestCLIOptions = {
         description: 'Directory of HTML coverage output to be served in UI mode and HTML reporter.',
         argument: '<path>',
       },
+      trackProcessAndWorker: {
+        description: 'Track coverage of the `node:child_process` and `node:worker_threads` spawned during test run. Supported only by `v8` provider. (default: false)',
+      },
     },
   },
   mode: {

packages/vitest/src/node/config/serializeConfig.ts

@@ -58,6 +58,7 @@ export function serializeConfig(project: TestProject): SerializedConfig {
           ? coverage.customProviderModule
           : undefined,
         htmlDir: coverage.htmlDir,
+        trackProcessAndWorker: coverage.trackProcessAndWorker ?? false,
       }
     })(config.coverage),
     fakeTimers: config.fakeTimers,

packages/vitest/src/node/coverage.ts

@@ -663,11 +663,11 @@ export class BaseCoverageProvider {
   // TODO: should this be abstracted in `project`/`vitest` instead?
   // if we decide to keep `viteModuleRunner: false`, we will need to abstract transformation in both main thread and tests
   // custom --import=module.registerHooks need to be transformed as well somehow
-  async transformFile(url: string, project: TestProject, viteEnvironment: string): Promise<TransformResult | null | undefined> {
+  async transformFile(url: string, project: TestProject, viteEnvironment: string, isTransformedByVite = true): Promise<TransformResult | null | undefined> {
     const config = project.config
 
     // vite is disabled, should transform manually if possible
-    if (config.experimental.viteModuleRunner === false) {
+    if (config.experimental.viteModuleRunner === false || !isTransformedByVite) {
       const pathname = url.split('?')[0]
       const filename = pathname.startsWith('file://') ? fileURLToPath(pathname) : pathname
       const extension = path.extname(filename)