rivetkit 2.0.42 → 2.1.0-rc.1

package/src/actor/conn/mod.ts

@@ -10,15 +10,23 @@ import {
 } from "@/schemas/client-protocol-zod/mod";
 import { bufferToArrayBuffer } from "@/utils";
 import type { AnyDatabaseProvider } from "../database";
-import { InternalError } from "../errors";
+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>;
+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");
@@ -34,10 +42,19 @@ export const CONN_SEND_MESSAGE_SYMBOL = Symbol("sendMessage");
  *
  * @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>;
+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> {
+	get [CONN_ACTOR_SYMBOL](): ActorInstance<S, CP, CS, V, I, DB, E, Q> {
 		return this.#actor;
 	}
 
@@ -111,7 +128,7 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 * If the underlying connection can hibernate.
 	 */
 	get isHibernatable(): boolean {
-		return this[CONN_DRIVER_SYMBOL]?.hibernatable !== undefined;
+		return this.#stateManager.hibernatableDataRaw !== undefined;
 	}
 
 	/**
@@ -122,7 +139,7 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 * @protected
 	 */
 	constructor(
-		actor: ActorInstance<S, CP, CS, V, I, DB>,
+		actor: ActorInstance<S, CP, CS, V, I, DB, E, Q>,
 		data: ConnDataInput<CP, CS>,
 	) {
 		this.#actor = actor;
@@ -159,6 +176,14 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 * @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]) {
@@ -169,13 +194,37 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			});
 		}
 
-		this.#actor.inspector.emitter.emit("eventFired", {
-			type: "event",
-			eventName,
-			args,
-			connId: this.id,
+		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 };
+		const eventData = { name: eventName, args: eventArgs };
 		this[CONN_SEND_MESSAGE_SYMBOL](
 			new CachedSerializer(
 				eventData,
@@ -223,14 +272,17 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 				});
 			}
 
-			this.#actor.connectionManager.connDisconnected(this);
+			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;
 		}
-
-		this[CONN_DRIVER_SYMBOL] = undefined;
 	}
 }

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

@@ -39,7 +39,7 @@ export type ConnData<CP, CS> =
  * Handles automatic state change detection for connection-specific state.
  */
 export class StateManager<CP, CS> {
-	#conn: Conn<any, CP, CS, any, any, any>;
+	#conn: Conn<any, CP, CS, any, any, any, any, any>;
 
 	/**
 	 * Data representing this connection.
@@ -50,7 +50,7 @@ export class StateManager<CP, CS> {
 	#data!: ConnData<CP, CS>;
 
 	constructor(
-		conn: Conn<any, CP, CS, any, any, any>,
+		conn: Conn<any, CP, CS, any, any, any, any, any>,
 		data: ConnDataInput<CP, CS>,
 	) {
 		this.#conn = conn;

package/src/actor/contexts/action.ts

@@ -2,6 +2,7 @@ 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";
 
 /**
@@ -14,26 +15,33 @@ export class ActionContext<
 	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
+	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,
-	any
->
-	? ActionContext<S, CP, CS, V, I, DB>
-	: never;
+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

@@ -5,9 +5,27 @@ 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 type { ActorInstance, SaveStateOptions } from "../../instance/mod";
+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
@@ -19,16 +37,33 @@ export class ActorContext<
 	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
+		TDatabase,
+		TEvents,
+		TQueues
 	>;
 	#kv: ActorKv | undefined;
+	#queue:
+		| ActorQueue<
+				TState,
+				TConnParams,
+				TConnState,
+				TVars,
+				TInput,
+				TDatabase,
+				TEvents,
+				TQueues
+		  >
+		| undefined;
 
 	constructor(
 		actor: ActorInstance<
@@ -37,10 +72,13 @@ export class ActorContext<
 			TConnState,
 			TVars,
 			TInput,
-			TDatabase
+			TDatabase,
+			TEvents,
+			TQueues
 		>,
 	) {
 		this.#actor = actor;
+		this[ACTOR_CONTEXT_INTERNAL_SYMBOL] = actor as AnyActorInstance;
 	}
 
 	/**
@@ -79,9 +117,46 @@ export class ActorContext<
 	 * @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 {
+	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);
-		return;
 	}
 
 	/**
@@ -91,6 +166,28 @@ export class ActorContext<
 		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.
 	 */
@@ -131,7 +228,16 @@ export class ActorContext<
 	 */
 	get conns(): Map<
 		ConnId,
-		Conn<TState, TConnParams, TConnState, TVars, TInput, TDatabase>
+		Conn<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>
 	> {
 		return this.#actor.conns;
 	}
@@ -173,6 +279,19 @@ export class ActorContext<
 		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.
 	 */
@@ -180,6 +299,15 @@ export class ActorContext<
 		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.
 	 *
@@ -211,7 +339,9 @@ export type ActorContextOf<AD extends AnyActorDefinition> =
 		infer V,
 		infer I,
 		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
 		any
 	>
-		? ActorContext<S, CP, CS, V, I, DB>
+		? ActorContext<S, CP, CS, V, I, DB, E, Q>
 		: never;

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

@@ -1,9 +1,7 @@
 import type { AnyDatabaseProvider } from "../../database";
-import type {
-	ActorDefinition,
-	AnyActorDefinition,
-} from "../../definition";
+import type { ActorDefinition, AnyActorDefinition } from "../../definition";
 import type { ActorInstance } from "../../instance/mod";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../../schema";
 import { ActorContext } from "./actor";
 
 /**
@@ -15,7 +13,18 @@ export abstract class ConnInitContext<
 	TVars,
 	TInput,
 	TDatabase extends AnyDatabaseProvider,
-> extends ActorContext<TState, never, never, TVars, TInput, TDatabase> {
+	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.
@@ -26,7 +35,16 @@ export abstract class ConnInitContext<
 	 * @internal
 	 */
 	constructor(
-		actor: ActorInstance<TState, any, any, TVars, TInput, TDatabase>,
+		actor: ActorInstance<
+			TState,
+			any,
+			any,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
 		request: Request | undefined,
 	) {
 		super(actor as any);
@@ -42,7 +60,9 @@ export type ConnInitContextOf<AD extends AnyActorDefinition> =
 		infer V,
 		infer I,
 		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
 		any
 	>
-		? ConnInitContext<S, V, I, DB>
+		? ConnInitContext<S, V, I, DB, E, Q>
 		: never;

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

@@ -1,10 +1,8 @@
 import type { Conn } from "../../conn/mod";
 import type { AnyDatabaseProvider } from "../../database";
-import type {
-	ActorDefinition,
-	AnyActorDefinition,
-} from "../../definition";
+import type { ActorDefinition, AnyActorDefinition } from "../../definition";
 import type { ActorInstance } from "../../instance/mod";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../../schema";
 import { ActorContext } from "./actor";
 
 /**
@@ -18,13 +16,17 @@ export abstract class ConnContext<
 	TVars,
 	TInput,
 	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
 > extends ActorContext<
 	TState,
 	TConnParams,
 	TConnState,
 	TVars,
 	TInput,
-	TDatabase
+	TDatabase,
+	TEvents,
+	TQueues
 > {
 	/**
 	 * @internal
@@ -36,7 +38,9 @@ export abstract class ConnContext<
 			TConnState,
 			TVars,
 			TInput,
-			TDatabase
+			TDatabase,
+			TEvents,
+			TQueues
 		>,
 		public readonly conn: Conn<
 			TState,
@@ -44,21 +48,26 @@ export abstract class ConnContext<
 			TConnState,
 			TVars,
 			TInput,
-			TDatabase
+			TDatabase,
+			TEvents,
+			TQueues
 		>,
 	) {
 		super(actor);
 	}
 }
 
-export type ConnContextOf<AD extends AnyActorDefinition> = AD extends ActorDefinition<
-	infer S,
-	infer CP,
-	infer CS,
-	infer V,
-	infer I,
-	infer DB extends AnyDatabaseProvider,
-	any
->
-	? ConnContext<S, CP, CS, V, I, DB>
-	: never;
+export type ConnContextOf<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
+	>
+		? ConnContext<S, CP, CS, V, I, DB, E, Q>
+		: never;

package/src/actor/contexts/before-action-response.ts

@@ -1,5 +1,6 @@
 import type { AnyDatabaseProvider } from "../database";
 import type { ActorDefinition, AnyActorDefinition } from "../definition";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
 import { ActorContext } from "./base/actor";
 
 /**
@@ -12,13 +13,17 @@ export class BeforeActionResponseContext<
 	TVars,
 	TInput,
 	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
 > extends ActorContext<
 	TState,
 	TConnParams,
 	TConnState,
 	TVars,
 	TInput,
-	TDatabase
+	TDatabase,
+	TEvents,
+	TQueues
 > {}
 
 export type BeforeActionResponseContextOf<AD extends AnyActorDefinition> =
@@ -29,7 +34,9 @@ export type BeforeActionResponseContextOf<AD extends AnyActorDefinition> =
 		infer V,
 		infer I,
 		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
 		any
 	>
-		? BeforeActionResponseContext<S, CP, CS, V, I, DB>
+		? BeforeActionResponseContext<S, CP, CS, V, I, DB, E, Q>
 		: never;

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

@@ -1,5 +1,6 @@
 import type { AnyDatabaseProvider } from "../database";
 import type { ActorDefinition, AnyActorDefinition } from "../definition";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
 import { ConnInitContext } from "./base/conn-init";
 
 /**
@@ -10,7 +11,9 @@ export class BeforeConnectContext<
 	TVars,
 	TInput,
 	TDatabase extends AnyDatabaseProvider,
-> extends ConnInitContext<TState, TVars, TInput, TDatabase> {}
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+> extends ConnInitContext<TState, TVars, TInput, TDatabase, TEvents, TQueues> {}
 
 export type BeforeConnectContextOf<AD extends AnyActorDefinition> =
 	AD extends ActorDefinition<
@@ -20,7 +23,9 @@ export type BeforeConnectContextOf<AD extends AnyActorDefinition> =
 		infer V,
 		infer I,
 		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
 		any
 	>
-		? BeforeConnectContext<S, V, I, DB>
+		? BeforeConnectContext<S, V, I, DB, E, Q>
 		: never;

package/src/actor/contexts/connect.ts

@@ -1,5 +1,6 @@
 import type { AnyDatabaseProvider } from "../database";
 import type { ActorDefinition, AnyActorDefinition } from "../definition";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
 import { ConnContext } from "./base/conn";
 
 /**
@@ -12,13 +13,17 @@ export class ConnectContext<
 	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
+	TDatabase,
+	TEvents,
+	TQueues
 > {}
 
 export type ConnectContextOf<AD extends AnyActorDefinition> =
@@ -29,7 +34,9 @@ export type ConnectContextOf<AD extends AnyActorDefinition> =
 		infer V,
 		infer I,
 		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
 		any
 	>
-		? ConnectContext<S, CP, CS, V, I, DB>
+		? ConnectContext<S, CP, CS, V, I, DB, E, Q>
 		: never;