fix: support WebSocket upgrades in dev server (#3823)

Vite's Connect-based dev server dispatches WebSocket upgrades as
'upgrade' events on the underlying http.Server rather than routing them
through middleware, so `Deno.upgradeWebSocket()` was never reached and
the client hung in CONNECTING. The plugin now attaches an upgrade
listener that hands the raw node socket off to Fresh's existing
`ctx.upgrade()` path.

The plumbing leans on denoland/deno#33342, which taught
`Deno.upgradeWebSocket` to accept a `{ socket, head }` option for
requests that didn't come in through `Deno.serve`. A
`WeakMap<Request, { socket, head }>` kept on `globalThis` acts as the
side channel: the plugin's upgrade listener populates it,
`ctx.upgrade()` reads and clears it before calling Deno. The map lives
on `globalThis` because Vite's SSR runner evaluates `context.ts`
separately from Deno's own load of the plugin, so each side would
otherwise hold its own module-local instance.

Fixes #3350

Author: bartlomieju

packages/fresh/src/context.ts

@@ -58,6 +58,28 @@ export interface WebSocketUpgradeOptions {
   protocol?: string;
 }
 
+/**
+ * Side channel used by `@fresh/plugin-vite` (and any other adapter sitting on
+ * top of `node:http`) to hand a raw socket + buffered head to `ctx.upgrade()`.
+ * `Deno.upgradeWebSocket()` normally pulls these off the request handled by
+ * `Deno.serve`, but a request synthesized from `node:http` carries neither, so
+ * the adapter stashes them here and `ctx.upgrade()` forwards them to Deno.
+ *
+ * Stored on `globalThis` so a single WeakMap is shared even when this module
+ * is evaluated more than once (e.g. once by Deno for the Vite plugin and once
+ * by Vite's SSR runner for the user's server code).
+ */
+const UPGRADE_SOURCE_KEY: unique symbol = Symbol.for(
+  "fresh.upgradeSourceMap",
+) as typeof UPGRADE_SOURCE_KEY;
+// deno-lint-ignore no-explicit-any
+type UpgradeSource = { socket: any; head: any };
+// deno-lint-ignore no-explicit-any
+const globalAny = globalThis as any;
+export const upgradeSourceMap: WeakMap<Request, UpgradeSource> =
+  globalAny[UPGRADE_SOURCE_KEY] ??
+    (globalAny[UPGRADE_SOURCE_KEY] = new WeakMap<Request, UpgradeSource>());
+
 /**
  * Duck-type check: treats the argument as managed-mode handlers when at least
  * one of the handler keys (`open`, `message`, `close`, `error`) is a
@@ -595,7 +617,16 @@ export class Context<State> {
       throw new HttpError(400, "Expected a WebSocket upgrade request");
     }
 
-    const { socket, response } = Deno.upgradeWebSocket(this.req, options);
+    const source = upgradeSourceMap.get(this.req);
+    if (source !== undefined) upgradeSourceMap.delete(this.req);
+    const upgradeOptions = source
+      ? { ...options, socket: source.socket, head: source.head }
+      : options;
+    const { socket, response } = Deno.upgradeWebSocket(
+      this.req,
+      // deno-lint-ignore no-explicit-any
+      upgradeOptions as any,
+    );
 
     if (handlers === undefined) {
       return { socket, response };

packages/fresh/src/internals.ts

@@ -4,3 +4,4 @@ export { setBuildCache, setErrorInterceptor } from "./app.ts";
 export { IslandPreparer, ProdBuildCache } from "./build_cache.ts";
 export { path };
 export { ASSET_CACHE_BUST_KEY } from "./constants.ts";
+export { upgradeSourceMap } from "./context.ts";

packages/plugin-vite/demo/routes/tests/ws_echo.ts

@@ -0,0 +1,11 @@
+import { define } from "../../utils.ts";
+
+export const handler = define.handlers({
+  GET(ctx) {
+    return ctx.upgrade({
+      message(socket, e) {
+        socket.send("echo: " + e.data);
+      },
+    });
+  },
+});

packages/plugin-vite/src/plugins/dev_server.ts

@@ -1,7 +1,7 @@
 import type { DevEnvironment, Plugin } from "vite";
 import * as path from "@std/path";
 import { contentType as getStdContentType } from "@std/media-types/content-type";
-import { ASSET_CACHE_BUST_KEY } from "fresh/internal";
+import { ASSET_CACHE_BUST_KEY, upgradeSourceMap } from "fresh/internal";
 import { createRequest, sendResponse } from "@remix-run/node-fetch-server";
 import { hashCode } from "../shared.ts";
 import type { ResolvedFreshViteConfig } from "../utils.ts";
@@ -28,6 +28,75 @@ export function devServer(freshConfig: ResolvedFreshViteConfig): Plugin[] {
           `^(${base})?/(@(vite|fs|id)|\\.vite)/`,
         );
 
+        // WebSocket upgrade requests don't flow through Connect middleware —
+        // node:http dispatches them as `upgrade` events on the http.Server.
+        // Vite's own HMR server identifies its connections via the
+        // `sec-websocket-protocol` header (`vite-hmr`/`vite-ping`), so we
+        // skip those and let Vite handle them. Everything else we route into
+        // the Fresh handler with the raw socket + head stashed in
+        // `upgradeSourceMap` so `ctx.upgrade()` can finish the handshake via
+        // `Deno.upgradeWebSocket(req, { socket, head })`.
+        server.httpServer?.on("upgrade", (nodeReq, nodeSocket, head) => {
+          const protocolHeader = nodeReq.headers["sec-websocket-protocol"];
+          if (
+            typeof protocolHeader === "string" &&
+            /\b(vite-hmr|vite-ping)\b/.test(protocolHeader)
+          ) {
+            return;
+          }
+
+          (async () => {
+            try {
+              const serverCfg = server.config.server;
+              const protocol = serverCfg.https ? "https" : "http";
+              const host = nodeReq.headers.host ?? "localhost";
+              const url = new URL(
+                nodeReq.url ?? "/",
+                `${protocol}://${host}`,
+              );
+
+              const headers = new Headers();
+              for (const [key, value] of Object.entries(nodeReq.headers)) {
+                if (value === undefined || value === null) continue;
+                if (Array.isArray(value)) {
+                  for (const v of value) {
+                    if (typeof v === "string") headers.append(key, v);
+                  }
+                } else if (typeof value === "string") {
+                  headers.set(key, value);
+                }
+              }
+
+              const req = new Request(url, {
+                method: nodeReq.method ?? "GET",
+                headers,
+              });
+
+              upgradeSourceMap.set(req, { socket: nodeSocket, head });
+
+              const mod = await server.ssrLoadModule("fresh:server_entry");
+              mod.setErrorInterceptor((err: unknown) => {
+                if (err instanceof Error) server.ssrFixStacktrace(err);
+              });
+
+              await mod.default.fetch(req);
+              // If the handler called ctx.upgrade(), the 101 has already been
+              // written to nodeSocket and the WebSocket is live. If it didn't,
+              // we have no usable channel left to report an error on — close
+              // the socket so the client doesn't hang.
+              if (!nodeSocket.destroyed && upgradeSourceMap.has(req)) {
+                nodeSocket.destroy();
+              }
+              upgradeSourceMap.delete(req);
+            } catch (err) {
+              if (err instanceof Error) server.ssrFixStacktrace(err);
+              // deno-lint-ignore no-console
+              console.error("[fresh] WebSocket upgrade failed:", err);
+              if (!nodeSocket.destroyed) nodeSocket.destroy();
+            }
+          })();
+        });
+
         server.middlewares.use(async (nodeReq, nodeRes, next) => {
           const serverCfg = server.config.server;
 

packages/plugin-vite/tests/dev_server_test.ts

@@ -27,6 +27,30 @@ integrationTest("vite dev - launches", async () => {
   expect(text).toContain("it works");
 });
 
+integrationTest("vite dev - upgrades WebSocket connections", async () => {
+  const address = demoServer.address();
+  const wsUrl = address.replace(/^http/, "ws") + "/tests/ws_echo";
+
+  const ws = new WebSocket(wsUrl);
+  const reply = await new Promise<string>((resolve, reject) => {
+    const timer = setTimeout(
+      () => reject(new Error("timed out waiting for echo")),
+      5000,
+    );
+    ws.onopen = () => ws.send("hello");
+    ws.onmessage = (e) => {
+      clearTimeout(timer);
+      resolve(typeof e.data === "string" ? e.data : "");
+    };
+    ws.onerror = () => {
+      clearTimeout(timer);
+      reject(new Error("websocket errored"));
+    };
+  });
+  ws.close();
+  expect(reply).toBe("echo: hello");
+});
+
 integrationTest("vite dev - serves static files", async () => {
   const res = await fetch(`${demoServer.address()}/test_static/foo.txt`);
   const text = await res.text();