fix: cancelling should stop current test immediately

docs/api/advanced/vitest.md

@@ -349,7 +349,7 @@ This makes this method very slow, unless you disable isolation before collecting
 function cancelCurrentRun(reason: CancelReason): Promise<void>
 ```
 
-This method will gracefully cancel all ongoing tests. It will wait for started tests to finish running and will not run tests that were scheduled to run but haven't started yet.
+This method will gracefully cancel all ongoing tests. It will stop the on-going tests and will not run tests that were scheduled to run but haven't started yet.
 
 ## setGlobalTestNamePattern
 

packages/runner/src/context.ts

@@ -109,6 +109,31 @@ export function withTimeout<T extends (...args: any[]) => any>(
   }) as T
 }
 
+export function withCancel<T extends (...args: any[]) => any>(
+  fn: T,
+  signal: AbortSignal,
+): T {
+  return (function runWithCancel(...args: T extends (...args: infer A) => any ? A : never) {
+    return new Promise((resolve, reject) => {
+      signal.addEventListener('abort', () => reject(signal.reason))
+
+      try {
+        const result = fn(...args) as PromiseLike<unknown>
+
+        if (typeof result === 'object' && result != null && typeof result.then === 'function') {
+          result.then(resolve, reject)
+        }
+        else {
+          resolve(result)
+        }
+      }
+      catch (error) {
+        reject(error)
+      }
+    })
+  }) as T
+}
+
 const abortControllers = new WeakMap<TestContext, AbortController>()
 
 export function abortIfTimeout([context]: [TestContext?], error: Error): void {

packages/runner/src/run.ts

@@ -792,6 +792,12 @@ function failTask(result: TaskResult, err: unknown, diffOptions: DiffOptions | u
     return
   }
 
+  if (err instanceof TestRunAbortError) {
+    result.state = 'skip'
+    result.note = err.message
+    return
+  }
+
   result.state = 'fail'
   const errors = Array.isArray(err) ? err : [err]
   for (const e of errors) {
@@ -814,6 +820,20 @@ function markTasksAsSkipped(suite: Suite, runner: VitestRunner) {
   })
 }
 
+function markPendingTasksAsSkipped(suite: Suite, runner: VitestRunner, note?: string) {
+  suite.tasks.forEach((t) => {
+    if (!t.result || t.result.state === 'run') {
+      t.mode = 'skip'
+      t.result = { ...t.result, state: 'skip', note }
+      updateTask('test-cancel', t, runner)
+    }
+
+    if (t.type === 'suite') {
+      markPendingTasksAsSkipped(t, runner, note)
+    }
+  })
+}
+
 export async function runSuite(suite: Suite, runner: VitestRunner): Promise<void> {
   await runner.onBeforeRunSuite?.(suite)
 
@@ -1028,8 +1048,10 @@ export async function startTests(specs: string[] | FileSpecification[], runner:
   runner.cancel = (reason) => {
     // We intentionally create only one error since there is only one test run that can be cancelled
     const error = new TestRunAbortError('The test run was aborted by the user.', reason)
-    getRunningTests().forEach(test =>
-      abortContextSignal(test.context, error),
+    getRunningTests().forEach((test) => {
+      abortContextSignal(test.context, error)
+      markPendingTasksAsSkipped(test.file, runner, error.message)
+    },
     )
     return cancel?.(reason)
   }

packages/runner/src/suite.ts

@@ -32,6 +32,7 @@ import {
   collectTask,
   createTestContext,
   runWithSuite,
+  withCancel,
   withTimeout,
 } from './context'
 import { configureProps, TestFixtures, withFixtures } from './fixture'
@@ -412,7 +413,7 @@ function createSuiteCollector(
       setFn(
         task,
         withTimeout(
-          withAwaitAsyncAssertions(withFixtures(handler, { context }), task),
+          withCancel(withAwaitAsyncAssertions(withFixtures(handler, { context }), task), task.context.signal),
           timeout,
           false,
           stackTraceError,

packages/runner/src/types/tasks.ts

@@ -259,6 +259,7 @@ export type TaskUpdateEvent
     | 'test-prepare'
     | 'test-finished'
     | 'test-retried'
+    | 'test-cancel'
     | 'suite-prepare'
     | 'suite-finished'
     | 'before-hook-start'

packages/vitest/src/node/test-run.ts

@@ -232,6 +232,11 @@ export class TestRun {
       return
     }
 
+    if (event === 'test-cancel' && entity.type === 'test') {
+      // This is used to just update state of the task
+      return
+    }
+
     if (event === 'test-prepare' && entity.type === 'test') {
       return await this.vitest.report('onTestCaseReady', entity)
     }

test/cli/fixtures/cancel-run/blocked-test-cases.test.ts

@@ -0,0 +1,26 @@
+import { afterEach, describe, test } from 'vitest'
+
+afterEach(async (context) => {
+  (context.task.meta as any).afterEachDone = true
+})
+
+describe('these should pass', () => {
+  test('one', async () => {})
+  test('two', async () => {})
+})
+
+test('this test starts and gets cancelled, its after each should be called', async ({ annotate }) => {
+  await annotate('Running long test, do the cancelling now!')
+
+  await new Promise(resolve => setTimeout(resolve, 100_000))
+})
+
+describe('these should not start but should be skipped', () => {
+  test('third, no after each expected', async () => {})
+
+  describe("nested", () => {
+    test('fourth, no after each expected', async () => {})
+  });
+})
+
+test('fifth, no after each expected', async () => {})

test/cli/fixtures/cancel-run/blocked-thread.test.ts

@@ -0,0 +1,7 @@
+import { execSync } from 'node:child_process'
+import { test } from 'vitest'
+
+test('block whole test runner thread/process', { timeout: 30_000 }, async () => {
+  // Note that this can also block the RPC before onTestCaseReady is emitted to main thread
+  execSync("sleep 40")
+})

test/cli/test/cancel-run.test.ts

@@ -1,21 +1,25 @@
+import type { TestModule } from 'vitest/node'
 import { Readable, Writable } from 'node:stream'
 import { stripVTControlCharacters } from 'node:util'
 import { createDefer } from '@vitest/utils/helpers'
 import { expect, onTestFinished, test, vi } from 'vitest'
 import { createVitest, registerConsoleShortcuts } from 'vitest/node'
 
-test('can force cancel a run', async () => {
+const CTRL_C = '\x03'
+
+test('can force cancel a run via CLI', async () => {
   const onExit = vi.fn<never>()
   const exit = process.exit
   onTestFinished(() => {
     process.exit = exit
   })
   process.exit = onExit
 
-  const onTestCaseReady = createDefer<void>()
+  const onTestModuleStart = createDefer<void>()
   const vitest = await createVitest('test', {
     root: 'fixtures/cancel-run',
-    reporters: [{ onTestCaseReady: () => onTestCaseReady.resolve() }],
+    include: ['blocked-thread.test.ts'],
+    reporters: [{ onTestModuleStart: () => onTestModuleStart.resolve() }],
   })
   onTestFinished(() => vitest.close())
 
@@ -27,17 +31,130 @@ test('can force cancel a run', async () => {
   const onLog = vi.spyOn(vitest.logger, 'log').mockImplementation(() => {})
   const promise = vitest.start()
 
-  await onTestCaseReady
+  await onTestModuleStart
 
   // First CTRL+c should log warning about graceful exit
-  stdin.emit('data', '\x03')
+  stdin.emit('data', CTRL_C)
+
+  // Let the test case start running
+  await new Promise(resolve => setTimeout(resolve, 100))
 
   const logs = onLog.mock.calls.map(log => stripVTControlCharacters(log[0] || '').trim())
   expect(logs).toContain('Cancelling test run. Press CTRL+c again to exit forcefully.')
 
   // Second CTRL+c should stop run
-  stdin.emit('data', '\x03')
+  stdin.emit('data', CTRL_C)
   await promise
 
   expect(onExit).toHaveBeenCalled()
 })
+
+test('cancelling test run stops test execution immediately', async () => {
+  const onTestRunEnd = createDefer<readonly TestModule[]>()
+  const onTestCaseReady = createDefer<void>()
+  const onTestCaseHooks: string[] = []
+
+  const vitest = await createVitest('test', {
+    root: 'fixtures/cancel-run',
+    include: ['blocked-test-cases.test.ts'],
+    reporters: [{
+      onTestCaseReady(testCase) {
+        onTestCaseHooks.push(`onTestCaseReady ${testCase.name}`)
+      },
+      onTestCaseResult(testCase) {
+        onTestCaseHooks.push(`onTestCaseResult ${testCase.name}`)
+        onTestCaseHooks.push('') // padding
+      },
+      onTestCaseAnnotate: (_, annotation) => {
+        if (annotation.message === 'Running long test, do the cancelling now!') {
+          onTestCaseReady.resolve()
+        }
+      },
+      onTestRunEnd(testModules) {
+        onTestRunEnd.resolve(testModules)
+      },
+    }],
+  })
+  onTestFinished(() => vitest.close())
+
+  const promise = vitest.start()
+
+  await onTestCaseReady
+  await vitest.cancelCurrentRun('keyboard-input')
+
+  const testModules = await onTestRunEnd
+  await Promise.all([vitest.close(), promise])
+
+  expect(testModules).toHaveLength(1)
+
+  const tests = Array.from(testModules[0].children.allTests().map(test => ({
+    name: test.name,
+    status: test.result().state,
+    note: (test.result() as any).note,
+    afterEachRun: (test.meta() as any).afterEachDone === true,
+  })))
+
+  expect(tests).toMatchInlineSnapshot(`
+    [
+      {
+        "afterEachRun": true,
+        "name": "one",
+        "note": undefined,
+        "status": "passed",
+      },
+      {
+        "afterEachRun": true,
+        "name": "two",
+        "note": undefined,
+        "status": "passed",
+      },
+      {
+        "afterEachRun": true,
+        "name": "this test starts and gets cancelled, its after each should be called",
+        "note": "The test run was aborted by the user.",
+        "status": "skipped",
+      },
+      {
+        "afterEachRun": false,
+        "name": "third, no after each expected",
+        "note": "The test run was aborted by the user.",
+        "status": "skipped",
+      },
+      {
+        "afterEachRun": false,
+        "name": "fourth, no after each expected",
+        "note": "The test run was aborted by the user.",
+        "status": "skipped",
+      },
+      {
+        "afterEachRun": false,
+        "name": "fifth, no after each expected",
+        "note": "The test run was aborted by the user.",
+        "status": "skipped",
+      },
+    ]
+  `)
+
+  expect(onTestCaseHooks).toMatchInlineSnapshot(`
+    [
+      "onTestCaseReady one",
+      "onTestCaseResult one",
+      "",
+      "onTestCaseReady two",
+      "onTestCaseResult two",
+      "",
+      "onTestCaseReady this test starts and gets cancelled, its after each should be called",
+      "onTestCaseResult this test starts and gets cancelled, its after each should be called",
+      "",
+      "onTestCaseReady third, no after each expected",
+      "onTestCaseResult third, no after each expected",
+      "",
+      "onTestCaseReady fourth, no after each expected",
+      "onTestCaseResult fourth, no after each expected",
+      "",
+      "onTestCaseReady fifth, no after each expected",
+      "onTestCaseResult fifth, no after each expected",
+      "",
+    ]
+  `)
+})