rivetkit 2.0.37 → 2.0.39

package/src/actor/instance/kv.ts

@@ -1,15 +1,241 @@
-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
-};
-
-// 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;
+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),
+		]);
+	}
 }

package/src/actor/instance/mod.ts

@@ -44,7 +44,7 @@ import {
 } from "../utils";
 import { ConnectionManager } from "./connection-manager";
 import { EventManager } from "./event-manager";
-import { KEYS } from "./kv";
+import { KEYS } from "./keys";
 import {
 	convertActorFromBarePersisted,
 	type PersistedActor,
@@ -203,7 +203,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	}
 
 	get actions(): string[] {
-		return Object.keys(this.#config.actions);
+		return Object.keys(this.#config.actions ?? {});
 	}
 
 	get config(): ActorConfig<S, CP, CS, V, I, DB> {
@@ -516,12 +516,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	): Promise<unknown> {
 		this.assertReady();
 
-		if (!(actionName in this.#config.actions)) {
+		const actions = this.#config.actions ?? {};
+		if (!(actionName in actions)) {
 			this.#rLog.warn({ msg: "action does not exist", actionName });
 			throw new errors.ActionNotFound(actionName);
 		}
 
-		const actionFunction = this.#config.actions[actionName];
+		const actionFunction = actions[actionName];
 		if (typeof actionFunction !== "function") {
 			this.#rLog.warn({
 				msg: "action is not a function",

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

@@ -12,7 +12,7 @@ import { convertConnToBarePersistedConn } from "../conn/persisted";
 import type { ActorDriver } from "../driver";
 import * as errors from "../errors";
 import { isConnStatePath, isStatePath } from "../utils";
-import { KEYS, makeConnKey } from "./kv";
+import { KEYS, makeConnKey } from "./keys";
 import type { ActorInstance } from "./mod";
 import { convertActorToBarePersisted, type PersistedActor } from "./persisted";
 

package/src/actor/mod.ts

@@ -76,7 +76,8 @@ export type { AnyConn, Conn } from "./conn/mod";
 export type { ActorDefinition, AnyActorDefinition } from "./definition";
 export { lookupInRegistry } from "./definition";
 export { UserError, type UserErrorOptions } from "./errors";
-export { KEYS as KV_KEYS } from "./instance/kv";
+export { KEYS as KV_KEYS } from "./instance/keys";
+export { ActorKv } from "./instance/kv";
 export type { AnyActorInstance } from "./instance/mod";
 export {
 	type ActorRouter,

package/src/actor/router-websocket-endpoints.ts

@@ -384,7 +384,8 @@ export function parseWebSocketProtocols(
 		}
 	}
 
-	const encoding = EncodingSchema.parse(encodingRaw);
+	// Default to "json" encoding for raw WebSocket connections without subprotocols
+	const encoding = EncodingSchema.parse(encodingRaw ?? "json");
 	const connParams = connParamsRaw ? JSON.parse(connParamsRaw) : undefined;
 
 	return { encoding, connParams };

package/src/client/actor-conn.ts

@@ -26,13 +26,14 @@ import { deserializeWithEncoding, serializeWithEncoding } from "@/serde";
 import { bufferToArrayBuffer, promiseWithResolvers } from "@/utils";
 import { getLogMessage } from "@/utils/env-vars";
 import type { ActorDefinitionActions } from "./actor-common";
-import { queryActor } from "./actor-query";
+import { checkForSchedulingError, queryActor } from "./actor-query";
 import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client";
 import * as errors from "./errors";
 import { logger } from "./log";
 import {
 	type WebSocketMessage as ConnMessage,
 	messageLength,
+	parseWebSocketCloseReason,
 	sendHttpRequest,
 } from "./utils";
 
@@ -387,9 +388,9 @@ enc
 				});
 			}
 		});
-		ws.addEventListener("close", (ev) => {
+		ws.addEventListener("close", async (ev) => {
 			try {
-				this.#handleOnClose(ev);
+				await this.#handleOnClose(ev);
 			} catch (err) {
 				logger().error({
 					msg: "error in websocket close handler",
@@ -519,27 +520,24 @@ enc
 					metadata,
 				});
 
-				// Check if this is an actor scheduling error
-				let errorToThrow: Error;
-				if (
-					group === "guard" &&
-					(code === "actor_ready_timeout" ||
-						code === "actor_runner_failed") &&
-					this.#actorId
-				) {
-					// Try to fetch actor details to get more specific error info
-					const schedulingError =
-						await this.#checkForSchedulingError();
-					errorToThrow =
-						schedulingError ??
-						new errors.ActorError(group, code, message, metadata);
-				} else {
-					errorToThrow = new errors.ActorError(
+				// Check if this is an actor scheduling error and try to get more details
+				let errorToThrow = new errors.ActorError(
+					group,
+					code,
+					message,
+					metadata,
+				);
+				if (errors.isSchedulingError(group, code) && this.#actorId) {
+					const schedulingError = await checkForSchedulingError(
 						group,
 						code,
-						message,
-						metadata,
+						this.#actorId,
+						this.#actorQuery,
+						this.#driver,
 					);
+					if (schedulingError) {
+						errorToThrow = schedulingError;
+					}
 				}
 
 				// If we have an onOpenPromise, reject it with the error
@@ -553,10 +551,7 @@ enc
 					this.#actionsInFlight.delete(id);
 				}
 
-				// Dispatch to error handler if it's an ActorError
-				if (errorToThrow instanceof errors.ActorError) {
-					this.#dispatchActorError(errorToThrow);
-				}
+				this.#dispatchActorError(errorToThrow);
 			}
 		} else if (response.body.tag === "ActionResponse") {
 			// Action response OK
@@ -587,7 +582,7 @@ enc
 	}
 
 	/** Called by the onclose event from drivers. */
-	#handleOnClose(event: Event | CloseEvent) {
+	async #handleOnClose(event: Event | CloseEvent) {
 		// We can't use `event instanceof CloseEvent` because it's not defined in NodeJS
 		const closeEvent = event as CloseEvent;
 		const wasClean = closeEvent.wasClean;
@@ -609,12 +604,55 @@ enc
 			this.#rejectPendingPromises(new errors.ActorConnDisposed(), true);
 		} else {
 			this.#setConnStatus("disconnected");
-			this.#rejectPendingPromises(
-				new Error(
-					`${wasClean ? "Connection closed" : "Connection lost"} (code: ${closeEvent.code}, reason: ${closeEvent.reason})`,
-				),
-				false,
-			);
+
+			// Build error from close event
+			let error: Error;
+			const reason = closeEvent.reason || "";
+			const parsed = parseWebSocketCloseReason(reason);
+
+			if (parsed) {
+				const { group, code } = parsed;
+
+				// Check if this is a scheduling error
+				if (errors.isSchedulingError(group, code) && this.#actorId) {
+					const schedulingError = await checkForSchedulingError(
+						group,
+						code,
+						this.#actorId,
+						this.#actorQuery,
+						this.#driver,
+					);
+					if (schedulingError) {
+						error = schedulingError;
+					} else {
+						error = new errors.ActorError(
+							group,
+							code,
+							`Connection closed: ${reason}`,
+							undefined,
+						);
+					}
+				} else {
+					error = new errors.ActorError(
+						group,
+						code,
+						`Connection closed: ${reason}`,
+						undefined,
+					);
+				}
+			} else {
+				// Default error for non-structured close reasons
+				error = new Error(
+					`${wasClean ? "Connection closed" : "Connection lost"} (code: ${closeEvent.code}, reason: ${reason})`,
+				);
+			}
+
+			this.#rejectPendingPromises(error, false);
+
+			// Dispatch to error handler if it's an ActorError
+			if (error instanceof errors.ActorError) {
+				this.#dispatchActorError(error);
+			}
 
 			// Automatically reconnect if we were connected
 			if (wasConnected) {
@@ -712,55 +750,6 @@ enc
 		}
 	}
 
-	async #checkForSchedulingError(): Promise<errors.ActorSchedulingError | null> {
-		if (!this.#actorId) return null;
-
-		// Extract name from the query (it's nested in one of the variant properties)
-		const query = this.#actorQuery;
-		const name =
-			"getForId" in query
-				? query.getForId.name
-				: "getForKey" in query
-					? query.getForKey.name
-					: "getOrCreateForKey" in query
-						? query.getOrCreateForKey.name
-						: "create" in query
-							? query.create.name
-							: null;
-
-		if (!name) return null;
-
-		try {
-			// Fetch actor details to check for scheduling errors
-			const actor = await this.#driver.getForId({
-				name,
-				actorId: this.#actorId,
-			});
-
-			if (actor?.error) {
-				logger().info({
-					msg: "found actor scheduling error",
-					actorId: this.#actorId,
-					error: actor.error,
-				});
-				return new errors.ActorSchedulingError(
-					this.#actorId,
-					actor.error,
-				);
-			}
-		} catch (err) {
-			// If we can't fetch actor details, just return null
-			// and fall back to the generic error
-			logger().warn({
-				msg: "failed to fetch actor details for scheduling error check",
-				actorId: this.#actorId,
-				error: stringifyError(err),
-			});
-		}
-
-		return null;
-	}
-
 	#addEventSubscription<Args extends Array<unknown>>(
 		eventName: string,
 		callback: (...args: Args) => void,

package/src/client/actor-handle.ts

@@ -25,9 +25,9 @@ import {
 import { bufferToArrayBuffer } from "@/utils";
 import type { ActorDefinitionActions } from "./actor-common";
 import { type ActorConn, ActorConnRaw } from "./actor-conn";
-import { queryActor } from "./actor-query";
+import { checkForSchedulingError, queryActor } from "./actor-query";
 import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client";
-import { ActorError } from "./errors";
+import { ActorError, isSchedulingError } from "./errors";
 import { logger } from "./log";
 import { rawHttpFetch, rawWebSocket } from "./raw-utils";
 import { sendHttpRequest } from "./utils";
@@ -81,22 +81,17 @@ export class ActorHandleRaw {
 		args: Args;
 		signal?: AbortSignal;
 	}): Promise<Response> {
-		// return await this.#driver.action<Args, Response>(
-		// 	undefined,
-		// 	this.#actorQuery,
-		// 	this.#encodingKind,
-		// 	this.#params,
-		// 	opts.name,
-		// 	opts.args,
-		// 	{ signal: opts.signal },
-		// );
+		// Track actorId for scheduling error lookups
+		let actorId: string | undefined;
+
 		try {
 			// Get the actor ID
-			const { actorId } = await queryActor(
+			const result = await queryActor(
 				undefined,
 				this.#actorQuery,
 				this.#driver,
 			);
+			actorId = result.actorId;
 			logger().debug({ msg: "found actor for action", actorId });
 			invariant(actorId, "Missing actor ID");
 
@@ -159,6 +154,21 @@ export class ActorHandleRaw {
 				{},
 				true,
 			);
+
+			// Check if this is a scheduling error and try to get more details
+			if (actorId && isSchedulingError(group, code)) {
+				const schedulingError = await checkForSchedulingError(
+					group,
+					code,
+					actorId,
+					this.#actorQuery,
+					this.#driver,
+				);
+				if (schedulingError) {
+					throw schedulingError;
+				}
+			}
+
 			throw new ActorError(group, code, message, metadata);
 		}
 	}
@@ -229,9 +239,7 @@ export class ActorHandleRaw {
 	}
 
 	/**
-	 * Resolves the actor to get its unique actor ID
-	 *
-	 * @returns {Promise<string>} - A promise that resolves to the actor's ID
+	 * Resolves the actor to get its unique actor ID.
 	 */
 	async resolve({ signal }: { signal?: AbortSignal } = {}): Promise<string> {
 		if (
@@ -267,6 +275,18 @@ export class ActorHandleRaw {
 			assertUnreachable(this.#actorQuery);
 		}
 	}
+
+	/**
+	 * Returns the raw URL for routing traffic to the actor.
+	 */
+	async getGatewayUrl(): Promise<string> {
+		const { actorId } = await queryActor(
+			undefined,
+			this.#actorQuery,
+			this.#driver,
+		);
+		return await this.#driver.buildGatewayUrl(actorId);
+	}
 }
 
 /**

package/src/client/actor-query.ts

@@ -1,7 +1,9 @@
 import type { Context as HonoContext } from "hono";
 import * as errors from "@/actor/errors";
+import { stringifyError } from "@/common/utils";
 import type { ManagerDriver } from "@/driver-helpers/mod";
 import type { ActorQuery } from "@/manager/protocol/query";
+import { ActorSchedulingError } from "./errors";
 import { logger } from "./log";
 
 /**
@@ -63,3 +65,48 @@ export async function queryActor(
 	logger().debug({ msg: "actor query result", actorId: actorOutput.actorId });
 	return { actorId: actorOutput.actorId };
 }
+
+/**
+ * Extract the actor name from a query.
+ */
+export function getActorNameFromQuery(query: ActorQuery): string {
+	if ("getForId" in query) return query.getForId.name;
+	if ("getForKey" in query) return query.getForKey.name;
+	if ("getOrCreateForKey" in query) return query.getOrCreateForKey.name;
+	if ("create" in query) return query.create.name;
+	throw new errors.InvalidRequest("Invalid query format");
+}
+
+/**
+ * Fetch actor details and check for scheduling errors.
+ */
+export async function checkForSchedulingError(
+	group: string,
+	code: string,
+	actorId: string,
+	query: ActorQuery,
+	driver: ManagerDriver,
+): Promise<ActorSchedulingError | null> {
+	const name = getActorNameFromQuery(query);
+
+	try {
+		const actor = await driver.getForId({ name, actorId });
+
+		if (actor?.error) {
+			logger().info({
+				msg: "found actor scheduling error",
+				actorId,
+				error: actor.error,
+			});
+			return new ActorSchedulingError(group, code, actorId, actor.error);
+		}
+	} catch (err) {
+		logger().warn({
+			msg: "failed to fetch actor details for scheduling error check",
+			actorId,
+			error: stringifyError(err),
+		});
+	}
+
+	return null;
+}