rivetkit 2.0.42 → 2.1.0-rc.1

package/src/actor/errors.ts

@@ -188,6 +188,111 @@ export class Unsupported extends ActorError {
 	}
 }
 
+export class QueueFull extends ActorError {
+	constructor(limit: number) {
+		super("queue", "full", `Queue is full. Limit is ${limit} messages.`, {
+			public: true,
+			metadata: { limit },
+		});
+	}
+}
+
+export class QueueMessageTooLarge extends ActorError {
+	constructor(size: number, limit: number) {
+		super(
+			"queue",
+			"message_too_large",
+			`Queue message too large (${size} bytes). Limit is ${limit} bytes.`,
+			{ public: true, metadata: { size, limit } },
+		);
+	}
+}
+
+export class QueueMessageInvalid extends ActorError {
+	constructor(path?: string) {
+		super(
+			"queue",
+			"message_invalid",
+			path
+				? `Queue message body contains unsupported type at ${path}.`
+				: "Queue message body contains unsupported type.",
+			{ public: true, metadata: path ? { path } : undefined },
+		);
+	}
+}
+
+export class EventPayloadInvalid extends ActorError {
+	constructor(name: string, issues?: unknown[]) {
+		super(
+			"event",
+			"invalid_payload",
+			`Event payload failed validation for '${name}'.`,
+			{ public: true, metadata: { name, issues } },
+		);
+	}
+}
+
+export class QueuePayloadInvalid extends ActorError {
+	constructor(name: string, issues?: unknown[]) {
+		super(
+			"queue",
+			"invalid_payload",
+			`Queue payload failed validation for '${name}'.`,
+			{ public: true, metadata: { name, issues } },
+		);
+	}
+}
+
+export class QueueCompletionPayloadInvalid extends ActorError {
+	constructor(name: string, issues?: unknown[]) {
+		super(
+			"queue",
+			"invalid_completion_payload",
+			`Queue completion payload failed validation for '${name}'.`,
+			{ public: true, metadata: { name, issues } },
+		);
+	}
+}
+
+export class QueueAlreadyCompleted extends ActorError {
+	constructor() {
+		super("queue", "already_completed", "Queue message was already completed.", {
+			public: true,
+		});
+	}
+}
+
+export class QueuePreviousMessageNotCompleted extends ActorError {
+	constructor() {
+		super(
+			"queue",
+			"previous_message_not_completed",
+			"Previous completable queue message is not completed. Call `message.complete(...)` before receiving the next message.",
+			{ public: true },
+		);
+	}
+}
+
+export class QueueCompleteNotConfigured extends ActorError {
+	constructor(name: string) {
+		super(
+			"queue",
+			"complete_not_configured",
+			`Queue '${name}' does not support completion responses.`,
+			{
+				public: true,
+				metadata: { name },
+			},
+		);
+	}
+}
+
+export class ActorAborted extends ActorError {
+	constructor() {
+		super("actor", "aborted", "Actor aborted.", { public: true });
+	}
+}
+
 /**
  * Options for the UserError class.
  */
@@ -344,6 +449,28 @@ export class InvalidRequestHandlerResponse extends ActorError {
 	}
 }
 
+export class InvalidCanSubscribeResponse extends ActorError {
+	constructor() {
+		super(
+			"handler",
+			"invalid_can_subscribe_response",
+			"Event canSubscribe hook must return a boolean value.",
+		);
+		this.statusCode = 500;
+	}
+}
+
+export class InvalidCanPublishResponse extends ActorError {
+	constructor() {
+		super(
+			"handler",
+			"invalid_can_publish_response",
+			"Queue canPublish hook must return a boolean value.",
+		);
+		this.statusCode = 500;
+	}
+}
+
 // Manager-specific errors
 export class MissingActorHeader extends ActorError {
 	constructor() {

package/src/actor/instance/connection-manager.ts

@@ -30,6 +30,7 @@ import {
 } from "../contexts";
 import type { AnyDatabaseProvider } from "../database";
 import { CachedSerializer } from "../protocol/serde";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
 import { deadline } from "../utils";
 import { makeConnKey } from "./keys";
 import type { ActorInstance } from "./mod";
@@ -44,22 +45,25 @@ export class ConnectionManager<
 	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>;
-	#connections = new Map<ConnId, Conn<S, CP, CS, V, I, DB>>();
+	#actor: ActorInstance<S, CP, CS, V, I, DB, E, Q>;
+	#connections = new Map<ConnId, Conn<S, CP, CS, V, I, DB, E, Q>>();
+	#pendingDisconnectCount = 0;
 
 	/** Connections that have had their state changed and need to be persisted. */
 	#connsWithPersistChanged = new Set<ConnId>();
 
-	constructor(actor: ActorInstance<S, CP, CS, V, I, DB>) {
+	constructor(actor: ActorInstance<S, CP, CS, V, I, DB, E, Q>) {
 		this.#actor = actor;
 	}
 
-	get connections(): Map<ConnId, Conn<S, CP, CS, V, I, DB>> {
+	get connections(): Map<ConnId, Conn<S, CP, CS, V, I, DB, E, Q>> {
 		return this.#connections;
 	}
 
-	getConnForId(id: string): Conn<S, CP, CS, V, I, DB> | undefined {
+	getConnForId(id: string): Conn<S, CP, CS, V, I, DB, E, Q> | undefined {
 		return this.#connections.get(id);
 	}
 
@@ -67,11 +71,15 @@ export class ConnectionManager<
 		return this.#connsWithPersistChanged;
 	}
 
+	get pendingDisconnectCount(): number {
+		return this.#pendingDisconnectCount;
+	}
+
 	clearConnWithPersistChanged() {
 		this.#connsWithPersistChanged.clear();
 	}
 
-	markConnWithPersistChanged(conn: Conn<S, CP, CS, V, I, DB>) {
+	markConnWithPersistChanged(conn: Conn<S, CP, CS, V, I, DB, E, Q>) {
 		invariant(
 			conn.isHibernatable,
 			"cannot mark non-hibernatable conn for persist",
@@ -100,7 +108,7 @@ export class ConnectionManager<
 		requestHeaders: Record<string, string> | undefined,
 		isHibernatable: boolean,
 		isRestoringHibernatable: boolean,
-	): Promise<Conn<S, CP, CS, V, I, DB>> {
+	): Promise<Conn<S, CP, CS, V, I, DB, E, Q>> {
 		this.#actor.assertReady();
 
 		// TODO: Add back
@@ -118,7 +126,13 @@ export class ConnectionManager<
 		// Create new connection
 		if (this.#actor.config.onBeforeConnect) {
 			const ctx = new BeforeConnectContext(this.#actor, request);
-			await this.#actor.config.onBeforeConnect(ctx, params);
+			await this.#actor.runInTraceSpan(
+				"actor.onBeforeConnect",
+				{
+					"rivet.conn.type": driver.type,
+				},
+				() => this.#actor.config.onBeforeConnect!(ctx, params),
+			);
 		}
 
 		// Create connection state if enabled
@@ -163,7 +177,7 @@ export class ConnectionManager<
 		}
 
 		// Create connection instance
-		const conn = new Conn<S, CP, CS, V, I, DB>(this.#actor, connData);
+		const conn = new Conn<S, CP, CS, V, I, DB, E, Q>(this.#actor, connData);
 		conn[CONN_DRIVER_SYMBOL] = driver;
 
 		return conn;
@@ -177,7 +191,7 @@ export class ConnectionManager<
 	 * be messed up and cause race conditions that can drop WebSocket messages.
 	 * So all async work in prepareConn.
 	 */
-	connectConn(conn: Conn<S, CP, CS, V, I, DB>) {
+	connectConn(conn: Conn<S, CP, CS, V, I, DB, E, Q>) {
 		invariant(!this.#connections.has(conn.id), "conn already connected");
 
 		this.#connections.set(conn.id, conn);
@@ -230,7 +244,9 @@ export class ConnectionManager<
 		}
 	}
 
-	#reconnectHibernatableConn(driver: ConnDriver): Conn<S, CP, CS, V, I, DB> {
+	#reconnectHibernatableConn(
+		driver: ConnDriver,
+	): Conn<S, CP, CS, V, I, DB, E, Q> {
 		invariant(driver.hibernatable, "missing requestIdBuf");
 		const existingConn = this.findHibernatableConn(
 			driver.hibernatable.gatewayId,
@@ -265,7 +281,7 @@ export class ConnectionManager<
 		return existingConn;
 	}
 
-	#disconnectExistingDriver(conn: Conn<S, CP, CS, V, I, DB>) {
+	#disconnectExistingDriver(conn: Conn<S, CP, CS, V, I, DB, E, Q>) {
 		const driver = conn[CONN_DRIVER_SYMBOL];
 		if (driver?.disconnect) {
 			driver.disconnect(
@@ -281,7 +297,7 @@ export class ConnectionManager<
 	 *
 	 * This is called by `Conn.disconnect`. This should not call `Conn.disconnect.`
 	 */
-	async connDisconnected(conn: Conn<S, CP, CS, V, I, DB>) {
+	async connDisconnected(conn: Conn<S, CP, CS, V, I, DB, E, Q>) {
 		// Remove from tracking
 		this.#connections.delete(conn.id);
 
@@ -296,55 +312,108 @@ export class ConnectionManager<
 			this.#actor.eventManager.removeSubscription(eventName, conn, true);
 		}
 
-		this.#actor.resetSleepTimer();
-
 		this.#actor.inspector.emitter.emit("connectionsUpdated");
+		this.#pendingDisconnectCount += 1;
+
+		const attributes = {
+			"rivet.conn.id": conn.id,
+			"rivet.conn.type": conn[CONN_DRIVER_SYMBOL]?.type,
+			"rivet.conn.hibernatable": conn.isHibernatable,
+		};
+		const span = this.#actor.startTraceSpan(
+			"actor.onDisconnect",
+			attributes,
+		);
 
-		// Trigger disconnect
-		if (this.#actor.config.onDisconnect) {
-			try {
-				const result = this.#actor.config.onDisconnect(
-					this.#actor.actorContext,
-					conn,
+		try {
+			if (this.#actor.config.onDisconnect) {
+				const result = this.#actor.traces.withSpan(span, () =>
+					this.#actor.config.onDisconnect!(
+						this.#actor.actorContext,
+						conn,
+					),
+				);
+				this.#actor.emitTraceEvent(
+					"connection.disconnect",
+					attributes,
+					span,
 				);
 				if (result instanceof Promise) {
-					result.catch((error) => {
-						this.#actor.rLog.error({
-							msg: "error in `onDisconnect`",
-							error: stringifyError(error),
-						});
+					await result;
+				}
+				this.#actor.endTraceSpan(span, { code: "OK" });
+			} else {
+				this.#actor.emitTraceEvent(
+					"connection.disconnect",
+					attributes,
+					span,
+				);
+				this.#actor.endTraceSpan(span, { code: "OK" });
+			}
+		} catch (error) {
+			this.#actor.endTraceSpan(span, {
+				code: "ERROR",
+				message: stringifyError(error),
+			});
+			this.#actor.rLog.error({
+				msg: "error in `onDisconnect`",
+				error: stringifyError(error),
+			});
+		} finally {
+			// Remove from connsWithPersistChanged after onDisconnect to handle any
+			// state changes made during the disconnect callback. Disconnected connections
+			// are removed from KV storage via kvBatchDelete below, not through the
+			// normal persist save flow, so they should not trigger persist saves.
+			this.#connsWithPersistChanged.delete(conn.id);
+
+			// Remove from KV storage.
+			if (conn.isHibernatable) {
+				const key = makeConnKey(conn.id);
+				try {
+					await this.#actor.driver.kvBatchDelete(this.#actor.id, [key]);
+					this.#actor.rLog.debug({
+						msg: "removed connection from KV",
+						connId: conn.id,
+					});
+				} catch (err) {
+					this.#actor.rLog.error({
+						msg: "kvBatchDelete failed for conn",
+						err: stringifyError(err),
 					});
 				}
-			} catch (error) {
-				this.#actor.rLog.error({
-					msg: "error in `onDisconnect`",
-					error: stringifyError(error),
-				});
 			}
+
+			this.#pendingDisconnectCount = Math.max(
+				0,
+				this.#pendingDisconnectCount - 1,
+			);
+			this.#actor.resetSleepTimer();
 		}
+	}
 
-		// Remove from connsWithPersistChanged after onDisconnect to handle any
-		// state changes made during the disconnect callback. Disconnected connections
-		// are removed from KV storage via kvBatchDelete below, not through the
-		// normal persist save flow, so they should not trigger persist saves.
-		this.#connsWithPersistChanged.delete(conn.id);
+	async cleanupPersistedHibernatableConnections(
+		reason?: string,
+	): Promise<number> {
+		const staleConnections = Array.from(this.#connections.values()).filter(
+			(conn) =>
+				conn.isHibernatable &&
+				conn[CONN_DRIVER_SYMBOL] === undefined,
+		);
+		if (staleConnections.length === 0) {
+			return 0;
+		}
 
-		// Remove from KV storage
-		if (conn.isHibernatable) {
-			const key = makeConnKey(conn.id);
-			try {
-				await this.#actor.driver.kvBatchDelete(this.#actor.id, [key]);
-				this.#actor.rLog.debug({
-					msg: "removed connection from KV",
-					connId: conn.id,
-				});
-			} catch (err) {
-				this.#actor.rLog.error({
-					msg: "kvBatchDelete failed for conn",
-					err: stringifyError(err),
-				});
-			}
+		this.#actor.rLog.info({
+			msg: "cleaning up persisted hibernatable connections",
+			reason: reason ?? "unknown",
+			count: staleConnections.length,
+		});
+
+		for (const conn of staleConnections) {
+			await this.connDisconnected(conn);
 		}
+
+		return staleConnections.length;
 	}
 
 	/**
@@ -356,7 +425,7 @@ export class ConnectionManager<
 		request: Request | undefined,
 		requestPath: string | undefined,
 		requestHeaders: Record<string, string> | undefined,
-	): Promise<Conn<S, CP, CS, V, I, DB>> {
+	): Promise<Conn<S, CP, CS, V, I, DB, E, Q>> {
 		const conn = await this.prepareConn(
 			driver,
 			params,
@@ -378,7 +447,7 @@ export class ConnectionManager<
 	restoreConnections(connections: PersistedConn<CP, CS>[]) {
 		for (const connPersist of connections) {
 			// Create connection instance
-			const conn = new Conn<S, CP, CS, V, I, DB>(this.#actor, {
+			const conn = new Conn<S, CP, CS, V, I, DB, E, Q>(this.#actor, {
 				hibernatable: connPersist,
 			});
 			this.#connections.set(conn.id, conn);
@@ -404,7 +473,7 @@ export class ConnectionManager<
 	findHibernatableConn(
 		gatewayIdBuf: ArrayBuffer,
 		requestIdBuf: ArrayBuffer,
-	): Conn<S, CP, CS, V, I, DB> | undefined {
+	): Conn<S, CP, CS, V, I, DB, E, Q> | undefined {
 		return Array.from(this.#connections.values()).find((conn) => {
 			const connStateManager = conn[CONN_STATE_MANAGER_SYMBOL];
 			const h = connStateManager.hibernatableDataRaw;
@@ -421,18 +490,22 @@ export class ConnectionManager<
 		request: Request | undefined,
 	): Promise<CS | undefined> {
 		if ("createConnState" in this.#actor.config) {
+			const createConnState = this.#actor.config.createConnState;
 			const ctx = new CreateConnStateContext(this.#actor, request);
-			const dataOrPromise = this.#actor.config.createConnState(
-				ctx,
-				params,
+			return await this.#actor.runInTraceSpan(
+				"actor.createConnState",
+				undefined,
+				() => {
+					const dataOrPromise = createConnState!(ctx, params);
+					if (dataOrPromise instanceof Promise) {
+						return deadline(
+							dataOrPromise,
+							this.#actor.config.options.createConnStateTimeout,
+						);
+					}
+					return dataOrPromise;
+				},
 			);
-			if (dataOrPromise instanceof Promise) {
-				return await deadline(
-					dataOrPromise,
-					this.#actor.config.options.createConnStateTimeout,
-				);
-			}
-			return dataOrPromise;
 		} else if ("connState" in this.#actor.config) {
 			return structuredClone(this.#actor.config.connState);
 		}
@@ -442,30 +515,60 @@ export class ConnectionManager<
 		);
 	}
 
-	#callOnConnect(conn: Conn<S, CP, CS, V, I, DB>) {
-		if (this.#actor.config.onConnect) {
-			try {
+	#callOnConnect(conn: Conn<S, CP, CS, V, I, DB, E, Q>) {
+		const attributes = {
+			"rivet.conn.id": conn.id,
+			"rivet.conn.type": conn[CONN_DRIVER_SYMBOL]?.type,
+			"rivet.conn.hibernatable": conn.isHibernatable,
+		};
+		const span = this.#actor.startTraceSpan("actor.onConnect", attributes);
+
+		try {
+			if (this.#actor.config.onConnect) {
 				const ctx = new ConnectContext(this.#actor, conn);
-				const result = this.#actor.config.onConnect(ctx, conn);
+				const result = this.#actor.traces.withSpan(span, () =>
+					this.#actor.config.onConnect!(ctx, conn),
+				);
+				this.#actor.emitTraceEvent(
+					"connection.connect",
+					attributes,
+					span,
+				);
 				if (result instanceof Promise) {
 					deadline(
 						result,
 						this.#actor.config.options.onConnectTimeout,
-					).catch((error) => {
-						this.#actor.rLog.error({
-							msg: "error in `onConnect`, closing socket",
-							error,
+					)
+						.then(() => {
+							this.#actor.endTraceSpan(span, { code: "OK" });
+						})
+						.catch((error) => {
+							this.#actor.endTraceSpan(span, {
+								code: "ERROR",
+								message: stringifyError(error),
+							});
+							this.#actor.rLog.error({
+								msg: "error in `onConnect`, closing socket",
+								error,
+							});
+							conn?.disconnect("`onConnect` failed");
 						});
-						conn?.disconnect("`onConnect` failed");
-					});
+					return;
 				}
-			} catch (error) {
-				this.#actor.rLog.error({
-					msg: "error in `onConnect`",
-					error: stringifyError(error),
-				});
-				conn?.disconnect("`onConnect` failed");
 			}
+
+			this.#actor.emitTraceEvent("connection.connect", attributes, span);
+			this.#actor.endTraceSpan(span, { code: "OK" });
+		} catch (error) {
+			this.#actor.endTraceSpan(span, {
+				code: "ERROR",
+				message: stringifyError(error),
+			});
+			this.#actor.rLog.error({
+				msg: "error in `onConnect`",
+				error: stringifyError(error),
+			});
+			conn?.disconnect("`onConnect` failed");
 		}
 	}
 }

package/src/actor/instance/event-manager.ts

@@ -18,17 +18,30 @@ import {
 import type { AnyDatabaseProvider } from "../database";
 import * as errors from "../errors";
 import { CachedSerializer } from "../protocol/serde";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
 import type { ActorInstance } from "./mod";
 
 /**
  * Manages event subscriptions and broadcasting for actor instances.
  * Handles subscription tracking and efficient message distribution to connected clients.
  */
-export class EventManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
-	#actor: ActorInstance<S, CP, CS, V, I, DB>;
-	#subscriptionIndex = new Map<string, Set<Conn<S, CP, CS, V, I, DB>>>();
+export class EventManager<
+	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>;
+	#subscriptionIndex = new Map<
+		string,
+		Set<Conn<S, CP, CS, V, I, DB, E, Q>>
+	>();
 
-	constructor(actor: ActorInstance<S, CP, CS, V, I, DB>) {
+	constructor(actor: ActorInstance<S, CP, CS, V, I, DB, E, Q>) {
 		this.#actor = actor;
 	}
 
@@ -43,7 +56,7 @@ export class EventManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 */
 	addSubscription(
 		eventName: string,
-		connection: Conn<S, CP, CS, V, I, DB>,
+		connection: Conn<S, CP, CS, V, I, DB, E, Q>,
 		fromPersist: boolean,
 	) {
 		// Check if already subscribed
@@ -94,7 +107,7 @@ export class EventManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 */
 	removeSubscription(
 		eventName: string,
-		connection: Conn<S, CP, CS, V, I, DB>,
+		connection: Conn<S, CP, CS, V, I, DB, E, Q>,
 		fromRemoveConn: boolean,
 	) {
 		// Check if subscription exists
@@ -154,13 +167,6 @@ export class EventManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	broadcast<Args extends Array<unknown>>(name: string, ...args: Args) {
 		this.#actor.assertReady();
 
-		// Emit to inspector
-		this.#actor.inspector.emitter.emit("eventFired", {
-			type: "broadcast",
-			eventName: name,
-			args,
-		});
-
 		// Get subscribers for this event
 		const subscribers = this.#subscriptionIndex.get(name);
 		if (!subscribers || subscribers.size === 0) {
@@ -171,6 +177,11 @@ export class EventManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			return;
 		}
 
+		this.#actor.emitTraceEvent("message.broadcast", {
+			"rivet.event.name": name,
+			"rivet.broadcast.subscribers": subscribers.size,
+		});
+
 		// Create serialized message
 		const eventData = { name, args };
 		const toClientSerializer = new CachedSerializer(
@@ -243,7 +254,7 @@ export class EventManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 */
 	getSubscribers(
 		eventName: string,
-	): Set<Conn<S, CP, CS, V, I, DB>> | undefined {
+	): Set<Conn<S, CP, CS, V, I, DB, E, Q>> | undefined {
 		return this.#subscriptionIndex.get(eventName);
 	}
 
@@ -266,7 +277,7 @@ export class EventManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	 *
 	 * @param connection - The connection to clear subscriptions for
 	 */
-	clearConnectionSubscriptions(connection: Conn<S, CP, CS, V, I, DB>) {
+	clearConnectionSubscriptions(connection: Conn<S, CP, CS, V, I, DB, E, Q>) {
 		for (const eventName of [...connection.subscriptions.values()]) {
 			this.removeSubscription(eventName, connection, true);
 		}