rivetkit 2.0.2 → 2.0.4-rc.1

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

@@ -0,0 +1,314 @@
+import * as cbor from "cbor-x";
+import type * as protocol from "@/schemas/client-protocol/mod";
+import {
+	CURRENT_VERSION as CLIENT_PROTOCOL_CURRENT_VERSION,
+	TO_CLIENT_VERSIONED,
+} from "@/schemas/client-protocol/versioned";
+import {
+	type ToClient as ToClientJson,
+	ToClientSchema,
+} from "@/schemas/client-protocol-zod/mod";
+import { bufferToArrayBuffer } from "@/utils";
+import {
+	CONN_SEND_MESSAGE_SYMBOL,
+	CONN_SPEAKS_RIVETKIT_SYMBOL,
+	CONN_STATE_MANAGER_SYMBOL,
+	type Conn,
+} from "../conn/mod";
+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,
+	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, E, Q>) {
+		this.#actor = actor;
+	}
+
+	// MARK: - Public API
+
+	/**
+	 * Adds a subscription for a connection to an event.
+	 *
+	 * @param eventName - The name of the event to subscribe to
+	 * @param connection - The connection subscribing to the event
+	 * @param fromPersist - Whether this subscription is being restored from persistence
+	 */
+	addSubscription(
+		eventName: string,
+		connection: Conn<S, CP, CS, V, I, DB, E, Q>,
+		fromPersist: boolean,
+	) {
+		// Check if already subscribed
+		if (connection.subscriptions.has(eventName)) {
+			this.#actor.rLog.debug({
+				msg: "connection already has subscription",
+				eventName,
+				connId: connection.id,
+			});
+			return;
+		}
+
+		// Update connection's subscription list
+		connection.subscriptions.add(eventName);
+
+		// Update subscription index
+		let subscribers = this.#subscriptionIndex.get(eventName);
+		if (!subscribers) {
+			subscribers = new Set();
+			this.#subscriptionIndex.set(eventName, subscribers);
+		}
+		subscribers.add(connection);
+
+		// Persist subscription if not restoring from persistence
+		if (!fromPersist) {
+			connection[CONN_STATE_MANAGER_SYMBOL].addSubscription({
+				eventName,
+			});
+
+			// Save state immediately
+			this.#actor.stateManager.saveState({ immediate: true });
+		}
+
+		this.#actor.rLog.debug({
+			msg: "subscription added",
+			eventName,
+			connId: connection.id,
+			totalSubscribers: subscribers.size,
+		});
+	}
+
+	/**
+	 * Removes a subscription for a connection from an event.
+	 *
+	 * @param eventName - The name of the event to unsubscribe from
+	 * @param connection - The connection unsubscribing from the event
+	 * @param fromRemoveConn - Whether this is being called as part of connection removal
+	 */
+	removeSubscription(
+		eventName: string,
+		connection: Conn<S, CP, CS, V, I, DB, E, Q>,
+		fromRemoveConn: boolean,
+	) {
+		// Check if subscription exists
+		if (!connection.subscriptions.has(eventName)) {
+			this.#actor.rLog.warn({
+				msg: "connection does not have subscription",
+				eventName,
+				connId: connection.id,
+			});
+			return;
+		}
+
+		// Remove from connection's subscription list
+		connection.subscriptions.delete(eventName);
+
+		// Update subscription index
+		const subscribers = this.#subscriptionIndex.get(eventName);
+		if (subscribers) {
+			subscribers.delete(connection);
+			if (subscribers.size === 0) {
+				this.#subscriptionIndex.delete(eventName);
+			}
+		}
+
+		// Update persistence if not part of connection removal
+		if (!fromRemoveConn) {
+			// Remove from persisted subscriptions
+			const removed = connection[
+				CONN_STATE_MANAGER_SYMBOL
+			].removeSubscription({ eventName });
+			if (!removed) {
+				this.#actor.rLog.warn({
+					msg: "subscription does not exist in persist",
+					eventName,
+					connId: connection.id,
+				});
+			}
+
+			// Save state immediately
+			this.#actor.stateManager.saveState({ immediate: true });
+		}
+
+		this.#actor.rLog.debug({
+			msg: "subscription removed",
+			eventName,
+			connId: connection.id,
+			remainingSubscribers: subscribers?.size || 0,
+		});
+	}
+
+	/**
+	 * Broadcasts an event to all subscribed connections.
+	 *
+	 * @param name - The name of the event to broadcast
+	 * @param args - The arguments to send with the event
+	 */
+	broadcast<Args extends Array<unknown>>(name: string, ...args: Args) {
+		this.#actor.assertReady();
+
+		// Get subscribers for this event
+		const subscribers = this.#subscriptionIndex.get(name);
+		if (!subscribers || subscribers.size === 0) {
+			this.#actor.rLog.debug({
+				msg: "no subscribers for event",
+				eventName: name,
+			});
+			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(
+			eventData,
+			TO_CLIENT_VERSIONED,
+			CLIENT_PROTOCOL_CURRENT_VERSION,
+			ToClientSchema,
+			// JSON: args is the raw value (array of arguments)
+			(value): ToClientJson => ({
+				body: {
+					tag: "Event" as const,
+					val: {
+						name: value.name,
+						args: value.args,
+					},
+				},
+			}),
+			// BARE/CBOR: args needs to be CBOR-encoded to ArrayBuffer
+			(value): protocol.ToClient => ({
+				body: {
+					tag: "Event" as const,
+					val: {
+						name: value.name,
+						args: bufferToArrayBuffer(cbor.encode(value.args)),
+					},
+				},
+			}),
+		);
+
+		// Send to all subscribers
+		let sentCount = 0;
+		for (const connection of subscribers) {
+			if (connection[CONN_SPEAKS_RIVETKIT_SYMBOL]) {
+				try {
+					connection[CONN_SEND_MESSAGE_SYMBOL](toClientSerializer);
+					sentCount++;
+				} catch (error) {
+					// Propagate message size errors to the call site so developers
+					// can handle them
+					if (error instanceof errors.OutgoingMessageTooLong) {
+						throw error;
+					}
+					// Log other errors (e.g., closed connections) and continue
+					this.#actor.rLog.error({
+						msg: "failed to send event to connection",
+						eventName: name,
+						connId: connection.id,
+						error:
+							error instanceof Error
+								? error.message
+								: String(error),
+					});
+				}
+			}
+		}
+
+		this.#actor.rLog.debug({
+			msg: "event broadcasted",
+			eventName: name,
+			subscriberCount: subscribers.size,
+			sentCount,
+		});
+	}
+
+	/**
+	 * Gets all subscribers for a specific event.
+	 *
+	 * @param eventName - The name of the event
+	 * @returns Set of connections subscribed to the event, or undefined if no subscribers
+	 */
+	getSubscribers(
+		eventName: string,
+	): Set<Conn<S, CP, CS, V, I, DB, E, Q>> | undefined {
+		return this.#subscriptionIndex.get(eventName);
+	}
+
+	/**
+	 * Gets all events and their subscriber counts.
+	 *
+	 * @returns Map of event names to subscriber counts
+	 */
+	getEventStats(): Map<string, number> {
+		const stats = new Map<string, number>();
+		for (const [eventName, subscribers] of this.#subscriptionIndex) {
+			stats.set(eventName, subscribers.size);
+		}
+		return stats;
+	}
+
+	/**
+	 * Clears all subscriptions for a connection.
+	 * Used during connection cleanup.
+	 *
+	 * @param connection - The connection to clear subscriptions for
+	 */
+	clearConnectionSubscriptions(connection: Conn<S, CP, CS, V, I, DB, E, Q>) {
+		for (const eventName of [...connection.subscriptions.values()]) {
+			this.removeSubscription(eventName, connection, true);
+		}
+	}
+
+	/**
+	 * Gets the total number of unique events being subscribed to.
+	 */
+	get eventCount(): number {
+		return this.#subscriptionIndex.size;
+	}
+
+	/**
+	 * Gets the total number of subscriptions across all events.
+	 */
+	get totalSubscriptionCount(): number {
+		let total = 0;
+		for (const subscribers of this.#subscriptionIndex.values()) {
+			total += subscribers.size;
+		}
+		return total;
+	}
+
+	/**
+	 * Checks if an event has any subscribers.
+	 *
+	 * @param eventName - The name of the event to check
+	 * @returns True if the event has at least one subscriber
+	 */
+	hasSubscribers(eventName: string): boolean {
+		const subscribers = this.#subscriptionIndex.get(eventName);
+		return subscribers !== undefined && subscribers.size > 0;
+	}
+}

package/src/actor/instance/keys.ts

@@ -0,0 +1,146 @@
+export const KEYS = {
+	PERSIST_DATA: Uint8Array.from([1]),
+	CONN_PREFIX: Uint8Array.from([2]), // Prefix for connection keys
+	INSPECTOR_TOKEN: Uint8Array.from([3]), // Inspector token key
+	KV: Uint8Array.from([4]), // Prefix for user-facing KV storage
+	QUEUE_PREFIX: Uint8Array.from([5]), // Prefix for queue storage
+	WORKFLOW_PREFIX: Uint8Array.from([6]), // Prefix for workflow storage
+	TRACES_PREFIX: Uint8Array.from([7]), // Prefix for traces storage
+	SQLITE_PREFIX: Uint8Array.from([8]), // Prefix for SQLite VFS data
+};
+
+export const STORAGE_VERSION = {
+	QUEUE: 1,
+	WORKFLOW: 1,
+	TRACES: 1,
+	SQLITE: 1,
+} as const;
+
+const STORAGE_VERSION_BYTES = {
+	QUEUE: Uint8Array.from([STORAGE_VERSION.QUEUE]),
+	WORKFLOW: Uint8Array.from([STORAGE_VERSION.WORKFLOW]),
+	TRACES: Uint8Array.from([STORAGE_VERSION.TRACES]),
+	SQLITE: Uint8Array.from([STORAGE_VERSION.SQLITE]),
+} as const;
+
+const QUEUE_NAMESPACE = {
+	METADATA: Uint8Array.from([1]),
+	MESSAGES: Uint8Array.from([2]),
+} as const;
+
+const QUEUE_ID_BYTES = 8;
+
+function concatPrefix(prefix: Uint8Array, suffix: Uint8Array): Uint8Array {
+	const merged = new Uint8Array(prefix.length + suffix.length);
+	merged.set(prefix, 0);
+	merged.set(suffix, prefix.length);
+	return merged;
+}
+
+const QUEUE_STORAGE_PREFIX = concatPrefix(
+	KEYS.QUEUE_PREFIX,
+	STORAGE_VERSION_BYTES.QUEUE,
+);
+const QUEUE_METADATA_KEY = concatPrefix(
+	QUEUE_STORAGE_PREFIX,
+	QUEUE_NAMESPACE.METADATA,
+);
+const QUEUE_MESSAGES_PREFIX = concatPrefix(
+	QUEUE_STORAGE_PREFIX,
+	QUEUE_NAMESPACE.MESSAGES,
+);
+const WORKFLOW_STORAGE_PREFIX = concatPrefix(
+	KEYS.WORKFLOW_PREFIX,
+	STORAGE_VERSION_BYTES.WORKFLOW,
+);
+const TRACES_STORAGE_PREFIX = concatPrefix(
+	KEYS.TRACES_PREFIX,
+	STORAGE_VERSION_BYTES.TRACES,
+);
+const SQLITE_STORAGE_PREFIX = concatPrefix(
+	KEYS.SQLITE_PREFIX,
+	STORAGE_VERSION_BYTES.SQLITE,
+);
+
+// Helper to create a prefixed key for user-facing KV storage
+export function makePrefixedKey(key: Uint8Array): Uint8Array {
+	const prefixed = new Uint8Array(KEYS.KV.length + key.length);
+	prefixed.set(KEYS.KV, 0);
+	prefixed.set(key, KEYS.KV.length);
+	return prefixed;
+}
+
+// Helper to remove the prefix from a key
+export function removePrefixFromKey(prefixedKey: Uint8Array): Uint8Array {
+	return prefixedKey.slice(KEYS.KV.length);
+}
+
+export function makeWorkflowKey(key: Uint8Array): Uint8Array {
+	return concatPrefix(WORKFLOW_STORAGE_PREFIX, key);
+}
+
+export function makeTracesKey(key: Uint8Array): Uint8Array {
+	return concatPrefix(TRACES_STORAGE_PREFIX, key);
+}
+
+export function workflowStoragePrefix(): Uint8Array {
+	return Uint8Array.from(WORKFLOW_STORAGE_PREFIX);
+}
+
+export function tracesStoragePrefix(): Uint8Array {
+	return Uint8Array.from(TRACES_STORAGE_PREFIX);
+}
+
+export function sqliteStoragePrefix(): Uint8Array {
+	return Uint8Array.from(SQLITE_STORAGE_PREFIX);
+}
+
+export function queueStoragePrefix(): Uint8Array {
+	return Uint8Array.from(QUEUE_STORAGE_PREFIX);
+}
+
+export function queueMetadataKey(): Uint8Array {
+	return Uint8Array.from(QUEUE_METADATA_KEY);
+}
+
+export function queueMessagesPrefix(): Uint8Array {
+	return Uint8Array.from(QUEUE_MESSAGES_PREFIX);
+}
+
+// Helper to create a connection key
+export function makeConnKey(connId: string): Uint8Array {
+	const encoder = new TextEncoder();
+	const connIdBytes = encoder.encode(connId);
+	const key = new Uint8Array(KEYS.CONN_PREFIX.length + connIdBytes.length);
+	key.set(KEYS.CONN_PREFIX, 0);
+	key.set(connIdBytes, KEYS.CONN_PREFIX.length);
+	return key;
+}
+
+// Helper to create a queue message key
+export function makeQueueMessageKey(id: bigint): Uint8Array {
+	const key = new Uint8Array(QUEUE_MESSAGES_PREFIX.length + QUEUE_ID_BYTES);
+	key.set(QUEUE_MESSAGES_PREFIX, 0);
+	const view = new DataView(key.buffer, key.byteOffset, key.byteLength);
+	view.setBigUint64(QUEUE_MESSAGES_PREFIX.length, id, false);
+	return key;
+}
+
+// Helper to decode a queue message key
+export function decodeQueueMessageKey(key: Uint8Array): bigint {
+	const offset = QUEUE_MESSAGES_PREFIX.length;
+	if (key.length < offset + QUEUE_ID_BYTES) {
+		throw new Error("Queue key is too short");
+	}
+	for (let i = 0; i < QUEUE_MESSAGES_PREFIX.length; i++) {
+		if (key[i] !== QUEUE_MESSAGES_PREFIX[i]) {
+			throw new Error("Queue key has invalid prefix");
+		}
+	}
+	const view = new DataView(
+		key.buffer,
+		key.byteOffset + offset,
+		QUEUE_ID_BYTES,
+	);
+	return view.getBigUint64(0, false);
+}

package/src/actor/instance/kv.ts

@@ -0,0 +1,241 @@
+import type { ActorDriver } from "../driver";
+import { makePrefixedKey, removePrefixFromKey } from "./keys";
+
+/**
+ * User-facing KV storage interface exposed on ActorContext.
+ */
+type KvValueType = "text" | "arrayBuffer" | "binary";
+type KvKeyType = "text" | "binary";
+type KvKey = Uint8Array | string;
+
+type KvValueTypeMap = {
+	text: string;
+	arrayBuffer: ArrayBuffer;
+	binary: Uint8Array;
+};
+
+type KvKeyTypeMap = {
+	text: string;
+	binary: Uint8Array;
+};
+
+type KvValueOptions<T extends KvValueType = "text"> = {
+	type?: T;
+};
+
+type KvListOptions<
+	T extends KvValueType = "text",
+	K extends KvKeyType = "text",
+> = KvValueOptions<T> & {
+	keyType?: K;
+};
+
+const textEncoder = new TextEncoder();
+const textDecoder = new TextDecoder();
+
+function encodeKey<K extends KvKeyType = KvKeyType>(
+	key: KvKeyTypeMap[K],
+	keyType?: K,
+): Uint8Array {
+	if (key instanceof Uint8Array) {
+		return key;
+	}
+	const resolvedKeyType = keyType ?? "text";
+	if (resolvedKeyType === "binary") {
+		throw new TypeError("Expected a Uint8Array when keyType is binary");
+	}
+	return textEncoder.encode(key);
+}
+
+function decodeKey<K extends KvKeyType = "text">(
+	key: Uint8Array,
+	keyType?: K,
+): KvKeyTypeMap[K] {
+	const resolvedKeyType = keyType ?? "text";
+	switch (resolvedKeyType) {
+		case "text":
+			return textDecoder.decode(key) as KvKeyTypeMap[K];
+		case "binary":
+			return key as KvKeyTypeMap[K];
+		default:
+			throw new TypeError("Invalid kv key type");
+	}
+}
+
+function resolveValueType(
+	value: string | Uint8Array | ArrayBuffer,
+): KvValueType {
+	if (typeof value === "string") {
+		return "text";
+	}
+	if (value instanceof Uint8Array) {
+		return "binary";
+	}
+	if (value instanceof ArrayBuffer) {
+		return "arrayBuffer";
+	}
+	throw new TypeError("Invalid kv value");
+}
+
+function encodeValue<T extends KvValueType = KvValueType>(
+	value: KvValueTypeMap[T],
+	options?: KvValueOptions<T>,
+): Uint8Array {
+	const type =
+		options?.type ??
+		resolveValueType(value as string | Uint8Array | ArrayBuffer);
+	switch (type) {
+		case "text":
+			if (typeof value !== "string") {
+				throw new TypeError("Expected a string when type is text");
+			}
+			return textEncoder.encode(value);
+		case "arrayBuffer":
+			if (!(value instanceof ArrayBuffer)) {
+				throw new TypeError("Expected an ArrayBuffer when type is arrayBuffer");
+			}
+			return new Uint8Array(value);
+		case "binary":
+			if (!(value instanceof Uint8Array)) {
+				throw new TypeError("Expected a Uint8Array when type is binary");
+			}
+			return value;
+		default:
+			throw new TypeError("Invalid kv value type");
+	}
+}
+
+function decodeValue<T extends KvValueType = "text">(
+	value: Uint8Array,
+	options?: KvValueOptions<T>,
+): KvValueTypeMap[T] {
+	const type = options?.type ?? "text";
+	switch (type) {
+		case "text":
+			return textDecoder.decode(value) as KvValueTypeMap[T];
+		case "arrayBuffer": {
+			const copy = new Uint8Array(value.byteLength);
+			copy.set(value);
+			return copy.buffer as KvValueTypeMap[T];
+		}
+		case "binary":
+			return value as KvValueTypeMap[T];
+		default:
+			throw new TypeError("Invalid kv value type");
+	}
+}
+
+export class ActorKv {
+	#driver: ActorDriver;
+	#actorId: string;
+
+	constructor(driver: ActorDriver, actorId: string) {
+		this.#driver = driver;
+		this.#actorId = actorId;
+	}
+
+	/**
+	 * Get a single value by key.
+	 */
+	async get<T extends KvValueType = "text">(
+		key: KvKey,
+		options?: KvValueOptions<T>,
+	): Promise<KvValueTypeMap[T] | null> {
+		const results = await this.#driver.kvBatchGet(this.#actorId, [
+			makePrefixedKey(encodeKey(key)),
+		]);
+		const result = results[0];
+		if (!result) {
+			return null;
+		}
+		return decodeValue(result, options);
+	}
+
+	/**
+	 * Get multiple values by keys.
+	 */
+	async getBatch<T extends KvValueType = "text">(
+		keys: KvKey[],
+		options?: KvValueOptions<T>,
+	): Promise<(KvValueTypeMap[T] | null)[]> {
+		const prefixedKeys = keys.map((key) =>
+			makePrefixedKey(encodeKey(key)),
+		);
+		const results = await this.#driver.kvBatchGet(
+			this.#actorId,
+			prefixedKeys,
+		);
+		return results.map((result) =>
+			result ? decodeValue(result, options) : null,
+		);
+	}
+
+	/**
+	 * Put a single key-value pair.
+	 */
+	async put<T extends KvValueType = KvValueType>(
+		key: KvKey,
+		value: KvValueTypeMap[T],
+		options?: KvValueOptions<T>,
+	): Promise<void> {
+		await this.#driver.kvBatchPut(this.#actorId, [
+			[makePrefixedKey(encodeKey(key)), encodeValue(value, options)],
+		]);
+	}
+
+	/**
+	 * Put multiple key-value pairs.
+	 */
+	async putBatch<T extends KvValueType = KvValueType>(
+		entries: [KvKey, KvValueTypeMap[T]][],
+		options?: KvValueOptions<T>,
+	): Promise<void> {
+		const prefixedEntries: [Uint8Array, Uint8Array][] = entries.map(
+			([key, value]) => [
+				makePrefixedKey(encodeKey(key)),
+				encodeValue(value, options),
+			],
+		);
+		await this.#driver.kvBatchPut(this.#actorId, prefixedEntries);
+	}
+
+	/**
+	 * Delete a single key.
+	 */
+	async delete(key: KvKey): Promise<void> {
+		await this.#driver.kvBatchDelete(this.#actorId, [
+			makePrefixedKey(encodeKey(key)),
+		]);
+	}
+
+	/**
+	 * Delete multiple keys.
+	 */
+	async deleteBatch(keys: KvKey[]): Promise<void> {
+		const prefixedKeys = keys.map((key) =>
+			makePrefixedKey(encodeKey(key)),
+		);
+		await this.#driver.kvBatchDelete(this.#actorId, prefixedKeys);
+	}
+
+	/**
+	 * List all keys with a given prefix.
+	 * Returns key-value pairs where keys have the user prefix removed.
+	 */
+	async list<T extends KvValueType = "text", K extends KvKeyType = "text">(
+		prefix: KvKeyTypeMap[K],
+		options?: KvListOptions<T, K>,
+	): Promise<[KvKeyTypeMap[K], KvValueTypeMap[T]][]> {
+		const prefixedPrefix = makePrefixedKey(
+			encodeKey(prefix, options?.keyType),
+		);
+		const results = await this.#driver.kvListPrefix(
+			this.#actorId,
+			prefixedPrefix,
+		);
+		return results.map(([key, value]) => [
+			decodeKey<K>(removePrefixFromKey(key), options?.keyType),
+			decodeValue<T>(value, options),
+		]);
+	}
+}