Use coupleWebSocket function from Miniflare (#8525)

Author: jamesopstad

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

@@ -277,11 +277,7 @@ export function cloudflare(pluginConfig: PluginConfig = {}): vite.Plugin[] {
 					{ alwaysCallNext: false }
 				);
 
-				handleWebSocket(
-					viteDevServer.httpServer,
-					entryWorker.fetch,
-					viteDevServer.config.logger
-				);
+				handleWebSocket(viteDevServer.httpServer, entryWorker.fetch);
 
 				return () => {
 					viteDevServer.middlewares.use((req, res, next) => {
@@ -306,11 +302,7 @@ export function cloudflare(pluginConfig: PluginConfig = {}): vite.Plugin[] {
 					{ alwaysCallNext: false }
 				);
 
-				handleWebSocket(
-					vitePreviewServer.httpServer,
-					miniflare.dispatchFetch,
-					vitePreviewServer.config.logger
-				);
+				handleWebSocket(vitePreviewServer.httpServer, miniflare.dispatchFetch);
 
 				return () => {
 					vitePreviewServer.middlewares.use((req, res, next) => {

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

@@ -1,3 +1,4 @@
+import { coupleWebSocket } from "miniflare";
 import { WebSocketServer } from "ws";
 import { UNKNOWN_HOST } from "./shared";
 import { nodeHeadersToWebHeaders } from "./utils";
@@ -12,8 +13,7 @@ import type * as vite from "vite";
  */
 export function handleWebSocket(
 	httpServer: vite.HttpServer,
-	fetcher: ReplaceWorkersTypes<Fetcher>["fetch"],
-	logger: vite.Logger
+	fetcher: ReplaceWorkersTypes<Fetcher>["fetch"]
 ) {
 	const nodeWebSocket = new WebSocketServer({ noServer: true });
 
@@ -44,41 +44,7 @@ export function handleWebSocket(
 				socket,
 				head,
 				async (clientWebSocket) => {
-					workerWebSocket.accept();
-
-					// Forward Worker events to client
-					workerWebSocket.addEventListener("message", (event) => {
-						clientWebSocket.send(event.data);
-					});
-					workerWebSocket.addEventListener("error", (event) => {
-						logger.error(
-							`WebSocket error:\n${event.error?.stack || event.error?.message}`,
-							{ error: event.error }
-						);
-					});
-					workerWebSocket.addEventListener("close", () => {
-						clientWebSocket.close();
-					});
-
-					// Forward client events to Worker
-					clientWebSocket.on("message", (data, isBinary) => {
-						workerWebSocket.send(
-							isBinary
-								? Array.isArray(data)
-									? Buffer.concat(data)
-									: data
-								: data.toString()
-						);
-					});
-					clientWebSocket.on("error", (error) => {
-						logger.error(`WebSocket error:\n${error.stack || error.message}`, {
-							error,
-						});
-					});
-					clientWebSocket.on("close", () => {
-						workerWebSocket.close();
-					});
-
+					coupleWebSocket(clientWebSocket, workerWebSocket);
 					nodeWebSocket.emit("connection", clientWebSocket, request);
 				}
 			);