From 05bf02e73f5f3262e17f6633c56bb27bfdaf61c8 Mon Sep 17 00:00:00 2001 From: talge-a11y Date: Thu, 15 Oct 2026 10:22:17 +0300 Subject: [PATCH 1/3] support stagehand in backend functions --- packages/functions-compiler/README.md | 15 +++ .../functions-compiler/src/deno-bundle.ts | 3 + .../src/esbuild/ws-native-virtual.ts | 101 ++++++++++++++++++ packages/functions-compiler/src/shim/entry.ts | 3 + .../test/package-matrix-data.ts | 5 + .../test/workerd-runtime.e2e.test.ts | 53 +++++++++ 6 files changed, 180 insertions(+) create mode 100644 packages/functions-compiler/src/esbuild/ws-native-virtual.ts diff --git a/packages/functions-compiler/README.md b/packages/functions-compiler/README.md index 263e500b8..cee640826 100644 --- a/packages/functions-compiler/README.md +++ b/packages/functions-compiler/README.md @@ -168,6 +168,21 @@ enforced without making the manifest user-importable. Release this package before bumping apper's compiler pin. Existing compiled Workers need rebuilding; the workspace toggle does not install new code. +## Workers compatibility for npm clients + +Two gaps in workerd break common Node-targeted npm packages, Stagehand among +them, and the compiler closes both for every function and actor: + +- **`ws`.** Its exports map sends the `browser` condition to a stub that throws + on use. `src/esbuild/ws-native-virtual.ts` resolves bare `ws` and + `npm:ws[@version]` to a `ws`-shaped client over workerd's native WebSocket + (`on`/`once`/`off`, `send`, `close`, `terminate`, `readyState`). Connection + options such as custom headers are ignored, and `WebSocketServer` throws: a + function can open sockets but not listen. +- **`AsyncLocalStorage.enterWith()`.** workerd omits it. The Deno shim makes it + a no-op, which suits libraries that also pass their context explicitly and + re-enter it with `run()`. + ## Layout | Path | What it is | diff --git a/packages/functions-compiler/src/deno-bundle.ts b/packages/functions-compiler/src/deno-bundle.ts index 2ddf1f868..297942b6c 100644 --- a/packages/functions-compiler/src/deno-bundle.ts +++ b/packages/functions-compiler/src/deno-bundle.ts @@ -22,6 +22,7 @@ import { privateDataSourcesVirtualPlugin } from "./esbuild/private-data-sources- import { runtimeContextVirtualPlugin } from "./esbuild/runtime-context-virtual.js"; import { runtimeVirtualPlugin } from "./esbuild/runtime-virtual.js"; import { USER_NAMESPACE, userFilesPlugin } from "./esbuild/user-files.js"; +import { wsNativeVirtualPlugin } from "./esbuild/ws-native-virtual.js"; import { lazyEsbuild } from "./lazy-deps.js"; import type { PreparedWorker } from "./worker-entry.js"; @@ -82,6 +83,8 @@ export async function bundleToModule( // esbuild from dumping diagnostics to the service's stderr. logLevel: "silent", plugins: [ + // `ws` → native-WebSocket client; claims the import before any resolver. + wsNativeVirtualPlugin(), // CJS require() of a node builtin → static re-export (workerd throws // on esbuild's lowered __require). Runs before the Deno resolver. nodeBuiltinRequirePlugin(), diff --git a/packages/functions-compiler/src/esbuild/ws-native-virtual.ts b/packages/functions-compiler/src/esbuild/ws-native-virtual.ts new file mode 100644 index 000000000..3ad2a28a1 --- /dev/null +++ b/packages/functions-compiler/src/esbuild/ws-native-virtual.ts @@ -0,0 +1,101 @@ +import type { Plugin } from "esbuild"; + +const NAMESPACE = "base44-ws-native"; + +// Bare `ws` from a dependency, or `npm:ws[@version]` from user code. Subpaths +// (`ws/lib/...`) are internals no client code reaches for, so they resolve as-is. +const WS_SPECIFIER = /^(npm:)?ws(@[^/]+)?$/; + +// The `ws` package's exports map sends the "browser" condition to a stub that +// throws on use, and its Node build needs node:net. workerd has a native +// WebSocket client, so serve a `ws`-shaped client over it. CommonJS so both +// `import WebSocket from "ws"` and `require("ws")` get the class. +const ADAPTER = `"use strict"; +const toBuffer = (value) => (globalThis.Buffer ? globalThis.Buffer.from(value) : value); +class WebSocket { + static CONNECTING = 0; + static OPEN = 1; + static CLOSING = 2; + static CLOSED = 3; + #socket; + #listeners = new Map(); + constructor(address, protocols) { + this.#socket = new globalThis.WebSocket(String(address), typeof protocols === "string" || Array.isArray(protocols) ? protocols : undefined); + try { this.#socket.binaryType = "arraybuffer"; } catch {} + this.#socket.addEventListener("open", () => this.emit("open")); + this.#socket.addEventListener("message", (event) => { + const isBinary = typeof event.data !== "string"; + this.emit("message", toBuffer(isBinary ? new Uint8Array(event.data) : event.data), isBinary); + }); + this.#socket.addEventListener("error", (event) => this.emit("error", event.error ?? new Error(event.message || "WebSocket error"))); + this.#socket.addEventListener("close", (event) => this.emit("close", event.code, toBuffer(event.reason ?? ""))); + } + get readyState() { return this.#socket.readyState; } + get url() { return this.#socket.url; } + get protocol() { return this.#socket.protocol; } + on(name, listener) { return this.#add(name, listener, false); } + addListener(name, listener) { return this.#add(name, listener, false); } + once(name, listener) { return this.#add(name, listener, true); } + off(name, listener) { return this.removeListener(name, listener); } + removeListener(name, listener) { + const listeners = this.#listeners.get(name); + if (listeners) this.#listeners.set(name, listeners.filter((entry) => entry.listener !== listener)); + return this; + } + removeAllListeners(name) { + if (name === undefined) this.#listeners.clear(); + else this.#listeners.delete(name); + return this; + } + emit(name, ...args) { + const listeners = this.#listeners.get(name) ?? []; + this.#listeners.set(name, listeners.filter((entry) => !entry.once)); + for (const entry of listeners) entry.listener.apply(this, args); + return listeners.length > 0; + } + send(data, options, callback) { + const done = typeof options === "function" ? options : callback; + try { + this.#socket.send(data); + done?.(); + } catch (error) { + if (!done) throw error; + done(error); + } + } + close(code, reason) { this.#socket.close(code, reason); } + terminate() { this.#socket.close(); } + #add(name, listener, once) { + const listeners = this.#listeners.get(name) ?? []; + listeners.push({ listener, once }); + this.#listeners.set(name, listeners); + return this; + } +} +class WebSocketServer { + constructor() { + throw new Error("ws WebSocketServer is not available in Base44 functions: a function can open WebSocket connections but cannot listen for them."); + } +} +module.exports = WebSocket; +module.exports.WebSocket = WebSocket; +module.exports.WebSocketServer = WebSocketServer; +module.exports.default = WebSocket; +`; + +export function wsNativeVirtualPlugin(): Plugin { + return { + name: "base44-ws-native", + setup(build) { + // Every namespace (user files and dependencies), one shared module path. + build.onResolve({ filter: WS_SPECIFIER }, () => ({ + path: "ws", + namespace: NAMESPACE, + })); + build.onLoad({ filter: /.*/, namespace: NAMESPACE }, () => ({ + contents: ADAPTER, + loader: "js", + })); + }, + }; +} diff --git a/packages/functions-compiler/src/shim/entry.ts b/packages/functions-compiler/src/shim/entry.ts index 4fa4351e7..a6c991adc 100644 --- a/packages/functions-compiler/src/shim/entry.ts +++ b/packages/functions-compiler/src/shim/entry.ts @@ -67,6 +67,9 @@ const deno = Object.freeze({ (globalThis as unknown as { Deno: unknown }).Deno = deno; (globalThis as unknown as { Buffer: typeof Buffer }).Buffer = Buffer; +// workerd omits enterWith(); libraries such as Stagehand call it to seed a +// context they also pass explicitly, so a no-op keeps them working. +AsyncLocalStorage.prototype.enterWith = () => {}; export function getRegisteredHandler(): DenoServeHandler | null { return registeredHandler; diff --git a/packages/functions-compiler/test/package-matrix-data.ts b/packages/functions-compiler/test/package-matrix-data.ts index 42cdf6026..185d9021d 100644 --- a/packages/functions-compiler/test/package-matrix-data.ts +++ b/packages/functions-compiler/test/package-matrix-data.ts @@ -157,4 +157,9 @@ export const PACKAGES: PkgCase[] = [ { name: "testJimp", pkg: "jimp", version: "0.16.13", expected: "pass", source: fn("jimp", "0.16.13", "import Jimp from 'npm:jimp@0.16.13';", " const img = new Jimp(1, 1, 0xFF0000FF); const base64 = await img.getBase64Async(Jimp.MIME_PNG); const output = { hasData: base64.length > 0 };") }, + // Constructing calls AsyncLocalStorage.enterWith() (no-op'd by the shim) and + // loading imports `ws` (served as the native-WebSocket adapter). + { name: "testStagehand", pkg: "@browserbasehq/stagehand", version: "3.7.3", expected: "pass", source: fn("@browserbasehq/stagehand", "3.7.3", + "import { Stagehand } from 'npm:@browserbasehq/stagehand@3.7.3';", + " const stagehand = new Stagehand({ env: 'LOCAL', localBrowserLaunchOptions: { cdpUrl: 'ws://127.0.0.1:9/devtools/browser/x' }, model: { modelName: 'openai/gpt-4.1-mini', apiKey: 'x', openaiEndpointFormat: 'chat' }, disablePino: true }); const output = { hasInit: typeof stagehand.init === 'function', hasAct: typeof stagehand.act === 'function' };") }, ]; diff --git a/packages/functions-compiler/test/workerd-runtime.e2e.test.ts b/packages/functions-compiler/test/workerd-runtime.e2e.test.ts index f5a24b1f7..9a79d7c89 100644 --- a/packages/functions-compiler/test/workerd-runtime.e2e.test.ts +++ b/packages/functions-compiler/test/workerd-runtime.e2e.test.ts @@ -1111,4 +1111,57 @@ export default { hasFunctionsInvoke: true, }); }); + + it("ws: the npm client runs over workerd's native WebSocket", async () => { + const echo = new Miniflare({ + modules: true, + script: `export default { + fetch() { + const [client, server] = Object.values(new WebSocketPair()); + server.accept(); + server.addEventListener("message", (e) => server.send("echo:" + e.data)); + return new Response(null, { status: 101, webSocket: client }); + }, + };`, + compatibilityDate: WFP_COMPAT_DATE, + port: 0, + }); + try { + const wsUrl = (await echo.ready).href.replace(/^http/, "ws"); + const m = await bundled(` + import WebSocket from "npm:ws@8.18.0"; + Deno.serve(async (req) => { + const ws = new WebSocket(new URL(req.url).searchParams.get("ws")); + const reply = await new Promise((resolve, reject) => { + ws.once("open", () => ws.send("hi")); + ws.once("message", (data, isBinary) => resolve(String(data) + "|" + isBinary)); + ws.once("error", reject); + }); + ws.close(); + return new Response(reply + "|" + (WebSocket.OPEN === 1)); + }); + `); + const { status, text } = await runInWorkerd(m, { + url: `http://localhost/?ws=${encodeURIComponent(wsUrl)}`, + }); + expect(status).toBe(200); + expect(text).toBe("echo:hi|false|true"); + } finally { + await echo.dispose(); + } + }); + + it("AsyncLocalStorage.enterWith() is a no-op instead of throwing", async () => { + const m = await bundled(` + import { AsyncLocalStorage } from "node:async_hooks"; + Deno.serve(() => { + const store = new AsyncLocalStorage(); + store.enterWith({ seeded: true }); + return new Response(String(store.getStore())); + }); + `); + const { status, text } = await runInWorkerd(m); + expect(status).toBe(200); + expect(text).toBe("undefined"); + }); }); From 05f10d92a0e8406b3401b1d55605ef8856a15bb0 Mon Sep 17 00:00:00 2001 From: talge-a11y Date: Thu, 15 Oct 2026 11:35:18 +0300 Subject: [PATCH 2/3] review fixes --- packages/functions-compiler/README.md | 16 ++-- .../src/esbuild/ws-native-virtual.ts | 93 +++++++++++++++---- packages/functions-compiler/src/shim/entry.ts | 12 ++- .../test/workerd-runtime.e2e.test.ts | 37 ++++++-- 4 files changed, 124 insertions(+), 34 deletions(-) diff --git a/packages/functions-compiler/README.md b/packages/functions-compiler/README.md index cee640826..7c172d6da 100644 --- a/packages/functions-compiler/README.md +++ b/packages/functions-compiler/README.md @@ -176,12 +176,16 @@ them, and the compiler closes both for every function and actor: - **`ws`.** Its exports map sends the `browser` condition to a stub that throws on use. `src/esbuild/ws-native-virtual.ts` resolves bare `ws` and `npm:ws[@version]` to a `ws`-shaped client over workerd's native WebSocket - (`on`/`once`/`off`, `send`, `close`, `terminate`, `readyState`). Connection - options such as custom headers are ignored, and `WebSocketServer` throws: a - function can open sockets but not listen. -- **`AsyncLocalStorage.enterWith()`.** workerd omits it. The Deno shim makes it - a no-op, which suits libraries that also pass their context explicitly and - re-enter it with `run()`. + (`on`/`once`/`off`, `send`, `close`, `terminate`, `readyState`). It covers + what Stagehand uses, not the whole `ws` API. When the options carry headers + beyond `User-Agent` (Stagehand's `cdpHeaders`), the socket is opened with a + `fetch` upgrade so they reach the server, and that connection follows the + function's egress routing. Other connection options are ignored, and + `WebSocketServer` throws: a function can open sockets but not listen. +- **`AsyncLocalStorage.enterWith()`.** workerd's version throws. The Deno shim + replaces it with a no-op, so the store is never seeded. That suits libraries + that also pass their context explicitly and re-enter it with `run()`. The + shim probes first and keeps a working native implementation. ## Layout diff --git a/packages/functions-compiler/src/esbuild/ws-native-virtual.ts b/packages/functions-compiler/src/esbuild/ws-native-virtual.ts index 3ad2a28a1..6a63d1b86 100644 --- a/packages/functions-compiler/src/esbuild/ws-native-virtual.ts +++ b/packages/functions-compiler/src/esbuild/ws-native-virtual.ts @@ -10,29 +10,45 @@ const WS_SPECIFIER = /^(npm:)?ws(@[^/]+)?$/; // throws on use, and its Node build needs node:net. workerd has a native // WebSocket client, so serve a `ws`-shaped client over it. CommonJS so both // `import WebSocket from "ws"` and `require("ws")` get the class. +// +// The native constructor cannot send request headers, so when the caller passes +// `options.headers` (Stagehand's `cdpHeaders`, e.g. a remote browser's auth +// token) the socket is opened with a fetch upgrade instead, which can. A lone +// User-Agent (Stagehand always sends one) is not worth leaving the native path. const ADAPTER = `"use strict"; const toBuffer = (value) => (globalThis.Buffer ? globalThis.Buffer.from(value) : value); +const READY_STATES = ["CONNECTING", "OPEN", "CLOSING", "CLOSED"]; +const upgradeHeaders = (options) => { + const headers = new Headers(options?.headers); + for (const name of headers.keys()) if (name !== "user-agent") return headers; + return null; +}; class WebSocket { static CONNECTING = 0; static OPEN = 1; static CLOSING = 2; static CLOSED = 3; - #socket; + #socket = null; + #state = 0; + #url; + #closeRequested = null; #listeners = new Map(); - constructor(address, protocols) { - this.#socket = new globalThis.WebSocket(String(address), typeof protocols === "string" || Array.isArray(protocols) ? protocols : undefined); - try { this.#socket.binaryType = "arraybuffer"; } catch {} - this.#socket.addEventListener("open", () => this.emit("open")); - this.#socket.addEventListener("message", (event) => { - const isBinary = typeof event.data !== "string"; - this.emit("message", toBuffer(isBinary ? new Uint8Array(event.data) : event.data), isBinary); - }); - this.#socket.addEventListener("error", (event) => this.emit("error", event.error ?? new Error(event.message || "WebSocket error"))); - this.#socket.addEventListener("close", (event) => this.emit("close", event.code, toBuffer(event.reason ?? ""))); + constructor(address, protocols, options) { + if (protocols && typeof protocols === "object" && !Array.isArray(protocols)) { + options = protocols; + protocols = undefined; + } + this.#url = String(address); + const headers = upgradeHeaders(options); + if (headers) { + this.#upgrade(headers, protocols); + } else { + this.#attach(new globalThis.WebSocket(this.#url, protocols)); + } } - get readyState() { return this.#socket.readyState; } - get url() { return this.#socket.url; } - get protocol() { return this.#socket.protocol; } + get readyState() { return this.#socket ? this.#socket.readyState : this.#state; } + get url() { return this.#url; } + get protocol() { return this.#socket ? this.#socket.protocol ?? "" : ""; } on(name, listener) { return this.#add(name, listener, false); } addListener(name, listener) { return this.#add(name, listener, false); } once(name, listener) { return this.#add(name, listener, true); } @@ -56,6 +72,9 @@ class WebSocket { send(data, options, callback) { const done = typeof options === "function" ? options : callback; try { + if (!this.#socket) { + throw new Error("WebSocket is not open: readyState " + this.#state + " (" + READY_STATES[this.#state] + ")"); + } this.#socket.send(data); done?.(); } catch (error) { @@ -63,8 +82,50 @@ class WebSocket { done(error); } } - close(code, reason) { this.#socket.close(code, reason); } - terminate() { this.#socket.close(); } + close(code, reason) { + if (this.#socket) this.#socket.close(code, reason); + else if (this.#state === 0) { + // Upgrade still in flight: close as soon as it lands. + this.#closeRequested = [code, reason]; + this.#state = 2; + } + } + terminate() { this.close(); } + #attach(socket) { + this.#socket = socket; + try { socket.binaryType = "arraybuffer"; } catch {} + socket.addEventListener("open", () => this.emit("open")); + socket.addEventListener("message", (event) => { + const isBinary = typeof event.data !== "string"; + this.emit("message", toBuffer(isBinary ? new Uint8Array(event.data) : event.data), isBinary); + }); + socket.addEventListener("error", (event) => this.emit("error", event.error ?? new Error(event.message || "WebSocket error"))); + socket.addEventListener("close", (event) => this.emit("close", event.code, toBuffer(event.reason ?? ""))); + } + #upgrade(headers, protocols) { + headers.set("Upgrade", "websocket"); + if (protocols) headers.set("Sec-WebSocket-Protocol", [].concat(protocols).join(", ")); + const httpUrl = this.#url.replace(/^ws(s?):/i, "http$1:"); + const fail = (error) => { + this.#state = 3; + this.emit("error", error); + this.emit("close", 1006, toBuffer("")); + }; + // Two-argument then(): an exception from an "open" listener must not be + // reported as a failed upgrade. + fetch(httpUrl, { headers }).then((response) => { + const socket = response.webSocket; + if (!socket) { + response.body?.cancel(); + return fail(new Error("Unexpected server response: " + response.status)); + } + socket.accept(); + this.#attach(socket); + if (this.#closeRequested) return socket.close(...this.#closeRequested); + // accept() opens the socket without dispatching an "open" event. + this.emit("open"); + }, fail); + } #add(name, listener, once) { const listeners = this.#listeners.get(name) ?? []; listeners.push({ listener, once }); diff --git a/packages/functions-compiler/src/shim/entry.ts b/packages/functions-compiler/src/shim/entry.ts index a6c991adc..b642a2efc 100644 --- a/packages/functions-compiler/src/shim/entry.ts +++ b/packages/functions-compiler/src/shim/entry.ts @@ -67,9 +67,15 @@ const deno = Object.freeze({ (globalThis as unknown as { Deno: unknown }).Deno = deno; (globalThis as unknown as { Buffer: typeof Buffer }).Buffer = Buffer; -// workerd omits enterWith(); libraries such as Stagehand call it to seed a -// context they also pass explicitly, so a no-op keeps them working. -AsyncLocalStorage.prototype.enterWith = () => {}; +// workerd's enterWith() is a stub that throws; libraries such as Stagehand +// call it to seed a context they also pass explicitly, so a no-op keeps them +// working. Probed on a throwaway instance so a future native implementation is +// left in place. +try { + new AsyncLocalStorage().enterWith(undefined); +} catch { + AsyncLocalStorage.prototype.enterWith = () => {}; +} export function getRegisteredHandler(): DenoServeHandler | null { return registeredHandler; diff --git a/packages/functions-compiler/test/workerd-runtime.e2e.test.ts b/packages/functions-compiler/test/workerd-runtime.e2e.test.ts index 9a79d7c89..e804bf3a7 100644 --- a/packages/functions-compiler/test/workerd-runtime.e2e.test.ts +++ b/packages/functions-compiler/test/workerd-runtime.e2e.test.ts @@ -1112,14 +1112,30 @@ export default { }); }); - it("ws: the npm client runs over workerd's native WebSocket", async () => { + // Without custom headers the adapter uses the native WebSocket constructor; + // with them (Stagehand's cdpHeaders) it opens the socket by fetch upgrade. + it.each([ + { label: "native constructor", options: "{}", auth: "null" }, + { + label: "fetch upgrade carrying headers", + options: '{ headers: { Authorization: "Bearer t", "User-Agent": "Stagehand/x" } }', + auth: "Bearer t", + }, + ])("ws: the npm client runs over workerd's WebSocket ($label)", async ({ + options, + auth, + }) => { const echo = new Miniflare({ modules: true, script: `export default { - fetch() { + fetch(req) { const [client, server] = Object.values(new WebSocketPair()); server.accept(); - server.addEventListener("message", (e) => server.send("echo:" + e.data)); + const authorization = req.headers.get("authorization"); + server.addEventListener("message", (e) => server.send("echo:" + e.data + ":" + authorization)); + server.addEventListener("close", (e) => { + try { server.close(e.code, "bye"); } catch {} + }); return new Response(null, { status: 101, webSocket: client }); }, };`, @@ -1131,21 +1147,24 @@ export default { const m = await bundled(` import WebSocket from "npm:ws@8.18.0"; Deno.serve(async (req) => { - const ws = new WebSocket(new URL(req.url).searchParams.get("ws")); + const ws = new WebSocket(new URL(req.url).searchParams.get("ws"), ${options}); const reply = await new Promise((resolve, reject) => { ws.once("open", () => ws.send("hi")); ws.once("message", (data, isBinary) => resolve(String(data) + "|" + isBinary)); ws.once("error", reject); }); - ws.close(); - return new Response(reply + "|" + (WebSocket.OPEN === 1)); + // Stagehand's shutdown closes and then waits for "close". + const closed = new Promise((resolve) => ws.once("close", (code) => resolve(code))); + ws.close(1000); + const code = await closed; + return new Response(reply + "|" + code + "|" + (ws.readyState === WebSocket.CLOSED)); }); `); const { status, text } = await runInWorkerd(m, { url: `http://localhost/?ws=${encodeURIComponent(wsUrl)}`, }); expect(status).toBe(200); - expect(text).toBe("echo:hi|false|true"); + expect(text).toBe(`echo:hi:${auth}|false|1000|true`); } finally { await echo.dispose(); } @@ -1157,11 +1176,11 @@ export default { Deno.serve(() => { const store = new AsyncLocalStorage(); store.enterWith({ seeded: true }); - return new Response(String(store.getStore())); + return new Response("ok"); }); `); const { status, text } = await runInWorkerd(m); expect(status).toBe(200); - expect(text).toBe("undefined"); + expect(text).toBe("ok"); }); }); From 9dfc12efc7784fca8a953975dba92687a94fe601 Mon Sep 17 00:00:00 2001 From: talge-a11y Date: Thu, 15 Oct 2026 12:02:13 +0300 Subject: [PATCH 3/3] align with vite plugin behavior --- .../src/esbuild/ws-native-virtual.ts | 27 ++++- .../test/workerd-runtime.e2e.test.ts | 101 ++++++++++++++++++ 2 files changed, 125 insertions(+), 3 deletions(-) diff --git a/packages/functions-compiler/src/esbuild/ws-native-virtual.ts b/packages/functions-compiler/src/esbuild/ws-native-virtual.ts index 6a63d1b86..43f655aed 100644 --- a/packages/functions-compiler/src/esbuild/ws-native-virtual.ts +++ b/packages/functions-compiler/src/esbuild/ws-native-virtual.ts @@ -18,6 +18,7 @@ const WS_SPECIFIER = /^(npm:)?ws(@[^/]+)?$/; const ADAPTER = `"use strict"; const toBuffer = (value) => (globalThis.Buffer ? globalThis.Buffer.from(value) : value); const READY_STATES = ["CONNECTING", "OPEN", "CLOSING", "CLOSED"]; +const later = (fn, ...args) => queueMicrotask(() => fn(...args)); const upgradeHeaders = (options) => { const headers = new Headers(options?.headers); for (const name of headers.keys()) if (name !== "user-agent") return headers; @@ -32,6 +33,7 @@ class WebSocket { #state = 0; #url; #closeRequested = null; + #terminated = false; #listeners = new Map(); constructor(address, protocols, options) { if (protocols && typeof protocols === "object" && !Array.isArray(protocols)) { @@ -46,7 +48,10 @@ class WebSocket { this.#attach(new globalThis.WebSocket(this.#url, protocols)); } } - get readyState() { return this.#socket ? this.#socket.readyState : this.#state; } + get readyState() { + if (this.#terminated) return 3; + return this.#socket ? this.#socket.readyState : this.#state; + } get url() { return this.#url; } get protocol() { return this.#socket ? this.#socket.protocol ?? "" : ""; } on(name, listener) { return this.#add(name, listener, false); } @@ -90,7 +95,20 @@ class WebSocket { this.#state = 2; } } - terminate() { this.close(); } + // \`ws\` tears the connection down at once. workerd has no abrupt close, so the + // socket is closed normally and "close" is reported now, not on the reply. + terminate() { + if (this.#terminated || this.readyState === 3) return; + this.#terminated = true; + if (this.#socket) { + try { this.#socket.close(); } catch {} + } else { + // Upgrade still in flight: close it on arrival, without an "open". + this.#closeRequested = []; + this.#state = 3; + } + later(() => this.emit("close", 1006, toBuffer(""))); + } #attach(socket) { this.#socket = socket; try { socket.binaryType = "arraybuffer"; } catch {} @@ -100,7 +118,9 @@ class WebSocket { this.emit("message", toBuffer(isBinary ? new Uint8Array(event.data) : event.data), isBinary); }); socket.addEventListener("error", (event) => this.emit("error", event.error ?? new Error(event.message || "WebSocket error"))); - socket.addEventListener("close", (event) => this.emit("close", event.code, toBuffer(event.reason ?? ""))); + socket.addEventListener("close", (event) => { + if (!this.#terminated) this.emit("close", event.code, toBuffer(event.reason ?? "")); + }); } #upgrade(headers, protocols) { headers.set("Upgrade", "websocket"); @@ -108,6 +128,7 @@ class WebSocket { const httpUrl = this.#url.replace(/^ws(s?):/i, "http$1:"); const fail = (error) => { this.#state = 3; + if (this.#terminated) return; this.emit("error", error); this.emit("close", 1006, toBuffer("")); }; diff --git a/packages/functions-compiler/test/workerd-runtime.e2e.test.ts b/packages/functions-compiler/test/workerd-runtime.e2e.test.ts index e804bf3a7..9b48ccb20 100644 --- a/packages/functions-compiler/test/workerd-runtime.e2e.test.ts +++ b/packages/functions-compiler/test/workerd-runtime.e2e.test.ts @@ -1170,6 +1170,107 @@ export default { } }); + // Close and terminate follow `ws`, within what workerd's WHATWG socket allows. + describe("ws: closing", () => { + async function runAgainstEcho(body: string): Promise { + const echo = new Miniflare({ + modules: true, + script: `export default { + fetch() { + const [client, server] = Object.values(new WebSocketPair()); + server.accept(); + server.addEventListener("close", (e) => { + try { server.close(e.code, "bye"); } catch {} + }); + return new Response(null, { status: 101, webSocket: client }); + }, + };`, + compatibilityDate: WFP_COMPAT_DATE, + port: 0, + }); + try { + const wsUrl = (await echo.ready).href.replace(/^http/, "ws"); + const m = await bundled(` + import WebSocket from "npm:ws@8.18.0"; + const settle = () => new Promise((resolve) => setTimeout(resolve, 200)); + Deno.serve(async (req) => { + const url = new URL(req.url).searchParams.get("ws"); + const closes = []; + ${body} + }); + `); + const { status, text } = await runInWorkerd(m, { + url: `http://localhost/?ws=${encodeURIComponent(wsUrl)}`, + }); + expect(status).toBe(200); + return text; + } finally { + await echo.dispose(); + } + } + + // workerd accepts and rejects the same close codes \`ws\` does, so codes pass + // through unchanged and an invalid one throws as it would in Node. + it("passes close codes through and rejects invalid ones like ws", async () => { + const text = await runAgainstEcho(` + const ws = new WebSocket(url); + await new Promise((resolve) => ws.once("open", resolve)); + let invalid = "accepted"; + try { ws.close(999); } catch { invalid = "threw"; } + const closed = new Promise((resolve) => ws.once("close", (code) => resolve(code))); + ws.close(1001); + return new Response(invalid + "|" + (await closed)); + `); + expect(text).toBe("threw|1001"); + }); + + it("terminate() reports one 1006 close at once and ignores the server's reply", async () => { + const text = await runAgainstEcho(` + const ws = new WebSocket(url); + await new Promise((resolve) => ws.once("open", resolve)); + ws.on("close", (code) => closes.push(code)); + ws.terminate(); + const stateAfter = ws.readyState; + await settle(); + return new Response(closes.join(",") + "|" + (stateAfter === WebSocket.CLOSED)); + `); + expect(text).toBe("1006|true"); + }); + + it("terminate() during a header upgrade never opens and closes once", async () => { + const text = await runAgainstEcho(` + const ws = new WebSocket(url, { headers: { Authorization: "Bearer t" } }); + let opened = false; + ws.on("open", () => { opened = true; }); + ws.on("close", (code) => closes.push(code)); + ws.terminate(); + await settle(); + return new Response(closes.join(",") + "|" + opened + "|" + ws.readyState); + `); + expect(text).toBe("1006|false|3"); + }); + + it("terminate() is a no-op when repeated or after the socket closed", async () => { + const text = await runAgainstEcho(` + const twice = new WebSocket(url); + await new Promise((resolve) => twice.once("open", resolve)); + twice.on("close", (code) => closes.push("twice:" + code)); + twice.terminate(); + twice.terminate(); + const after = new WebSocket(url); + await new Promise((resolve) => after.once("open", resolve)); + const closed = new Promise((resolve) => after.once("close", resolve)); + after.close(1000); + await closed; + after.on("close", (code) => closes.push("after:" + code)); + after.terminate(); + await settle(); + return new Response(closes.join(",")); + `); + expect(text).toBe("twice:1006"); + }); + }); + it("AsyncLocalStorage.enterWith() is a no-op instead of throwing", async () => { const m = await bundled(` import { AsyncLocalStorage } from "node:async_hooks";