[vite-plugin] Route v1 DevTools WebSocket upgrades explicitly

.changeset/witty-pandas-serve.md

@@ -2,6 +2,6 @@
 "@cloudflare/vite-plugin": patch
 ---
 
-fix(vite-plugin): don't destroy WebSockets owned by other upgrade listeners
+Fix Vite DevTools WebSocket connections closing unexpectedly
 
-Node dispatches `upgrade` to every registered listener. The dev-server handler destroyed the raw socket when the Worker returned no route, even if another listener (such as Vite DevTools at `/__devtools/__ws`) had already completed the handshake. It now leaves the socket alone once another listener has written a response.
+Leave upgrades at `/__devtools/__ws` to Vite DevTools instead of forwarding them to the Worker. Sandbox preview requests continue to reach the Worker.

packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts

@@ -1,3 +1,4 @@
+import { once } from "node:events";
 import http from "node:http";
 import net from "node:net";
 import { DeferredPromise, Miniflare, Response } from "miniflare";
@@ -11,10 +12,9 @@ import {
 	test,
 	vi,
 } from "vitest";
-import { WebSocketServer } from "ws";
+import { WebSocket, WebSocketServer } from "ws";
 import { handleWebSocket } from "../websockets";
 import type { AddressInfo } from "node:net";
-import type { Duplex } from "node:stream";
 
 describe("handleWebSocket", () => {
 	let httpServer: http.Server;
@@ -186,7 +186,8 @@ describe("handleWebSocket", () => {
 	async function dispatchUpgrade(
 		expect: ExpectStatic,
 		mf: Miniflare,
-		headers: Record<string, string>
+		headers: Record<string, string>,
+		pathname = "/"
 	) {
 		const deferred = new DeferredPromise<Response>();
 		const mockedDispatchFetch = vi
@@ -199,7 +200,7 @@ describe("handleWebSocket", () => {
 			.map(([k, v]) => `${k}: ${v}`)
 			.join("\r\n");
 		socket.write(
-			"GET / HTTP/1.1\r\n" +
+			`GET ${pathname} HTTP/1.1\r\n` +
 				`${headerLines}\r\n` +
 				"Upgrade: websocket\r\n" +
 				"Connection: Upgrade\r\n" +
@@ -456,72 +457,114 @@ describe("handleWebSocket", () => {
 	});
 
 	// https://github.com/cloudflare/workers-sdk/issues/15654
-	test("does not destroy a WebSocket another upgrade listener has already claimed", async ({
-		expect,
-	}) => {
+	test.for(["/__devtools/__ws", "/__devtools/__ws?token=testing"])(
+		"leaves %s to Vite DevTools without dispatching to the Worker",
+		async (pathname, { expect }) => {
+			const mf = await listen();
+			const dispatchFetch = vi.spyOn(mf, "dispatchFetch");
+
+			const owner = new WebSocketServer({ noServer: true });
+			onTestFinished(() => {
+				for (const client of owner.clients) {
+					client.terminate();
+				}
+				owner.close();
+			});
+			httpServer.on("upgrade", (request, socket, head) => {
+				if (request.url === pathname) {
+					owner.handleUpgrade(request, socket, head, (ws) => {
+						ws.on("message", (message) => ws.send(message));
+					});
+				}
+			});
+
+			const socket = await connect();
+			const client = new WebSocket(`ws://127.0.0.1:${port}${pathname}`, {
+				createConnection: () => socket,
+			});
+			onTestFinished(() => client.terminate());
+			await once(client, "open");
+			const received = once(client, "message");
+			client.send("DevTools connected");
+			const [message] = await received;
+			expect(message.toString()).toBe("DevTools connected");
+			expect(dispatchFetch).not.toHaveBeenCalled();
+		}
+	);
+
+	test.for(["/__devtools/__ws", "/__devtools/__ws?token=testing"])(
+		"forwards Sandbox requests for %s to the Worker",
+		async (pathname, { expect }) => {
+			const mf = await listen();
+			const host = `4567-my-sandbox-sup3rs3cr3t.localhost:${port}`;
+			const url = await dispatchUpgrade(expect, mf, { Host: host }, pathname);
+			expect(`${url}`).toBe(`http://${host}${pathname}`);
+		}
+	);
+
+	test.for([
+		"/__devtools/__ws-extra",
+		"/__devtools/__ws/child",
+		"/__devtools/other",
+	])("forwards %s to the Worker", async (pathname, { expect }) => {
 		const mf = await listen();
-		const deferred = new DeferredPromise<Response>();
-		vi.spyOn(mf, "dispatchFetch").mockReturnValue(deferred);
+		const host = `127.0.0.1:${port}`;
+		const url = await dispatchUpgrade(expect, mf, { Host: host }, pathname);
+		expect(`${url}`).toBe(`http://${host}${pathname}`);
+	});
 
-		// The Worker has no route for the DevTools socket, so the plugin's
-		// dispatch resolves without a WebSocket after the owner has upgraded.
-		const owner = new WebSocketServer({ noServer: true });
-		let serverSocket: Duplex | undefined;
-		onTestFinished(() => {
-			for (const client of owner.clients) {
-				client.terminate();
-			}
-			owner.close();
+	test("upgrades a Worker WebSocket after an earlier pipelined HTTP response", async ({
+		expect,
+	}) => {
+		const pageResponse = new DeferredPromise<http.ServerResponse>();
+		httpServer.removeAllListeners("request");
+		httpServer.on("request", (_request, response) => {
+			pageResponse.resolve(response);
 		});
-		httpServer.on("upgrade", (request, socket, head) => {
-			if (request.url === "/__devtools/__ws") {
-				serverSocket = socket;
-				owner.handleUpgrade(request, socket, head, (ws) => {
-					owner.emit("connection", ws, request);
-				});
-			}
+		const mf = await listen();
+		const dispatchFetch = mf.dispatchFetch.bind(mf);
+		const dispatchStarted = new DeferredPromise<void>();
+		const resumeDispatch = new DeferredPromise<void>();
+		onTestFinished(() => resumeDispatch.resolve());
+		vi.spyOn(mf, "dispatchFetch").mockImplementation(async (...args) => {
+			dispatchStarted.resolve();
+			await resumeDispatch;
+			return dispatchFetch(...args);
 		});
 
 		const socket = await connect();
 		const chunks: Buffer[] = [];
-		const upgraded = new Promise<void>((resolve) => {
-			socket.on("data", (chunk) => {
-				chunks.push(chunk);
-				if (Buffer.concat(chunks).toString("utf8").includes("HTTP/1.1 101")) {
-					resolve();
-				}
-			});
+		const upgraded = new DeferredPromise<void>();
+		socket.on("data", (chunk) => {
+			chunks.push(chunk);
+			if (Buffer.concat(chunks).toString("utf8").includes("HTTP/1.1 101")) {
+				upgraded.resolve();
+			}
 		});
-
 		socket.write(
-			"GET /__devtools/__ws HTTP/1.1\r\n" +
+			"GET /page HTTP/1.1\r\n" +
+				`Host: 127.0.0.1:${port}\r\n\r\n` +
+				"GET /ws HTTP/1.1\r\n" +
 				`Host: 127.0.0.1:${port}\r\n` +
 				"Upgrade: websocket\r\n" +
 				"Connection: Upgrade\r\n" +
 				"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" +
 				"Sec-WebSocket-Version: 13\r\n\r\n"
 		);
-		await upgraded;
-
-		if (!serverSocket) {
-			throw new Error("expected the DevTools listener to claim the socket");
-		}
-		const destroy = vi.fn();
-		const originalDestroy = serverSocket.destroy.bind(serverSocket);
-		vi.spyOn(serverSocket, "destroy").mockImplementation((...args) => {
-			destroy();
-			return originalDestroy(...args);
-		});
 
-		deferred.resolve(new Response("not found", { status: 404 }));
-		// Let the plugin handler resume and make its teardown decision.
-		await new Promise<void>((resolve) => setImmediate(resolve));
+		// Finish the earlier HTTP response after the upgrade handler starts,
+		// but before the Worker returns its WebSocket.
+		await dispatchStarted;
+		const response = await pageResponse;
+		const finished = once(response, "finish");
+		response.end("OK");
+		await finished;
+		resumeDispatch.resolve();
+		await upgraded;
 
-		// Assert on the server socket directly; the client may not observe the
-		// close until a later poll phase.
-		expect(destroy).not.toHaveBeenCalled();
-		expect(socket.destroyed).toBe(false);
-		expect(Buffer.concat(chunks).toString("utf8")).toContain("HTTP/1.1 101");
+		const raw = Buffer.concat(chunks).toString("utf8");
+		expect(raw).toMatch(/^HTTP\/1\.1 200/);
+		expect(raw).toContain("OKHTTP/1.1 101 Switching Protocols");
 		socket.destroy();
 	});
 

packages/vite-plugin-cloudflare/src/websockets.ts

@@ -4,7 +4,6 @@ import { UNKNOWN_HOST } from "./shared";
 import { createHeadersForIncomingMessage, getForwardedProto } from "./utils";
 import type { Headers, Miniflare } from "miniflare";
 import type { IncomingMessage } from "node:http";
-import type { Socket } from "node:net";
 import type { Duplex } from "node:stream";
 import type * as vite from "vite";
 
@@ -39,26 +38,12 @@ export function handleWebSocket(
 		}
 	);
 
-	// Prepend so the write-count baseline below is captured before another
-	// `upgrade` listener can answer the same request. Node dispatches `upgrade`
-	// to every listener, not just the first.
-	httpServer.prependListener(
+	httpServer.on(
 		"upgrade",
 		async (request: IncomingMessage, socket: Duplex, head: Buffer) => {
 			// Socket errors crash Node.js if unhandled
 			socket.on("error", () => socket.destroy());
 
-			// Another listener may already have completed the handshake (for
-			// example Vite DevTools at `/__devtools/__ws`, which the Worker does
-			// not route). `bytesWritten` is cumulative across keep-alive
-			// connections, so compare against a baseline rather than zero.
-			// See https://github.com/cloudflare/workers-sdk/issues/15654
-			const bytesWrittenAtStart =
-				(socket as unknown as Socket).bytesWritten ?? 0;
-			const isClaimed = () =>
-				socket.destroyed ||
-				((socket as unknown as Socket).bytesWritten ?? 0) > bytesWrittenAtStart;
-
 			try {
 				const rawHost = request.headers.host ?? UNKNOWN_HOST;
 				// Honor `X-Forwarded-Proto` so that the upgrade URL reflects the
@@ -72,10 +57,12 @@ export function handleWebSocket(
 
 				const isViteRequest =
 					request.headers["sec-websocket-protocol"]?.startsWith("vite");
+				const isDevToolsRequest = url.pathname === "/__devtools/__ws";
 				const isSandboxRequest = hasSandboxOrigin(url.origin);
 
-				// Ignore Vite HMR WebSockets but forward on all sandbox requests.
-				if (isViteRequest && !isSandboxRequest) {
+				// Leave Vite HMR and DevTools upgrades to their own handlers, but
+				// forward all sandbox requests to the Worker.
+				if ((isViteRequest || isDevToolsRequest) && !isSandboxRequest) {
 					return;
 				}
 
@@ -92,17 +79,11 @@ export function handleWebSocket(
 				const workerWebSocket = response.webSocket;
 
 				if (!workerWebSocket) {
-					// This listener does not own the socket. Destroy it only if no
-					// other listener has written a response in the meantime.
-					if (!isClaimed()) {
-						socket.destroy();
-					}
+					socket.destroy();
 					return;
 				}
 
-				// A listener claimed the socket while `dispatchFetch` was in
-				// flight, so a second upgrade would corrupt the connection.
-				if (isClaimed()) {
+				if (socket.destroyed) {
 					return;
 				}
 
@@ -129,14 +110,9 @@ export function handleWebSocket(
 				// This listener is `async`, so an uncaught rejection here escapes as
 				// an unhandled rejection — which terminates the Node.js process on
 				// modern versions and leaks the client socket. Tear the socket down
-				// instead. Defer the teardown one tick so another listener can claim
-				// the socket first, mirroring the `!workerWebSocket` path above.
+				// instead.
 				workerResponseHeaders.delete(request);
-				setImmediate(() => {
-					if (!isClaimed()) {
-						socket.destroy();
-					}
-				});
+				socket.destroy();
 			}
 		}
 	);