feat(experimental): add message option to vi.waitFor and vi.waitUntil (#11754)

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: hi-ogawa

docs/api/expect.md

@@ -118,7 +118,7 @@ The callback receives an `AbortSignal` that is aborted when the poll timeout is
 
 If an error is thrown inside the `expect.poll` callback, Vitest will retry again until the timeout runs out.
 
-The `message` option describes what the poll waits for. If the poll fails, Vitest adds the message to the start of the error.
+The `message` option describes what the poll waits for. If the poll fails, Vitest adds the message to the start of the error. The message also names the step in the [trace timeline](/config/experimental#experimental-tracetimeline).
 
 ```ts
 await expect.poll(() => queue.count(), { message: 'queue drained' }).toBe(0)

docs/api/vi.md

@@ -1362,6 +1362,13 @@ test('Element exists in a DOM', async () => {
 })
 ```
 
+The `message` option describes what the callback waits for. If the wait fails, Vitest adds the message to the start of the error. The message also names the step in the [trace timeline](/config/experimental#experimental-tracetimeline).
+
+```ts
+await vi.waitFor(() => server.ping(), { message: 'server started' })
+// Error: server started: connect ECONNREFUSED
+```
+
 If `vi.useFakeTimers` is used, `vi.waitFor` automatically calls `vi.advanceTimersByTime(interval)` in every check callback. See [Automatic Advancing](/guide/mocking/timers#automatic-advancing) for details.
 
 ### vi.waitUntil {#vi-waituntil}
@@ -1394,6 +1401,13 @@ test('Element render correctly', async () => {
 })
 ```
 
+Like `vi.waitFor`, `vi.waitUntil` accepts a `message` option that prefixes the error and names the step in the [trace timeline](/config/experimental#experimental-tracetimeline).
+
+```ts
+await vi.waitUntil(() => db.users.findByName('alice'), { message: 'user saved' })
+// Error: user saved: Timed out in waitUntil!
+```
+
 If `vi.useFakeTimers` is used, `vi.waitUntil` automatically calls `vi.advanceTimersByTime(interval)` in every check callback. See [Automatic Advancing](/guide/mocking/timers#automatic-advancing) for details.
 
 ### vi.hoisted {#vi-hoisted}

docs/config/experimental.md

@@ -485,6 +485,8 @@ Steps are reported only for failed attempts, so the option adds almost no overhe
 
 Step locations are resolved like error stacks, so a step in a helper file points at that file, and a step in a [`vi.defineHelper`](/api/vi#vi-definehelper) helper points at the line that called the helper.
 
+The `message` option of `expect.poll`, `vi.waitFor`, and `vi.waitUntil` names the step, as with `user saved (vi.waitUntil)` in the report above.
+
 This option has no effect in browser mode, which prints the same timeline with [`browser.traceView`](/config/browser/traceview).
 
 ::: warning

docs/config/snippets/trace-timeline-failure.ansi

@@ -9,7 +9,7 @@
 
  ❯ src/search.test.ts:8:50
       6|   await api.post('/users', { name: 'alice' })
-      7|   await vi.waitUntil(() => db.users.findByName('alice'))
+      7|   await vi.waitUntil(() => db.users.findByName('alice'), { message: 'user saved' })
       8|   await expect.poll(() => search.query('alice')).toHaveLength(1)
        |                                                  ^
       9| })
@@ -19,5 +19,5 @@
 Trace:
   step                           start  duration  location
   vi.waitFor()                    +2ms     120ms  src/search.test.ts:5:12
-  vi.waitUntil()                +171ms      42ms  src/search.test.ts:7:12
+  user saved (vi.waitUntil)     +171ms      42ms  src/search.test.ts:7:12
   expect.poll().toHaveLength()  +214ms    1001ms  src/search.test.ts:8:50

packages/vitest/src/integrations/chai/poll.ts

@@ -265,7 +265,9 @@ function setTraceHandlers(assertion: Assertion, test: Test) {
   // so this and `getTraceName` in packages/browser/src/client/tester/expect-element.ts can share it
   const getName = (assertion: Assertion) => {
     const negated = chai.util.flag(assertion, 'negate') ? 'not.' : ''
-    return `expect.poll().${negated}${chai.util.flag(assertion, '_name')}()`
+    const message = chai.util.flag(assertion, 'message')
+    const name = `expect.poll().${negated}${chai.util.flag(assertion, '_name')}`
+    return message ? `${message} (${name})` : `${name}()`
   }
   const onStart: PollOnStart = ({ assertion }) => {
     recordTraceEntry(

packages/vitest/src/integrations/wait.ts

@@ -18,17 +18,47 @@ export interface WaitForOptions {
    * @default 1000ms
    */
   timeout?: number
+  /**
+   * @description Message that describes what is being waited for. It prefixes the error and names the trace step
+   */
+  message?: string
+}
+
+function prefixMessage(error: unknown, message: string | undefined) {
+  if (message && error instanceof Error) {
+    const original = error.message
+    try {
+      error.message = `${message}: ${original}`
+    } catch {
+      // message can be read-only, e.g. DOMException
+      return error
+    }
+    if (error.stack !== undefined) {
+      error.stack = error.stack.replace(original, error.message)
+    }
+  }
+  return error
+}
+
+function getTraceName(name: string, message: string | undefined) {
+  return message ? `${message} (${name})` : `${name}()`
 }
 
 export function waitFor<T>(
   callback: WaitForCallback<T>,
   options: number | WaitForOptions = {},
 ): Promise<T> {
   const { setTimeout, setInterval, clearTimeout, clearInterval } = getSafeTimers()
-  const { interval = 50, timeout = 1000 } =
-    typeof options === 'number' ? { timeout: options } : options
+  const {
+    interval = 50,
+    timeout = 1000,
+    message,
+  } = typeof options === 'number' ? { timeout: options } : options
   const STACK_TRACE_ERROR = new Error('STACK_TRACE_ERROR')
-  const trackPromise = startPromiseTraceRange('vi.waitFor()', STACK_TRACE_ERROR)
+  const trackPromise = startPromiseTraceRange(
+    getTraceName('vi.waitFor', message),
+    STACK_TRACE_ERROR,
+  )
 
   const promise = new Promise<T>((resolve, reject) => {
     let lastError: unknown
@@ -56,7 +86,7 @@ export function waitFor<T>(
         error = copyStackTrace(new Error('Timed out in waitFor!'), STACK_TRACE_ERROR)
       }
 
-      reject(error)
+      reject(prefixMessage(error, message))
     }
 
     const checkCallback = () => {
@@ -107,7 +137,10 @@ export function waitFor<T>(
 
 export type WaitUntilCallback<T> = () => T | Promise<T>
 
-export interface WaitUntilOptions extends Pick<WaitForOptions, 'interval' | 'timeout'> {}
+export interface WaitUntilOptions extends Pick<
+  WaitForOptions,
+  'interval' | 'timeout' | 'message'
+> {}
 
 type Truthy<T> = T extends false | '' | 0 | null | undefined ? never : T
 
@@ -116,10 +149,16 @@ export function waitUntil<T>(
   options: number | WaitUntilOptions = {},
 ): Promise<Truthy<T>> {
   const { setTimeout, setInterval, clearTimeout, clearInterval } = getSafeTimers()
-  const { interval = 50, timeout = 1000 } =
-    typeof options === 'number' ? { timeout: options } : options
+  const {
+    interval = 50,
+    timeout = 1000,
+    message,
+  } = typeof options === 'number' ? { timeout: options } : options
   const STACK_TRACE_ERROR = new Error('STACK_TRACE_ERROR')
-  const trackPromise = startPromiseTraceRange('vi.waitUntil()', STACK_TRACE_ERROR)
+  const trackPromise = startPromiseTraceRange(
+    getTraceName('vi.waitUntil', message),
+    STACK_TRACE_ERROR,
+  )
 
   const promise = new Promise<Truthy<T>>((resolve, reject) => {
     let promiseStatus: 'idle' | 'pending' | 'resolved' | 'rejected' = 'idle'
@@ -133,7 +172,7 @@ export function waitUntil<T>(
       if (!error) {
         error = copyStackTrace(new Error('Timed out in waitUntil!'), STACK_TRACE_ERROR)
       }
-      reject(error)
+      reject(prefixMessage(error, message))
     }
 
     const onResolve = (result: T) => {

test/e2e/test/trace-timeline.test.ts

@@ -117,6 +117,45 @@ test('prints vi.waitFor and vi.waitUntil steps on failure', async () => {
   `)
 })
 
+test('names steps with message', async () => {
+  const { stderr, errorTree } = await runInlineTests(
+    {
+      'basic.test.ts': /* ts */ `
+      import { expect, test, vi } from 'vitest'
+
+      test('fail', async () => {
+        await vi.waitFor(() => {}, { message: 'server ready' })
+        await vi.waitUntil(() => true, { message: 'user saved' })
+        await expect.poll(() => 1, { timeout: 100, message: 'user indexed' }).toBe(10)
+      })
+    `,
+    },
+    {
+      reporters: 'default',
+      experimental: { traceTimeline: true },
+    },
+  )
+
+  expect(errorTree()).toMatchInlineSnapshot(`
+    {
+      "basic.test.ts": {
+        "fail": [
+          "user indexed: expected 1 to be 10 // Object.is equality",
+        ],
+      },
+    }
+  `)
+  expect(extractTimelines(stderr)).toMatchInlineSnapshot(`
+    [
+      "Trace:
+    step                               start    duration  location
+    server ready (vi.waitFor)          +<time>  <time>    basic.test.ts:5:18
+    user saved (vi.waitUntil)          +<time>  <time>    basic.test.ts:6:18
+    user indexed (expect.poll().toBe)  +<time>  <time>    basic.test.ts:7:79",
+    ]
+  `)
+})
+
 test('skips waits while tests run concurrently', async () => {
   const { stderr, errorTree } = await runInlineTests(
     {

test/unit/test/wait.test.ts

@@ -28,6 +28,31 @@ describe('waitFor', () => {
 
       expect(callback).toHaveBeenCalledTimes(2)
     })
+
+    test('message', async () => {
+      await expect(
+        vi.waitFor(
+          () => {
+            throw new Error('server not ready')
+          },
+          { timeout: 50, message: 'server started' },
+        ),
+      ).rejects.toThrowErrorMatchingInlineSnapshot(`[Error: server started: server not ready]`)
+
+      await expect(
+        vi.waitFor(() => new Promise(() => {}), { timeout: 50, message: 'server started' }),
+      ).rejects.toThrowErrorMatchingInlineSnapshot(`[Error: server started: Timed out in waitFor!]`)
+
+      // read-only message is left as is
+      await expect(
+        vi.waitFor(
+          () => {
+            throw new DOMException('aborted', 'AbortError')
+          },
+          { timeout: 50, message: 'server started' },
+        ),
+      ).rejects.toThrow(expect.objectContaining({ name: 'AbortError', message: 'aborted' }))
+    })
   })
 
   test('basic', async () => {
@@ -149,6 +174,31 @@ describe('waitUntil', () => {
 
       expect(callback).toHaveBeenCalledTimes(2)
     })
+
+    test('message', async () => {
+      await expect(
+        vi.waitUntil(() => false, { timeout: 50, message: 'user saved' }),
+      ).rejects.toThrowErrorMatchingInlineSnapshot(`[Error: user saved: Timed out in waitUntil!]`)
+
+      await expect(
+        vi.waitUntil(
+          () => {
+            throw new Error('db closed')
+          },
+          { timeout: 50, message: 'user saved' },
+        ),
+      ).rejects.toThrowErrorMatchingInlineSnapshot(`[Error: user saved: db closed]`)
+
+      // read-only message is left as is
+      await expect(
+        vi.waitUntil(
+          () => {
+            throw new DOMException('aborted', 'AbortError')
+          },
+          { timeout: 50, message: 'user saved' },
+        ),
+      ).rejects.toThrow(expect.objectContaining({ name: 'AbortError', message: 'aborted' }))
+    })
   })
 
   test('basic', async () => {