rivetkit 2.0.2 → 2.0.4-rc.1

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

@@ -0,0 +1,196 @@
+import type { HibernatingWebSocketMetadata } from "@rivetkit/engine-runner";
+import onChange from "@rivetkit/on-change";
+import * as cbor from "cbor-x";
+import invariant from "invariant";
+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, 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, 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,47 @@
+import type { Conn } from "../conn/mod";
+import type { AnyDatabaseProvider } from "../database";
+import type { ActorDefinition, AnyActorDefinition } from "../definition";
+import type { ActorInstance } from "../instance/mod";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
+import { ConnContext } from "./base/conn";
+
+/**
+ * Context for a remote procedure call.
+ */
+export class ActionContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+> extends ConnContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase,
+	TEvents,
+	TQueues
+> {}
+
+/**
+ * Extracts the ActionContext type from an ActorDefinition.
+ */
+export type ActionContextOf<AD extends AnyActorDefinition> =
+	AD extends ActorDefinition<
+		infer S,
+		infer CP,
+		infer CS,
+		infer V,
+		infer I,
+		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
+		any
+	>
+		? ActionContext<S, CP, CS, V, I, DB, E, Q>
+		: never;

package/src/actor/contexts/base/actor.ts

@@ -0,0 +1,347 @@
+import type { ActorKey } from "@/actor/mod";
+import type { Client } from "@/client/client";
+import type { Logger } from "@/common/log";
+import type { Registry } from "@/registry";
+import type { Conn, ConnId } from "../../conn/mod";
+import type { AnyDatabaseProvider, InferDatabaseClient } from "../../database";
+import type { ActorDefinition, AnyActorDefinition } from "../../definition";
+import * as errors from "../../errors";
+import { ActorKv } from "../../instance/kv";
+import type {
+	ActorInstance,
+	AnyActorInstance,
+	SaveStateOptions,
+} from "../../instance/mod";
+import { ActorQueue } from "../../instance/queue";
+import type { Schedule } from "../../schedule";
+import {
+	type EventSchemaConfig,
+	type InferEventArgs,
+	type InferSchemaMap,
+	type QueueSchemaConfig,
+	hasSchemaConfigKey,
+	validateSchemaSync,
+} from "../../schema";
+
+export const ACTOR_CONTEXT_INTERNAL_SYMBOL = Symbol.for(
+	"rivetkit.actorContextInternal",
+);
+
+/**
+ * ActorContext class that provides access to actor methods and state
+ */
+export class ActorContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+> {
+	[ACTOR_CONTEXT_INTERNAL_SYMBOL]!: AnyActorInstance;
+	#actor: ActorInstance<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+	#kv: ActorKv | undefined;
+	#queue:
+		| ActorQueue<
+				TState,
+				TConnParams,
+				TConnState,
+				TVars,
+				TInput,
+				TDatabase,
+				TEvents,
+				TQueues
+		  >
+		| undefined;
+
+	constructor(
+		actor: ActorInstance<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	) {
+		this.#actor = actor;
+		this[ACTOR_CONTEXT_INTERNAL_SYMBOL] = actor as AnyActorInstance;
+	}
+
+	/**
+	 * Gets the KV storage interface.
+	 */
+	get kv(): ActorKv {
+		if (!this.#kv) {
+			this.#kv = new ActorKv(this.#actor.driver, this.#actor.id);
+		}
+		return this.#kv;
+	}
+
+	/**
+	 * Get the actor state
+	 *
+	 * @remarks
+	 * This property is not available in `createState` since the state hasn't been created yet.
+	 */
+	get state(): TState extends never ? never : TState {
+		return this.#actor.state as TState extends never ? never : TState;
+	}
+
+	/**
+	 * Get the actor variables
+	 *
+	 * @remarks
+	 * This property is not available in `createVars` since the variables haven't been created yet.
+	 * Variables are only available if you define `vars` or `createVars` in your actor config.
+	 */
+	get vars(): TVars extends never ? never : TVars {
+		return this.#actor.vars as TVars extends never ? never : TVars;
+	}
+
+	/**
+	 * Broadcasts an event to all connected clients.
+	 * @param name - The name of the event.
+	 * @param args - The arguments to send with the event.
+	 */
+	broadcast<K extends keyof TEvents & string>(
+		name: K,
+		...args: InferEventArgs<InferSchemaMap<TEvents>[K]>
+	): void;
+	broadcast(
+		name: keyof TEvents extends never ? string : never,
+		...args: Array<unknown>
+	): void;
+	broadcast(name: string, ...args: Array<unknown>): void {
+		if (
+			this.#actor.config.events !== undefined &&
+			!hasSchemaConfigKey(this.#actor.config.events, name)
+		) {
+			this.#actor.rLog.warn({
+				msg: "broadcasting event not defined in actor events config",
+				eventName: name,
+			});
+		}
+
+		const payload = args.length === 1 ? args[0] : args;
+		const result = validateSchemaSync(
+			this.#actor.config.events,
+			name as keyof TEvents & string,
+			payload,
+		);
+		if (!result.success) {
+			throw new errors.EventPayloadInvalid(name, result.issues);
+		}
+		if (args.length === 1) {
+			this.#actor.eventManager.broadcast(name, result.data);
+			return;
+		}
+		if (Array.isArray(result.data)) {
+			this.#actor.eventManager.broadcast(
+				name,
+				...(result.data as unknown[]),
+			);
+			return;
+		}
+		this.#actor.eventManager.broadcast(name, ...args);
+	}
+
+	/**
+	 * Gets the logger instance.
+	 */
+	get log(): Logger {
+		return this.#actor.log;
+	}
+
+	/**
+	 * Access to queue receive helpers.
+	 */
+	get queue(): ActorQueue<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	> {
+		if (!this.#queue) {
+			this.#queue = new ActorQueue(
+				this.#actor.queueManager,
+				this.#actor.abortSignal,
+			);
+		}
+		return this.#queue;
+	}
+
+	/**
+	 * 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,
+			TDatabase,
+			TEvents,
+			TQueues
+		>
+	> {
+		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
+	 * @remarks
+	 * This property is only available if you define a `db` provider in your actor config.
+	 * @throws {DatabaseNotEnabled} If the database is not enabled.
+	 */
+	get db(): TDatabase extends never ? never : InferDatabaseClient<TDatabase> {
+		return this.#actor.db as TDatabase extends never
+			? never
+			: InferDatabaseClient<TDatabase>;
+	}
+
+	/**
+	 * Forces the state to get saved.
+	 *
+	 * @param opts - Options for saving the state.
+	 */
+	async saveState(opts: SaveStateOptions): Promise<void> {
+		return this.#actor.stateManager.saveState(opts);
+	}
+
+	/**
+	 * Prevents the actor from sleeping until promise is complete.
+	 */
+	waitUntil(promise: Promise<void>): void {
+		this.#actor.waitUntil(promise);
+	}
+
+	/**
+	 * Prevents the actor from sleeping while the given promise is running.
+	 *
+	 * Use this when performing async operations in the `run` handler or other
+	 * background contexts where you need to ensure the actor stays awake.
+	 *
+	 * Returns the resolved value and resets the sleep timer on completion.
+	 * Errors are propagated to the caller.
+	 */
+	keepAwake<T>(promise: Promise<T>): Promise<T> {
+		return this.#actor.keepAwake(promise);
+	}
+
+	/**
+	 * AbortSignal that fires when the actor is stopping.
+	 */
+	get abortSignal(): AbortSignal {
+		return this.#actor.abortSignal;
+	}
+
+	/**
+	 * True when the actor is stopping.
+	 *
+	 * Alias for `c.abortSignal.aborted`.
+	 */
+	get aborted(): boolean {
+		return this.#actor.abortSignal.aborted;
+	}
+
+	/**
+	 * Forces the actor to sleep.
+	 *
+	 * Not supported on all drivers.
+	 *
+	 * @experimental
+	 */
+	sleep() {
+		this.#actor.startSleep();
+	}
+
+	/**
+	 * Forces the actor to destroy.
+	 *
+	 * This will return immediately, then call `onStop` and `onDestroy`.
+	 *
+	 * @experimental
+	 */
+	destroy() {
+		this.#actor.startDestroy();
+	}
+}
+
+export type ActorContextOf<AD extends AnyActorDefinition> =
+	AD extends ActorDefinition<
+		infer S,
+		infer CP,
+		infer CS,
+		infer V,
+		infer I,
+		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
+		any
+	>
+		? ActorContext<S, CP, CS, V, I, DB, E, Q>
+		: never;

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

@@ -0,0 +1,68 @@
+import type { AnyDatabaseProvider } from "../../database";
+import type { ActorDefinition, AnyActorDefinition } from "../../definition";
+import type { ActorInstance } from "../../instance/mod";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../../schema";
+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,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+> extends ActorContext<
+	TState,
+	never,
+	never,
+	TVars,
+	TInput,
+	TDatabase,
+	TEvents,
+	TQueues
+> {
+	/**
+	 * 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,
+			TEvents,
+			TQueues
+		>,
+		request: Request | undefined,
+	) {
+		super(actor as any);
+		this.request = request;
+	}
+}
+
+export type ConnInitContextOf<AD extends AnyActorDefinition> =
+	AD extends ActorDefinition<
+		infer S,
+		any,
+		any,
+		infer V,
+		infer I,
+		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
+		any
+	>
+		? ConnInitContext<S, V, I, DB, E, Q>
+		: never;