diff --git a/docs/guide/advanced/pool.md b/docs/guide/advanced/pool.md
index ea632de7f..635d04e34 100644
--- a/docs/guide/advanced/pool.md
+++ b/docs/guide/advanced/pool.md
@@ -120,7 +120,7 @@ Your `CustomPoolRunner` will be controlling how your custom test runner worker l
 In your worker file, you can import helper utilities from `vitest/worker`:
 
 ```ts [my-worker.ts]
-import { init, runBaseTests } from 'vitest/worker'
+import { init, runBaseTests, setupEnvironment } from 'vitest/worker'
 
 init({
   post: (response) => {
@@ -141,7 +141,8 @@ init({
   deserialize: (value) => {
     // Optional, provide custom deserializer for `on` callbacks
   },
-  runTests: state => runBaseTests('run', state),
-  collectTests: state => runBaseTests('collect', state),
+  runTests: (state, traces) => runBaseTests('run', state, traces),
+  collectTests: (state, traces) => runBaseTests('collect', state, traces),
+  setup: setupEnvironment,
 })
 ```
diff --git a/packages/vitest/src/public/worker.ts b/packages/vitest/src/public/worker.ts
index f80a2e3f1..786ad6bbf 100644
--- a/packages/vitest/src/public/worker.ts
+++ b/packages/vitest/src/public/worker.ts
@@ -1,2 +1,2 @@
-export { runBaseTests } from '../runtime/workers/base'
+export { runBaseTests, setupEnvironment } from '../runtime/workers/base'
 export { init } from '../runtime/workers/init'
diff --git a/packages/vitest/src/runtime/workers/base.ts b/packages/vitest/src/runtime/workers/base.ts
index 810a9da4f..e0dc6af90 100644
--- a/packages/vitest/src/runtime/workers/base.ts
+++ b/packages/vitest/src/runtime/workers/base.ts
@@ -30,6 +30,7 @@ function startModuleRunner(options: ContextModuleRunnerOptions) {
 let _currentEnvironment!: Environment
 let _environmentTime: number
 
+/** @experimental */
 export async function setupEnvironment(context: WorkerSetupContext): Promise<() => Promise<void>> {
   const startTime = performance.now()
   const {
diff --git a/test/core/test/exports.test.ts b/test/core/test/exports.test.ts
index 5f6ab03f6..849371ed2 100644
--- a/test/core/test/exports.test.ts
+++ b/test/core/test/exports.test.ts
@@ -340,6 +340,7 @@ it('exports snapshot', async ({ skip, task }) => {
           "./worker": {
             "init": "function",
             "runBaseTests": "function",
+            "setupEnvironment": "function",
           },
         }
       `)