feat(watch): test run cancelling

packages/runner/src/run.ts

@@ -105,6 +105,14 @@ async function callCleanupHooks(cleanups: HookCleanupCallback[]) {
 }
 
 export async function runTest(test: Test, runner: VitestRunner) {
+  // @ts-expect-error untyped global
+  if (globalThis.__vitest_worker__.isCancelling) {
+    test.mode = 'skip'
+    test.result = { ...test.result, state: 'skip' }
+    updateTask(test, runner)
+    return
+  }
+
   await runner.onBeforeRunTest?.(test)
 
   if (test.mode !== 'run')
@@ -227,6 +235,15 @@ function markTasksAsSkipped(suite: Suite, runner: VitestRunner) {
 }
 
 export async function runSuite(suite: Suite, runner: VitestRunner) {
+  // @ts-expect-error untyped global
+  if (globalThis.__vitest_worker__.isCancelling) {
+    suite.result = { state: 'skip' }
+    suite.mode = 'skip'
+    markTasksAsSkipped(suite, runner)
+    updateTask(suite, runner)
+    return
+  }
+
   await runner.onBeforeRunSuite?.(suite)
 
   if (suite.result?.state === 'fail') {

packages/vitest/src/node/core.ts

@@ -46,6 +46,7 @@ export class Vitest {
   filenamePattern?: string
   runningPromise?: Promise<void>
   closingPromise?: Promise<void>
+  isCancelling = false
 
   isFirstRun = true
   restartsCount = 0
@@ -64,6 +65,7 @@ export class Vitest {
 
   private _onRestartListeners: OnServerRestartHandler[] = []
   private _onSetServer: OnServerRestartHandler[] = []
+  private _onCancelListeners: (() => Promise<void> | void)[] = []
 
   async setServer(options: UserConfig, server: ViteDevServer, cliOptions: UserConfig) {
     this.unregisterWatcher?.()
@@ -394,9 +396,10 @@ export class Vitest {
 
   async runFiles(paths: WorkspaceSpec[]) {
     const filepaths = paths.map(([, file]) => file)
-
     this.state.collectPaths(filepaths)
 
+    this.isCancelling = false
+
     await this.report('onPathsCollected', filepaths)
 
     // previous run
@@ -436,6 +439,13 @@ export class Vitest {
     return await this.runningPromise
   }
 
+  async cancelCurrentRun() {
+    this.isCancelling = true
+    this.logger.log('Cancelling current run...')
+
+    await Promise.all(this._onCancelListeners.splice(0).map(listener => listener()))
+  }
+
   async rerunFiles(files: string[] = this.state.getFilepaths(), trigger?: string) {
     if (this.filenamePattern) {
       const filteredFiles = await this.globTestFiles([this.filenamePattern])
@@ -765,4 +775,8 @@ export class Vitest {
   onAfterSetServer(fn: OnServerRestartHandler) {
     this._onSetServer.push(fn)
   }
+
+  onCancel(fn: () => void) {
+    this._onCancelListeners.push(fn)
+  }
 }

packages/vitest/src/node/pools/browser.ts

@@ -24,6 +24,9 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
     const isolate = project.config.isolate
     if (isolate) {
       for (const path of paths) {
+        if (ctx.isCancelling)
+          break
+
         const url = new URL('/', origin)
         url.searchParams.append('path', path)
         url.searchParams.set('id', path)

packages/vitest/src/node/pools/child.ts

@@ -4,7 +4,7 @@ import { fork } from 'node:child_process'
 import { fileURLToPath, pathToFileURL } from 'node:url'
 import { createBirpc } from 'birpc'
 import { resolve } from 'pathe'
-import type { ContextTestEnvironment, ResolvedConfig, RuntimeRPC, Vitest } from '../../types'
+import type { ContextTestEnvironment, ResolvedConfig, RunnerRPC, RuntimeRPC, Vitest } from '../../types'
 import type { ChildContext } from '../../types/child'
 import type { PoolProcessOptions, ProcessPool, WorkspaceSpec } from '../pool'
 import { distDir } from '../../paths'
@@ -16,7 +16,7 @@ import { createMethodsRPC } from './rpc'
 const childPath = fileURLToPath(pathToFileURL(resolve(distDir, './child.js')).href)
 
 function setupChildProcessChannel(project: WorkspaceProject, fork: ChildProcess): void {
-  createBirpc<{}, RuntimeRPC>(
+  const rpc = createBirpc<RunnerRPC, RuntimeRPC>(
     createMethodsRPC(project),
     {
       serialize: v8.serialize,
@@ -29,6 +29,8 @@ function setupChildProcessChannel(project: WorkspaceProject, fork: ChildProcess)
       },
     },
   )
+
+  project.ctx.onCancel(() => rpc.onCancel())
 }
 
 function stringifyRegex(input: RegExp | string): string {

packages/vitest/src/node/pools/threads.ts

@@ -1,12 +1,13 @@
 import { MessageChannel } from 'node:worker_threads'
 import { cpus } from 'node:os'
 import { pathToFileURL } from 'node:url'
+import EventEmitter from 'node:events'
 import { createBirpc } from 'birpc'
 import { resolve } from 'pathe'
 import type { Options as TinypoolOptions } from 'tinypool'
 import Tinypool from 'tinypool'
 import { distDir } from '../../paths'
-import type { ContextTestEnvironment, ResolvedConfig, RuntimeRPC, Vitest, WorkerContext } from '../../types'
+import type { ContextTestEnvironment, ResolvedConfig, RunnerRPC, RuntimeRPC, Vitest, WorkerContext } from '../../types'
 import type { PoolProcessOptions, ProcessPool, RunWithFiles } from '../pool'
 import { envsOrder, groupFilesByEnv } from '../../utils/test-helpers'
 import { groupBy } from '../../utils/base'
@@ -20,7 +21,7 @@ function createWorkerChannel(project: WorkspaceProject) {
   const port = channel.port2
   const workerPort = channel.port1
 
-  createBirpc<{}, RuntimeRPC>(
+  const rpc = createBirpc<RunnerRPC, RuntimeRPC>(
     createMethodsRPC(project),
     {
       post(v) {
@@ -32,6 +33,8 @@ function createWorkerChannel(project: WorkspaceProject) {
     },
   )
 
+  project.ctx.onCancel(() => rpc.onCancel())
+
   return { workerPort, port }
 }
 
@@ -71,6 +74,10 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt
 
   const pool = new Tinypool(options)
 
+  const signal = new EventEmitter()
+  signal.setMaxListeners(ctx.state.getFiles().length)
+  ctx.onCancel(() => signal.emit('abort'))
+
   const runWithFiles = (name: string): RunWithFiles => {
     let id = 0
 
@@ -87,14 +94,28 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt
         workerId,
       }
       try {
-        await pool.run(data, { transferList: [workerPort], name })
+        await pool.run(data, { transferList: [workerPort], name, signal })
       }
       catch (error) {
         // Worker got stuck and won't terminate - this may cause process to hang
-        if (error instanceof Error && /Failed to terminate worker/.test(error.message))
+        if (error instanceof Error && /Failed to terminate worker/.test(error.message)) {
           ctx.state.addProcessTimeoutCause(`Failed to terminate worker while running ${files.join(', ')}.`)
-        else
+        }
+        else if (ctx.isCancelling && error instanceof Error && /The task has been aborted/.test(error.message)) {
+          // Intentionally cancelled
+          // TODO: This could be a "ctx.state.cancelFiles(files: string[])" instead
+          ctx.state.collectFiles(files.map(filepath => ({
+            filepath,
+            id: filepath,
+            mode: 'skip',
+            name: filepath,
+            tasks: [],
+            type: 'suite',
+          })))
+        }
+        else {
           throw error
+        }
       }
       finally {
         port.close()

packages/vitest/src/node/reporters/base.ts

@@ -13,6 +13,7 @@ const HELP_QUITE = `${c.dim('press ')}${c.bold('q')}${c.dim(' to quit')}`
 
 const WAIT_FOR_CHANGE_PASS = `\n${c.bold(c.inverse(c.green(' PASS ')))}${c.green(' Waiting for file changes...')}`
 const WAIT_FOR_CHANGE_FAIL = `\n${c.bold(c.inverse(c.red(' FAIL ')))}${c.red(' Tests failed. Watching for file changes...')}`
+const WAIT_FOR_CHANGE_CANCELLED = `\n${c.bold(c.inverse(c.red(' CANCELLED ')))}${c.red(' Test run cancelled. Watching for file changes...')}`
 
 const LAST_RUN_LOG_TIMEOUT = 1_500
 
@@ -100,9 +101,13 @@ export abstract class BaseReporter implements Reporter {
   async onWatcherStart(files = this.ctx.state.getFiles(), errors = this.ctx.state.getUnhandledErrors()) {
     this.resetLastRunLog()
 
+    const cancelled = this.ctx.isCancelling
     const failed = errors.length > 0 || hasFailed(files)
     const failedSnap = hasFailedSnapshot(files)
-    if (failed)
+
+    if (cancelled)
+      this.ctx.logger.log(WAIT_FOR_CHANGE_CANCELLED)
+    else if (failed)
       this.ctx.logger.log(WAIT_FOR_CHANGE_FAIL)
     else
       this.ctx.logger.log(WAIT_FOR_CHANGE_PASS)

packages/vitest/src/node/stdin.ts

@@ -37,12 +37,14 @@ export function registerConsoleShortcuts(ctx: Vitest) {
       return
     }
 
-    // is running, ignore keypress
-    if (ctx.runningPromise)
-      return
-
     const name = key?.name
 
+    if (ctx.runningPromise) {
+      if (['space', 'c'].includes(name) || keys.map(key => key[0]).includes(name))
+        await ctx.cancelCurrentRun()
+      return
+    }
+
     // quit
     if (name === 'q')
       return ctx.exit(true)

packages/vitest/src/runtime/child.ts

@@ -2,8 +2,9 @@ import v8 from 'node:v8'
 import { createBirpc } from 'birpc'
 import { parseRegexp } from '@vitest/utils'
 import type { ResolvedConfig } from '../types'
-import type { RuntimeRPC } from '../types/rpc'
+import type { RunnerRPC, RuntimeRPC } from '../types/rpc'
 import type { ChildContext } from '../types/child'
+import { getWorkerState } from '../utils'
 import { mockMap, moduleCache, startViteNode } from './execute'
 import { rpcDone } from './rpc'
 import { setupInspect } from './inspector'
@@ -26,8 +27,12 @@ function init(ctx: ChildContext) {
       environment: 0,
       prepare: performance.now(),
     },
-    rpc: createBirpc<RuntimeRPC>(
-      {},
+    rpc: createBirpc<RuntimeRPC, RunnerRPC>(
+      {
+        onCancel: () => {
+          getWorkerState().isCancelling = true
+        },
+      },
       {
         eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'],
         serialize: v8.serialize,

packages/vitest/src/runtime/entry.ts

@@ -96,6 +96,10 @@ export async function run(files: string[], config: ResolvedConfig, environment:
     workerState.durations.environment = performance.now() - workerState.durations.environment
 
     for (const file of files) {
+      if (workerState.isCancelling) {
+        // TODO: Mark as skipped and break
+      }
+
       // it doesn't matter if running with --threads
       // if running with --no-threads, we usually want to reset everything before running a test
       // but we have --isolate option to disable this

packages/vitest/src/runtime/worker.ts

@@ -1,6 +1,6 @@
 import { createBirpc } from 'birpc'
 import { workerId as poolId } from 'tinypool'
-import type { RuntimeRPC, WorkerContext } from '../types'
+import type { RunnerRPC, RuntimeRPC, WorkerContext } from '../types'
 import { getWorkerState } from '../utils/global'
 import { mockMap, moduleCache, startViteNode } from './execute'
 import { setupInspect } from './inspector'
@@ -28,8 +28,12 @@ function init(ctx: WorkerContext) {
       environment: 0,
       prepare: performance.now(),
     },
-    rpc: createBirpc<RuntimeRPC>(
-      {},
+    rpc: createBirpc<RuntimeRPC, RunnerRPC>(
+      {
+        onCancel: () => {
+          getWorkerState().isCancelling = true
+        },
+      },
       {
         eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'],
         post(v) { port.postMessage(v) },

packages/vitest/src/types/rpc.ts

@@ -24,6 +24,11 @@ export interface RuntimeRPC {
   resolveSnapshotPath: (testPath: string) => string
 }
 
+export interface RunnerRPC {
+  // TODO: This could be "(reason: 'INPUT' | 'BAIL') => void" instead
+  onCancel: () => void
+}
+
 export interface ContextTestEnvironment {
   name: VitestEnvironment
   options: EnvironmentOptions | null

packages/vitest/src/types/worker.ts

@@ -24,6 +24,7 @@ export interface WorkerGlobalState {
   current?: Test
   filepath?: string
   environmentTeardownRun?: boolean
+  isCancelling?: boolean
   moduleCache: ModuleCacheMap
   mockMap: MockMap
   durations: {