fix(vitest-pool-workers): dispose remote proxy sessions on pool close (#12682)

Co-authored-by: Dario Piotrowicz <dario.piotrowicz@gmail.com>

Author: hiendv

.changeset/wicked-wombats-decide.md

@@ -0,0 +1,5 @@
+---
+"@cloudflare/vitest-pool-workers": patch
+---
+
+Fix resource leak where remote proxy sessions were not disposed during pool shutdown, causing vitest processes to hang.

packages/vitest-pool-workers/src/pool/config.ts

@@ -199,7 +199,7 @@ function filterTails(
 }
 
 /** Map that maps worker configPaths to their existing remote proxy session data (if any) */
-const remoteProxySessionsDataMap = new Map<
+export const remoteProxySessionsDataMap = new Map<
 	string,
 	{
 		session: RemoteProxySession;

packages/vitest-pool-workers/src/pool/index.ts

@@ -29,7 +29,11 @@ import { experimental_readRawConfig } from "wrangler";
 import { workerdBuiltinModules } from "../shared/builtin-modules";
 import { createChunkingSocket } from "../shared/chunking-socket";
 import { CompatibilityFlagAssertions } from "./compatibility-flag-assertions";
-import { OPTIONS_PATH, parseProjectOptions } from "./config";
+import {
+	OPTIONS_PATH,
+	parseProjectOptions,
+	remoteProxySessionsDataMap,
+} from "./config";
 import { guessWorkerExports } from "./guess-exports";
 import {
 	getProjectPath,
@@ -1284,6 +1288,14 @@ export default function (ctx: Vitest): ProcessPool {
 				}
 			}
 			allProjects.clear();
+			// Dispose remote proxy sessions to prevent handle leaks
+			log.debug("Disposing remote proxy sessions...");
+			for (const sessionData of remoteProxySessionsDataMap.values()) {
+				if (sessionData?.session?.dispose) {
+					promises.push(sessionData.session.dispose());
+				}
+			}
+			remoteProxySessionsDataMap.clear();
 			await Promise.all(promises);
 		},
 	};

packages/vitest-pool-workers/test/remote-proxy-cleanup.test.ts

@@ -0,0 +1,62 @@
+import dedent from "ts-dedent";
+import { test } from "./helpers";
+
+// This test requires CLOUDFLARE_ACCOUNT_ID and CLOUDFLARE_API_TOKEN
+// environment variables to be set, as it exercises remote proxy sessions
+// that connect to the Cloudflare API.
+test.skipIf(
+	!process.env.CLOUDFLARE_ACCOUNT_ID || !process.env.CLOUDFLARE_API_TOKEN
+)(
+	"disposes remote proxy sessions on pool close",
+	async ({ expect, seed, vitestRun }) => {
+		await seed({
+			"vitest.config.mts": dedent`
+				import { defineWorkersConfig } from "@cloudflare/vitest-pool-workers/config";
+				export default defineWorkersConfig({
+					test: {
+						reporters: ["hanging-process", "verbose"],
+						poolOptions: {
+							workers: {
+								wrangler: { configPath: "./wrangler.jsonc" },
+							},
+						},
+					}
+				});
+			`,
+			"wrangler.jsonc": dedent`
+				{
+					"name": "test-worker",
+					"main": "src/index.ts",
+					"compatibility_date": "2025-06-01",
+					"compatibility_flags": ["nodejs_compat"],
+					"ai": { "binding": "AI" },
+					"account_id": "${process.env.CLOUDFLARE_ACCOUNT_ID}"
+				}
+			`,
+			"src/index.ts": dedent`
+				export default {
+					async fetch(request, env, ctx) {
+						return new Response("Hello");
+					}
+				}
+			`,
+			"src/index.test.ts": dedent`
+				import { SELF } from "cloudflare:test";
+				import { it, expect } from "vitest";
+				it("responds with Hello", async () => {
+					const response = await SELF.fetch("http://localhost/");
+					expect(await response.text()).toBe("Hello");
+				});
+			`,
+		});
+
+		const result = await vitestRun({
+			flags: ["--reporter=hanging-process", "--reporter=verbose"],
+		});
+
+		expect(result.stderr).not.toContain(
+			"something prevents Vite server from exiting"
+		);
+	},
+	20_000
+);

packages/vitest-pool-workers/turbo.json

@@ -4,6 +4,9 @@
 	"tasks": {
 		"build": {
 			"outputs": ["dist/**"]
+		},
+		"test": {
+			"env": ["CLOUDFLARE_ACCOUNT_ID", "CLOUDFLARE_API_TOKEN"]
 		}
 	}
 }