From 8e1bf01f7be238922c9ab4cf3ef0dea7e39e9905 Mon Sep 17 00:00:00 2001 From: Kevin Jain Date: Thu, 20 Aug 2026 15:42:23 -0700 Subject: [PATCH] fix(vitest-plugin): dispose shared remote proxy sessions only after the last pool worker stops ``` 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. ``` This fixes issue: https://github.com/cloudflare/workers-sdk/issues/13306 --- ...test-plugin-shared-remote-proxy-session.md | 7 ++ .../src/pool/cloudflare-pool-worker.ts | 22 +++--- packages/vitest-plugin/src/pool/config.ts | 14 ++++ packages/vitest-plugin/src/pool/pages.ts | 5 +- .../test/remote-proxy-sessions.test.ts | 77 +++++++++++++++++++ 5 files changed, 113 insertions(+), 12 deletions(-) create mode 100644 .changeset/vitest-plugin-shared-remote-proxy-session.md create mode 100644 packages/vitest-plugin/test/remote-proxy-sessions.test.ts diff --git a/.changeset/vitest-plugin-shared-remote-proxy-session.md b/.changeset/vitest-plugin-shared-remote-proxy-session.md new file mode 100644 index 00000000000..f7019dee53b --- /dev/null +++ b/.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. diff --git a/packages/vitest-plugin/src/pool/cloudflare-pool-worker.ts b/packages/vitest-plugin/src/pool/cloudflare-pool-worker.ts index 6da37b98b66..c5e048f3bb1 100644 --- a/packages/vitest-plugin/src/pool/cloudflare-pool-worker.ts +++ b/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 { diff --git a/packages/vitest-plugin/src/pool/config.ts b/packages/vitest-plugin/src/pool/config.ts index b078ad48a52..af7e5cb55e4 100644 --- a/packages/vitest-plugin/src/pool/config.ts +++ b/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 { + const sessions = [...remoteProxySessionsDataMap.values()]; + remoteProxySessionsDataMap.clear(); + await Promise.all(sessions.map((data) => data?.session.dispose())); +} + /** * Normalise the `experimental.newConfig` option into its resolved form. * diff --git a/packages/vitest-plugin/src/pool/pages.ts b/packages/vitest-plugin/src/pool/pages.ts index 56b41df8acf..e6b4818cd11 100644 --- a/packages/vitest-plugin/src/pool/pages.ts +++ b/packages/vitest-plugin/src/pool/pages.ts @@ -17,7 +17,8 @@ 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; @@ -25,7 +26,9 @@ export function poolWorkerStopped(): void { ac.abort(); } registeredControllers.clear(); + return true; } + return false; } export async function buildPagesASSETSBinding( diff --git a/packages/vitest-plugin/test/remote-proxy-sessions.test.ts b/packages/vitest-plugin/test/remote-proxy-sessions.test.ts new file mode 100644 index 00000000000..9970f165c39 --- /dev/null +++ b/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): 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); + }); +});