fix(rivetkit): fix race condition with websocket open events

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/raw-websocket.ts

@@ -51,16 +51,17 @@ export const rawWebSocketActor = actor({
 							}),
 						);
 					} else if (parsed.type === "getRequestInfo") {
-						throw "TODO";
-						// Send back the request URL info
-						// websocket.send(
-						// 	JSON.stringify({
-						// 		type: "requestInfo",
-						// 		url: opts.request.url,
-						// 		pathname: new URL(opts.request.url).pathname,
-						// 		search: new URL(opts.request.url).search,
-						// 	}),
-						// );
+						// Send back the request URL info if available
+						const url = ctx.request?.url || "ws://actor/websocket";
+						const urlObj = new URL(url);
+						websocket.send(
+							JSON.stringify({
+								type: "requestInfo",
+								url: url,
+								pathname: urlObj.pathname,
+								search: urlObj.search,
+							}),
+						);
 					} else {
 						// Echo back
 						websocket.send(data);

rivetkit-typescript/packages/rivetkit/src/actor/config.ts

@@ -330,7 +330,14 @@ interface BaseActorConfig<
 	 * @returns Void or a Promise that resolves when connection handling is complete
 	 */
 	onConnect?: (
-		c: OnConnectContext<TState, TVars, TInput, TDatabase>,
+		c: OnConnectContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase
+		>,
 		conn: Conn<TState, TConnParams, TConnState, TVars, TInput, TDatabase>,
 	) => void | Promise<void>;
 

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

@@ -1,7 +1,8 @@
 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";
+import { loggerWithoutContext } from "../../log";
+import { type ConnDriver, DriverReadyState } from "../driver";
 
 /**
  * Creates a raw WebSocket connection driver.
@@ -15,10 +16,11 @@ export function createRawWebSocketSocket(
 	requestId: string,
 	requestIdBuf: ArrayBuffer | undefined,
 	hibernatable: boolean,
-	websocket: UniversalWebSocket,
 	closePromise: Promise<void>,
-): ConnDriver {
-	return {
+): { driver: ConnDriver; setWebSocket(ws: UniversalWebSocket): void } {
+	let websocket: UniversalWebSocket | undefined;
+
+	const driver: ConnDriver = {
 		type: "raw-websocket",
 		requestId,
 		requestIdBuf,
@@ -32,6 +34,13 @@ export function createRawWebSocketSocket(
 			_conn: AnyConn,
 			reason?: string,
 		) => {
+			if (!websocket) {
+				loggerWithoutContext().warn(
+					"disconnecting raw ws without websocket",
+				);
+				return;
+			}
+
 			// Close socket
 			websocket.close(1000, reason);
 
@@ -40,14 +49,21 @@ export function createRawWebSocketSocket(
 		},
 
 		terminate: () => {
-			(websocket as any).terminate?.();
+			(websocket as any)?.terminate?.();
 		},
 
 		getConnectionReadyState: (
 			_actor: AnyActorInstance,
 			_conn: AnyConn,
 		): DriverReadyState | undefined => {
-			return websocket.readyState;
+			return websocket?.readyState ?? DriverReadyState.CONNECTING;
+		},
+	};
+
+	return {
+		driver,
+		setWebSocket(ws) {
+			websocket = ws;
 		},
 	};
 }

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

@@ -3,6 +3,7 @@ import type { AnyConn } from "@/actor/conn/mod";
 import type { AnyActorInstance } from "@/actor/instance/mod";
 import type { CachedSerializer, Encoding } from "@/actor/protocol/serde";
 import type * as protocol from "@/schemas/client-protocol/mod";
+import { loggerWithoutContext } from "../../log";
 import { type ConnDriver, DriverReadyState } from "../driver";
 
 export type ConnDriverWebSocketState = Record<never, never>;
@@ -12,10 +13,12 @@ export function createWebSocketSocket(
 	requestIdBuf: ArrayBuffer | undefined,
 	hibernatable: boolean,
 	encoding: Encoding,
-	websocket: WSContext,
 	closePromise: Promise<void>,
-): ConnDriver {
-	return {
+): { driver: ConnDriver; setWebSocket(ws: WSContext): void } {
+	// Wait for WS to open
+	let websocket: WSContext | undefined;
+
+	const driver: ConnDriver = {
 		type: "websocket",
 		requestId,
 		requestIdBuf,
@@ -25,6 +28,13 @@ export function createWebSocketSocket(
 			conn: AnyConn,
 			message: CachedSerializer<any, any, any>,
 		) => {
+			if (!websocket) {
+				actor.rLog.warn({
+					msg: "websocket not open",
+					connId: conn.id,
+				});
+				return;
+			}
 			if (websocket.readyState !== DriverReadyState.OPEN) {
 				actor.rLog.warn({
 					msg: "attempting to send message to closed websocket, this is likely a bug in RivetKit",
@@ -83,6 +93,13 @@ export function createWebSocketSocket(
 			_conn: AnyConn,
 			reason?: string,
 		) => {
+			if (!websocket) {
+				loggerWithoutContext().warn(
+					"disconnecting ws without websocket",
+				);
+				return;
+			}
+
 			// Close socket
 			websocket.close(1000, reason);
 
@@ -98,7 +115,14 @@ export function createWebSocketSocket(
 			_actor: AnyActorInstance,
 			_conn: AnyConn,
 		): DriverReadyState | undefined => {
-			return websocket.readyState;
+			return websocket?.readyState ?? DriverReadyState.CONNECTING;
+		},
+	};
+
+	return {
+		driver,
+		setWebSocket(ws) {
+			websocket = ws;
 		},
 	};
 }

rivetkit-typescript/packages/rivetkit/src/actor/conn/mod.ts

@@ -7,10 +7,8 @@ import {
 } from "@/schemas/client-protocol-zod/mod";
 import { arrayBuffersEqual, bufferToArrayBuffer } from "@/utils";
 import type { AnyDatabaseProvider } from "../database";
-import {
-	ACTOR_INSTANCE_PERSIST_SYMBOL,
-	type ActorInstance,
-} from "../instance/mod";
+import { InternalError } from "../errors";
+import type { ActorInstance } from "../instance/mod";
 import type { PersistedConn } from "../instance/persisted";
 import { CachedSerializer } from "../protocol/serde";
 import type { ConnDriver } from "./driver";
@@ -24,6 +22,7 @@ export type ConnId = string;
 
 export type AnyConn = Conn<any, any, any, any, any, any>;
 
+export const CONN_CONNECTED_SYMBOL = Symbol("connected");
 export const CONN_PERSIST_SYMBOL = Symbol("persist");
 export const CONN_DRIVER_SYMBOL = Symbol("driver");
 export const CONN_ACTOR_SYMBOL = Symbol("actor");
@@ -60,6 +59,16 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		return this.#actor;
 	}
 
+	/** Connections exist before being connected to an actor. If true, this connection has been connected. */
+	[CONN_CONNECTED_SYMBOL] = false;
+
+	#assertConnected() {
+		if (!this[CONN_CONNECTED_SYMBOL])
+			throw new InternalError(
+				"Connection not connected yet. This happens when trying to use the connection in onBeforeConnect or createConnState.",
+			);
+	}
+
 	get [CONN_PERSIST_SYMBOL](): PersistedConn<CP, CS> {
 		return this.#stateManager.persist;
 	}
@@ -109,9 +118,7 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			return false;
 		}
 		return (
-			(this.#actor as any)[
-				ACTOR_INSTANCE_PERSIST_SYMBOL
-			].hibernatableConns.findIndex((conn: any) =>
+			this.#actor.persist.hibernatableConns.findIndex((conn: any) =>
 				arrayBuffersEqual(
 					conn.hibernatableRequestId,
 					hibernatableRequestId,
@@ -191,6 +198,8 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 * @see {@link https://rivet.dev/docs/events|Events Documentation}
 	 */
 	send(eventName: string, ...args: unknown[]) {
+		this.#assertConnected();
+
 		this.#actor.inspector.emitter.emit("eventFired", {
 			type: "event",
 			eventName,
@@ -244,7 +253,7 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 				});
 			}
 
-			this.#actor.connDisconnected(this, true);
+			this.#actor.connectionManager.connDisconnected(this);
 		} else {
 			this.#actor.rLog.warn({
 				msg: "missing connection driver state for disconnect",

rivetkit-typescript/packages/rivetkit/src/actor/conn/state-manager.ts

@@ -134,6 +134,8 @@ export class StateManager<CP, CS> {
 		});
 
 		// Notify actor that this connection has changed
-		this.#conn[CONN_ACTOR_SYMBOL].markConnChanged(this.#conn);
+		this.#conn[CONN_ACTOR_SYMBOL].connectionManager.markConnChanged(
+			this.#conn,
+		);
 	}
 }

rivetkit-typescript/packages/rivetkit/src/actor/contexts/actor.ts

@@ -60,7 +60,7 @@ export class ActorContext<
 	 * @param args - The arguments to send with the event.
 	 */
 	broadcast<Args extends Array<unknown>>(name: string, ...args: Args): void {
-		this.#actor.broadcast(name, ...args);
+		this.#actor.eventManager.broadcast(name, ...args);
 		return;
 	}
 

rivetkit-typescript/packages/rivetkit/src/actor/contexts/on-connect.ts

@@ -1,13 +1,22 @@
 import type { AnyDatabaseProvider } from "../database";
-import { ConnInitContext } from "./conn-init";
+import { ConnContext } from "./conn";
 
 /**
  * Context for the onConnect lifecycle hook.
  * Called when a connection is successfully established.
  */
 export class OnConnectContext<
 	TState,
+	TConnParams,
+	TConnState,
 	TVars,
 	TInput,
 	TDatabase extends AnyDatabaseProvider,
-> extends ConnInitContext<TState, TVars, TInput, TDatabase> {}
+> extends ConnContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase
+> {}