rivetkit 2.0.2 → 2.0.3

package/src/actor/connection.ts

@@ -0,0 +1,260 @@
+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 { bufferToArrayBuffer } from "@/utils";
+import type { AnyDatabaseProvider } from "./database";
+import { type ConnDriver, ConnectionReadyState } from "./driver";
+import * as errors from "./errors";
+import type { ActorInstance } from "./instance";
+import { logger } from "./log";
+import type { PersistedConn } from "./persisted";
+import { CachedSerializer } from "./protocol/serde";
+import { generateSecureToken } from "./utils";
+
+export function generateConnId(): string {
+	return crypto.randomUUID();
+}
+
+export function generateConnToken(): string {
+	return generateSecureToken(32);
+}
+
+export type ConnId = string;
+
+export type AnyConn = Conn<any, any, any, any, any, any, any>;
+
+export const CONNECTION_DRIVER_WEBSOCKET = "webSocket";
+export const CONNECTION_DRIVER_SSE = "sse";
+export const CONNECTION_DRIVER_HTTP = "http";
+
+export type ConnectionDriver =
+	| typeof CONNECTION_DRIVER_WEBSOCKET
+	| typeof CONNECTION_DRIVER_SSE
+	| typeof CONNECTION_DRIVER_HTTP;
+
+export type ConnectionStatus = "connected" | "reconnecting";
+
+export const CONNECTION_CHECK_LIVENESS_SYMBOL = Symbol("checkLiveness");
+
+/**
+ * Represents a client connection to a actor.
+ *
+ * Manages connection-specific data and controls the connection lifecycle.
+ *
+ * @see {@link https://rivet.gg/docs/connections|Connection Documentation}
+ */
+export class Conn<S, CP, CS, V, I, AD, DB extends AnyDatabaseProvider> {
+	subscriptions: Set<string> = new Set<string>();
+
+	#stateEnabled: boolean;
+
+	// TODO: Remove this cyclical reference
+	#actor: ActorInstance<S, CP, CS, V, I, AD, DB>;
+
+	#status: ConnectionStatus = "connected";
+
+	/**
+	 * The proxied state that notifies of changes automatically.
+	 *
+	 * Any data that should be stored indefinitely should be held within this object.
+	 */
+	__persist: PersistedConn<CP, CS>;
+
+	/**
+	 * Driver used to manage realtime connection communication.
+	 *
+	 * @protected
+	 */
+	#driver: ConnDriver;
+
+	public get params(): CP {
+		return this.__persist.params;
+	}
+
+	public get auth(): AD {
+		return this.__persist.authData as AD;
+	}
+
+	public get driver(): ConnectionDriver {
+		return this.__persist.connDriver as ConnectionDriver;
+	}
+
+	public get _stateEnabled() {
+		return this.#stateEnabled;
+	}
+
+	/**
+	 * Gets the current state of the connection.
+	 *
+	 * Throws an error if the state is not enabled.
+	 */
+	public get state(): CS {
+		this.#validateStateEnabled();
+		if (!this.__persist.state) throw new Error("state should exists");
+		return this.__persist.state;
+	}
+
+	/**
+	 * Sets the state of the connection.
+	 *
+	 * Throws an error if the state is not enabled.
+	 */
+	public set state(value: CS) {
+		this.#validateStateEnabled();
+		this.__persist.state = value;
+	}
+
+	/**
+	 * Unique identifier for the connection.
+	 */
+	public get id(): ConnId {
+		return this.__persist.connId;
+	}
+
+	/**
+	 * Token used to authenticate this request.
+	 */
+	public get _token(): string {
+		return this.__persist.token;
+	}
+
+	/**
+	 * Status of the connection.
+	 */
+	public get status(): ConnectionStatus {
+		return this.#status;
+	}
+
+	/**
+	 * Timestamp of the last time the connection was seen, i.e. the last time the connection was active and checked for liveness.
+	 */
+	public get lastSeen(): number {
+		return this.__persist.lastSeen;
+	}
+
+	/**
+	 * Initializes a new instance of the Connection class.
+	 *
+	 * This should only be constructed by {@link Actor}.
+	 *
+	 * @protected
+	 */
+	public constructor(
+		actor: ActorInstance<S, CP, CS, V, I, AD, DB>,
+		persist: PersistedConn<CP, CS>,
+		driver: ConnDriver,
+		stateEnabled: boolean,
+	) {
+		this.#actor = actor;
+		this.__persist = persist;
+		this.#driver = driver;
+		this.#stateEnabled = stateEnabled;
+	}
+
+	#validateStateEnabled() {
+		if (!this.#stateEnabled) {
+			throw new errors.ConnStateNotEnabled();
+		}
+	}
+
+	/**
+	 * Sends a WebSocket message to the client.
+	 *
+	 * @param message - The message to send.
+	 *
+	 * @protected
+	 */
+	public _sendMessage(message: CachedSerializer<protocol.ToClient>) {
+		this.#driver.sendMessage?.(
+			this.#actor,
+			this,
+			this.__persist.connDriverState,
+			message,
+		);
+	}
+
+	/**
+	 * 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.gg/docs/events|Events Documentation}
+	 */
+	public send(eventName: string, ...args: unknown[]) {
+		this.#actor.inspector.emitter.emit("eventFired", {
+			type: "event",
+			eventName,
+			args,
+			connId: this.id,
+		});
+		this._sendMessage(
+			new CachedSerializer<protocol.ToClient>(
+				{
+					body: {
+						tag: "Event",
+						val: {
+							name: eventName,
+							args: bufferToArrayBuffer(cbor.encode(args)),
+						},
+					},
+				},
+				TO_CLIENT_VERSIONED,
+			),
+		);
+	}
+
+	/**
+	 * Disconnects the client with an optional reason.
+	 *
+	 * @param reason - The reason for disconnection.
+	 */
+	public async disconnect(reason?: string) {
+		this.#status = "reconnecting";
+		await this.#driver.disconnect(
+			this.#actor,
+			this,
+			this.__persist.connDriverState,
+			reason,
+		);
+	}
+
+	/**
+	 * This method checks the connection's liveness by querying the driver for its ready state.
+	 * If the connection is not closed, it updates the last liveness timestamp and returns `true`.
+	 * Otherwise, it returns `false`.
+	 * @internal
+	 */
+	[CONNECTION_CHECK_LIVENESS_SYMBOL]() {
+		const readyState = this.#driver.getConnectionReadyState(this.#actor, this);
+
+		const isConnectionClosed =
+			readyState === ConnectionReadyState.CLOSED ||
+			readyState === ConnectionReadyState.CLOSING ||
+			readyState === undefined;
+
+		const newLastSeen = Date.now();
+		const newStatus = isConnectionClosed ? "reconnecting" : "connected";
+
+		logger().debug("liveness probe for connection", {
+			connId: this.id,
+			actorId: this.#actor.id,
+			readyState,
+
+			status: this.#status,
+			newStatus,
+
+			lastSeen: this.__persist.lastSeen,
+			currentTs: newLastSeen,
+		});
+
+		if (!isConnectionClosed) {
+			this.__persist.lastSeen = newLastSeen;
+		}
+
+		this.#status = newStatus;
+		return {
+			status: this.#status,
+			lastSeen: this.__persist.lastSeen,
+		};
+	}
+}

package/src/actor/context.ts

@@ -0,0 +1,171 @@
+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 "./connection";
+import type { AnyDatabaseProvider, InferDatabaseClient } from "./database";
+import type { ActorInstance, SaveStateOptions } from "./instance";
+import type { Schedule } from "./schedule";
+
+/**
+ * ActorContext class that provides access to actor methods and state
+ */
+export class ActorContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TAuthData,
+	TDatabase extends AnyDatabaseProvider,
+> {
+	#actor: ActorInstance<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TAuthData,
+		TDatabase
+	>;
+
+	constructor(
+		actor: ActorInstance<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TAuthData,
+			TDatabase
+		>,
+	) {
+		this.#actor = actor;
+	}
+
+	/**
+	 * Get the actor state
+	 */
+	get state(): TState {
+		return this.#actor.state;
+	}
+
+	/**
+	 * Get the actor variables
+	 */
+	get vars(): TVars {
+		return this.#actor.vars;
+	}
+
+	/**
+	 * Broadcasts an event to all connected clients.
+	 * @param name - The name of the event.
+	 * @param args - The arguments to send with the event.
+	 */
+	broadcast<Args extends Array<unknown>>(name: string, ...args: Args): void {
+		this.#actor._broadcast(name, ...args);
+		return;
+	}
+
+	/**
+	 * Gets the logger instance.
+	 */
+	get log(): Logger {
+		return this.#actor.log;
+	}
+
+	/**
+	 * Gets actor ID.
+	 */
+	get actorId(): string {
+		return this.#actor.id;
+	}
+
+	/**
+	 * Gets the actor name.
+	 */
+	get name(): string {
+		return this.#actor.name;
+	}
+
+	/**
+	 * Gets the actor key.
+	 */
+	get key(): ActorKey {
+		return this.#actor.key;
+	}
+
+	/**
+	 * Gets the region.
+	 */
+	get region(): string {
+		return this.#actor.region;
+	}
+
+	/**
+	 * Gets the scheduler.
+	 */
+	get schedule(): Schedule {
+		return this.#actor.schedule;
+	}
+
+	/**
+	 * Gets the map of connections.
+	 */
+	get conns(): Map<
+		ConnId,
+		Conn<TState, TConnParams, TConnState, TVars, TInput, TAuthData, TDatabase>
+	> {
+		return this.#actor.conns;
+	}
+
+	/**
+	 * Returns the client for the given registry.
+	 */
+	client<R extends Registry<any>>(): Client<R> {
+		return this.#actor.inlineClient as Client<R>;
+	}
+
+	/**
+	 * Gets the database.
+	 * @experimental
+	 * @throws {DatabaseNotEnabled} If the database is not enabled.
+	 */
+	get db(): InferDatabaseClient<TDatabase> {
+		return this.#actor.db;
+	}
+
+	/**
+	 * Forces the state to get saved.
+	 *
+	 * @param opts - Options for saving the state.
+	 */
+	async saveState(opts: SaveStateOptions): Promise<void> {
+		return this.#actor.saveState(opts);
+	}
+
+	/**
+	 * Prevents the actor from sleeping until promise is complete.
+	 */
+	waitUntil(promise: Promise<void>): void {
+		this.#actor._waitUntil(promise);
+	}
+
+	/**
+	 * AbortSignal that fires when the actor is stopping.
+	 */
+	get abortSignal(): AbortSignal {
+		return this.#actor.abortSignal;
+	}
+
+	/**
+	 * Forces the actor to sleep.
+	 *
+	 * Not supported on all drivers.
+	 *
+	 * @experimental
+	 */
+	sleep() {
+		this.#actor._sleep();
+	}
+}

package/src/actor/database.ts

@@ -0,0 +1,23 @@
+export type InferDatabaseClient<DBProvider extends AnyDatabaseProvider> =
+	DBProvider extends DatabaseProvider<any>
+		? Awaited<ReturnType<DBProvider["createClient"]>>
+		: never;
+
+export type AnyDatabaseProvider = DatabaseProvider<any> | undefined;
+
+export type DatabaseProvider<DB extends { execute: (query: string) => any }> = {
+	/**
+	 * Creates a new database client for the actor.
+	 * The result is passed to the actor context as `c.db`.
+	 * @experimental
+	 */
+	createClient: (ctx: {
+		getDatabase: () => Promise<string | unknown>;
+	}) => Promise<DB>;
+	/**
+	 * Runs before the actor has started.
+	 * Use this to run migrations or other setup tasks.
+	 * @experimental
+	 */
+	onMigrate: (client: DB) => void | Promise<void>;
+};

package/src/actor/definition.ts

@@ -0,0 +1,86 @@
+import type { RegistryConfig } from "@/registry/config";
+import type { ActionContext } from "./action";
+import type { Actions, ActorConfig } from "./config";
+import type { ActorContext } from "./context";
+import type { AnyDatabaseProvider } from "./database";
+import { ActorInstance } from "./instance";
+
+export type AnyActorDefinition = ActorDefinition<
+	any,
+	any,
+	any,
+	any,
+	any,
+	any,
+	any,
+	any
+>;
+
+/**
+ * Extracts the context type from an ActorDefinition
+ */
+export type ActorContextOf<AD extends AnyActorDefinition> =
+	AD extends ActorDefinition<
+		infer S,
+		infer CP,
+		infer CS,
+		infer V,
+		infer I,
+		infer AD,
+		infer DB,
+		any
+	>
+		? ActorContext<S, CP, CS, V, I, AD, DB>
+		: never;
+
+/**
+ * Extracts the context type from an ActorDefinition
+ */
+export type ActionContextOf<AD extends AnyActorDefinition> =
+	AD extends ActorDefinition<
+		infer S,
+		infer CP,
+		infer CS,
+		infer V,
+		infer I,
+		infer AD,
+		infer DB,
+		any
+	>
+		? ActionContext<S, CP, CS, V, I, AD, DB>
+		: never;
+
+export class ActorDefinition<
+	S,
+	CP,
+	CS,
+	V,
+	I,
+	AD,
+	DB extends AnyDatabaseProvider,
+	R extends Actions<S, CP, CS, V, I, AD, DB>,
+> {
+	#config: ActorConfig<S, CP, CS, V, I, AD, DB>;
+
+	constructor(config: ActorConfig<S, CP, CS, V, I, AD, DB>) {
+		this.#config = config;
+	}
+
+	get config(): ActorConfig<S, CP, CS, V, I, AD, DB> {
+		return this.#config;
+	}
+
+	instantiate(): ActorInstance<S, CP, CS, V, I, AD, DB> {
+		return new ActorInstance(this.#config);
+	}
+}
+
+export function lookupInRegistry(
+	registryConfig: RegistryConfig,
+	name: string,
+): AnyActorDefinition {
+	// Build actor
+	const definition = registryConfig.use[name];
+	if (!definition) throw new Error(`no actor in registry for name ${name}`);
+	return definition;
+}

package/src/actor/driver.ts

@@ -0,0 +1,84 @@
+import type { CachedSerializer } from "@/actor/protocol/serde";
+import type { AnyClient } from "@/client/client";
+import type { ManagerDriver } from "@/manager/driver";
+import type { RegistryConfig } from "@/registry/config";
+import type { RunConfig } from "@/registry/run-config";
+import type * as protocol from "@/schemas/client-protocol/mod";
+import type { AnyConn, ConnectionDriver } from "./connection";
+import type { GenericConnGlobalState } from "./generic-conn-driver";
+import type { AnyActorInstance } from "./instance";
+
+export type ConnectionDriversMap = Record<ConnectionDriver, ConnDriver>;
+
+export type ActorDriverBuilder = (
+	registryConfig: RegistryConfig,
+	runConfig: RunConfig,
+	managerDriver: ManagerDriver,
+	inlineClient: AnyClient,
+) => ActorDriver;
+
+export interface ActorDriver {
+	//load(): Promise<LoadOutput>;
+
+	loadActor(actorId: string): Promise<AnyActorInstance>;
+
+	getGenericConnGlobalState(actorId: string): GenericConnGlobalState;
+
+	getContext(actorId: string): unknown;
+
+	readPersistedData(actorId: string): Promise<Uint8Array | undefined>;
+
+	/** ActorInstance ensure that only one instance of writePersistedData is called in parallel at a time. */
+	writePersistedData(actorId: string, data: Uint8Array): Promise<void>;
+
+	// Schedule
+	/** ActorInstance ensure that only one instance of setAlarm is called in parallel at a time. */
+	setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void>;
+
+	// Database
+	/**
+	 * @experimental
+	 * This is an experimental API that may change in the future.
+	 */
+	getDatabase(actorId: string): Promise<unknown | undefined>;
+
+	sleep?(actorId: string): Promise<void>;
+
+	shutdown?(immediate: boolean): Promise<void>;
+}
+
+export enum ConnectionReadyState {
+	UNKNOWN = -1,
+	CONNECTING = 0,
+	OPEN = 1,
+	CLOSING = 2,
+	CLOSED = 3,
+}
+
+export interface ConnDriver<ConnDriverState = unknown> {
+	sendMessage?(
+		actor: AnyActorInstance,
+		conn: AnyConn,
+		state: ConnDriverState,
+		message: CachedSerializer<protocol.ToClient>,
+	): void;
+
+	/**
+	 * 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,
+		state: ConnDriverState,
+		reason?: string,
+	): Promise<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,
+	): ConnectionReadyState | undefined;
+}