rivetkit 2.1.4 → 2.1.6-rc.1

package/src/driver-test-suite/tests/actor-workflow.ts

@@ -1,5 +1,7 @@
-import { describe, expect, test } from "vitest";
+import { describe, expect, test, vi } from "vitest";
+import type { ActorError } from "@/client/mod";
 import {
+	WORKFLOW_NESTED_QUEUE_NAME,
 	WORKFLOW_QUEUE_NAME,
 } from "../../../fixtures/driver-test-suite/workflow";
 import type { DriverTestConfig } from "../mod";
@@ -32,7 +34,9 @@ export function runActorWorkflowTests(driverTestConfig: DriverTestConfig) {
 
 		test("consumes queue messages via workflow queue.next", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
-			const actor = client.workflowQueueActor.getOrCreate(["workflow-queue"]);
+			const actor = client.workflowQueueActor.getOrCreate([
+				"workflow-queue",
+			]);
 
 			await actor.send(WORKFLOW_QUEUE_NAME, {
 				hello: "world",
@@ -56,6 +60,143 @@ export function runActorWorkflowTests(driverTestConfig: DriverTestConfig) {
 			});
 		});
 
+		for (const testCase of [
+			{
+				name: "nested loops",
+				key: "loop" as const,
+				getActor: (
+					client: Awaited<
+						ReturnType<typeof setupDriverTest>
+					>["client"],
+				) => client.workflowNestedLoopActor,
+				firstItems: ["a", "b"],
+				secondItems: ["c"],
+				expected: ["a", "b", "c"],
+			},
+			{
+				name: "nested joins",
+				key: "join" as const,
+				getActor: (
+					client: Awaited<
+						ReturnType<typeof setupDriverTest>
+					>["client"],
+				) => client.workflowNestedJoinActor,
+				firstItems: ["a", "b"],
+				secondItems: ["c"],
+				expected: ["a", "b", "c"],
+			},
+			{
+				name: "nested races",
+				key: "race" as const,
+				getActor: (
+					client: Awaited<
+						ReturnType<typeof setupDriverTest>
+					>["client"],
+				) => client.workflowNestedRaceActor,
+				firstItems: ["a"],
+				secondItems: ["b"],
+				expected: ["a", "b"],
+			},
+			]) {
+				test(
+					`replays ${testCase.name} across workflow queue iterations`,
+				async (c) => {
+					const { client } = await setupDriverTest(
+						c,
+						driverTestConfig,
+					);
+					const actor = testCase.getActor(client).getOrCreate([
+						`workflow-nested-${testCase.key}`,
+					]);
+
+					const first = await actor.send(
+						WORKFLOW_NESTED_QUEUE_NAME,
+						{
+							items: testCase.firstItems,
+						},
+						{
+							wait: true,
+							timeout: 1_000,
+						},
+					);
+					expect(first).toEqual({
+						status: "completed",
+						response: {
+							processed: testCase.firstItems.length,
+						},
+					});
+
+					const second = await actor.send(
+						WORKFLOW_NESTED_QUEUE_NAME,
+						{
+							items: testCase.secondItems,
+						},
+						{
+							wait: true,
+							timeout: 1_000,
+						},
+					);
+					expect(second).toEqual({
+						status: "completed",
+						response: {
+							processed: testCase.secondItems.length,
+						},
+					});
+
+					const state = await actor.getState();
+					expect(state.processed).toEqual(testCase.expected);
+				},
+				);
+			}
+
+			test(
+				"starts child workflows created inside workflow steps",
+				async (c) => {
+					const { client } = await setupDriverTest(
+						c,
+						driverTestConfig,
+					);
+					const parent = client.workflowSpawnParentActor.getOrCreate([
+						"workflow-spawn-parent",
+					]);
+
+					expect(await parent.triggerSpawn("child-1")).toEqual({
+						queued: true,
+					});
+
+					let parentState = await parent.getState();
+					for (
+						let i = 0;
+						i < 30 && parentState.results.length === 0;
+						i++
+					) {
+						await waitFor(driverTestConfig, 100);
+						parentState = await parent.getState();
+					}
+
+					expect(parentState.results).toEqual([
+						{
+							key: "child-1",
+							result: {
+								status: "completed",
+								response: { ok: true },
+							},
+							error: null,
+						},
+					]);
+
+					const child = client.workflowSpawnChildActor.getOrCreate([
+						"child-1",
+					]);
+					const childState = await child.getState();
+					expect(childState).toEqual({
+						label: "child-1",
+						started: true,
+						processed: ["hello"],
+					});
+				},
+			);
+
 		test("db and client are step-only in workflow context", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
 			const actor = client.workflowAccessActor.getOrCreate([
@@ -80,7 +221,9 @@ export function runActorWorkflowTests(driverTestConfig: DriverTestConfig) {
 
 		test("sleeps and resumes between ticks", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
-			const actor = client.workflowSleepActor.getOrCreate(["workflow-sleep"]);
+			const actor = client.workflowSleepActor.getOrCreate([
+				"workflow-sleep",
+			]);
 
 			const initial = await actor.getState();
 			await waitFor(driverTestConfig, 200);
@@ -89,6 +232,247 @@ export function runActorWorkflowTests(driverTestConfig: DriverTestConfig) {
 			expect(next.ticks).toBeGreaterThan(initial.ticks);
 		});
 
+		test("workflow onError reports retry metadata", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const actor = client.workflowErrorHookActor.getOrCreate([
+				"workflow-error-hook",
+			]);
+
+			let state = await actor.getErrorState();
+			for (
+				let i = 0;
+				i < 80 && (state.attempts < 2 || state.events.length === 0);
+				i++
+			) {
+				await waitFor(driverTestConfig, 50);
+				state = await actor.getErrorState();
+			}
+
+			expect(state.attempts).toBe(2);
+			expect(state.events).toHaveLength(1);
+			expect(state.events[0]).toEqual(
+				expect.objectContaining({
+					step: expect.objectContaining({
+						stepName: "flaky",
+						attempt: 1,
+						willRetry: true,
+						retryDelay: 1,
+						error: expect.objectContaining({
+							name: "Error",
+							message: "workflow hook failed",
+						}),
+					}),
+				}),
+			);
+		});
+
+		test("workflow onError can update actor state", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const actor = client.workflowErrorHookEffectsActor.getOrCreate([
+				"workflow-error-state",
+			]);
+
+			await actor.startWorkflow();
+
+			let state = await actor.getErrorState();
+			for (
+				let i = 0;
+				i < 80 &&
+				(state.attempts < 2 ||
+					state.lastError === null ||
+					state.errorCount === 0);
+				i++
+			) {
+				await waitFor(driverTestConfig, 50);
+				state = await actor.getErrorState();
+			}
+
+			expect(state.attempts).toBe(2);
+			expect(state.errorCount).toBe(1);
+			expect(state.lastError).toEqual(
+				expect.objectContaining({
+					step: expect.objectContaining({
+						stepName: "flaky",
+						attempt: 1,
+						willRetry: true,
+						retryDelay: 1,
+						error: expect.objectContaining({
+							name: "Error",
+							message: "workflow hook failed",
+						}),
+					}),
+				}),
+			);
+		});
+
+		test("workflow onError can broadcast actor events", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const actor = client.workflowErrorHookEffectsActor
+				.getOrCreate(["workflow-error-broadcast"])
+				.connect();
+
+			try {
+				const eventPromise = new Promise((resolve) => {
+					actor.once("workflowError", resolve);
+				});
+
+				await actor.startWorkflow();
+
+				const event = await eventPromise;
+				expect(event).toEqual(
+					expect.objectContaining({
+						step: expect.objectContaining({
+							stepName: "flaky",
+							attempt: 1,
+							willRetry: true,
+							retryDelay: 1,
+							error: expect.objectContaining({
+								name: "Error",
+								message: "workflow hook failed",
+							}),
+						}),
+					}),
+				);
+			} finally {
+				await actor.dispose();
+			}
+		});
+
+		test("workflow onError can enqueue actor messages", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const actor = client.workflowErrorHookEffectsActor.getOrCreate([
+				"workflow-error-queue",
+			]);
+
+			await actor.startWorkflow();
+
+			const queuedError = await actor.receiveQueuedError();
+			expect(queuedError).toEqual(
+				expect.objectContaining({
+					step: expect.objectContaining({
+						stepName: "flaky",
+						attempt: 1,
+						willRetry: true,
+						retryDelay: 1,
+						error: expect.objectContaining({
+							name: "Error",
+							message: "workflow hook failed",
+						}),
+					}),
+				}),
+			);
+		});
+
+		test.skipIf(driverTestConfig.skip?.sleep)(
+			"completed workflows sleep instead of destroying the actor",
+			async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+				const actor = client.workflowCompleteActor.getOrCreate([
+					"workflow-complete",
+				]);
+
+				let state = await actor.getState();
+				for (let i = 0; i < 10 && state.sleepCount === 0; i++) {
+					await waitFor(driverTestConfig, 100);
+					state = await actor.getState();
+				}
+				expect(state.runCount).toBeGreaterThan(0);
+				expect(state.sleepCount).toBeGreaterThan(0);
+				expect(state.startCount).toBeGreaterThan(1);
+			},
+		);
+
+		test("workflow steps can destroy the actor", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const actorKey = "workflow-destroy";
+			const observer = client.destroyObserver.getOrCreate(["observer"]);
+			await observer.reset();
+
+			const actor = client.workflowDestroyActor.getOrCreate([actorKey]);
+			const actorId = await actor.resolve();
+
+			await vi.waitFor(async () => {
+				const wasDestroyed = await observer.wasDestroyed(actorKey);
+				expect(wasDestroyed, "actor onDestroy not called").toBeTruthy();
+			});
+
+			await vi.waitFor(async () => {
+				let actorRunning = false;
+				try {
+					await client.workflowDestroyActor.getForId(actorId).resolve();
+					actorRunning = true;
+				} catch (err) {
+					expect((err as ActorError).group).toBe("actor");
+					expect((err as ActorError).code).toBe("not_found");
+				}
+
+				expect(actorRunning, "actor still running").toBeFalsy();
+			});
+		});
+
+		test.skipIf(driverTestConfig.skip?.sleep)(
+			"failed workflow steps sleep instead of surfacing as run errors",
+			async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+				const actor = client.workflowFailedStepActor.getOrCreate([
+					"workflow-failed-step",
+				]);
+
+				let state = await actor.getState();
+				for (let i = 0; i < 10 && state.sleepCount === 0; i++) {
+					await waitFor(driverTestConfig, 100);
+					state = await actor.getState();
+				}
+				expect(state.runCount).toBeGreaterThan(0);
+				expect(state.sleepCount).toBeGreaterThan(0);
+				expect(state.startCount).toBeGreaterThan(1);
+			},
+		);
+
+		test.skipIf(driverTestConfig.skip?.sleep)(
+			"workflow onError is not reported again after sleep and wake",
+			async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+				const actor = client.workflowErrorHookSleepActor.getOrCreate([
+					"workflow-error-hook-sleep",
+				]);
+
+				let state = await actor.getErrorState();
+				for (
+					let i = 0;
+					i < 80 &&
+					(state.attempts < 2 || state.events.length === 0);
+					i++
+				) {
+					await waitFor(driverTestConfig, 50);
+					state = await actor.getErrorState();
+				}
+
+				expect(state.attempts).toBe(2);
+				expect(state.events).toHaveLength(1);
+				expect(state.wakeCount).toBe(1);
+
+				await actor.triggerSleep();
+				await waitFor(driverTestConfig, 250);
+
+				let resumedState = await actor.getErrorState();
+				for (
+					let i = 0;
+					i < 40 &&
+					(resumedState.wakeCount < 2 || resumedState.sleepCount < 1);
+					i++
+				) {
+					await waitFor(driverTestConfig, 50);
+					resumedState = await actor.getErrorState();
+				}
+
+				expect(resumedState.sleepCount).toBeGreaterThanOrEqual(1);
+				expect(resumedState.wakeCount).toBeGreaterThanOrEqual(2);
+				expect(resumedState.attempts).toBe(2);
+				expect(resumedState.events).toHaveLength(1);
+			},
+		);
+
 		test.skipIf(driverTestConfig.skip?.sleep)(
 			"workflow run teardown does not wait for runStopTimeout",
 			async (c) => {

package/src/driver-test-suite/tests/conn-error-serialization.ts

@@ -2,7 +2,9 @@ import { describe, expect, test } from "vitest";
 import type { DriverTestConfig } from "../mod";
 import { setupDriverTest } from "../utils";
 
-export function runConnErrorSerializationTests(driverTestConfig: DriverTestConfig) {
+export function runConnErrorSerializationTests(
+	driverTestConfig: DriverTestConfig,
+) {
 	describe("Connection Error Serialization Tests", () => {
 		test("error thrown in createConnState preserves group and code through WebSocket serialization", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);

package/src/driver-test-suite/tests/raw-websocket.ts

@@ -484,7 +484,11 @@ export function runRawWebSocketTests(driverTestConfig: DriverTestConfig) {
 				ws.addEventListener("open", () => resolve(), { once: true });
 				ws.addEventListener("error", reject);
 				ws.addEventListener("close", (evt: any) => {
-					reject(new Error(`WebSocket closed: code=${evt.code} reason=${evt.reason}`));
+					reject(
+						new Error(
+							`WebSocket closed: code=${evt.code} reason=${evt.reason}`,
+						),
+					);
 				});
 			});
 

package/src/drivers/default.ts

@@ -7,9 +7,7 @@ import { DriverConfig, RegistryConfig } from "@/registry/config";
 /**
  * Chooses the appropriate driver based on the run configuration.
  */
-export function chooseDefaultDriver(
-	config: RegistryConfig,
-): DriverConfig {
+export function chooseDefaultDriver(config: RegistryConfig): DriverConfig {
 	if (config.endpoint && config.driver) {
 		throw new UserError(
 			"Cannot specify both 'endpoint' and 'driver' in configuration",

package/src/drivers/engine/actor-driver.ts

@@ -3,7 +3,8 @@ import type {
 	RunnerConfig as EngineRunnerConfig,
 	HibernatingWebSocketMetadata,
 } from "@rivetkit/engine-runner";
-import type { SqliteVfs } from "@rivetkit/sqlite-vfs";
+import type { ISqliteVfs } from "@rivetkit/sqlite-vfs";
+import { SqliteVfsPoolManager } from "@/driver-helpers/sqlite-pool";
 import { idToStr, Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
 import type { Context as HonoContext } from "hono";
@@ -84,9 +85,22 @@ export class EngineActorDriver implements ActorDriver {
 	#runner: Runner;
 	#actors: Map<string, ActorHandler> = new Map();
 	#actorRouter: ActorRouter;
+	#sqlitePool: SqliteVfsPoolManager;
 
-	#runnerStarted: PromiseWithResolvers<undefined> = promiseWithResolvers((reason) => logger().warn({ msg: "unhandled runner started promise rejection", reason }));
-	#runnerStopped: PromiseWithResolvers<undefined> = promiseWithResolvers((reason) => logger().warn({ msg: "unhandled runner stopped promise rejection", reason }));
+	#runnerStarted: PromiseWithResolvers<undefined> = promiseWithResolvers(
+		(reason) =>
+			logger().warn({
+				msg: "unhandled runner started promise rejection",
+				reason,
+			}),
+	);
+	#runnerStopped: PromiseWithResolvers<undefined> = promiseWithResolvers(
+		(reason) =>
+			logger().warn({
+				msg: "unhandled runner stopped promise rejection",
+				reason,
+			}),
+	);
 	#isRunnerStopped: boolean = false;
 
 	// HACK: Track actor stop intent locally since the runner protocol doesn't
@@ -124,6 +138,7 @@ export class EngineActorDriver implements ActorDriver {
 		this.#config = config;
 		this.#managerDriver = managerDriver;
 		this.#inlineClient = inlineClient;
+		this.#sqlitePool = new SqliteVfsPoolManager(config);
 
 		// HACK: Override inspector token (which are likely to be
 		// removed later on) with token from x-rivet-token header
@@ -148,7 +163,6 @@ export class EngineActorDriver implements ActorDriver {
 			namespace: config.namespace,
 			totalSlots: config.runner.totalSlots,
 			runnerName: config.runner.runnerName,
-			runnerKey: config.runner.runnerKey ?? crypto.randomUUID(),
 			metadata: {
 				rivetkit: { version: VERSION },
 			},
@@ -169,6 +183,9 @@ export class EngineActorDriver implements ActorDriver {
 			onActorStart: this.#runnerOnActorStart.bind(this),
 			onActorStop: this.#runnerOnActorStop.bind(this),
 			logger: getLogger("engine-runner"),
+			debugLatencyMs: process.env._RIVET_DEBUG_LATENCY_MS
+				? Number.parseInt(process.env._RIVET_DEBUG_LATENCY_MS, 10)
+				: undefined,
 		};
 
 		// Create and start runner
@@ -257,6 +274,14 @@ export class EngineActorDriver implements ActorDriver {
 		await this.#runner.kvDelete(actorId, keys);
 	}
 
+	async kvDeleteRange(
+		actorId: string,
+		start: Uint8Array,
+		end: Uint8Array,
+	): Promise<void> {
+		await this.#runner.kvDeleteRange(actorId, start, end);
+	}
+
 	async kvList(actorId: string): Promise<Uint8Array[]> {
 		const entries = await this.#runner.kvListPrefix(
 			actorId,
@@ -275,8 +300,16 @@ export class EngineActorDriver implements ActorDriver {
 	async kvListPrefix(
 		actorId: string,
 		prefix: Uint8Array,
+		options?: {
+			reverse?: boolean;
+			limit?: number;
+		},
 	): Promise<[Uint8Array, Uint8Array][]> {
-		const result = await this.#runner.kvListPrefix(actorId, prefix);
+		const result = await this.#runner.kvListPrefix(
+			actorId,
+			prefix,
+			options,
+		);
 		logger().info({
 			msg: "kvListPrefix called",
 			actorId,
@@ -287,14 +320,27 @@ export class EngineActorDriver implements ActorDriver {
 		return result;
 	}
 
+	async kvListRange(
+		actorId: string,
+		start: Uint8Array,
+		end: Uint8Array,
+		options?: {
+			reverse?: boolean;
+			limit?: number;
+		},
+	): Promise<[Uint8Array, Uint8Array][]> {
+		return await this.#runner.kvListRange(
+			actorId,
+			start,
+			end,
+			true,
+			options,
+		);
+	}
+
 	/** Creates a SQLite VFS instance for creating KV-backed databases */
-	async createSqliteVfs(): Promise<SqliteVfs> {
-		// Dynamic import keeps @rivetkit/sqlite out of the main entrypoint bundle.
-		// Returning a fresh SqliteVfs gives each actor an isolated sqlite module
-		// instance, avoiding async re-entrancy across actors.
-		const specifier = "@rivetkit/" + "sqlite-vfs";
-		const { SqliteVfs } = await import(specifier);
-		return new SqliteVfs();
+	async createSqliteVfs(actorId: string): Promise<ISqliteVfs> {
+		return await this.#sqlitePool.acquire(actorId);
 	}
 
 	// MARK: - Actor Lifecycle
@@ -364,11 +410,16 @@ export class EngineActorDriver implements ActorDriver {
 		await Promise.all(stopPromises);
 		logger().debug({ msg: "all actors stopped" });
 
+		await this.#sqlitePool.shutdown();
+
 		try {
 			await this.#runner.shutdown(immediate);
 		} catch (error) {
-			const message = error instanceof Error ? error.message : String(error);
-			if (message.includes("WebSocket connection closed during shutdown")) {
+			const message =
+				error instanceof Error ? error.message : String(error);
+			if (
+				message.includes("WebSocket connection closed during shutdown")
+			) {
 				logger().debug({
 					msg: "ignoring shutdown websocket close race",
 					error: message,
@@ -465,7 +516,12 @@ export class EngineActorDriver implements ActorDriver {
 			// async operations to avoid race conditions where multiple calls might try to
 			// create the same handler simultaneously.
 			handler = {
-				actorStartPromise: promiseWithResolvers((reason) => logger().warn({ msg: "unhandled actor start promise rejection", reason })),
+				actorStartPromise: promiseWithResolvers((reason) =>
+					logger().warn({
+						msg: "unhandled actor start promise rejection",
+						reason,
+					}),
+				),
 			};
 			this.#actors.set(actorId, handler);
 		}
@@ -508,12 +564,19 @@ export class EngineActorDriver implements ActorDriver {
 					protocolMetadata,
 				});
 
-				const stopThresholdMax = Math.max(Number(protocolMetadata.actorStopThreshold) - 1000, 0);
+				const stopThresholdMax = Math.max(
+					Number(protocolMetadata.actorStopThreshold) - 1000,
+					0,
+				);
 				handler.actor.overrides.onSleepTimeout = stopThresholdMax;
 				handler.actor.overrides.onDestroyTimeout = stopThresholdMax;
 
 				if (protocolMetadata.serverlessDrainGracePeriod) {
-					const drainMax = Math.max(Number(protocolMetadata.serverlessDrainGracePeriod) - 1000, 0);
+					const drainMax = Math.max(
+						Number(protocolMetadata.serverlessDrainGracePeriod) -
+						1000,
+						0,
+					);
 					handler.actor.overrides.runStopTimeout = drainMax;
 					handler.actor.overrides.waitUntilTimeout = drainMax;
 				}
@@ -537,7 +600,9 @@ export class EngineActorDriver implements ActorDriver {
 						`Failed to start actor ${actorId}: ${innerError.message}`,
 						{ cause: innerError },
 					)
-					: new Error(`Failed to start actor ${actorId}: ${String(innerError)}`);
+					: new Error(
+						`Failed to start actor ${actorId}: ${String(innerError)}`,
+					);
 			handler.actor = undefined;
 			handler.actorStartError = error;
 			handler.actorStartPromise?.reject(error);
@@ -580,13 +645,21 @@ export class EngineActorDriver implements ActorDriver {
 
 		const handler = this.#actors.get(actorId);
 		if (!handler) {
-			logger().debug({ msg: "no runner actor handler to stop", actorId, reason });
+			logger().debug({
+				msg: "no runner actor handler to stop",
+				actorId,
+				reason,
+			});
 			return;
 		}
 
 		if (handler.actorStartPromise) {
 			try {
-				logger().debug({ msg: "runner actor stopping before it started, waiting", actorId, generation });
+				logger().debug({
+					msg: "runner actor stopping before it started, waiting",
+					actorId,
+					generation,
+				});
 				await handler.actorStartPromise.promise;
 			} catch (err) {
 				// Start failed, but we still want to clean up the handler
@@ -1056,4 +1129,4 @@ export class EngineActorDriver implements ActorDriver {
 			entry.bufferedMessageSize = 0;
 		}
 	}
-}
+}