rivetkit 2.0.42 → 2.1.0-rc.1

package/src/actor/instance/queue.ts

@@ -0,0 +1,287 @@
+import * as errors from "../errors";
+import type { AnyDatabaseProvider } from "../database";
+import type {
+	EventSchemaConfig,
+	InferQueueCompleteMap,
+	InferSchemaMap,
+	QueueSchemaConfig,
+} from "../schema";
+import { joinAbortSignals } from "../utils";
+import type { QueueManager, QueueMessage } from "./queue-manager";
+
+export type QueueMessageOf<Name extends string, Body> = Omit<
+	QueueMessage,
+	"name" | "body"
+> & {
+	name: Name;
+	body: Body;
+};
+
+export type QueueName<TQueues extends QueueSchemaConfig> = keyof TQueues & string;
+export type QueueFilterName<TQueues extends QueueSchemaConfig> =
+	keyof TQueues extends never ? string : QueueName<TQueues>;
+
+type QueueMessageForName<
+	TQueues extends QueueSchemaConfig,
+	TName extends QueueFilterName<TQueues>,
+> = keyof TQueues extends never
+	? QueueMessage
+	: TName extends QueueName<TQueues>
+		? QueueMessageOf<TName, InferSchemaMap<TQueues>[TName]>
+		: never;
+
+type QueueCompleteArgs<T> = undefined extends T
+	? [response?: T]
+	: [response: T];
+
+type QueueCompleteArgsForName<
+	TQueues extends QueueSchemaConfig,
+	TName extends QueueFilterName<TQueues>,
+> = keyof TQueues extends never
+	? [response?: unknown]
+	: TName extends QueueName<TQueues>
+		? [InferQueueCompleteMap<TQueues>[TName]] extends [never]
+			? [response?: unknown]
+			: QueueCompleteArgs<InferQueueCompleteMap<TQueues>[TName]>
+		: [response?: unknown];
+
+type QueueCompletableMessageForName<
+	TQueues extends QueueSchemaConfig,
+	TName extends QueueFilterName<TQueues>,
+> = QueueMessageForName<TQueues, TName> & {
+	complete(
+		...args: QueueCompleteArgsForName<TQueues, TName>
+	): Promise<void>;
+};
+
+export type QueueResultMessageForName<
+	TQueues extends QueueSchemaConfig,
+	TName extends QueueFilterName<TQueues>,
+	TCompletable extends boolean,
+> = TCompletable extends true
+	? QueueCompletableMessageForName<TQueues, TName>
+	: QueueMessageForName<TQueues, TName>;
+
+/** Options for receiving queue messages. */
+export interface QueueNextOptions<
+	TName extends string = string,
+	TCompletable extends boolean = boolean,
+> {
+	/** Queue names to receive from. If omitted, reads from all queue names. */
+	names?: readonly TName[];
+	/** Maximum number of messages to receive. Defaults to 1. */
+	count?: number;
+	/** Timeout in milliseconds. Omit to wait indefinitely. */
+	timeout?: number;
+	/** Optional abort signal for this receive call. */
+	signal?: AbortSignal;
+	/** Whether to return completable messages. */
+	completable?: TCompletable;
+}
+
+/** Options for non-blocking queue reads. */
+export interface QueueTryNextOptions<
+	TName extends string = string,
+	TCompletable extends boolean = boolean,
+> {
+	/** Queue names to receive from. If omitted, reads from all queue names. */
+	names?: readonly TName[];
+	/** Maximum number of messages to receive. Defaults to 1. */
+	count?: number;
+	/** Whether to return completable messages. */
+	completable?: TCompletable;
+}
+
+/** Options for queue async iteration. */
+export interface QueueIterOptions<
+	TName extends string = string,
+	TCompletable extends boolean = boolean,
+> {
+	/** Queue names to receive from. If omitted, reads from all queue names. */
+	names?: readonly TName[];
+	/** Optional abort signal for this iterator. */
+	signal?: AbortSignal;
+	/** Whether to return completable messages. */
+	completable?: TCompletable;
+}
+
+/** User-facing queue interface exposed on ActorContext. */
+export class ActorQueue<
+	S,
+	CP,
+	CS,
+	V,
+	I,
+	DB extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+> {
+	#queueManager: QueueManager<S, CP, CS, V, I, DB, TEvents, TQueues>;
+	#abortSignal: AbortSignal;
+	#pendingCompletableMessageIds = new Set<string>();
+
+	constructor(
+		queueManager: QueueManager<S, CP, CS, V, I, DB, TEvents, TQueues>,
+		abortSignal: AbortSignal,
+	) {
+		this.#queueManager = queueManager;
+		this.#abortSignal = abortSignal;
+	}
+
+	async next<
+		const TName extends QueueFilterName<TQueues>,
+		const TCompletable extends boolean = false,
+	>(
+		opts?: QueueNextOptions<TName, TCompletable>,
+	): Promise<Array<QueueResultMessageForName<TQueues, TName, TCompletable>>> {
+		const resolvedOpts = (opts ?? {}) as QueueNextOptions<
+			TName,
+			TCompletable
+		>;
+		const completable = resolvedOpts.completable === true;
+
+		if (this.#pendingCompletableMessageIds.size > 0) {
+			throw new errors.QueuePreviousMessageNotCompleted();
+		}
+
+		const names = this.#normalizeNames(resolvedOpts.names);
+		const count = Math.max(1, resolvedOpts.count ?? 1);
+		const { signal, cleanup } = joinAbortSignals(
+			this.#abortSignal,
+			resolvedOpts.signal,
+		);
+		const messages = await this.#queueManager
+			.receive(
+				names,
+				count,
+				resolvedOpts.timeout,
+				signal,
+				completable,
+			)
+			.finally(cleanup);
+		if (!completable) {
+			return messages as Array<
+				QueueResultMessageForName<TQueues, TName, TCompletable>
+			>;
+		}
+		return messages.map((message) => this.#makeCompletableMessage(message)) as unknown as Array<
+			QueueResultMessageForName<TQueues, TName, TCompletable>
+		>;
+	}
+
+	async tryNext<
+		const TName extends QueueFilterName<TQueues>,
+		const TCompletable extends boolean = false,
+	>(
+		opts?: QueueTryNextOptions<TName, TCompletable>,
+	): Promise<Array<QueueResultMessageForName<TQueues, TName, TCompletable>>> {
+		const resolvedOpts = (opts ?? {}) as QueueTryNextOptions<
+			TName,
+			TCompletable
+		>;
+		if (resolvedOpts.completable === true) {
+			return (await this.next<TName, true>({
+				names: resolvedOpts.names,
+				count: resolvedOpts.count,
+				timeout: 0,
+				completable: true,
+			})) as Array<QueueResultMessageForName<TQueues, TName, TCompletable>>;
+		}
+		return (await this.next<TName, false>({
+			names: resolvedOpts.names,
+			count: resolvedOpts.count,
+			timeout: 0,
+		})) as Array<QueueResultMessageForName<TQueues, TName, TCompletable>>;
+	}
+
+	async *iter<
+		const TName extends QueueFilterName<TQueues>,
+		const TCompletable extends boolean = false,
+	>(
+		opts?: QueueIterOptions<TName, TCompletable>,
+	): AsyncIterableIterator<
+		QueueResultMessageForName<TQueues, TName, TCompletable>
+	> {
+		const resolvedOpts = (opts ?? {}) as QueueIterOptions<
+			TName,
+			TCompletable
+		>;
+		while (!this.#abortSignal.aborted) {
+			try {
+				const messages = resolvedOpts.completable === true
+					? await this.next<TName, true>({
+							names: resolvedOpts.names,
+							count: 1,
+							signal: resolvedOpts.signal,
+							completable: true,
+						})
+					: await this.next<TName, false>({
+							names: resolvedOpts.names,
+							count: 1,
+							signal: resolvedOpts.signal,
+						});
+				if (messages.length === 0) {
+					continue;
+				}
+				yield messages[0] as QueueResultMessageForName<
+					TQueues,
+					TName,
+					TCompletable
+				>;
+			} catch (error) {
+				if (error instanceof errors.ActorAborted) {
+					return;
+				}
+				throw error;
+			}
+		}
+	}
+
+	/** Sends a message to the specified queue. */
+	send<K extends keyof TQueues & string>(
+		name: K,
+		body: InferSchemaMap<TQueues>[K],
+	): Promise<QueueMessage>;
+	send(
+		name: keyof TQueues extends never ? string : never,
+		body: unknown,
+	): Promise<QueueMessage>;
+	async send(name: string, body: unknown): Promise<QueueMessage> {
+		return await this.#queueManager.enqueue(name, body);
+	}
+
+	#normalizeNames(names: readonly string[] | undefined): string[] | undefined {
+		if (!names || names.length === 0) {
+			return undefined;
+		}
+		return [...new Set(names)];
+	}
+
+	#makeCompletableMessage(
+		message: QueueMessage,
+	): QueueMessage & {
+		complete: (response?: unknown) => Promise<void>;
+	} {
+		const messageId = message.id.toString();
+		this.#pendingCompletableMessageIds.add(messageId);
+
+		let completed = false;
+		const completableMessage = {
+			...message,
+			complete: async (response?: unknown) => {
+				if (completed) {
+					throw new errors.QueueAlreadyCompleted();
+				}
+				completed = true;
+				try {
+					await this.#queueManager.completeMessage(message, response);
+					this.#pendingCompletableMessageIds.delete(messageId);
+				} catch (error) {
+					completed = false;
+					throw error;
+				}
+			},
+		};
+		return completableMessage;
+	}
+}

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

@@ -6,6 +6,7 @@ import {
 } from "@/utils";
 import type { AnyDatabaseProvider } from "../database";
 import type { ActorDriver } from "../driver";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
 import type { ActorInstance } from "./mod";
 import type { PersistedScheduleEvent } from "./persisted";
 
@@ -13,15 +14,24 @@ import type { PersistedScheduleEvent } from "./persisted";
  * Manages scheduled events and alarms for actor instances.
  * Handles event scheduling, alarm triggers, and automatic event execution.
  */
-export class ScheduleManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
-	#actor: ActorInstance<S, CP, CS, V, I, DB>;
+export class ScheduleManager<
+	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>;
 	#actorDriver: ActorDriver;
 	#alarmWriteQueue = new SinglePromiseQueue();
 	#config: any; // ActorConfig type
 	#persist: any; // Reference to PersistedActor
 
 	constructor(
-		actor: ActorInstance<S, CP, CS, V, I, DB>,
+		actor: ActorInstance<S, CP, CS, V, I, DB, E, Q>,
 		actorDriver: ActorDriver,
 		config: any,
 	) {
@@ -59,6 +69,12 @@ export class ScheduleManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			args: bufferToArrayBuffer(cbor.encode(args)),
 		};
 
+		this.#actor.emitTraceEvent("schedule.created", {
+			"rivet.schedule.event_id": newEvent.eventId,
+			"rivet.schedule.action": newEvent.action,
+			"rivet.schedule.timestamp_ms": newEvent.timestamp,
+		});
+
 		await this.#scheduleEventInner(newEvent);
 	}
 
@@ -252,7 +268,25 @@ export class ScheduleManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 
 	async #executeDueEvents(events: PersistedScheduleEvent[]): Promise<void> {
 		for (const event of events) {
+			const span = this.#actor.startTraceSpan(
+				`actor.action.${event.action}`,
+				{
+					"rivet.action.name": event.action,
+					"rivet.action.scheduled": true,
+					"rivet.schedule.event_id": event.eventId,
+					"rivet.schedule.timestamp_ms": event.timestamp,
+				},
+			);
 			try {
+				this.#actor.emitTraceEvent(
+					"schedule.triggered",
+					{
+						"rivet.schedule.event_id": event.eventId,
+						"rivet.schedule.action": event.action,
+						"rivet.schedule.timestamp_ms": event.timestamp,
+					},
+					span,
+				);
 				this.#actor.log.info({
 					msg: "executing scheduled event",
 					eventId: event.eventId,
@@ -281,10 +315,8 @@ export class ScheduleManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 					? cbor.decode(new Uint8Array(event.args))
 					: [];
 
-				const result = fn.call(
-					undefined,
-					this.#actor.actorContext,
-					...args,
+				const result = this.#actor.traces.withSpan(span, () =>
+					fn.call(undefined, this.#actor.actorContext, ...args),
 				);
 
 				// Handle async actions
@@ -292,12 +324,22 @@ export class ScheduleManager<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 					await result;
 				}
 
+				this.#actor.endTraceSpan(span, { code: "OK" });
 				this.#actor.log.debug({
 					msg: "scheduled event completed",
 					eventId: event.eventId,
 					action: event.action,
 				});
 			} catch (error) {
+				this.#actor.traces.setAttributes(span, {
+					"error.message": stringifyError(error),
+					"error.type":
+						error instanceof Error ? error.name : typeof error,
+				});
+				this.#actor.endTraceSpan(span, {
+					code: "ERROR",
+					message: stringifyError(error),
+				});
 				this.#actor.log.error({
 					msg: "error executing scheduled event",
 					error: stringifyError(error),

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

@@ -7,10 +7,12 @@ import {
 	CONN_VERSIONED,
 } from "@/schemas/actor-persist/versioned";
 import { promiseWithResolvers, SinglePromiseQueue } from "@/utils";
+import { loggerWithoutContext } from "@/actor/log";
 import { type AnyConn, CONN_STATE_MANAGER_SYMBOL } from "../conn/mod";
 import { convertConnToBarePersistedConn } from "../conn/persisted";
 import type { ActorDriver } from "../driver";
 import * as errors from "../errors";
+import type { EventSchemaConfig, QueueSchemaConfig } from "../schema";
 import { isConnStatePath, isStatePath } from "../utils";
 import { KEYS, makeConnKey } from "./keys";
 import type { ActorInstance } from "./mod";
@@ -35,8 +37,15 @@ export interface SaveStateOptions {
  * Manages actor state persistence, proxying, and synchronization.
  * Handles automatic state change detection and throttled persistence to KV storage.
  */
-export class StateManager<S, CP, CS, I> {
-	#actor: ActorInstance<S, CP, CS, any, I, any>;
+export class StateManager<
+	S,
+	CP,
+	CS,
+	I,
+	E extends EventSchemaConfig = Record<never, never>,
+	Q extends QueueSchemaConfig = Record<never, never>,
+> {
+	#actor: ActorInstance<S, CP, CS, any, I, any, E, Q>;
 	#actorDriver: ActorDriver;
 
 	// State tracking
@@ -57,7 +66,7 @@ export class StateManager<S, CP, CS, I> {
 	#stateSaveInterval: number;
 
 	constructor(
-		actor: ActorInstance<S, CP, CS, any, I, any>,
+		actor: ActorInstance<S, CP, CS, any, I, any, E, Q>,
 		actorDriver: ActorDriver,
 		config: any,
 	) {
@@ -108,9 +117,14 @@ export class StateManager<S, CP, CS, I> {
 				this.#actor.rLog.info({ msg: "actor state initializing" });
 
 				if ("createState" in this.#config) {
-					stateData = await this.#config.createState(
-						this.#actor.actorContext,
-						persistData.input!,
+					stateData = await this.#actor.runInTraceSpan(
+						"actor.createState",
+						undefined,
+						() =>
+							this.#config.createState!(
+								this.#actor.actorContext,
+								persistData.input!,
+							),
 					);
 				} else if ("state" in this.#config) {
 					stateData = structuredClone(this.#config.state);
@@ -205,7 +219,7 @@ export class StateManager<S, CP, CS, I> {
 			} else {
 				// Create promise for waiting
 				if (!this.#onPersistSavedPromise) {
-					this.#onPersistSavedPromise = promiseWithResolvers();
+					this.#onPersistSavedPromise = promiseWithResolvers((reason) => loggerWithoutContext().warn({ msg: "unhandled persist saved promise rejection", reason }));
 				}
 
 				// Save throttled
@@ -349,13 +363,23 @@ export class StateManager<S, CP, CS, I> {
 			this.#actor.isReady() &&
 			!this.#isInOnStateChange
 		) {
+			const span = this.#actor.startTraceSpan("actor.onStateChange", {
+				"rivet.state.path": path,
+			});
 			try {
 				this.#isInOnStateChange = true;
-				this.#config.onStateChange(
-					this.#actor.actorContext,
-					this.#persistRaw.state,
+				this.#actor.traces.withSpan(span, () =>
+					this.#config.onStateChange!(
+						this.#actor.actorContext,
+						this.#persistRaw.state,
+					),
 				);
+				this.#actor.endTraceSpan(span, { code: "OK" });
 			} catch (error) {
+				this.#actor.endTraceSpan(span, {
+					code: "ERROR",
+					message: stringifyError(error),
+				});
 				this.#actor.rLog.error({
 					msg: "error in `_onStateChange`",
 					error: stringifyError(error),
@@ -440,7 +464,7 @@ export class StateManager<S, CP, CS, I> {
 						this.#actor.rLog.info({
 							msg: "persisting connection",
 							connId,
-							gatewayId: idToStr(hibernatableDataRaw.requestId),
+							gatewayId: idToStr(hibernatableDataRaw.gatewayId),
 							requestId: idToStr(hibernatableDataRaw.requestId),
 							serverMessageIndex:
 								hibernatableDataRaw.serverMessageIndex,

package/src/actor/instance/traces-driver.ts

@@ -0,0 +1,128 @@
+import type { TracesDriver } from "@rivetkit/traces";
+import type { ActorDriver } from "../driver";
+import { tracesStoragePrefix } from "./keys";
+
+function concatPrefix(prefix: Uint8Array, key: Uint8Array): Uint8Array {
+	const merged = new Uint8Array(prefix.length + key.length);
+	merged.set(prefix, 0);
+	merged.set(key, prefix.length);
+	return merged;
+}
+
+function stripPrefix(prefix: Uint8Array, key: Uint8Array): Uint8Array {
+	return key.slice(prefix.length);
+}
+
+function compareBytes(a: Uint8Array, b: Uint8Array): number {
+	const len = Math.min(a.length, b.length);
+	for (let i = 0; i < len; i++) {
+		if (a[i] !== b[i]) {
+			return a[i] - b[i];
+		}
+	}
+	return a.length - b.length;
+}
+
+export class ActorTracesDriver implements TracesDriver {
+	#driver: ActorDriver;
+	#actorId: string;
+	#prefix: Uint8Array;
+
+	constructor(driver: ActorDriver, actorId: string) {
+		this.#driver = driver;
+		this.#actorId = actorId;
+		this.#prefix = tracesStoragePrefix();
+	}
+
+	async get(key: Uint8Array): Promise<Uint8Array | null> {
+		const [value] = await this.#driver.kvBatchGet(this.#actorId, [
+			concatPrefix(this.#prefix, key),
+		]);
+		return value ?? null;
+	}
+
+	async set(key: Uint8Array, value: Uint8Array): Promise<void> {
+		await this.#driver.kvBatchPut(this.#actorId, [
+			[concatPrefix(this.#prefix, key), value],
+		]);
+	}
+
+	async delete(key: Uint8Array): Promise<void> {
+		await this.#driver.kvBatchDelete(this.#actorId, [
+			concatPrefix(this.#prefix, key),
+		]);
+	}
+
+	async deletePrefix(prefix: Uint8Array): Promise<void> {
+		const fullPrefix = concatPrefix(this.#prefix, prefix);
+		const entries = await this.#driver.kvListPrefix(
+			this.#actorId,
+			fullPrefix,
+		);
+		if (entries.length === 0) {
+			return;
+		}
+		await this.#driver.kvBatchDelete(
+			this.#actorId,
+			entries.map(([key]) => key),
+		);
+	}
+
+	async list(
+		prefix: Uint8Array,
+	): Promise<Array<{ key: Uint8Array; value: Uint8Array }>> {
+		const fullPrefix = concatPrefix(this.#prefix, prefix);
+		const entries = await this.#driver.kvListPrefix(
+			this.#actorId,
+			fullPrefix,
+		);
+		return entries.map(([key, value]) => ({
+			key: stripPrefix(this.#prefix, key),
+			value,
+		}));
+	}
+
+	async listRange(
+		start: Uint8Array,
+		end: Uint8Array,
+		options?: { reverse?: boolean; limit?: number },
+	): Promise<Array<{ key: Uint8Array; value: Uint8Array }>> {
+		const fullStart = concatPrefix(this.#prefix, start);
+		const fullEnd = concatPrefix(this.#prefix, end);
+		const entries = await this.#driver.kvListPrefix(
+			this.#actorId,
+			this.#prefix,
+		);
+		const filtered = entries
+			.filter(([key]) => {
+				return (
+					compareBytes(key, fullStart) >= 0 &&
+					compareBytes(key, fullEnd) < 0
+				);
+			})
+			.sort(([keyA], [keyB]) => compareBytes(keyA, keyB));
+		if (options?.reverse) {
+			filtered.reverse();
+		}
+		const limited = options?.limit
+			? filtered.slice(0, options.limit)
+			: filtered;
+		return limited.map(([key, value]) => ({
+			key: stripPrefix(this.#prefix, key),
+			value,
+		}));
+	}
+
+	async batch(writes: Array<{ key: Uint8Array; value: Uint8Array }>): Promise<void> {
+		if (writes.length === 0) {
+			return;
+		}
+		await this.#driver.kvBatchPut(
+			this.#actorId,
+			writes.map(({ key, value }) => [
+				concatPrefix(this.#prefix, key),
+				value,
+			]),
+		);
+	}
+}

package/src/actor/mod.ts

@@ -7,6 +7,8 @@ import {
 } from "./config";
 import type { AnyDatabaseProvider } from "./database";
 import { ActorDefinition } from "./definition";
+import { event as schemaEvent, queue as schemaQueue } from "./schema";
+import type { EventSchemaConfig, QueueSchemaConfig } from "./schema";
 
 export function actor<
 	TState,
@@ -15,13 +17,26 @@ export function actor<
 	TVars,
 	TInput,
 	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
 	TActions extends Actions<
 		TState,
 		TConnParams,
 		TConnState,
 		TVars,
 		TInput,
-		TDatabase
+		TDatabase,
+		TEvents,
+		TQueues
+	> = Actions<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
 	>,
 >(
 	input: ActorConfigInput<
@@ -31,6 +46,8 @@ export function actor<
 		TVars,
 		TInput,
 		TDatabase,
+		TEvents,
+		TQueues,
 		TActions
 	>,
 ): ActorDefinition<
@@ -40,6 +57,8 @@ export function actor<
 	TVars,
 	TInput,
 	TDatabase,
+	TEvents,
+	TQueues,
 	TActions
 > {
 	const config = ActorConfigSchema.parse(input) as ActorConfig<
@@ -48,7 +67,9 @@ export function actor<
 		TConnState,
 		TVars,
 		TInput,
-		TDatabase
+		TDatabase,
+		TEvents,
+		TQueues
 	>;
 	return new ActorDefinition(config);
 }
@@ -84,3 +105,6 @@ export {
 	createActorRouter,
 } from "./router";
 export { routeWebSocket } from "./router-websocket-endpoints";
+export type { Type } from "./schema";
+export const event = schemaEvent;
+export const queue = schemaQueue;