rivetkit 2.0.24-rc.1 → 2.0.24

package/src/actor/conn/mod.ts

@@ -0,0 +1,232 @@
+import * as cbor from "cbor-x";
+import type * as protocol from "@/schemas/client-protocol/mod";
+import { 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 { InternalError } from "../errors";
+import type { ActorInstance } from "../instance/mod";
+import { CachedSerializer } from "../protocol/serde";
+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>;
+
+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> {
+	#actor: ActorInstance<S, CP, CS, V, I, DB>;
+
+	get [CONN_ACTOR_SYMBOL](): ActorInstance<S, CP, CS, V, I, DB> {
+		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[CONN_DRIVER_SYMBOL]?.hibernatable !== 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>,
+		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(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,
+			});
+		}
+
+		this.#actor.inspector.emitter.emit("eventFired", {
+			type: "event",
+			eventName,
+			args,
+			connId: this.id,
+		});
+		const eventData = { name: eventName, args };
+		this[CONN_SEND_MESSAGE_SYMBOL](
+			new CachedSerializer(
+				eventData,
+				TO_CLIENT_VERSIONED,
+				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,
+				});
+			}
+
+			this.#actor.connectionManager.connDisconnected(this);
+		} 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),
+	};
+}

package/src/actor/conn/state-manager.ts

@@ -0,0 +1,196 @@
+import type { HibernatingWebSocketMetadata } from "@rivetkit/engine-runner";
+import * as cbor from "cbor-x";
+import invariant from "invariant";
+import onChange from "on-change";
+import { isCborSerializable } from "@/common/utils";
+import * as errors from "../errors";
+import { assertUnreachable } from "../utils";
+import { CONN_ACTOR_SYMBOL, type Conn } from "./mod";
+import type { PersistedConn } from "./persisted";
+
+/** Pick a subset of persisted data used to represent ephemeral connections */
+export type EphemeralConn<CP, CS> = Pick<
+	PersistedConn<CP, CS>,
+	"id" | "parameters" | "state"
+>;
+
+export type ConnDataInput<CP, CS> =
+	| { ephemeral: EphemeralConn<CP, CS> }
+	| { hibernatable: PersistedConn<CP, CS> };
+
+export type ConnData<CP, CS> =
+	| {
+			ephemeral: {
+				/** In-memory data representing this connection */
+				data: EphemeralConn<CP, CS>;
+			};
+	  }
+	| {
+			hibernatable: {
+				/** Persisted data with on-change proxy */
+				data: PersistedConn<CP, CS>;
+				/** Raw persisted data without proxy */
+				dataRaw: PersistedConn<CP, CS>;
+			};
+	  };
+
+/**
+ * Manages connection state persistence, proxying, and change tracking.
+ * Handles automatic state change detection for connection-specific state.
+ */
+export class StateManager<CP, CS> {
+	#conn: Conn<any, CP, CS, any, any, any>;
+
+	/**
+	 * Data representing this connection.
+	 *
+	 * This is stored as a struct for both ephemeral and hibernatable conns in
+	 * order to keep the separation clear between the two.
+	 */
+	#data!: ConnData<CP, CS>;
+
+	constructor(
+		conn: Conn<any, CP, CS, any, any, any>,
+		data: ConnDataInput<CP, CS>,
+	) {
+		this.#conn = conn;
+
+		if ("ephemeral" in data) {
+			this.#data = { ephemeral: { data: data.ephemeral } };
+		} else if ("hibernatable" in data) {
+			// Listen for changes to the object
+			const persistRaw = data.hibernatable;
+			const persist = onChange(
+				persistRaw,
+				(
+					path: string,
+					value: any,
+					_previousValue: any,
+					_applyData: any,
+				) => {
+					this.#handleChange(path, value);
+				},
+				{ ignoreDetached: true },
+			);
+			this.#data = {
+				hibernatable: { data: persist, dataRaw: persistRaw },
+			};
+		} else {
+			assertUnreachable(data);
+		}
+	}
+
+	/**
+	 * Returns the ephemeral or persisted data for this connectioned.
+	 *
+	 * This property is used to be able to treat both memory & persist conns
+	 * identical by looking up the correct underlying data structure.
+	 */
+	get ephemeralData(): EphemeralConn<CP, CS> {
+		if ("hibernatable" in this.#data) {
+			return this.#data.hibernatable.data;
+		} else if ("ephemeral" in this.#data) {
+			return this.#data.ephemeral.data;
+		} else {
+			return assertUnreachable(this.#data);
+		}
+	}
+
+	get hibernatableData(): PersistedConn<CP, CS> | undefined {
+		if ("hibernatable" in this.#data) {
+			return this.#data.hibernatable.data;
+		} else {
+			return undefined;
+		}
+	}
+
+	hibernatableDataOrError(): PersistedConn<CP, CS> {
+		const hibernatable = this.hibernatableData;
+		invariant(hibernatable, "missing hibernatable data");
+		return hibernatable;
+	}
+
+	get hibernatableDataRaw(): PersistedConn<CP, CS> | undefined {
+		if ("hibernatable" in this.#data) {
+			return this.#data.hibernatable.dataRaw;
+		} else {
+			return undefined;
+		}
+	}
+
+	get stateEnabled(): boolean {
+		return this.#conn[CONN_ACTOR_SYMBOL].connStateEnabled;
+	}
+
+	get state(): CS {
+		this.#validateStateEnabled();
+		const state = this.ephemeralData.state;
+		if (!state) throw new Error("state should exists");
+		return state;
+	}
+
+	set state(value: CS) {
+		this.#validateStateEnabled();
+		this.ephemeralData.state = value;
+	}
+
+	#validateStateEnabled() {
+		if (!this.#conn[CONN_ACTOR_SYMBOL].connStateEnabled) {
+			throw new errors.ConnStateNotEnabled();
+		}
+	}
+
+	#handleChange(path: string, value: any) {
+		// NOTE: This will only be called for hibernatable conns since only
+		// hibernatable conns have the on-change proxy
+
+		// Validate CBOR serializability for state changes
+		if (path.startsWith("state")) {
+			let invalidPath = "";
+			if (
+				!isCborSerializable(
+					value,
+					(invalidPathPart: string) => {
+						invalidPath = invalidPathPart;
+					},
+					"",
+				)
+			) {
+				throw new errors.InvalidStateType({
+					path: path + (invalidPath ? `.${invalidPath}` : ""),
+				});
+			}
+		}
+
+		this.#conn[CONN_ACTOR_SYMBOL].rLog.debug({
+			msg: "conn onChange triggered",
+			connId: this.#conn.id,
+			path,
+		});
+
+		// Notify actor that this connection has changed
+		this.#conn[
+			CONN_ACTOR_SYMBOL
+		].connectionManager.markConnWithPersistChanged(this.#conn);
+	}
+
+	addSubscription({ eventName }: { eventName: string }) {
+		const hibernatable = this.hibernatableData;
+		if (!hibernatable) return;
+		hibernatable.subscriptions.push({
+			eventName,
+		});
+	}
+
+	removeSubscription({ eventName }: { eventName: string }) {
+		const hibernatable = this.hibernatableData;
+		if (!hibernatable) return;
+		const subIdx = hibernatable.subscriptions.findIndex(
+			(s) => s.eventName === eventName,
+		);
+		if (subIdx !== -1) {
+			hibernatable.subscriptions.splice(subIdx, 1);
+		}
+		return subIdx !== -1;
+	}
+}

package/src/actor/contexts/action.ts

@@ -0,0 +1,23 @@
+import type { Conn } from "../conn/mod";
+import type { AnyDatabaseProvider } from "../database";
+import type { ActorInstance } from "../instance/mod";
+import { ConnContext } from "./conn";
+
+/**
+ * Context for a remote procedure call.
+ */
+export class ActionContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+> extends ConnContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase
+> {}

package/src/actor/{context.ts → contexts/actor.ts}

@@ -2,10 +2,10 @@ import type { ActorKey } from "@/actor/mod";
 import type { Client } from "@/client/client";
 import type { Logger } from "@/common/log";
 import type { Registry } from "@/registry/mod";
-import type { Conn, ConnId } from "./conn";
-import type { AnyDatabaseProvider, InferDatabaseClient } from "./database";
-import type { ActorInstance, SaveStateOptions } from "./instance";
-import type { Schedule } from "./schedule";
+import type { Conn, ConnId } from "../conn/mod";
+import type { AnyDatabaseProvider, InferDatabaseClient } from "../database";
+import type { ActorInstance, SaveStateOptions } from "../instance/mod";
+import type { Schedule } from "../schedule";
 
 /**
  * ActorContext class that provides access to actor methods and state
@@ -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;
 	}
 
@@ -138,14 +138,14 @@ export class ActorContext<
 	 * @param opts - Options for saving the state.
 	 */
 	async saveState(opts: SaveStateOptions): Promise<void> {
-		return this.#actor.saveState(opts);
+		return this.#actor.stateManager.saveState(opts);
 	}
 
 	/**
 	 * Prevents the actor from sleeping until promise is complete.
 	 */
 	waitUntil(promise: Promise<void>): void {
-		this.#actor._waitUntil(promise);
+		this.#actor.waitUntil(promise);
 	}
 
 	/**
@@ -163,6 +163,17 @@ export class ActorContext<
 	 * @experimental
 	 */
 	sleep() {
-		this.#actor._startSleep();
+		this.#actor.startSleep();
+	}
+
+	/**
+	 * Forces the actor to destroy.
+	 *
+	 * This will return immediately, then call `onStop` and `onDestroy`.
+	 *
+	 * @experimental
+	 */
+	destroy() {
+		this.#actor.startDestroy();
 	}
 }

package/src/actor/contexts/conn-init.ts

@@ -0,0 +1,31 @@
+import type { AnyDatabaseProvider } from "../database";
+import type { ActorInstance } from "../instance/mod";
+import { ActorContext } from "./actor";
+
+/**
+ * Base context for connection initialization handlers.
+ * Extends ActorContext with request-specific functionality for connection lifecycle events.
+ */
+export abstract class ConnInitContext<
+	TState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+> extends ActorContext<TState, undefined, undefined, TVars, TInput, TDatabase> {
+	/**
+	 * The incoming request that initiated the connection.
+	 * May be undefined for connections initiated without a direct HTTP request.
+	 */
+	public readonly request: Request | undefined;
+
+	/**
+	 * @internal
+	 */
+	constructor(
+		actor: ActorInstance<TState, any, any, TVars, TInput, TDatabase>,
+		request: Request | undefined,
+	) {
+		super(actor);
+		this.request = request;
+	}
+}

package/src/actor/contexts/conn.ts

@@ -0,0 +1,48 @@
+import type { Conn } from "../conn/mod";
+import type { AnyDatabaseProvider } from "../database";
+import type { ActorInstance } from "../instance/mod";
+import { ActorContext } from "./actor";
+
+/**
+ * Base context for connection-based handlers.
+ * Extends ActorContext with connection-specific functionality.
+ */
+export abstract class ConnContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+> extends ActorContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase
+> {
+	/**
+	 * @internal
+	 */
+	constructor(
+		actor: ActorInstance<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase
+		>,
+		public readonly conn: Conn<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase
+		>,
+	) {
+		super(actor);
+	}
+}

package/src/actor/contexts/create-conn-state.ts

@@ -0,0 +1,13 @@
+import type { AnyDatabaseProvider } from "../database";
+import { ConnInitContext } from "./conn-init";
+
+/**
+ * Context for the createConnState lifecycle hook.
+ * Called to initialize connection-specific state when a connection is created.
+ */
+export class CreateConnStateContext<
+	TState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+> extends ConnInitContext<TState, TVars, TInput, TDatabase> {}

package/src/actor/contexts/on-before-connect.ts

@@ -0,0 +1,13 @@
+import type { AnyDatabaseProvider } from "../database";
+import { ConnInitContext } from "./conn-init";
+
+/**
+ * Context for the onBeforeConnect lifecycle hook.
+ * Called before a connection is established, allowing for validation and early rejection.
+ */
+export class OnBeforeConnectContext<
+	TState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+> extends ConnInitContext<TState, TVars, TInput, TDatabase> {}