fix(vitest-plugin): dispose shared remote proxy sessions only after the last pool worker stops

.changeset/vitest-plugin-shared-remote-proxy-session.md

@@ -0,0 +1,7 @@
+---
+"@cloudflare/vitest-plugin": patch
+---
+
+Fix "Network connection lost" when multiple test files use remote bindings
+
+Remote proxy sessions are shared across pool workers by Wrangler config path, but were disposed during each test file's teardown. Because Vitest starts the next file's worker before the previous one finishes stopping, later files reused an already-disposed session and failed with "Network connection lost". Sessions are now disposed only once the last pool worker stops.

packages/vitest-plugin/src/pool/cloudflare-pool-worker.ts

@@ -4,7 +4,7 @@ import util from "node:util";
 import { compileModuleRules, testRegExps } from "miniflare";
 import { type ProvidedContext } from "vitest";
 import { workerdBuiltinModules } from "../shared/builtin-modules";
-import { parseProjectOptions, remoteProxySessionsDataMap } from "./config";
+import { disposeAllRemoteProxySessions, parseProjectOptions } from "./config";
 import { poolWorkerStarted, poolWorkerStopped } from "./pages";
 import { type WorkerPoolOptionsContext } from "./plugin";
 import {
@@ -120,20 +120,20 @@ export class CloudflarePoolWorker implements PoolWorker {
 		});
 		this.mf = undefined;
 
-		if (this.parsedPoolOptions?.resolvedConfig) {
-			const session = remoteProxySessionsDataMap.get(
-				this.parsedPoolOptions.resolvedConfig.path
-			)?.session;
-			await session?.dispose?.()?.catch((err) => {
-				this.debug("remote proxy session dispose rejected: %O", err);
-			});
-		}
-
 		// Decrement the active worker count. When the last worker stops, this
 		// closes file watchers created by buildPagesASSETSBinding() during config
 		// evaluation — they're registered globally because vitest evaluates all
 		// project configs at startup, even for projects that won't run.
-		poolWorkerStopped();
+		const wasLastWorker = poolWorkerStopped();
+
+		// Remote proxy sessions are shared by all pool workers using the same
+		// Wrangler config, and consecutive workers overlap, so only dispose them
+		// once the last worker stops.
+		if (wasLastWorker) {
+			await disposeAllRemoteProxySessions().catch((err) => {
+				this.debug("remote proxy session dispose rejected: %O", err);
+			});
+		}
 	}
 
 	send(message: WorkerRequest): void {

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

@@ -289,6 +289,20 @@ export const remoteProxySessionsDataMap = new Map<
 	RemoteProxySessionData | null
 >();
 
+/**
+ * Disposes every remote proxy session and clears the map.
+ *
+ * Sessions are shared across pool workers by Wrangler config path and
+ * consecutive workers overlap, so this is only safe to call once the last
+ * pool worker has stopped — calling it earlier would dispose sessions that
+ * later workers still depend on.
+ */
+export async function disposeAllRemoteProxySessions(): Promise<void> {
+	const sessions = [...remoteProxySessionsDataMap.values()];
+	remoteProxySessionsDataMap.clear();
+	await Promise.all(sessions.map((data) => data?.session.dispose()));
+}
+
 /**
  * Normalise the `experimental.newConfig` option into its resolved form.
  *

packages/vitest-plugin/src/pool/pages.ts

@@ -17,15 +17,18 @@ export function poolWorkerStarted(): void {
 	activeWorkers++;
 }
 
-export function poolWorkerStopped(): void {
+// Returns whether the stopped worker was the last active one.
+export function poolWorkerStopped(): boolean {
 	activeWorkers--;
 	if (activeWorkers <= 0) {
 		activeWorkers = 0;
 		for (const ac of registeredControllers) {
 			ac.abort();
 		}
 		registeredControllers.clear();
+		return true;
 	}
+	return false;
 }
 
 export async function buildPagesASSETSBinding(

packages/vitest-plugin/test/remote-proxy-sessions.test.ts

@@ -0,0 +1,77 @@
+import util from "node:util";
+import { beforeEach, describe, it, vi } from "vitest";
+import { CloudflarePoolWorker } from "../src/pool/cloudflare-pool-worker";
+import {
+	disposeAllRemoteProxySessions,
+	remoteProxySessionsDataMap,
+} from "../src/pool/config";
+import { poolWorkerStarted } from "../src/pool/pages";
+import type { RemoteProxySessionData } from "@cloudflare/remote-bindings";
+import type { RemoteProxyConnectionString } from "miniflare";
+
+function fakeSessionData(dispose: () => Promise<void>): RemoteProxySessionData {
+	return {
+		session: {
+			ready: Promise.resolve(),
+			dispose,
+			updateBindings: vi.fn(),
+			remoteProxyConnectionString: new URL(
+				"http://localhost"
+			) as RemoteProxyConnectionString,
+		},
+		remoteBindings: {},
+	};
+}
+
+// Bypasses the constructor's version check; start() is never called so
+// socket/miniflare are undefined and stop() exercises only session disposal.
+function createPoolWorker(): CloudflarePoolWorker {
+	const worker = Object.create(
+		CloudflarePoolWorker.prototype
+	) as CloudflarePoolWorker;
+	Object.defineProperty(worker, "debug", {
+		value: util.debuglog("vitest-plugin"),
+	});
+	return worker;
+}
+
+describe("remote proxy session disposal", () => {
+	beforeEach(() => {
+		remoteProxySessionsDataMap.clear();
+	});
+
+	it("disposes every session and clears the map", async ({ expect }) => {
+		const a = vi.fn(async () => {});
+		const b = vi.fn(async () => {});
+		remoteProxySessionsDataMap.set("/a/wrangler.toml", fakeSessionData(a));
+		remoteProxySessionsDataMap.set("/b/wrangler.toml", fakeSessionData(b));
+
+		await disposeAllRemoteProxySessions();
+
+		expect(a).toHaveBeenCalledTimes(1);
+		expect(b).toHaveBeenCalledTimes(1);
+		expect(remoteProxySessionsDataMap.size).toBe(0);
+	});
+
+	it("keeps the shared session alive until the last worker using it stops", async ({
+		expect,
+	}) => {
+		const dispose = vi.fn(async () => {});
+		const configPath = "/shared/wrangler.toml";
+		remoteProxySessionsDataMap.set(configPath, fakeSessionData(dispose));
+
+		// Two overlapping pool workers share one session.
+		poolWorkerStarted();
+		poolWorkerStarted();
+
+		const workerA = createPoolWorker();
+		const workerB = createPoolWorker();
+
+		await workerA.stop();
+		expect(dispose).not.toHaveBeenCalled();
+
+		await workerB.stop();
+		expect(dispose).toHaveBeenCalledTimes(1);
+		expect(remoteProxySessionsDataMap.has(configPath)).toBe(false);
+	});
+});