fix(core): update websocket to use event listeners

Author: NathanFlurry

examples/rivet/package.json

@@ -14,7 +14,7 @@
     "typescript": "^5.5.2"
   },
   "dependencies": {
-    "@rivetkit/actor": "workspace:*"
+    "@rivetkit/actor": "https://pkg.pr.new/rivet-gg/rivetkit/@rivetkit/actor@40cc7ac"
   },
   "stableVersion": "0.8.0"
 }

examples/rivet/src/server.ts

@@ -1,3 +1,4 @@
 import { registry } from "./registry";
 
+console.log("new version");
 registry.runServer();

packages/core/src/client/actor-conn.ts

@@ -255,18 +255,18 @@ enc
 			signal ? { signal } : undefined,
 		);
 		this.#transport = { websocket: ws };
-		ws.onopen = () => {
+		ws.addEventListener("open", () => {
 			logger().debug("websocket open");
-		};
-		ws.onmessage = async (ev) => {
+		});
+		ws.addEventListener("message", async (ev) => {
 			this.#handleOnMessage(ev.data);
-		};
-		ws.onclose = (ev) => {
+		});
+		ws.addEventListener("close", (ev) => {
 			this.#handleOnClose(ev);
-		};
-		ws.onerror = (ev) => {
+		});
+		ws.addEventListener("error", (ev) => {
 			this.#handleOnError();
-		};
+		});
 	}
 
 	async #connectSse({ signal }: { signal?: AbortSignal } = {}) {

packages/core/src/drivers/rivet/ws-proxy.ts

@@ -26,7 +26,7 @@ export async function createWebSocketProxy(
 			targetWs = new WebSocket(targetUrl, { headers });
 
 			// Set up target websocket handlers
-			targetWs.onopen = () => {
+			targetWs.addEventListener("open", () => {
 				invariant(targetWs, "targetWs does not exist");
 
 				// Process any queued messages once connected
@@ -37,13 +37,13 @@ export async function createWebSocketProxy(
 					// Clear the queue after sending
 					messageQueue.length = 0;
 				}
-			};
+			});
 
-			targetWs.onmessage = (event: any) => {
+			targetWs.addEventListener("message", (event: any) => {
 				wsContext.send(event.data as any);
-			};
+			});
 
-			targetWs.onclose = (event: any) => {
+			targetWs.addEventListener("close", (event: any) => {
 				logger().debug("target websocket closed", {
 					code: event.code,
 					reason: event.reason,
@@ -53,17 +53,17 @@ export async function createWebSocketProxy(
 					// Forward the close code and reason from target to client
 					wsContext.close(event.code, event.reason);
 				}
-			};
+			});
 
-			targetWs.onerror = () => {
+			targetWs.addEventListener("error", () => {
 				logger().warn("target websocket error");
 
 				if (wsContext.readyState === WebSocket.OPEN) {
 					// Use standard WebSocket error code: 1006 - Abnormal Closure
 					// The connection was closed abnormally, e.g., without sending or receiving a Close control frame
 					wsContext.close(1006, "Error in target connection");
 				}
-			};
+			});
 		},
 
 		// Handle messages from client to target

packages/core/src/inline-client-driver/fake-websocket.ts

@@ -22,50 +22,15 @@ export class FakeWebSocket {
 	#wsContext: WSContext;
 	#readyState: 0 | 1 | 2 | 3 = 0; // Start in CONNECTING state
 	#queuedMessages: Array<string | ArrayBuffer | Uint8Array> = [];
-	// Event buffering is needed since onopen/onmessage events can be fired
-	// before JavaScript has a chance to assign handlers (e.g. within the same tick)
+	// Event buffering is needed since events can be fired
+	// before JavaScript has a chance to add event listeners (e.g. within the same tick)
 	#bufferedEvents: Array<{
-		type: "open" | "close" | "error" | "message";
+		type: string;
 		event: any;
 	}> = [];
 
-	// Event handlers with buffering
-	#onopen: ((ev: any) => void) | null = null;
-	#onclose: ((ev: any) => void) | null = null;
-	#onerror: ((ev: any) => void) | null = null;
-	#onmessage: ((ev: any) => void) | null = null;
-
-	get onopen() {
-		return this.#onopen;
-	}
-	set onopen(handler: ((ev: any) => void) | null) {
-		this.#onopen = handler;
-		if (handler) this.#flushBufferedEvents("open");
-	}
-
-	get onclose() {
-		return this.#onclose;
-	}
-	set onclose(handler: ((ev: any) => void) | null) {
-		this.#onclose = handler;
-		if (handler) this.#flushBufferedEvents("close");
-	}
-
-	get onerror() {
-		return this.#onerror;
-	}
-	set onerror(handler: ((ev: any) => void) | null) {
-		this.#onerror = handler;
-		if (handler) this.#flushBufferedEvents("error");
-	}
-
-	get onmessage() {
-		return this.#onmessage;
-	}
-	set onmessage(handler: ((ev: any) => void) | null) {
-		this.#onmessage = handler;
-		if (handler) this.#flushBufferedEvents("message");
-	}
+	// Event listeners with buffering
+	#eventListeners: Map<string, ((ev: any) => void)[]> = new Map();
 
 	constructor(handler: ConnectWebSocketOutput) {
 		this.#handler = handler;
@@ -303,13 +268,7 @@ export class FakeWebSocket {
 		} as unknown as MessageEvent;
 
 		// Dispatch the event
-		if (this.onmessage) {
-			logger().debug("dispatching message to onmessage handler");
-			this.onmessage(event);
-		} else {
-			logger().debug("no onmessage handler registered, buffering message");
-			this.#bufferedEvents.push({ type: "message", event });
-		}
+		this.#dispatchEvent("message", event);
 	}
 
 	#handleClose(code: number, reason: string): void {
@@ -328,14 +287,47 @@ export class FakeWebSocket {
 		} as unknown as CloseEvent;
 
 		// Dispatch the event
-		if (this.onclose) {
-			this.onclose(event);
+		this.#dispatchEvent("close", event);
+	}
+
+	addEventListener(type: string, listener: (ev: any) => void): void {
+		if (!this.#eventListeners.has(type)) {
+			this.#eventListeners.set(type, []);
+		}
+		this.#eventListeners.get(type)!.push(listener);
+		
+		// Flush any buffered events for this type
+		this.#flushBufferedEvents(type);
+	}
+
+	removeEventListener(type: string, listener: (ev: any) => void): void {
+		const listeners = this.#eventListeners.get(type);
+		if (listeners) {
+			const index = listeners.indexOf(listener);
+			if (index !== -1) {
+				listeners.splice(index, 1);
+			}
+		}
+	}
+
+	#dispatchEvent(type: string, event: any): void {
+		const listeners = this.#eventListeners.get(type);
+		if (listeners && listeners.length > 0) {
+			logger().debug(`dispatching ${type} event to ${listeners.length} listeners`);
+			for (const listener of listeners) {
+				try {
+					listener(event);
+				} catch (err) {
+					logger().error(`error in ${type} event listener`, { error: err });
+				}
+			}
 		} else {
-			this.#bufferedEvents.push({ type: "close", event });
+			logger().debug(`no ${type} listeners registered, buffering event`);
+			this.#bufferedEvents.push({ type, event });
 		}
 	}
 
-	#flushBufferedEvents(type: "open" | "close" | "error" | "message"): void {
+	#flushBufferedEvents(type: string): void {
 		const eventsToFlush = this.#bufferedEvents.filter(
 			(buffered) => buffered.type === type,
 		);
@@ -344,24 +336,7 @@ export class FakeWebSocket {
 		);
 
 		for (const { event } of eventsToFlush) {
-			try {
-				switch (type) {
-					case "open":
-						this.#onopen?.(event);
-						break;
-					case "close":
-						this.#onclose?.(event);
-						break;
-					case "error":
-						this.#onerror?.(event);
-						break;
-					case "message":
-						this.#onmessage?.(event);
-						break;
-				}
-			} catch (err) {
-				logger().error(`error in buffered ${type} handler`, { error: err });
-			}
+			this.#dispatchEvent(type, event);
 		}
 	}
 
@@ -374,25 +349,17 @@ export class FakeWebSocket {
 				currentTarget: this,
 			} as unknown as Event;
 
-			if (this.onopen) {
-				this.onopen(event);
-			} else {
-				this.#bufferedEvents.push({ type: "open", event });
-			}
+			this.#dispatchEvent("open", event);
 		} catch (err) {
-			logger().error("error in onopen handler", { error: err });
+			logger().error("error in open event", { error: err });
 		}
 	}
 
 	#fireClose(event: CloseEvent): void {
 		try {
-			if (this.onclose) {
-				this.onclose(event);
-			} else {
-				this.#bufferedEvents.push({ type: "close", event });
-			}
+			this.#dispatchEvent("close", event);
 		} catch (err) {
-			logger().error("error in onclose handler", { error: err });
+			logger().error("error in close event", { error: err });
 		}
 	}
 
@@ -407,13 +374,9 @@ export class FakeWebSocket {
 				message: error instanceof Error ? error.message : String(error),
 			} as unknown as Event;
 
-			if (this.onerror) {
-				this.onerror(event);
-			} else {
-				this.#bufferedEvents.push({ type: "error", event });
-			}
+			this.#dispatchEvent("error", event);
 		} catch (err) {
-			logger().error("error in onerror handler", { error: err });
+			logger().error("error in error event", { error: err });
 		}
 
 		// Log the error

packages/core/src/manager/router.ts

@@ -442,34 +442,37 @@ export function createManagerRouter(
 								clientWs = await clientWsPromise;
 
 								// Add message handler to forward messages from client to server
-								clientWs.onmessage = (clientEvt: MessageEvent) => {
-									logger().debug("test websocket connection message");
-
-									if (serverWs.readyState === 1) {
-										// OPEN
-										serverWs.send(clientEvt.data as any);
-									}
-								};
+								clientWs.addEventListener(
+									"message",
+									(clientEvt: MessageEvent) => {
+										logger().debug("test websocket connection message");
+
+										if (serverWs.readyState === 1) {
+											// OPEN
+											serverWs.send(clientEvt.data as any);
+										}
+									},
+								);
 
 								// Add close handler to close server when client closes
-								clientWs.onclose = (clientEvt: CloseEvent) => {
+								clientWs.addEventListener("close", (clientEvt: CloseEvent) => {
 									logger().debug("test websocket connection closed");
 
 									if (serverWs.readyState !== 3) {
 										// Not CLOSED
 										serverWs.close(clientEvt.code, clientEvt.reason);
 									}
-								};
+								});
 
 								// Add error handler
-								clientWs.onerror = () => {
+								clientWs.addEventListener("error", () => {
 									logger().debug("test websocket connection error");
 
 									if (serverWs.readyState !== 3) {
 										// Not CLOSED
 										serverWs.close(1011, "Error in client websocket");
 									}
-								};
+								});
 							} catch (error) {
 								logger().error(
 									"failed to establish client websocket connection",

packages/core/src/registry/run-config.ts

@@ -2,7 +2,7 @@ import type { ActorDriver } from "@/actor/driver";
 import { createMemoryDriver } from "@/drivers/memory/mod";
 import type { ManagerDriver } from "@/manager/driver";
 import type { CoordinateDriver } from "@/topologies/coordinate/driver";
-import type { UpgradeWebSocket } from "@/utils";
+import { createDefaultDriver, type UpgradeWebSocket } from "@/utils";
 import type { Hono } from "hono";
 import type { cors } from "hono/cors";
 import { z } from "zod";
@@ -60,7 +60,7 @@ export type DriverConfig = z.infer<typeof DriverConfigSchema>;
 /** Base config used for the actor config across all platforms. */
 export const RunConfigSchema = z
 	.object({
-		driver: DriverConfigSchema.optional().default(() => createMemoryDriver()),
+		driver: DriverConfigSchema.optional().default(() => createDefaultDriver()),
 
 		// This is dynamic since NodeJS requires a reference to the router to initialize WebSockets
 		getUpgradeWebSocket: z.custom<GetUpgradeWebSocket>().optional(),

packages/core/src/utils.ts

@@ -37,7 +37,7 @@ export type UpgradeWebSocket = (
 /**
  * Determines which driver to use if none is provided.
  */
-export function getDefaultDriver(): DriverConfig {
+export function createDefaultDriver(): DriverConfig {
 	const driver = getEnvUniversal("RIVETKIT_DRIVER");
 	console.log("driver", driver);
 	if (!driver || driver === "memory") {