rivetkit 2.0.2 → 2.0.4-rc.1

package/src/actor/conn/driver.ts

@@ -0,0 +1,61 @@
+import type { AnyConn } from "@/actor/conn/mod";
+import type { AnyActorInstance } from "@/actor/instance/mod";
+import type { CachedSerializer } from "@/actor/protocol/serde";
+
+export enum DriverReadyState {
+	UNKNOWN = -1,
+	CONNECTING = 0,
+	OPEN = 1,
+	CLOSING = 2,
+	CLOSED = 3,
+}
+
+export interface ConnDriver {
+	/** The type of driver. Used for debug purposes only. */
+	type: string;
+
+	/**
+	 * If defined, this connection driver talks the RivetKit client driver (see
+	 * schemas/client-protocol/).
+	 *
+	 * If enabled, events like `Init`, subscription events, etc. will be sent
+	 * to this connection.
+	 */
+	rivetKitProtocol?: {
+		/** Sends a RivetKit client message. */
+		sendMessage(
+			actor: AnyActorInstance,
+			conn: AnyConn,
+			message: CachedSerializer<any, any, any>,
+		): void;
+	};
+
+	/**
+	 * If the connection can be hibernated. If true, this will allow the actor to go to sleep while the connection is still active.
+	 **/
+	hibernatable?: {
+		gatewayId: ArrayBuffer;
+		requestId: ArrayBuffer;
+	};
+
+	/**
+	 * This returns a promise since we commonly disconnect at the end of a program, and not waiting will cause the socket to not close cleanly.
+	 */
+	disconnect(
+		actor: AnyActorInstance,
+		conn: AnyConn,
+		reason?: string,
+	): Promise<void>;
+
+	/** Terminates the connection without graceful handling. */
+	terminate?(actor: AnyActorInstance, conn: AnyConn): void;
+
+	/**
+	 * Returns the ready state of the connection.
+	 * This is used to determine if the connection is ready to send messages, or if the connection is stale.
+	 */
+	getConnectionReadyState(
+		actor: AnyActorInstance,
+		conn: AnyConn,
+	): DriverReadyState | undefined;
+}

package/src/actor/conn/drivers/http.ts

@@ -0,0 +1,17 @@
+import { type ConnDriver, DriverReadyState } from "../driver";
+
+export type ConnHttpState = Record<never, never>;
+
+export function createHttpDriver(): ConnDriver {
+	return {
+		type: "http",
+		getConnectionReadyState(_actor, _conn) {
+			// TODO: This might not be the correct logic
+			return DriverReadyState.OPEN;
+		},
+		disconnect: async () => {
+			// Noop
+			// TODO: Configure with abort signals to abort the request
+		},
+	};
+}

package/src/actor/conn/drivers/raw-request.ts

@@ -0,0 +1,24 @@
+import type { ConnDriver } from "../driver";
+import { DriverReadyState } from "../driver";
+
+/**
+ * Creates a raw HTTP connection driver.
+ *
+ * This driver is used for raw HTTP connections that don't use the RivetKit protocol.
+ * Unlike the standard HTTP driver, this provides connection lifecycle management
+ * for tracking the HTTP request through the actor's onRequest handler.
+ */
+export function createRawRequestDriver(): ConnDriver {
+	return {
+		type: "raw-request",
+
+		disconnect: async () => {
+			// Noop
+		},
+
+		getConnectionReadyState: (): DriverReadyState | undefined => {
+			// HTTP connections are always considered open until the request completes
+			return DriverReadyState.OPEN;
+		},
+	};
+}

package/src/actor/conn/drivers/raw-websocket.ts

@@ -0,0 +1,65 @@
+import type { AnyConn } from "@/actor/conn/mod";
+import type { AnyActorInstance } from "@/actor/instance/mod";
+import type { UniversalWebSocket } from "@/common/websocket-interface";
+import { loggerWithoutContext } from "../../log";
+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 createRawWebSocketDriver(
+	hibernatable: ConnDriver["hibernatable"],
+	closePromise: Promise<void>,
+): { driver: ConnDriver; setWebSocket(ws: UniversalWebSocket): void } {
+	let websocket: UniversalWebSocket | undefined;
+
+	const driver: ConnDriver = {
+		type: "raw-websocket",
+		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,
+		) => {
+			if (!websocket) {
+				loggerWithoutContext().warn(
+					"disconnecting raw ws without websocket",
+				);
+				return;
+			}
+
+			// 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 ?? DriverReadyState.CONNECTING;
+		},
+	};
+
+	return {
+		driver,
+		setWebSocket(ws) {
+			websocket = ws;
+		},
+	};
+}

package/src/actor/conn/drivers/websocket.ts

@@ -0,0 +1,144 @@
+import type { WSContext } from "hono/ws";
+import type { AnyConn } from "@/actor/conn/mod";
+import type { AnyActorInstance } from "@/actor/instance/mod";
+import type { CachedSerializer, Encoding } from "@/actor/protocol/serde";
+import * as errors from "@/actor/errors";
+import { loggerWithoutContext } from "../../log";
+import { type ConnDriver, DriverReadyState } from "../driver";
+import { RegistryConfig } from "@/registry/config";
+
+export type ConnDriverWebSocketState = Record<never, never>;
+
+export function createWebSocketDriver(
+	hibernatable: ConnDriver["hibernatable"],
+	encoding: Encoding,
+	closePromise: Promise<void>,
+	config: RegistryConfig,
+): { driver: ConnDriver; setWebSocket(ws: WSContext): void } {
+	loggerWithoutContext().debug({
+		msg: "createWebSocketDriver creating driver",
+		hibernatable,
+	});
+	// Wait for WS to open
+	let websocket: WSContext | undefined;
+
+	const driver: ConnDriver = {
+		type: "websocket",
+		hibernatable,
+		rivetKitProtocol: {
+			sendMessage: (
+				actor: AnyActorInstance,
+				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",
+						connId: conn.id,
+						wsReadyState: websocket.readyState,
+					});
+					return;
+				}
+
+				const serialized = message.serialize(encoding);
+
+				actor.rLog.debug({
+					msg: "sending websocket message",
+					encoding: encoding,
+					dataType: typeof serialized,
+					isUint8Array: serialized instanceof Uint8Array,
+					isArrayBuffer: serialized instanceof ArrayBuffer,
+					dataLength:
+						(serialized as any).byteLength ||
+						(serialized as any).length,
+				});
+
+				// Check outgoing message size
+				const messageSize =
+					(serialized as any).byteLength || (serialized as any).length;
+				if (messageSize > config.maxOutgoingMessageSize) {
+					actor.rLog.error({
+						msg: "outgoing message exceeds maxOutgoingMessageSize",
+						messageSize,
+						maxOutgoingMessageSize: config.maxOutgoingMessageSize,
+					});
+					throw new errors.OutgoingMessageTooLong();
+				}
+
+				// Convert Uint8Array to ArrayBuffer for proper transmission
+				if (serialized instanceof Uint8Array) {
+					const buffer = serialized.buffer.slice(
+						serialized.byteOffset,
+						serialized.byteOffset + serialized.byteLength,
+					);
+					// Handle SharedArrayBuffer case
+					if (buffer instanceof SharedArrayBuffer) {
+						const arrayBuffer = new ArrayBuffer(buffer.byteLength);
+						new Uint8Array(arrayBuffer).set(new Uint8Array(buffer));
+						actor.rLog.debug({
+							msg: "converted SharedArrayBuffer to ArrayBuffer",
+							byteLength: arrayBuffer.byteLength,
+						});
+						websocket.send(arrayBuffer);
+					} else {
+						actor.rLog.debug({
+							msg: "sending ArrayBuffer",
+							byteLength: buffer.byteLength,
+						});
+						websocket.send(buffer);
+					}
+				} else {
+					actor.rLog.debug({
+						msg: "sending string data",
+						length: (serialized as string).length,
+					});
+					websocket.send(serialized);
+				}
+			},
+		},
+
+		disconnect: async (
+			_actor: AnyActorInstance,
+			_conn: AnyConn,
+			reason?: string,
+		) => {
+			if (!websocket) {
+				loggerWithoutContext().warn(
+					"disconnecting ws without websocket",
+				);
+				return;
+			}
+
+			// Close socket
+			websocket.close(1000, reason);
+
+			// Create promise to wait for socket to close gracefully
+			await closePromise;
+		},
+
+		terminate: () => {
+			(websocket as any).terminate();
+		},
+
+		getConnectionReadyState: (
+			_actor: AnyActorInstance,
+			_conn: AnyConn,
+		): DriverReadyState | undefined => {
+			return websocket?.readyState ?? DriverReadyState.CONNECTING;
+		},
+	};
+
+	return {
+		driver,
+		setWebSocket(ws) {
+			websocket = ws;
+		},
+	};
+}

package/src/actor/conn/mod.ts

@@ -0,0 +1,288 @@
+import * as cbor from "cbor-x";
+import type * as protocol from "@/schemas/client-protocol/mod";
+import {
+	CURRENT_VERSION as CLIENT_PROTOCOL_CURRENT_VERSION,
+	TO_CLIENT_VERSIONED,
+} from "@/schemas/client-protocol/versioned";
+import {
+	type ToClient as ToClientJson,
+	ToClientSchema,
+} from "@/schemas/client-protocol-zod/mod";
+import { bufferToArrayBuffer } from "@/utils";
+import type { AnyDatabaseProvider } from "../database";
+import { EventPayloadInvalid, InternalError } from "../errors";
+import type { ActorInstance } from "../instance/mod";
+import { CachedSerializer } from "../protocol/serde";
+import {
+	type EventSchemaConfig,
+	hasSchemaConfigKey,
+	type InferEventArgs,
+	type InferSchemaMap,
+	type QueueSchemaConfig,
+	validateSchemaSync,
+} from "../schema";
+import type { ConnDriver } from "./driver";
+import { type ConnDataInput, StateManager } from "./state-manager";
+
+export type ConnId = string;
+
+export type AnyConn = Conn<any, any, any, any, any, any, any, any>;
+
+export const CONN_CONNECTED_SYMBOL = Symbol("connected");
+export const CONN_SPEAKS_RIVETKIT_SYMBOL = Symbol("speaksRivetKit");
+export const CONN_DRIVER_SYMBOL = Symbol("driver");
+export const CONN_ACTOR_SYMBOL = Symbol("actor");
+export const CONN_STATE_MANAGER_SYMBOL = Symbol("stateManager");
+export const CONN_SEND_MESSAGE_SYMBOL = Symbol("sendMessage");
+
+/**
+ * Represents a client connection to a actor.
+ *
+ * Manages connection-specific data and controls the connection lifecycle.
+ *
+ * @see {@link https://rivet.dev/docs/connections|Connection Documentation}
+ */
+export class Conn<
+	S,
+	CP,
+	CS,
+	V,
+	I,
+	DB extends AnyDatabaseProvider,
+	E extends EventSchemaConfig = Record<never, never>,
+	Q extends QueueSchemaConfig = Record<never, never>,
+> {
+	#actor: ActorInstance<S, CP, CS, V, I, DB, E, Q>;
+
+	get [CONN_ACTOR_SYMBOL](): ActorInstance<S, CP, CS, V, I, DB, E, Q> {
+		return this.#actor;
+	}
+
+	#stateManager!: StateManager<CP, CS>;
+
+	get [CONN_STATE_MANAGER_SYMBOL]() {
+		return this.#stateManager;
+	}
+
+	/**
+	 * Connections exist before being connected to an actor. If true, this
+	 * connection has been connected.
+	 **/
+	[CONN_CONNECTED_SYMBOL] = false;
+
+	/**
+	 * If undefined, then no socket is connected to this conn
+	 */
+	[CONN_DRIVER_SYMBOL]?: ConnDriver;
+
+	/**
+	 * If this connection is speaking the RivetKit protocol. If false, this is
+	 * a raw connection for WebSocket or fetch or inspector.
+	 **/
+	get [CONN_SPEAKS_RIVETKIT_SYMBOL](): boolean {
+		return this[CONN_DRIVER_SYMBOL]?.rivetKitProtocol !== undefined;
+	}
+
+	subscriptions: Set<string> = new Set<string>();
+
+	#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.",
+			);
+	}
+
+	// MARK: - Public Getters
+	get params(): CP {
+		return this.#stateManager.ephemeralData.parameters;
+	}
+
+	/**
+	 * Gets the current state of the connection.
+	 *
+	 * Throws an error if the state is not enabled.
+	 */
+	get state(): CS {
+		return this.#stateManager.state;
+	}
+
+	/**
+	 * Sets the state of the connection.
+	 *
+	 * Throws an error if the state is not enabled.
+	 */
+	set state(value: CS) {
+		this.#stateManager.state = value;
+	}
+
+	/**
+	 * Unique identifier for the connection.
+	 */
+	get id(): ConnId {
+		return this.#stateManager.ephemeralData.id;
+	}
+
+	/**
+	 * @experimental
+	 *
+	 * If the underlying connection can hibernate.
+	 */
+	get isHibernatable(): boolean {
+		return this.#stateManager.hibernatableDataRaw !== undefined;
+	}
+
+	/**
+	 * Initializes a new instance of the Connection class.
+	 *
+	 * This should only be constructed by {@link Actor}.
+	 *
+	 * @protected
+	 */
+	constructor(
+		actor: ActorInstance<S, CP, CS, V, I, DB, E, Q>,
+		data: ConnDataInput<CP, CS>,
+	) {
+		this.#actor = actor;
+		this.#stateManager = new StateManager(this, data);
+	}
+
+	/**
+	 * Sends a raw message to the underlying connection.
+	 */
+	[CONN_SEND_MESSAGE_SYMBOL](message: CachedSerializer<any, any, any>) {
+		if (this[CONN_DRIVER_SYMBOL]) {
+			const driver = this[CONN_DRIVER_SYMBOL];
+
+			if (driver.rivetKitProtocol) {
+				driver.rivetKitProtocol.sendMessage(this.#actor, this, message);
+			} else {
+				this.#actor.rLog.warn({
+					msg: "attempting to send RivetKit protocol message to connection that does not support it",
+					conn: this.id,
+				});
+			}
+		} else {
+			this.#actor.rLog.warn({
+				msg: "missing connection driver state for send message",
+				conn: this.id,
+			});
+		}
+	}
+
+	/**
+	 * Sends an event with arguments to the client.
+	 *
+	 * @param eventName - The name of the event.
+	 * @param args - The arguments for the event.
+	 * @see {@link https://rivet.dev/docs/events|Events Documentation}
+	 */
+	send<K extends keyof E & string>(
+		eventName: K,
+		...args: InferEventArgs<InferSchemaMap<E>[K]>
+	): void;
+	send(
+		eventName: keyof E extends never ? string : never,
+		...args: unknown[]
+	): void;
+	send(eventName: string, ...args: unknown[]) {
+		this.#assertConnected();
+		if (!this[CONN_SPEAKS_RIVETKIT_SYMBOL]) {
+			this.#actor.rLog.warn({
+				msg: "cannot send messages to this connection type",
+				connId: this.id,
+				connType: this[CONN_DRIVER_SYMBOL]?.type,
+			});
+		}
+
+		if (
+			this.#actor.config.events !== undefined &&
+			!hasSchemaConfigKey(this.#actor.config.events, eventName)
+		) {
+			this.#actor.rLog.warn({
+				msg: "sending event not defined in actor events config",
+				eventName,
+				connId: this.id,
+			});
+		}
+
+		const payload = args.length === 1 ? args[0] : args;
+		const result = validateSchemaSync(
+			this.#actor.config.events,
+			eventName as keyof E & string,
+			payload,
+		);
+		if (!result.success) {
+			throw new EventPayloadInvalid(eventName, result.issues);
+		}
+		const eventArgs =
+			args.length === 1
+				? [result.data]
+				: Array.isArray(result.data)
+					? (result.data as unknown[])
+					: args;
+		this.#actor.emitTraceEvent("message.send", {
+			"rivet.event.name": eventName,
+			"rivet.conn.id": this.id,
+		});
+		const eventData = { name: eventName, args: eventArgs };
+		this[CONN_SEND_MESSAGE_SYMBOL](
+			new CachedSerializer(
+				eventData,
+				TO_CLIENT_VERSIONED,
+				CLIENT_PROTOCOL_CURRENT_VERSION,
+				ToClientSchema,
+				// JSON: args is the raw value (array of arguments)
+				(value): ToClientJson => ({
+					body: {
+						tag: "Event" as const,
+						val: {
+							name: value.name,
+							args: value.args,
+						},
+					},
+				}),
+				// BARE/CBOR: args needs to be CBOR-encoded to ArrayBuffer
+				(value): protocol.ToClient => ({
+					body: {
+						tag: "Event" as const,
+						val: {
+							name: value.name,
+							args: bufferToArrayBuffer(cbor.encode(value.args)),
+						},
+					},
+				}),
+			),
+		);
+	}
+
+	/**
+	 * Disconnects the client with an optional reason.
+	 *
+	 * @param reason - The reason for disconnection.
+	 */
+	async disconnect(reason?: string) {
+		if (this[CONN_DRIVER_SYMBOL]) {
+			const driver = this[CONN_DRIVER_SYMBOL];
+			if (driver.disconnect) {
+				driver.disconnect(this.#actor, this, reason);
+			} else {
+				this.#actor.rLog.debug({
+					msg: "no disconnect handler for conn driver",
+					conn: this.id,
+				});
+			}
+
+			try {
+				await this.#actor.connectionManager.connDisconnected(this);
+			} finally {
+				this[CONN_DRIVER_SYMBOL] = undefined;
+			}
+		} else {
+			this.#actor.rLog.warn({
+				msg: "missing connection driver state for disconnect",
+				conn: this.id,
+			});
+			this[CONN_DRIVER_SYMBOL] = undefined;
+		}
+	}
+}

package/src/actor/conn/persisted.ts

@@ -0,0 +1,81 @@
+/**
+ * Persisted data structures for connections.
+ *
+ * Keep this file in sync with the Connection section of rivetkit-typescript/packages/rivetkit/schemas/actor-persist/
+ */
+
+import * as cbor from "cbor-x";
+import type * as persistSchema from "@/schemas/actor-persist/mod";
+import { bufferToArrayBuffer } from "@/utils";
+
+export type GatewayId = ArrayBuffer;
+export type RequestId = ArrayBuffer;
+
+export type Cbor = ArrayBuffer;
+
+// MARK: Connection
+/** Event subscription for connection */
+export interface PersistedSubscription {
+	eventName: string;
+}
+
+/** Connection associated with hibernatable WebSocket that should persist across lifecycles */
+export interface PersistedConn<CP, CS> {
+	/** Connection ID generated by RivetKit */
+	id: string;
+	parameters: CP;
+	state: CS;
+	subscriptions: PersistedSubscription[];
+	gatewayId: GatewayId;
+	requestId: RequestId;
+	serverMessageIndex: number;
+	clientMessageIndex: number;
+	requestPath: string;
+	requestHeaders: Record<string, string>;
+}
+
+/**
+ * Converts persisted connection data to BARE schema format for serialization.
+ * @throws {Error} If the connection is ephemeral (not hibernatable)
+ */
+export function convertConnToBarePersistedConn<CP, CS>(
+	persist: PersistedConn<CP, CS>,
+): persistSchema.Conn {
+	return {
+		id: persist.id,
+		parameters: bufferToArrayBuffer(cbor.encode(persist.parameters)),
+		state: bufferToArrayBuffer(cbor.encode(persist.state)),
+		subscriptions: persist.subscriptions.map((sub) => ({
+			eventName: sub.eventName,
+		})),
+		gatewayId: persist.gatewayId,
+		requestId: persist.requestId,
+		serverMessageIndex: persist.serverMessageIndex,
+		clientMessageIndex: persist.clientMessageIndex,
+		requestPath: persist.requestPath,
+		requestHeaders: new Map(Object.entries(persist.requestHeaders)),
+	};
+}
+
+/**
+ * Converts BARE schema format to persisted connection data.
+ * @throws {Error} If the connection is ephemeral (not hibernatable)
+ */
+export function convertConnFromBarePersistedConn<CP, CS>(
+	bareData: persistSchema.Conn,
+): PersistedConn<CP, CS> {
+	return {
+		id: bareData.id,
+		parameters: cbor.decode(new Uint8Array(bareData.parameters)),
+		state: cbor.decode(new Uint8Array(bareData.state)),
+		subscriptions: bareData.subscriptions.map((sub) => ({
+			eventName: sub.eventName,
+		})),
+		gatewayId: bareData.gatewayId,
+		requestId: bareData.requestId,
+		serverMessageIndex: bareData.serverMessageIndex,
+		clientMessageIndex: bareData.clientMessageIndex,
+		requestPath: bareData.requestPath,
+		requestHeaders: Object.fromEntries(bareData.requestHeaders),
+	};
+}