refactor: create a single safe rpc

packages/browser/src/client/main.ts

@@ -1,4 +1,3 @@
-import type { VitestClient } from '@vitest/ws-client'
 import { createClient } from '@vitest/ws-client'
 // eslint-disable-next-line no-restricted-imports
 import type { ResolvedConfig } from 'vitest'
@@ -7,7 +6,7 @@ import { createBrowserRunner } from './runner'
 import { BrowserSnapshotEnvironment } from './snapshot'
 import { importId } from './utils'
 import { setupConsoleLogSpy } from './logger'
-import { rpc, rpcDone } from './rpc'
+import { createSafeRpc, rpc, rpcDone } from './rpc'
 
 // @ts-expect-error mocking some node apis
 globalThis.process = { env: {}, argv: [], cwd: () => '/', stdout: { write: () => {} }, nextTick: cb => cb() }
@@ -54,12 +53,16 @@ async function loadConfig() {
 ws.addEventListener('open', async () => {
   await loadConfig()
 
+  const { getSafeTimers } = await importId('vitest/utils') as typeof import('vitest/utils')
+  const safeRpc = createSafeRpc(client, getSafeTimers)
+
   // @ts-expect-error mocking vitest apis
   globalThis.__vitest_worker__ = {
     config,
     browserHashMap,
     moduleCache: new Map(),
     rpc: client.rpc,
+    safeRpc,
   }
 
   // @ts-expect-error mocking vitest apis
@@ -70,11 +73,11 @@ ws.addEventListener('open', async () => {
   iFrame.setAttribute('src', '/__vitest__/')
 
   await setupConsoleLogSpy()
-  await runTests(paths, config, client)
+  await runTests(paths, config)
 })
 
 let hasSnapshot = false
-async function runTests(paths: string[], config: any, client: VitestClient) {
+async function runTests(paths: string[], config: any) {
   // need to import it before any other import, otherwise Vite optimizer will hang
   const viteClientPath = '/@vite/client'
   await import(viteClientPath)
@@ -84,11 +87,11 @@ async function runTests(paths: string[], config: any, client: VitestClient) {
   if (!runner) {
     const { VitestTestRunner } = await importId('vitest/runners') as typeof import('vitest/runners')
     const BrowserRunner = createBrowserRunner(VitestTestRunner)
-    runner = new BrowserRunner({ config, client, browserHashMap })
+    runner = new BrowserRunner({ config, browserHashMap })
   }
 
   if (!hasSnapshot) {
-    setupSnapshotEnvironment(new BrowserSnapshotEnvironment(client))
+    setupSnapshotEnvironment(new BrowserSnapshotEnvironment())
     hasSnapshot = true
   }
 

packages/browser/src/client/rpc.ts

@@ -1,13 +1,13 @@
-import {
+import type {
   getSafeTimers,
 } from '@vitest/utils'
 import type { VitestClient } from '@vitest/ws-client'
 
 const { get } = Reflect
 const safeRandom = Math.random
 
-function withSafeTimers(fn: () => void) {
-  const { setTimeout, clearTimeout, nextTick, setImmediate, clearImmediate } = getSafeTimers()
+function withSafeTimers(getTimers: typeof getSafeTimers, fn: () => void) {
+  const { setTimeout, clearTimeout, nextTick, setImmediate, clearImmediate } = getTimers()
 
   const currentSetTimeout = globalThis.setTimeout
   const currentClearTimeout = globalThis.clearTimeout
@@ -48,13 +48,11 @@ export const rpcDone = async () => {
   return Promise.all(awaitable)
 }
 
-export const rpc = (): VitestClient['rpc'] => {
-  // @ts-expect-error not typed global
-  const { rpc } = globalThis.__vitest_worker__
-  return new Proxy(rpc, {
+export const createSafeRpc = (client: VitestClient, getTimers: () => any): VitestClient['rpc'] => {
+  return new Proxy(client.rpc, {
     get(target, p, handler) {
       const sendCall = get(target, p, handler)
-      const safeSendCall = (...args: any[]) => withSafeTimers(async () => {
+      const safeSendCall = (...args: any[]) => withSafeTimers(getTimers, async () => {
         const result = sendCall(...args)
         promises.add(result)
         try {
@@ -69,3 +67,8 @@ export const rpc = (): VitestClient['rpc'] => {
     },
   })
 }
+
+export const rpc = (): VitestClient['rpc'] => {
+  // @ts-expect-error not typed global
+  return globalThis.__vitest_worker__.safeRpc
+}

packages/browser/src/client/runner.ts

@@ -1,25 +1,21 @@
 import type { File, TaskResult, Test } from '@vitest/runner'
-import type { VitestClient } from '@vitest/ws-client'
 import { rpc } from './rpc'
 import type { ResolvedConfig } from '#types'
 
 interface BrowserRunnerOptions {
   config: ResolvedConfig
-  client: VitestClient
   browserHashMap: Map<string, string>
 }
 
 export function createBrowserRunner(original: any) {
   return class BrowserTestRunner extends original {
     public config: ResolvedConfig
     hashMap = new Map<string, string>()
-    client: VitestClient
 
     constructor(options: BrowserRunnerOptions) {
       super(options.config)
       this.config = options.config
       this.hashMap = options.browserHashMap
-      this.client = options.client
     }
 
     async onAfterRunTest(task: Test) {

packages/browser/src/client/snapshot.ts

@@ -1,10 +1,7 @@
-import type { VitestClient } from '@vitest/ws-client'
 import { rpc } from './rpc'
 import type { SnapshotEnvironment } from '#types'
 
 export class BrowserSnapshotEnvironment implements SnapshotEnvironment {
-  constructor(private client: VitestClient) {}
-
   readSnapshotFile(filepath: string): Promise<string | null> {
     return rpc().readFile(filepath)
   }