rivetkit 2.1.5 → 2.1.6-rc.1

package/src/actor/config.ts

@@ -212,13 +212,14 @@ export const ActorConfigSchema = z
 				onDestroyTimeout: z.number().positive().default(5000),
 				stateSaveInterval: z.number().positive().default(10_000),
 				actionTimeout: z.number().positive().default(60_000),
-				// Max time to wait for waitUntil background promises during shutdown
+				// Deprecated timeout for legacy background shutdown tasks
 				waitUntilTimeout: z.number().positive().default(15_000),
 				// Max time to wait for run handler to stop during shutdown
 				runStopTimeout: z.number().positive().default(15_000),
 				connectionLivenessTimeout: z.number().positive().default(2500),
 				connectionLivenessInterval: z.number().positive().default(5000),
-				noSleep: z.boolean().default(false),
+				/** @deprecated Use `c.setPreventSleep(true)` in `onWake` instead. Will be removed in 2.2. */
+			noSleep: z.boolean().default(false),
 				sleepTimeout: z.number().positive().default(30_000),
 				maxQueueSize: z.number().positive().default(1000),
 				maxQueueMessageSize: z
@@ -286,11 +287,11 @@ type CreateState<
 > =
 	| { state: TState }
 	| {
-		createState: (
-			c: CreateContext<TState, TInput, TDatabase, TEvents, TQueues>,
-			input: TInput,
-		) => TState | Promise<TState>;
-	}
+			createState: (
+				c: CreateContext<TState, TInput, TDatabase, TEvents, TQueues>,
+				input: TInput,
+			) => TState | Promise<TState>;
+	  }
 	| Record<never, never>;
 
 // Creates connection state config
@@ -310,18 +311,18 @@ type CreateConnState<
 > =
 	| { connState: TConnState }
 	| {
-		createConnState: (
-			c: CreateConnStateContext<
-				TState,
-				TVars,
-				TInput,
-				TDatabase,
-				TEvents,
-				TQueues
-			>,
-			params: TConnParams,
-		) => TConnState | Promise<TConnState>;
-	}
+			createConnState: (
+				c: CreateConnStateContext<
+					TState,
+					TVars,
+					TInput,
+					TDatabase,
+					TEvents,
+					TQueues
+				>,
+				params: TConnParams,
+			) => TConnState | Promise<TConnState>;
+	  }
 	| Record<never, never>;
 
 // Creates vars config
@@ -341,26 +342,26 @@ type CreateVars<
 	TQueues extends QueueSchemaConfig,
 > =
 	| {
-		/**
-		 * @experimental
-		 */
-		vars: TVars;
-	}
+			/**
+			 * @experimental
+			 */
+			vars: TVars;
+	  }
 	| {
-		/**
-		 * @experimental
-		 */
-		createVars: (
-			c: CreateVarsContext<
-				TState,
-				TInput,
-				TDatabase,
-				TEvents,
-				TQueues
-			>,
-			driverCtx: any,
-		) => TVars | Promise<TVars>;
-	}
+			/**
+			 * @experimental
+			 */
+			createVars: (
+				c: CreateVarsContext<
+					TState,
+					TInput,
+					TDatabase,
+					TEvents,
+					TQueues
+				>,
+				driverCtx: any,
+			) => TVars | Promise<TVars>;
+	  }
 	| Record<never, never>;
 
 export interface Actions<
@@ -498,43 +499,48 @@ interface BaseActorConfig<
 	 * - Custom workflow logic
 	 *
 	 * **Important:** The actor may go to sleep at any time during the `run`
-	 * handler. Use `c.keepAwake(promise)` to wrap async operations that should
-	 * not be interrupted by sleep.
+	 * handler. Use `c.setPreventSleep(true)` while work is active, then clear
+	 * it with `c.setPreventSleep(false)` once the actor can sleep again.
 	 *
 	 * The handler receives an abort signal via `c.abortSignal` and a
 	 * `c.aborted` alias for loop checks. Use these to gracefully exit.
 	 *
-	 * If this handler exits or throws, the actor will crash and reschedule.
+	 * If this handler exits, the actor will follow the normal idle sleep timeout
+	 * once it becomes idle.
+	 * If this handler throws, the actor logs the error and then sleeps once it
+	 * becomes idle.
+	 * Call `c.destroy()` explicitly if a run handler should destroy the actor.
+	 *
 	 * On shutdown, the actor waits for this handler to complete with a
 	 * configurable timeout (options.runStopTimeout, default 15s).
 	 *
 	 * Can be either a function or a RunConfig object with optional name/icon metadata.
 	 *
-	 * @returns Void or a Promise. If the promise exits, the actor crashes.
+	 * @returns Void or a Promise.
 	 */
 	run?:
-	| ((
-		c: RunContext<
-			TState,
-			TConnParams,
-			TConnState,
-			TVars,
-			TInput,
-			TDatabase,
-			TEvents,
-			TQueues
-		>,
-	) => void | Promise<void>)
-	| RunConfig<
-		TState,
-		TConnParams,
-		TConnState,
-		TVars,
-		TInput,
-		TDatabase,
-		TEvents,
-		TQueues
-	>;
+		| ((
+				c: RunContext<
+					TState,
+					TConnParams,
+					TConnState,
+					TVars,
+					TInput,
+					TDatabase,
+					TEvents,
+					TQueues
+				>,
+		  ) => void | Promise<void>)
+		| RunConfig<
+				TState,
+				TConnParams,
+				TConnState,
+				TVars,
+				TInput,
+				TDatabase,
+				TEvents,
+				TQueues
+		  >;
 
 	/**
 	 * Called when the actor's state changes.
@@ -739,11 +745,11 @@ interface BaseActorConfig<
 
 type ActorDatabaseConfig<TDatabase extends AnyDatabaseProvider> =
 	| {
-		/**
-		 * @experimental
-		 */
-		db: TDatabase;
-	}
+			/**
+			 * @experimental
+			 */
+			db: TDatabase;
+	  }
 	| Record<never, never>;
 
 // 1. Infer schema
@@ -1041,7 +1047,7 @@ export const DocActorOptionsSchema = z
 			.number()
 			.optional()
 			.describe(
-				"Max time in ms to wait for waitUntil background promises during shutdown. Default: 15000",
+				"Deprecated. Max time in ms to wait for legacy background shutdown tasks. Default: 15000",
 			),
 		runStopTimeout: z
 			.number()
@@ -1064,7 +1070,7 @@ export const DocActorOptionsSchema = z
 		noSleep: z
 			.boolean()
 			.optional()
-			.describe("If true, the actor will never sleep. Default: false"),
+			.describe("Deprecated: use c.setPreventSleep(true) in onWake instead. Will be removed in 2.2. If true, the actor will never sleep. Default: false"),
 		sleepTimeout: z
 			.number()
 			.optional()
@@ -1160,24 +1166,24 @@ export const DocActorConfigSchema = z
 			.unknown()
 			.optional()
 			.describe(
-				"Called after actor starts. Does not block startup. Use for background tasks like queue processing or tick loops. If it exits or throws, the actor crashes.",
+				"Called after actor starts. Does not block startup. Use for background tasks like queue processing or tick loops. If it exits, the actor follows the normal idle sleep timeout once idle. If it throws, the actor logs the error and then follows the normal idle sleep timeout once idle.",
+			),
+		onStateChange: z
+			.unknown()
+			.optional()
+			.describe(
+				"Called when the actor's state changes. State changes within this hook won't trigger recursion.",
 			),
-			onStateChange: z
-				.unknown()
-				.optional()
-				.describe(
-					"Called when the actor's state changes. State changes within this hook won't trigger recursion.",
-				),
-			onBeforeConnect: z
-				.unknown()
-				.optional()
-				.describe(
-					"Called before a client connects. Throw an error to reject the connection.",
-				),
-			onConnect: z
-				.unknown()
-				.optional()
-				.describe("Called when a client successfully connects."),
+		onBeforeConnect: z
+			.unknown()
+			.optional()
+			.describe(
+				"Called before a client connects. Throw an error to reject the connection.",
+			),
+		onConnect: z
+			.unknown()
+			.optional()
+			.describe("Called when a client successfully connects."),
 		onDisconnect: z
 			.unknown()
 			.optional()

package/src/actor/conn/drivers/websocket.ts

@@ -62,7 +62,8 @@ export function createWebSocketDriver(
 
 				// Check outgoing message size
 				const messageSize =
-					(serialized as any).byteLength || (serialized as any).length;
+					(serialized as any).byteLength ||
+					(serialized as any).length;
 				if (messageSize > config.maxOutgoingMessageSize) {
 					actor.rLog.error({
 						msg: "outgoing message exceeds maxOutgoingMessageSize",

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

@@ -245,7 +245,7 @@ export class ActorContext<
 	/**
 	 * Returns the client for the given registry.
 	 */
-	client<R extends Registry<any>>(): Client<R> {
+	client<R extends Registry<any> = Registry<any>>(): Client<R> {
 		return this.#actor.inlineClient as Client<R>;
 	}
 
@@ -274,19 +274,42 @@ export class ActorContext<
 
 	/**
 	 * Prevents the actor from sleeping until promise is complete.
+	 *
+	 * @deprecated Use `onSleep` for shutdown or flush work, or
+	 * `c.setPreventSleep(true)` while work is active if the actor must stay
+	 * awake until it finishes.
 	 */
 	waitUntil(promise: Promise<void>): void {
 		this.#actor.waitUntil(promise);
 	}
 
 	/**
-	 * Prevents the actor from sleeping while the given promise is running.
+	 * Prevents the actor from automatically sleeping until cleared.
 	 *
-	 * Use this when performing async operations in the `run` handler or other
-	 * background contexts where you need to ensure the actor stays awake.
+	 * @experimental
+	 */
+	setPreventSleep(prevent: boolean): void {
+		this.#actor.setPreventSleep(prevent);
+	}
+
+	/**
+	 * True when the actor is explicitly blocking automatic sleep.
+	 *
+	 * @experimental
+	 */
+	get preventSleep(): boolean {
+		return this.#actor.preventSleep;
+	}
+
+	/**
+	 * Prevents the actor from sleeping while the given promise is running.
 	 *
 	 * Returns the resolved value and resets the sleep timer on completion.
 	 * Errors are propagated to the caller.
+	 *
+	 * @deprecated Use `c.setPreventSleep(true)` while work is active, or move
+	 * shutdown and flush work to `onSleep` if it can wait until the actor is
+	 * sleeping.
 	 */
 	keepAwake<T>(promise: Promise<T>): Promise<T> {
 		return this.#actor.keepAwake(promise);

package/src/actor/database.ts

@@ -10,4 +10,9 @@ export type InferDatabaseClient<DBProvider extends AnyDatabaseProvider> =
 		? Awaited<ReturnType<DBProvider["createClient"]>>
 		: never;
 
-export type { AnyDatabaseProvider, DatabaseProvider, RawDatabaseClient, DrizzleDatabaseClient };
+export type {
+	AnyDatabaseProvider,
+	DatabaseProvider,
+	RawDatabaseClient,
+	DrizzleDatabaseClient,
+};

package/src/actor/driver.ts

@@ -4,10 +4,8 @@ import type { ManagerDriver } from "@/manager/driver";
 import { type AnyConn } from "./conn/mod";
 import type { AnyActorInstance } from "./instance/mod";
 import type { RegistryConfig } from "@/registry/config";
-import type {
-	RawDatabaseClient,
-} from "@/db/config";
-import type { SqliteVfs } from "@rivetkit/sqlite-vfs";
+import type { RawDatabaseClient } from "@/db/config";
+import type { ISqliteVfs } from "@rivetkit/sqlite-vfs";
 import { BaseSQLiteDatabase } from "drizzle-orm/sqlite-core";
 
 export type ActorDriverBuilder = (
@@ -39,10 +37,32 @@ export interface ActorDriver {
 	/** Batch delete multiple keys. */
 	kvBatchDelete(actorId: string, keys: Uint8Array[]): Promise<void>;
 
+	/** Delete all keys in the half-open range [start, end). */
+	kvDeleteRange(
+		actorId: string,
+		start: Uint8Array,
+		end: Uint8Array,
+	): Promise<void>;
+
 	/** List all keys with a given prefix. */
 	kvListPrefix(
 		actorId: string,
 		prefix: Uint8Array,
+		options?: {
+			reverse?: boolean;
+			limit?: number;
+		},
+	): Promise<[Uint8Array, Uint8Array][]>;
+
+	/** List all keys in the half-open range [start, end). */
+	kvListRange(
+		actorId: string,
+		start: Uint8Array,
+		end: Uint8Array,
+		options?: {
+			reverse?: boolean;
+			limit?: number;
+		},
 	): Promise<[Uint8Array, Uint8Array][]>;
 
 	// Schedule
@@ -66,19 +86,18 @@ export interface ActorDriver {
 	 */
 	overrideDrizzleDatabaseClient?(
 		actorId: string,
-	): Promise<BaseSQLiteDatabase<any,any,any,any> | undefined>;
+	): Promise<BaseSQLiteDatabase<any, any, any, any> | undefined>;
 
 	/**
 	 * Creates a SQLite VFS instance for creating KV-backed databases.
 	 * If not provided, the database provider will need an override.
 	 *
-	 * @rivetkit/sqlite's async build is not re-entrant per module instance. Drivers
-	 * should return a new instance per call for actor-level isolation.
+	 * Drivers may use the actorId for pool-based instance assignment.
 	 *
 	 * This is a method (not a property) so drivers can use dynamic imports,
 	 * keeping the core driver tree-shakeable from @rivetkit/sqlite.
 	 */
-	createSqliteVfs?(): SqliteVfs | Promise<SqliteVfs>;
+	createSqliteVfs?(actorId: string): ISqliteVfs | Promise<ISqliteVfs>;
 
 	/**
 	 * Requests the actor to go to sleep.

package/src/actor/errors.ts

@@ -138,7 +138,7 @@ export class IncomingMessageTooLong extends ActorError {
 		super(
 			"message",
 			"incoming_too_long",
-			"Incoming message too long. This can be configured with: `registry.start({ maxIncomingMessageSize: ... })`",
+			"Incoming message too long. This can be configured with: `setup({ maxIncomingMessageSize: ... })`",
 			{ public: true },
 		);
 	}
@@ -149,7 +149,7 @@ export class OutgoingMessageTooLong extends ActorError {
 		super(
 			"message",
 			"outgoing_too_long",
-			"Outgoing message too long. This can be configured with: `registry.start({ maxOutgoingMessageSize: ... })`",
+			"Outgoing message too long. This can be configured with: `setup({ maxOutgoingMessageSize: ... })`",
 			{ public: true },
 		);
 	}
@@ -256,9 +256,14 @@ export class QueueCompletionPayloadInvalid extends ActorError {
 
 export class QueueAlreadyCompleted extends ActorError {
 	constructor() {
-		super("queue", "already_completed", "Queue message was already completed.", {
-			public: true,
-		});
+		super(
+			"queue",
+			"already_completed",
+			"Queue message was already completed.",
+			{
+				public: true,
+			},
+		);
 	}
 }
 

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

@@ -370,7 +370,9 @@ export class ConnectionManager<
 			if (conn.isHibernatable) {
 				const key = makeConnKey(conn.id);
 				try {
-					await this.#actor.driver.kvBatchDelete(this.#actor.id, [key]);
+					await this.#actor.driver.kvBatchDelete(this.#actor.id, [
+						key,
+					]);
 					this.#actor.rLog.debug({
 						msg: "removed connection from KV",
 						connId: conn.id,
@@ -396,8 +398,7 @@ export class ConnectionManager<
 	): Promise<number> {
 		const staleConnections = Array.from(this.#connections.values()).filter(
 			(conn) =>
-				conn.isHibernatable &&
-				conn[CONN_DRIVER_SYMBOL] === undefined,
+				conn.isHibernatable && conn[CONN_DRIVER_SYMBOL] === undefined,
 		);
 		if (staleConnections.length === 0) {
 			return 0;

package/src/actor/instance/kv.ts

@@ -28,6 +28,8 @@ type KvListOptions<
 	K extends KvKeyType = "text",
 > = KvValueOptions<T> & {
 	keyType?: K;
+	reverse?: boolean;
+	limit?: number;
 };
 
 const textEncoder = new TextEncoder();
@@ -92,12 +94,16 @@ function encodeValue<T extends KvValueType = KvValueType>(
 			return textEncoder.encode(value);
 		case "arrayBuffer":
 			if (!(value instanceof ArrayBuffer)) {
-				throw new TypeError("Expected an ArrayBuffer when type is 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");
+				throw new TypeError(
+					"Expected a Uint8Array when type is binary",
+				);
 			}
 			return value;
 		default:
@@ -158,9 +164,7 @@ export class ActorKv {
 		keys: KvKey[],
 		options?: KvValueOptions<T>,
 	): Promise<(KvValueTypeMap[T] | null)[]> {
-		const prefixedKeys = keys.map((key) =>
-			makePrefixedKey(encodeKey(key)),
-		);
+		const prefixedKeys = keys.map((key) => makePrefixedKey(encodeKey(key)));
 		const results = await this.#driver.kvBatchGet(
 			this.#actorId,
 			prefixedKeys,
@@ -212,12 +216,21 @@ export class ActorKv {
 	 * Delete multiple keys.
 	 */
 	async deleteBatch(keys: KvKey[]): Promise<void> {
-		const prefixedKeys = keys.map((key) =>
-			makePrefixedKey(encodeKey(key)),
-		);
+		const prefixedKeys = keys.map((key) => makePrefixedKey(encodeKey(key)));
 		await this.#driver.kvBatchDelete(this.#actorId, prefixedKeys);
 	}
 
+	/**
+	 * Delete all keys in the half-open range [start, end).
+	 */
+	async deleteRange(start: KvKey, end: KvKey): Promise<void> {
+		await this.#driver.kvDeleteRange(
+			this.#actorId,
+			makePrefixedKey(encodeKey(start)),
+			makePrefixedKey(encodeKey(end)),
+		);
+	}
+
 	/**
 	 * List all keys with a given prefix.
 	 * Returns key-value pairs where keys have the user prefix removed.
@@ -232,10 +245,40 @@ export class ActorKv {
 		const results = await this.#driver.kvListPrefix(
 			this.#actorId,
 			prefixedPrefix,
+			{
+				reverse: options?.reverse,
+				limit: options?.limit,
+			},
 		);
 		return results.map(([key, value]) => [
 			decodeKey<K>(removePrefixFromKey(key), options?.keyType),
 			decodeValue<T>(value, options),
 		]);
 	}
-}
+
+	/**
+	 * List all key-value pairs in the half-open range [start, end).
+	 */
+	async listRange<
+		T extends KvValueType = "text",
+		K extends KvKeyType = "text",
+	>(
+		start: KvKeyTypeMap[K],
+		end: KvKeyTypeMap[K],
+		options?: KvListOptions<T, K>,
+	): Promise<[KvKeyTypeMap[K], KvValueTypeMap[T]][]> {
+		const results = await this.#driver.kvListRange(
+			this.#actorId,
+			makePrefixedKey(encodeKey(start, options?.keyType)),
+			makePrefixedKey(encodeKey(end, options?.keyType)),
+			{
+				reverse: options?.reverse,
+				limit: options?.limit,
+			},
+		);
+		return results.map(([key, value]) => [
+			decodeKey<K>(removePrefixFromKey(key), options?.keyType),
+			decodeValue<T>(value, options),
+		]);
+	}
+}