fix: add handshake timeout to iframe communication

packages/browser/src/client/channel.ts

@@ -70,6 +70,8 @@ export type IframeChannelEvent
   = | IframeChannelIncomingEvent
     | IframeChannelOutgoingEvent
 
+export type IframeReceivedEvent = IframeChannelEvent & { messageId: number }
+
 export const channel: BroadcastChannel = new BroadcastChannel(
   `vitest:${getBrowserState().sessionId}`,
 )

packages/browser/src/client/orchestrator.ts

@@ -1,5 +1,5 @@
 import type { Context as OTELContext } from '@opentelemetry/api'
-import type { GlobalChannelIncomingEvent, IframeChannelEvent, IframeChannelOutgoingEvent, IframeViewportDoneEvent, IframeViewportFailEvent } from '@vitest/browser/client'
+import type { GlobalChannelIncomingEvent, IframeChannelEvent, IframeChannelOutgoingEvent, IframeReceivedEvent, IframeViewportDoneEvent, IframeViewportFailEvent } from '@vitest/browser/client'
 import type { BrowserTesterOptions, SerializedConfig } from 'vitest'
 import type { FileSpecification } from 'vitest/internal/browser'
 import { channel, client, globalChannel } from '@vitest/browser/client'
@@ -17,7 +17,8 @@ export class IframeOrchestrator {
   private recreateNonIsolatedIframe = false
   private iframes = new Map<string, HTMLIFrameElement>()
   private readyIframes = new Set<string>()
-  private readyWaiters = new Map<string, () => void>()
+  private readyWaiters = new Map<string, { resolve: () => void; reject: (error: Error) => void }>()
+  private messageId = 0
 
   public eventTarget: EventTarget = new EventTarget()
 
@@ -290,7 +291,7 @@ export class IframeOrchestrator {
     const waiter = this.readyWaiters.get(iframeId)
     if (waiter) {
       this.readyWaiters.delete(iframeId)
-      waiter()
+      waiter.resolve()
     }
   }
 
@@ -299,16 +300,41 @@ export class IframeOrchestrator {
       return Promise.resolve()
     }
 
-    return new Promise((resolve) => {
-      this.readyWaiters.set(iframeId, resolve)
+    return new Promise<void>((resolve, reject) => {
+      const timeout = getIframeTimeout()
+      // the tester reports readiness as soon as its module evaluates; if it
+      // never does (e.g. it threw during bootstrap), don't wait forever
+      const timer = setTimeout(() => {
+        this.readyWaiters.delete(iframeId)
+        reject(new Error(
+          `The iframe "${iframeId}" did not become ready within ${timeout}ms. `
+          + `The tester likely failed to initialize, check the browser console for errors.`,
+        ))
+      }, timeout)
+
+      this.readyWaiters.set(iframeId, {
+        resolve: () => {
+          clearTimeout(timer)
+          resolve()
+        },
+        reject: (error) => {
+          clearTimeout(timer)
+          reject(error)
+        },
+      })
     })
   }
 
   private removeIframe(iframeId: string) {
     const iframe = this.iframes.get(iframeId)
     this.iframes.delete(iframeId)
     this.readyIframes.delete(iframeId)
-    this.readyWaiters.delete(iframeId)
+    const waiter = this.readyWaiters.get(iframeId)
+    if (waiter) {
+      this.readyWaiters.delete(iframeId)
+      // surface an error instead of silently abandoning whoever awaits readiness
+      waiter.reject(new Error(`The iframe "${iframeId}" was removed before it became ready.`))
+    }
     iframe?.remove()
   }
 
@@ -432,10 +458,11 @@ export class IframeOrchestrator {
         break
       }
       default: {
-        // ignore responses
+        // ignore acknowledgements and responses to events we sent
+        const event = e.data.event
         if (
-          typeof e.data.event === 'string'
-          && (e.data.event as string).startsWith('response:')
+          typeof event === 'string'
+          && (event.startsWith('response:') || event.startsWith('ack:'))
         ) {
           break
         }
@@ -465,25 +492,51 @@ export class IframeOrchestrator {
     }
     events.add(event.event)
 
-    channel.postMessage(event)
+    const messageId = this.messageId++
+    channel.postMessage({ ...event, messageId } satisfies IframeReceivedEvent)
+
     return new Promise<void>((resolve, reject) => {
+      let ackTimer: ReturnType<typeof setTimeout>
+
       const cleanupEvents = () => {
+        clearTimeout(ackTimer)
         channel.removeEventListener('message', onReceived)
         this.eventTarget.removeEventListener('iframeerror', onError)
+        events!.delete(event.event)
       }
 
+      // The tester acknowledges the message as soon as it receives it, then
+      // sends the actual response once the work is done. We only time out
+      // waiting for the acknowledgement: it proves the tester is alive, after
+      // which the work (e.g. running a whole test file) may take any amount of
+      // time, so there is intentionally no deadline on the response itself.
+      const timeout = getIframeTimeout()
+      ackTimer = setTimeout(() => {
+        cleanupEvents()
+        reject(new Error(
+          `The iframe "${event.iframeId}" did not acknowledge the "${event.event}" message within ${timeout}ms. `
+          + `The tester might have crashed, been removed, or be blocked by a long synchronous task.`,
+        ))
+      }, timeout)
+
       function onReceived(e: MessageEvent) {
-        if (e.data.iframeId === event.iframeId && e.data.event === `response:${event.event}`) {
-          resolve()
+        if (e.data.iframeId !== event.iframeId || e.data.messageId !== messageId) {
+          return
+        }
+        if (e.data.event === `ack:${event.event}`) {
+          // alive and processing: wait for the response without a deadline
+          clearTimeout(ackTimer)
+          return
+        }
+        if (e.data.event === `response:${event.event}`) {
           cleanupEvents()
-          events!.delete(event.event)
+          resolve()
         }
       }
 
       function onError(e: Event) {
-        reject((e as CustomEvent).detail)
         cleanupEvents()
-        events!.delete(event.event)
+        reject((e as CustomEvent).detail)
       }
 
       this.eventTarget.addEventListener('iframeerror', onError)
@@ -551,3 +604,10 @@ function debug(...args: unknown[]) {
     client.rpc.debug(...args.map(String))
   }
 }
+
+// Liveness timeout for tester iframes (readiness and message acknowledgement),
+// not a timeout for the test work itself. Overridable via the `VITEST_BROWSER_IFRAME_TIMEOUT`
+// env in case a tester legitimately needs longer to boot or acknowledge.
+function getIframeTimeout(): number {
+  return Number(getConfig().env.VITEST_BROWSER_IFRAME_TIMEOUT) || 60_000
+}

packages/browser/src/client/tester/tester.ts

@@ -1,4 +1,4 @@
-import type { BrowserRPC, IframeChannelEvent } from '@vitest/browser/client'
+import type { BrowserRPC, IframeReceivedEvent } from '@vitest/browser/client'
 import type { FileSpecification } from 'vitest/internal/browser'
 import { channel, client, onCancel, registerPageMarkHandler } from '@vitest/browser/client'
 import { parse } from 'flatted'
@@ -35,12 +35,10 @@ let rootTesterSpan: ReturnType<Traces['startContextSpan']> | undefined
 getBrowserState().traces = traces
 
 channel.addEventListener('message', async (e) => {
-  await client.waitForConnection()
-
   const data = e.data
-  debug?.('event from orchestrator', JSON.stringify(e.data))
 
   if (!isEvent(data)) {
+    await client.waitForConnection()
     const error = new Error(`Unknown message: ${JSON.stringify(e.data)}`)
     unhandledError(error, 'Unknown Iframe Message')
     return
@@ -51,60 +49,83 @@ channel.addEventListener('message', async (e) => {
     return
   }
 
-  switch (data.event) {
-    case 'execute': {
-      const { method, files, context, concurrencyId, workerId } = data
-      const state = getWorkerState()
-      const parsedContext = parse(context)
-
-      state.ctx.concurrencyId = concurrencyId
-      state.ctx.workerId = workerId
-      state.ctx.providedContext = parsedContext
-      state.providedContext = parsedContext
-      state.metaEnv.VITEST_POOL_ID = String(concurrencyId)
-      state.metaEnv.VITEST_WORKER_ID = String(workerId)
-
-      if (method === 'collect') {
-        await executeTests('collect', files).catch(err => unhandledError(err, 'Collect Error'))
+  // tell the orchestrator we received the event before doing any work (which
+  // may be long-running or gated on the connection), so it can tell a busy
+  // tester apart from a crashed one. See `sendEventToIframe` in orchestrator.ts.
+  channel.postMessage({
+    event: `ack:${data.event}`,
+    iframeId: data.iframeId,
+    messageId: data.messageId,
+  })
+
+  await client.waitForConnection()
+  debug?.('event from orchestrator', JSON.stringify(e.data))
+
+  try {
+    switch (data.event) {
+      case 'execute': {
+        const { method, files, context, concurrencyId, workerId } = data
+        const state = getWorkerState()
+        const parsedContext = parse(context)
+
+        state.ctx.concurrencyId = concurrencyId
+        state.ctx.workerId = workerId
+        state.ctx.providedContext = parsedContext
+        state.providedContext = parsedContext
+        state.metaEnv.VITEST_POOL_ID = String(concurrencyId)
+        state.metaEnv.VITEST_WORKER_ID = String(workerId)
+
+        if (method === 'collect') {
+          await executeTests('collect', files).catch(err => unhandledError(err, 'Collect Error'))
+        }
+        else {
+          await executeTests('run', files).catch(err => unhandledError(err, 'Run Error'))
+        }
+        break
       }
-      else {
-        await executeTests('run', files).catch(err => unhandledError(err, 'Run Error'))
+      case 'cleanup': {
+        await cleanup().catch(err => unhandledError(err, 'Cleanup Error'))
+        rootTesterSpan?.span.end()
+        await traces.finish()
+        break
+      }
+      case 'prepare': {
+        await traces.waitInit()
+        const tracesContext = traces.getContextFromCarrier(data.otelCarrier)
+        traces.recordInitSpan(tracesContext)
+        rootTesterSpan = traces.startContextSpan(
+          `vitest.browser.tester.run`,
+          tracesContext,
+        )
+        traces.bind(rootTesterSpan.context)
+        await prepare(data).catch(err => unhandledError(err, 'Prepare Error'))
+        break
+      }
+      case 'viewport:done':
+      case 'viewport:fail':
+      case 'viewport': {
+        break
+      }
+      default: {
+        const error = new Error(`Unknown event: ${(data as any).event}`)
+        unhandledError(error, 'Unknown Event')
       }
-      break
-    }
-    case 'cleanup': {
-      await cleanup().catch(err => unhandledError(err, 'Cleanup Error'))
-      rootTesterSpan?.span.end()
-      await traces.finish()
-      break
-    }
-    case 'prepare': {
-      await traces.waitInit()
-      const tracesContext = traces.getContextFromCarrier(data.otelCarrier)
-      traces.recordInitSpan(tracesContext)
-      rootTesterSpan = traces.startContextSpan(
-        `vitest.browser.tester.run`,
-        tracesContext,
-      )
-      traces.bind(rootTesterSpan.context)
-      await prepare(data).catch(err => unhandledError(err, 'Prepare Error'))
-      break
-    }
-    case 'viewport:done':
-    case 'viewport:fail':
-    case 'viewport': {
-      break
-    }
-    default: {
-      const error = new Error(`Unknown event: ${(data as any).event}`)
-      unhandledError(error, 'Unknown Event')
     }
   }
-
-  channel.postMessage({
-    event: `response:${data.event}`,
-    iframeId: getBrowserState().iframeId!,
-  })
+  catch (error: any) {
+    // errors not handled by the cases above (e.g. tracing setup/teardown) must
+    // not stop us from responding, otherwise the orchestrator would wait forever
+    await unhandledError(error, 'Tester Error')
+  }
+  finally {
+    // always let the orchestrator know the event was handled so its
+    // `sendEventToIframe` promise resolves, even if the work above threw
+    channel.postMessage({
+      event: `response:${data.event}`,
+      iframeId: data.iframeId,
+      messageId: data.messageId,
+    })
+  }
 })
 
 const url = new URL(location.href)
@@ -307,6 +328,6 @@ function unhandledError(e: Error, type: string) {
     stack: e.stack,
   }, type).catch(() => {})
 }
-function isEvent(data: unknown): data is IframeChannelEvent {
+function isEvent(data: unknown): data is IframeReceivedEvent {
   return typeof data === 'object' && !!data && 'event' in data
 }