feat(rivetkit): refactor handleRawWebSocketHandler to use actor.createConn

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/src/actor/conn/drivers/raw-websocket.ts

@@ -0,0 +1,52 @@
+import type { AnyConn } from "@/actor/conn/mod";
+import type { AnyActorInstance } from "@/actor/instance/mod";
+import type { UniversalWebSocket } from "@/common/websocket-interface";
+import type { ConnDriver, DriverReadyState } from "../driver";
+
+/**
+ * Creates a raw WebSocket connection driver.
+ *
+ * This driver is used for raw WebSocket connections that don't use the RivetKit protocol.
+ * Unlike the standard WebSocket driver, this doesn't have sendMessage since raw WebSockets
+ * don't handle messages from the RivetKit protocol - they handle messages directly in the
+ * actor's onWebSocket handler.
+ */
+export function createRawWebSocketSocket(
+	requestId: string,
+	requestIdBuf: ArrayBuffer | undefined,
+	hibernatable: boolean,
+	websocket: UniversalWebSocket,
+	closePromise: Promise<void>,
+): ConnDriver {
+	return {
+		requestId,
+		requestIdBuf,
+		hibernatable,
+
+		// No sendMessage implementation since this is a raw WebSocket that doesn't
+		// handle messages from the RivetKit protocol
+
+		disconnect: async (
+			_actor: AnyActorInstance,
+			_conn: AnyConn,
+			reason?: string,
+		) => {
+			// Close socket
+			websocket.close(1000, reason);
+
+			// Wait for socket to close gracefully
+			await closePromise;
+		},
+
+		terminate: () => {
+			(websocket as any).terminate?.();
+		},
+
+		getConnectionReadyState: (
+			_actor: AnyActorInstance,
+			_conn: AnyConn,
+		): DriverReadyState | undefined => {
+			return websocket.readyState;
+		},
+	};
+}

rivetkit-typescript/packages/rivetkit/src/actor/router-endpoints.ts

@@ -36,6 +36,7 @@ import {
 	promiseWithResolvers,
 } from "@/utils";
 import { createHttpSocket } from "./conn/drivers/http";
+import { createRawWebSocketSocket } from "./conn/drivers/raw-websocket";
 import { createWebSocketSocket } from "./conn/drivers/websocket";
 import type { ActorDriver } from "./driver";
 import { loggerWithoutContext } from "./log";
@@ -383,12 +384,19 @@ export async function handleRawWebSocketHandler(
 ): Promise<UpgradeWebSocketArgs> {
 	const actor = await actorDriver.loadActor(actorId);
 
+	// Promise used to wait for the websocket close in `disconnect`
+	const closePromiseResolvers = promiseWithResolvers<void>();
+
+	// Track connection outside of scope for cleanup
+	let createdConn: AnyConn | undefined;
+
 	// Return WebSocket event handlers
 	return {
-		onOpen: (evt: any, ws: any) => {
+		onOpen: async (evt: any, ws: any) => {
 			// Extract rivetRequestId provided by engine runner
 			const rivetRequestId = evt?.rivetRequestId;
 			const isHibernatable =
+				!!rivetRequestId &&
 				actor[
 					ACTOR_INSTANCE_PERSIST_SYMBOL
 				].hibernatableConns.findIndex((conn) =>
@@ -424,10 +432,36 @@ export async function handleRawWebSocketHandler(
 				toUrl: newRequest.url,
 			});
 
-			// Call the actor's onWebSocket handler with the adapted WebSocket
-			actor.handleWebSocket(adapter, {
-				request: newRequest,
-			});
+			try {
+				// Create connection using actor.createConn - this handles deduplication for hibernatable connections
+				const requestId = rivetRequestId
+					? String(rivetRequestId)
+					: crypto.randomUUID();
+				const conn = await actor.createConn(
+					createRawWebSocketSocket(
+						requestId,
+						rivetRequestId,
+						isHibernatable,
+						adapter,
+						closePromiseResolvers.promise,
+					),
+					{}, // No parameters for raw WebSocket
+					newRequest,
+				);
+
+				createdConn = conn;
+
+				// Call the actor's onWebSocket handler with the adapted WebSocket
+				actor.handleWebSocket(adapter, {
+					request: newRequest,
+				});
+			} catch (error) {
+				actor.rLog.error({
+					msg: "failed to create raw WebSocket connection",
+					error: String(error),
+				});
+				ws.close(1011, "Failed to create connection");
+			}
 		},
 		onMessage: (event: any, ws: any) => {
 			// Find the adapter for this WebSocket
@@ -442,6 +476,15 @@ export async function handleRawWebSocketHandler(
 			if (adapter) {
 				adapter._handleClose(evt?.code || 1006, evt?.reason || "");
 			}
+
+			// Resolve the close promise
+			closePromiseResolvers.resolve();
+
+			// Clean up the connection
+			if (createdConn) {
+				const wasClean = evt?.wasClean || evt?.code === 1000;
+				actor.connDisconnected(createdConn, wasClean);
+			}
 		},
 		onError: (error: any, ws: any) => {
 			// Find the adapter for this WebSocket