rivetkit 2.1.4 → 2.1.6-rc.1

package/src/utils/endpoint-parser.test.ts

@@ -87,7 +87,9 @@ describe("tryParseEndpoint", () => {
 			});
 			expect(result).toBeUndefined();
 			expect(issues).toHaveLength(1);
-			expect(issues[0]?.message).toBe("endpoint cannot contain a query string");
+			expect(issues[0]?.message).toBe(
+				"endpoint cannot contain a query string",
+			);
 		});
 
 		test("adds issue for fragment", () => {
@@ -97,7 +99,9 @@ describe("tryParseEndpoint", () => {
 			});
 			expect(result).toBeUndefined();
 			expect(issues).toHaveLength(1);
-			expect(issues[0]?.message).toBe("endpoint cannot contain a fragment");
+			expect(issues[0]?.message).toBe(
+				"endpoint cannot contain a fragment",
+			);
 		});
 
 		test("adds issue for token without namespace", () => {

package/src/utils/endpoint-parser.ts

@@ -39,7 +39,12 @@ export function tryParseEndpoint(
 	ctx: z.RefinementCtx,
 	options: TryParseEndpointOptions,
 ): ParsedEndpoint | undefined {
-	const { endpoint, path = ["endpoint"], namespace: configNamespace, token: configToken } = options;
+	const {
+		endpoint,
+		path = ["endpoint"],
+		namespace: configNamespace,
+		token: configToken,
+	} = options;
 	// Parse the URL
 	let url: URL;
 	try {
@@ -121,4 +126,3 @@ export function tryParseEndpoint(
 		token,
 	};
 }
-

package/src/utils/env-vars.ts

@@ -20,8 +20,6 @@ export const getRivetTotalSlots = (): number | undefined => {
 	const value = getEnvUniversal("RIVET_TOTAL_SLOTS");
 	return value !== undefined ? parseInt(value, 10) : undefined;
 };
-export const getRivetRunnerKey = (): string | undefined =>
-	getEnvUniversal("RIVET_RUNNER_KEY");
 export const getRivetRunEngine = (): boolean =>
 	getEnvUniversal("RIVET_RUN_ENGINE") === "1";
 export const getRivetRunEngineVersion = (): string | undefined =>

package/src/utils/node.ts

@@ -175,4 +175,4 @@ export function getNodeUrl(): typeof import("node:url") {
 		);
 	}
 	return nodeUrl;
-}
+}

package/src/utils/serve.ts

@@ -9,8 +9,9 @@ import getPort from "get-port";
 const DEFAULT_PORT = 6420;
 export type ServeStatic =
 	typeof import("@hono/node-server/serve-static").serveStatic;
-const serveStaticLoaderPromises: Partial<Record<Runtime, Promise<ServeStatic>>> =
-	{};
+const serveStaticLoaderPromises: Partial<
+	Record<Runtime, Promise<ServeStatic>>
+> = {};
 
 /**
  * Finds a free port starting from the given port.
@@ -39,7 +40,7 @@ export async function crossPlatformServe(
 	managerPort: number,
 	app: Hono<any>,
 	runtime: Runtime = detectRuntime(),
-): Promise<{ upgradeWebSocket: any }> {
+): Promise<{ upgradeWebSocket: any; closeServer?: () => void }> {
 	logger().debug({ msg: "detected runtime for serve", runtime });
 
 	switch (runtime) {
@@ -88,7 +89,7 @@ async function serveNode(
 	config: RegistryConfig,
 	managerPort: number,
 	app: Hono<any>,
-): Promise<{ upgradeWebSocket: any }> {
+): Promise<{ upgradeWebSocket: any; closeServer: () => void }> {
 	// Import @hono/node-server using string variable to prevent static analysis
 	const nodeServerModule = "@hono/node-server";
 	let serve: any;
@@ -136,7 +137,11 @@ async function serveNode(
 	);
 	injectWebSocket(server);
 
-	return { upgradeWebSocket };
+	const closeServer = () => {
+		server.close();
+	};
+
+	return { upgradeWebSocket, closeServer };
 }
 
 async function serveDeno(

package/src/utils.ts

@@ -272,7 +272,12 @@ export class SinglePromiseQueue {
 
 		// Ensure a shared resolver exists for all callers in this cycle
 		if (!this.#pending) {
-			this.#pending = promiseWithResolvers<void>((reason) => logger().warn({ msg: "unhandled single promise queue rejection", reason }));
+			this.#pending = promiseWithResolvers<void>((reason) =>
+				logger().warn({
+					msg: "unhandled single promise queue rejection",
+					reason,
+				}),
+			);
 		}
 
 		const waitForThisCycle = this.#pending.promise;

package/src/workflow/constants.ts

@@ -1,2 +1 @@
-export const WORKFLOW_GUARD_KV_KEY =
-	"__rivet_actor_workflow_guard_triggered";
+export const WORKFLOW_GUARD_KV_KEY = "__rivet_actor_workflow_guard_triggered";

package/src/workflow/context.ts

@@ -2,7 +2,10 @@ import type { RunContext } from "@/actor/contexts/run";
 import type { Client } from "@/client/client";
 import type { Registry } from "@/registry";
 import type { ActorDefinition, AnyActorDefinition } from "@/actor/definition";
-import type { AnyDatabaseProvider, InferDatabaseClient } from "@/actor/database";
+import type {
+	AnyDatabaseProvider,
+	InferDatabaseClient,
+} from "@/actor/database";
 import type {
 	QueueFilterName,
 	QueueNextBatchOptions,
@@ -42,9 +45,8 @@ type WorkflowActorQueueNextBatchOptions<
 	TCompletable extends boolean,
 > = Omit<QueueNextBatchOptions<TName, TCompletable>, "signal">;
 
-type WorkflowActorQueueNextBatchOptionsFallback<
-	TCompletable extends boolean,
-> = Omit<QueueNextBatchOptions<string, TCompletable>, "signal">;
+type WorkflowActorQueueNextBatchOptionsFallback<TCompletable extends boolean> =
+	Omit<QueueNextBatchOptions<string, TCompletable>, "signal">;
 
 type ActorWorkflowLoopConfig<
 	S,
@@ -182,7 +184,9 @@ export class ActorWorkflowContext<
 		>(
 			name: string,
 			opts?: WorkflowActorQueueNextBatchOptions<TName, TCompletable>,
-		): Promise<Array<QueueResultMessageForName<TQueues, TName, TCompletable>>>;
+		): Promise<
+			Array<QueueResultMessageForName<TQueues, TName, TCompletable>>
+		>;
 		function nextBatch<const TCompletable extends boolean = false>(
 			name: string,
 			opts?: WorkflowActorQueueNextBatchOptionsFallback<TCompletable>,
@@ -200,7 +204,9 @@ export class ActorWorkflowContext<
 			opts?: WorkflowActorQueueNextBatchOptions<string, boolean>,
 		): Promise<Array<WorkflowQueueMessage<unknown>>> {
 			const messages = await self.#inner.queue.nextBatch(name, opts);
-			return messages.map((message) => self.#toActorQueueMessage(message));
+			return messages.map((message) =>
+				self.#toActorQueueMessage(message),
+			);
 		}
 
 		function send<K extends keyof TQueues & string>(
@@ -232,7 +238,9 @@ export class ActorWorkflowContext<
 				throw new Error("Step run function missing");
 			}
 			return await this.#wrapActive(() =>
-				this.#inner.step(nameOrConfig, () => this.#withActorAccess(run)),
+				this.#inner.step(nameOrConfig, () =>
+					this.#withActorAccess(run),
+				),
 			);
 		}
 		const stepConfig = nameOrConfig as StepConfig<T>;
@@ -334,6 +342,11 @@ export class ActorWorkflowContext<
 		return this.#inner.sleepUntil(name, timestampMs);
 	}
 
+	destroy(): void {
+		this.#ensureActorAccess("destroy");
+		this.#runCtx.destroy();
+	}
+
 	async rollbackCheckpoint(name: string): Promise<void> {
 		await this.#wrapActive(() => this.#inner.rollbackCheckpoint(name));
 	}
@@ -429,7 +442,7 @@ export class ActorWorkflowContext<
 		return this.#runCtx.vars as TVars extends never ? never : TVars;
 	}
 
-	client<R extends Registry<any>>(): Client<R> {
+	client<R extends Registry<any> = Registry<any>>(): Client<R> {
 		this.#ensureActorAccess("client");
 		return this.#runCtx.client<R>();
 	}
@@ -445,11 +458,31 @@ export class ActorWorkflowContext<
 		return this.#runCtx.log;
 	}
 
+	setPreventSleep(prevent: boolean): void {
+		this.#ensureActorAccess("setPreventSleep");
+		this.#runCtx.setPreventSleep(prevent);
+	}
+
+	get preventSleep(): boolean {
+		this.#ensureActorAccess("preventSleep");
+		return this.#runCtx.preventSleep;
+	}
+
+	/**
+	 * @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> {
 		this.#ensureActorAccess("keepAwake");
 		return this.#runCtx.keepAwake(promise);
 	}
 
+	/**
+	 * @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.#ensureActorAccess("waitUntil");
 		this.#runCtx.waitUntil(promise);
@@ -471,7 +504,7 @@ export class ActorWorkflowContext<
 		this.#ensureActorAccess("broadcast");
 		this.#runCtx.broadcast(
 			name as never,
-			...((args as unknown[]) as never[]),
+			...(args as unknown[] as never[]),
 		);
 	}
 

package/src/workflow/driver.ts

@@ -15,6 +15,17 @@ function stripWorkflowKey(prefixed: Uint8Array): Uint8Array {
 	return prefixed.slice(WORKFLOW_STORAGE_PREFIX.length);
 }
 
+function computeUpperBound(prefix: Uint8Array): Uint8Array | null {
+	const upperBound = prefix.slice();
+	for (let i = upperBound.length - 1; i >= 0; i--) {
+		if (upperBound[i] !== 0xff) {
+			upperBound[i]++;
+			return upperBound.slice(0, i + 1);
+		}
+	}
+	return null;
+}
+
 class ActorWorkflowMessageDriver implements WorkflowMessageDriver {
 	#actor: AnyActorInstance;
 	#runCtx: RunContext<any, any, any, any, any, any, any, any>;
@@ -40,7 +51,9 @@ class ActorWorkflowMessageDriver implements WorkflowMessageDriver {
 	}): Promise<Message[]> {
 		const messages = await this.#runCtx.keepAwake(
 			this.#actor.queueManager.receive(
-				opts.names && opts.names.length > 0 ? [...opts.names] : undefined,
+				opts.names && opts.names.length > 0
+					? [...opts.names]
+					: undefined,
 				opts.count,
 				0,
 				undefined,
@@ -54,20 +67,23 @@ class ActorWorkflowMessageDriver implements WorkflowMessageDriver {
 			sentAt: message.createdAt,
 			...(opts.completable
 				? {
-						complete: async (response?: unknown) => {
-							await this.#runCtx.keepAwake(
-								this.#actor.queueManager.completeMessage(
-									message,
-									response,
-								),
-							);
-						},
-					}
+					complete: async (response?: unknown) => {
+						await this.#runCtx.keepAwake(
+							this.#actor.queueManager.completeMessage(
+								message,
+								response,
+							),
+						);
+					},
+				}
 				: {}),
 		}));
 	}
 
-	async completeMessage(messageId: string, response?: unknown): Promise<void> {
+	async completeMessage(
+		messageId: string,
+		response?: unknown,
+	): Promise<void> {
 		let parsedId: bigint;
 		try {
 			parsedId = BigInt(messageId);
@@ -122,19 +138,34 @@ export class ActorWorkflowDriver implements EngineDriver {
 	}
 
 	async deletePrefix(prefix: Uint8Array): Promise<void> {
-		const entries = await this.#runCtx.keepAwake(
-			this.#actor.driver.kvListPrefix(
-				this.#actor.id,
-				makeWorkflowKey(prefix),
-			),
-		);
-		if (entries.length === 0) {
-			return;
+		const start = makeWorkflowKey(prefix);
+		const end = computeUpperBound(start);
+		if (end) {
+			await this.#runCtx.keepAwake(
+				this.#actor.driver.kvDeleteRange(this.#actor.id, start, end),
+			);
+		} else {
+			const entries = await this.#runCtx.keepAwake(
+				this.#actor.driver.kvListPrefix(this.#actor.id, start),
+			);
+			if (entries.length === 0) {
+				return;
+			}
+			await this.#runCtx.keepAwake(
+				this.#actor.driver.kvBatchDelete(
+					this.#actor.id,
+					entries.map(([key]) => key),
+				),
+			);
 		}
+	}
+
+	async deleteRange(start: Uint8Array, end: Uint8Array): Promise<void> {
 		await this.#runCtx.keepAwake(
-			this.#actor.driver.kvBatchDelete(
+			this.#actor.driver.kvDeleteRange(
 				this.#actor.id,
-				entries.map(([key]) => key),
+				makeWorkflowKey(start),
+				makeWorkflowKey(end),
 			),
 		);
 	}
@@ -161,7 +192,10 @@ export class ActorWorkflowDriver implements EngineDriver {
 			Promise.all([
 				this.#actor.driver.kvBatchPut(
 					this.#actor.id,
-					writes.map(({ key, value }) => [makeWorkflowKey(key), value]),
+					writes.map(({ key, value }) => [
+						makeWorkflowKey(key),
+						value,
+					]),
 				),
 				this.#actor.stateManager.saveState({
 					immediate: true,
@@ -191,4 +225,4 @@ export class ActorWorkflowDriver implements EngineDriver {
 			abortSignal,
 		);
 	}
-}
+}

package/src/workflow/inspector.ts

@@ -62,9 +62,7 @@ function toU64(value: number): bigint {
 	return BigInt(Math.max(0, Math.floor(value)));
 }
 
-function toWorkflowLocation(
-	location: Location,
-): transport.WorkflowLocation {
+function toWorkflowLocation(location: Location): transport.WorkflowLocation {
 	return location.map((segment) => {
 		if (typeof segment === "number") {
 			return { tag: "WorkflowNameIndex", val: segment };
@@ -79,9 +77,7 @@ function toWorkflowLocation(
 	});
 }
 
-function toWorkflowEntryKind(
-	kind: EntryKind,
-): transport.WorkflowEntryKind {
+function toWorkflowEntryKind(kind: EntryKind): transport.WorkflowEntryKind {
 	switch (kind.type) {
 		case "step":
 			return {
@@ -124,7 +120,9 @@ function toWorkflowEntryKind(
 		case "join":
 			return {
 				tag: "WorkflowJoinEntry",
-				val: { branches: toWorkflowBranchStatusMap(kind.data.branches) },
+				val: {
+					branches: toWorkflowBranchStatusMap(kind.data.branches),
+				},
 			};
 		case "race":
 			return {
@@ -147,9 +145,7 @@ function toWorkflowEntryKind(
 	}
 }
 
-function toWorkflowEntry(
-	entry: WorkflowHistoryEntry,
-): transport.WorkflowEntry {
+function toWorkflowEntry(entry: WorkflowHistoryEntry): transport.WorkflowEntry {
 	return {
 		id: entry.id,
 		location: toWorkflowLocation(entry.location),
@@ -176,9 +172,7 @@ function toWorkflowEntryStatus(
 	}
 }
 
-function toWorkflowSleepState(
-	state: SleepState,
-): transport.WorkflowSleepState {
+function toWorkflowSleepState(state: SleepState): transport.WorkflowSleepState {
 	switch (state) {
 		case "pending":
 			return transport.WorkflowSleepState.PENDING;

package/src/workflow/mod.ts

@@ -5,13 +5,28 @@ import type { AnyActorInstance } from "@/actor/instance/mod";
 import type { EventSchemaConfig, QueueSchemaConfig } from "@/actor/schema";
 import { RUN_FUNCTION_CONFIG_SYMBOL } from "@/actor/config";
 import { stringifyError } from "@/utils";
-import { runWorkflow } from "@rivetkit/workflow-engine";
+import {
+	CriticalError,
+	EntryInProgressError,
+	HistoryDivergedError,
+	JoinError,
+	RaceError,
+	RollbackCheckpointError,
+	RollbackError,
+	runWorkflow,
+	StepExhaustedError,
+	type WorkflowErrorEvent,
+} from "@rivetkit/workflow-engine";
 import invariant from "invariant";
 import { ActorWorkflowContext } from "./context";
 import { ActorWorkflowDriver } from "./driver";
 import { createWorkflowInspectorAdapter } from "./inspector";
 
 export { Loop } from "@rivetkit/workflow-engine";
+export type {
+	WorkflowError,
+	WorkflowErrorEvent,
+} from "@rivetkit/workflow-engine";
 export {
 	ActorWorkflowContext,
 	type WorkflowBranchContextOf,
@@ -20,6 +35,53 @@ export {
 	type WorkflowStepContextOf,
 } from "./context";
 
+function shouldRethrowWorkflowError(error: unknown): boolean {
+	if (
+		error instanceof CriticalError ||
+		error instanceof JoinError ||
+		error instanceof RaceError ||
+		error instanceof RollbackError ||
+		error instanceof StepExhaustedError
+	) {
+		return false;
+	}
+
+	if (
+		error instanceof EntryInProgressError ||
+		error instanceof HistoryDivergedError ||
+		error instanceof RollbackCheckpointError
+	) {
+		return true;
+	}
+
+	return true;
+}
+
+export interface WorkflowOptions<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+> {
+	onError?: (
+		ctx: RunContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+		event: WorkflowErrorEvent,
+	) => void | Promise<void>;
+}
+
 export function workflow<
 	TState,
 	TConnParams,
@@ -42,6 +104,16 @@ export function workflow<
 			TQueues
 		>,
 	) => Promise<unknown>,
+	options: WorkflowOptions<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	> = {},
 ): (
 	c: RunContext<
 		TState,
@@ -55,6 +127,7 @@ export function workflow<
 	>,
 ) => Promise<void> {
 	const workflowInspector = createWorkflowInspectorAdapter();
+	const onError = options.onError;
 
 	async function run(
 		runCtx: RunContext<
@@ -86,6 +159,9 @@ export function workflow<
 				mode: "live",
 				logger: runCtx.log,
 				onHistoryUpdated: workflowInspector.update,
+				onError: onError
+					? async (event) => await onError(runCtx, event)
+					: undefined,
 			},
 		);
 
@@ -103,11 +179,22 @@ export function workflow<
 		try {
 			await handle.result;
 		} catch (error) {
-			runCtx.log.error({
-				msg: "workflow run failed",
+			if (runCtx.abortSignal.aborted) {
+				return;
+			}
+
+			if (shouldRethrowWorkflowError(error)) {
+				runCtx.log.error({
+					msg: "workflow run failed",
+					error: stringifyError(error),
+				});
+				throw error;
+			}
+
+			runCtx.log.warn({
+				msg: "workflow failed and will sleep until woken",
 				error: stringifyError(error),
 			});
-			throw error;
 		} finally {
 			runCtx.abortSignal.removeEventListener("abort", onAbort);
 		}