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 0000000000..f7019dee53 --- /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 6da37b98b6..c5e048f3bb 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 b078ad48a5..af7e5cb55e 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 56b41df8ac..e6b4818cd1 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 0000000000..9970f165c3 --- /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); + }); +});