[vitest-pool-workers] Treat webSocket* as optional Durable Object handlers

.changeset/olive-jokes-shave.md

@@ -0,0 +1,9 @@
+---
+"@cloudflare/vitest-pool-workers": patch
+---
+
+Treat `webSocketMessage()`, `webSocketClose()` and `webSocketError()` as optional Durable Object handlers
+
+The pool wraps each Durable Object class and installs a prototype method for every default handler before user code is loaded, so `workerd` always sees a handler and always dispatches. When the wrapped class didn't actually define one, the wrapper threw ``<ClassName> exported by <path> does not define a `webSocketClose()` method``, even though deployed Workers silently ignore these events for classes that omit them. A hibernatable Durable Object defining only `webSocketMessage()` would log an uncaught `TypeError` on every close.
+
+These three handlers now no-op when absent, matching deployed behaviour. `alarm()` is unchanged and still reports a missing handler, since `workerd` rejects `setAlarm()` up front on a class without one.

fixtures/vitest-pool-workers-examples/durable-objects/src/index.ts

@@ -48,10 +48,6 @@ export class Counter extends DurableObject {
 		this.#webSocketMessages.push(value);
 		ws.send(value);
 	}
-
-	webSocketClose() {}
-
-	webSocketError() {}
 }
 
 export class SQLiteDurableObject extends DurableObject {

packages/vitest-pool-workers/src/worker/entrypoints.ts

@@ -247,6 +247,17 @@ const DURABLE_OBJECT_KEYS = [
 	"webSocketClose",
 	"webSocketError",
 ] as const;
+// Handlers `workerd` treats as optional: when a Durable Object doesn't define
+// one, the corresponding event is silently dropped. We can't leave them off the
+// wrapper's prototype to get that behaviour for free — the prototype is built
+// before user code is loaded, so `workerd` always sees a handler and always
+// dispatches. The no-op has to be reproduced here instead.
+//
+// `alarm()` deliberately isn't in this set: `workerd` rejects `setAlarm()` up
+// front on a class with no `alarm()` handler, so a missing one is a real error.
+const OPTIONAL_DURABLE_OBJECT_KEYS: ReadonlySet<
+	(typeof DURABLE_OBJECT_KEYS)[number]
+> = new Set(["webSocketMessage", "webSocketClose", "webSocketError"] as const);
 
 // This type will grab the keys from T and remove "branded" keys
 type UnbrandedKeys<T> = Exclude<keyof T, `__${string}_BRAND`>;
@@ -555,6 +566,8 @@ export function createDurableObjectWrapper(
 					const maybeFn = instance[key];
 					if (typeof maybeFn === "function") {
 						return (maybeFn as (...a: unknown[]) => void).apply(instance, args);
+					} else if (OPTIONAL_DURABLE_OBJECT_KEYS.has(key)) {
+						return;
 					} else {
 						const message = `${className} exported by ${mainPath} does not define a \`${key}()\` method`;
 						throw new TypeError(message);

packages/vitest-pool-workers/test/bindings.test.ts

@@ -52,6 +52,103 @@ test("hello_world support", async ({ expect, seed, vitestRun }) => {
 	await expect(result.exitCode).resolves.toBe(0);
 });
 
+test("Durable Objects may omit optional WebSocket handlers", async ({
+	expect,
+	seed,
+	vitestRun,
+}) => {
+	await seed({
+		"vitest.config.mts": vitestConfig({
+			wrangler: { configPath: "./wrangler.jsonc" },
+		}),
+		"wrangler.jsonc": dedent`
+			{
+				"name": "test-worker",
+				"main": "./index.ts",
+				"compatibility_date": "2025-12-02",
+				"compatibility_flags": ["nodejs_compat"],
+				"durable_objects": {
+					"bindings": [
+						{ "name": "OPTIONAL_WS", "class_name": "OptionalWebSocketHandlers" }
+					]
+				},
+				"migrations": [
+					{ "tag": "v1", "new_sqlite_classes": ["OptionalWebSocketHandlers"] }
+				]
+			}
+		`,
+		"index.ts": dedent /* javascript */ `
+			import { DurableObject } from "cloudflare:workers";
+
+			// Defines webSocketMessage() but deliberately omits webSocketClose()
+			// and webSocketError(), which workerd treats as optional handlers.
+			export class OptionalWebSocketHandlers extends DurableObject {
+				fetch(request) {
+					if (request.headers.get("Upgrade") !== "websocket") {
+						return new Response("ok");
+					}
+					const { 0: client, 1: server } = new WebSocketPair();
+					this.ctx.acceptWebSocket(server);
+					return new Response(null, { status: 101, webSocket: client });
+				}
+
+				webSocketMessage(ws, message) {
+					ws.send("echo:" + message);
+				}
+
+				socketCount() {
+					return this.ctx.getWebSockets().length;
+				}
+			}
+
+			export default {
+				async fetch() { return new Response("ok"); },
+			};
+		`,
+		"index.test.ts": dedent /* javascript */ `
+			import { env } from "cloudflare:workers";
+			import { it, vi } from "vitest";
+
+			it("echoes a message and closes without a webSocketClose() handler", async ({ expect }) => {
+				const stub = env.OPTIONAL_WS.get(env.OPTIONAL_WS.idFromName("ws"));
+				const response = await stub.fetch("https://example.com", {
+					headers: { Upgrade: "websocket" },
+				});
+				const socket = response.webSocket;
+				if (!socket) { throw new Error("Expected WebSocket response"); }
+
+				const message = new Promise((resolve) => {
+					socket.addEventListener("message", (event) => resolve(event.data));
+				});
+				socket.accept();
+				socket.send("hello");
+				expect(await message).toBe("echo:hello");
+
+				// Dispatches webSocketClose() on the Durable Object, which doesn't define it
+				socket.close(1000, "done");
+
+				// Wait until the runtime has actually processed the close, so the
+				// dispatch has definitely happened before the run ends
+				await vi.waitFor(
+					async () => {
+						expect(await stub.socketCount()).toBe(0);
+					},
+					{ timeout: 5_000, interval: 100 }
+				);
+			});
+		`,
+	});
+
+	const result = await vitestRun();
+
+	await expect(result.exitCode).resolves.toBe(0);
+	// Dispatching to the absent handlers must be a no-op. Previously the wrapper
+	// threw "<ClassName> exported by <path> does not define a `webSocketClose()`
+	// method", surfacing as an uncaught exception from the Durable Object.
+	expect(result.stderr).not.toMatch("does not define");
+	expect(result.stdout).not.toMatch("does not define");
+});
+
 test("adminSecretsStore seeds and reads secrets", async ({
 	expect,
 	seed,