[vite-plugin] Dispose remote proxy sessions when the server closes

.changeset/spicy-donkeys-listen.md

@@ -0,0 +1,7 @@
+---
+"@cloudflare/vite-plugin": patch
+---
+
+Fix `vite build` hanging when `remoteBindings` is enabled
+
+With `remoteBindings` enabled, `vite build` produced all of its output but then never exited, so builds had to be killed manually and could not complete in CI. Builds using remote bindings now finish and exit as expected.

packages/vite-plugin-cloudflare/src/__tests__/preview-server.spec.ts

@@ -1,11 +1,18 @@
 import { fileURLToPath } from "node:url";
+import { maybeStartOrUpdateRemoteProxySession } from "@cloudflare/remote-bindings";
 import { Miniflare } from "miniflare";
 import { createBuilder, preview } from "vite";
 import { afterEach, describe, test, vi } from "vitest";
 import { cloudflare } from "../index";
+import type { RemoteProxySession } from "@cloudflare/remote-bindings";
 
 vi.mock("@cloudflare/workers-utils");
 
+vi.mock("@cloudflare/remote-bindings", async (importOriginal) => ({
+	...(await importOriginal<typeof import("@cloudflare/remote-bindings")>()),
+	maybeStartOrUpdateRemoteProxySession: vi.fn(),
+}));
+
 const fixturesPath = fileURLToPath(new URL("./fixtures", import.meta.url));
 
 describe("preview server", () => {
@@ -38,4 +45,120 @@ describe("preview server", () => {
 		await previewServer.close();
 		expect(disposeSpy).toHaveBeenCalled();
 	});
+
+	test("disposes the remote proxy session when preview server is closed", async ({
+		expect,
+	}) => {
+		const dispose = vi.fn(async () => {});
+
+		vi.mocked(maybeStartOrUpdateRemoteProxySession).mockResolvedValue({
+			session: {
+				ready: Promise.resolve(),
+				dispose,
+				updateBindings: async () => {},
+				remoteProxyConnectionString:
+					"http://127.0.0.1:1234" as unknown as RemoteProxySession["remoteProxyConnectionString"],
+			},
+			remoteBindings: {},
+		});
+
+		const builder = await createBuilder({
+			root: fixturesPath,
+			logLevel: "silent",
+			plugins: [
+				cloudflare({
+					inspectorPort: false,
+					persistState: false,
+					remoteBindings: true,
+				}),
+			],
+		});
+
+		// Build the worker
+		await builder.buildApp();
+
+		// Start a preview server
+		const previewServer = await preview({
+			root: fixturesPath,
+			logLevel: "silent",
+			preview: { port: 0 },
+			plugins: [
+				cloudflare({
+					inspectorPort: false,
+					persistState: false,
+					remoteBindings: true,
+				}),
+			],
+		});
+
+		expect(dispose).not.toHaveBeenCalled();
+		// The session holds a listening server handle, so leaving it open keeps
+		// the event loop alive and `vite build` never exits
+		await previewServer.close();
+		expect(dispose).toHaveBeenCalled();
+	});
+
+	test("retries remote proxy session dispose after a failed close", async ({
+		expect,
+	}) => {
+		const dispose = vi
+			.fn()
+			.mockRejectedValueOnce(new Error("dispose failed"))
+			.mockResolvedValue(undefined);
+
+		vi.mocked(maybeStartOrUpdateRemoteProxySession).mockResolvedValue({
+			session: {
+				ready: Promise.resolve(),
+				dispose,
+				updateBindings: async () => {},
+				remoteProxyConnectionString:
+					"http://127.0.0.1:1234" as unknown as RemoteProxySession["remoteProxyConnectionString"],
+			},
+			remoteBindings: {},
+		});
+
+		const plugins = [
+			cloudflare({
+				inspectorPort: false,
+				persistState: false,
+				remoteBindings: true,
+			}),
+		];
+
+		const builder = await createBuilder({
+			root: fixturesPath,
+			logLevel: "silent",
+			plugins,
+		});
+		await builder.buildApp();
+
+		const firstPreview = await preview({
+			root: fixturesPath,
+			logLevel: "silent",
+			preview: { port: 0 },
+			plugins,
+		});
+
+		await expect(firstPreview.close()).resolves.toBeUndefined();
+		expect(dispose).toHaveBeenCalledTimes(1);
+
+		const secondPreview = await preview({
+			root: fixturesPath,
+			logLevel: "silent",
+			preview: { port: 0 },
+			plugins,
+		});
+
+		const lastStart = vi
+			.mocked(maybeStartOrUpdateRemoteProxySession)
+			.mock.calls.at(-1);
+		expect(lastStart?.[1]).toEqual(
+			expect.objectContaining({
+				session: expect.objectContaining({ dispose }),
+			})
+		);
+
+		await secondPreview.close();
+		expect(dispose).toHaveBeenCalledTimes(2);
+	});
 });

packages/vite-plugin-cloudflare/src/miniflare-options.ts

@@ -115,6 +115,28 @@ const remoteProxySessionsDataMap = new Map<
 	RemoteProxySessionData | null
 >();
 
+/**
+ * Disposes every remote proxy session that has been started.
+ *
+ * Each session runs a listening server, which keeps the event loop alive, so
+ * the sessions must be disposed for a `vite build` or a programmatic server
+ * close to be able to exit.
+ */
+export async function disposeRemoteProxySessions(): Promise<void> {
+	await Promise.all(
+		[...remoteProxySessionsDataMap.entries()].map(
+			async ([configPath, remoteProxySessionData]) => {
+				try {
+					await remoteProxySessionData?.session.dispose();
+					remoteProxySessionsDataMap.delete(configPath);
+				} catch (error) {
+					debuglog("Failed to dispose remote proxy session:", error);
+				}
+			}
+		)
+	);
+}
+
 function createRemoteBindingsLogger(logger: vite.Logger): RemoteBindingsLogger {
 	const write = (
 		level: "info" | "warn" | "error",

packages/vite-plugin-cloudflare/src/plugins/dev.ts

@@ -21,7 +21,10 @@ import {
 	compareWorkerNameToExportTypesMaps,
 	getCurrentWorkerNameToExportTypesMap,
 } from "../export-types";
-import { getDevMiniflareOptions } from "../miniflare-options";
+import {
+	disposeRemoteProxySessions,
+	getDevMiniflareOptions,
+} from "../miniflare-options";
 import { UNKNOWN_HOST } from "../shared";
 import {
 	createPlugin,
@@ -84,6 +87,11 @@ export const devPlugin = createPlugin("dev", (ctx) => {
 						} catch (error) {
 							debuglog("Failed to dispose Miniflare instance:", error);
 						}
+						try {
+							await disposeRemoteProxySessions();
+						} catch (error) {
+							debuglog("Failed to dispose remote proxy sessions:", error);
+						}
 					}
 				}
 			};

packages/vite-plugin-cloudflare/src/plugins/preview.ts

@@ -8,8 +8,11 @@ import { buildPublicUrl, Request as MiniflareRequest } from "miniflare";
 import colors from "picocolors";
 import { configureContainerPull, getDockerPath } from "../containers";
 import { assertIsPreview } from "../context";
-import { getPreviewMiniflareOptions } from "../miniflare-options";
-import { createPlugin, createRequestHandler } from "../utils";
+import {
+	disposeRemoteProxySessions,
+	getPreviewMiniflareOptions,
+} from "../miniflare-options";
+import { createPlugin, createRequestHandler, debuglog } from "../utils";
 import { handleWebSocket } from "../websockets";
 import { rewriteLegacyMiniflarePath } from "./trigger-handlers";
 
@@ -27,11 +30,18 @@ export const previewPlugin = createPlugin("preview", (ctx) => {
 		async configurePreviewServer(vitePreviewServer) {
 			assertIsPreview(ctx);
 
-			// Ensure Miniflare is disposed when the preview server is closed during prerendering
+			// Ensure Miniflare and any remote proxy sessions are disposed when the
+			// preview server is closed during prerendering
 			const closePreviewServer =
 				vitePreviewServer.close.bind(vitePreviewServer);
 			vitePreviewServer.close = async () => {
-				await Promise.all([ctx.disposeMiniflare(), closePreviewServer()]);
+				await Promise.all([
+					ctx.disposeMiniflare(),
+					disposeRemoteProxySessions().catch((error) => {
+						debuglog("Failed to dispose remote proxy sessions:", error);
+					}),
+					closePreviewServer(),
+				]);
 			};
 
 			const { miniflareOptions, containerTagToOptionsMap } =