rivetkit 2.0.2 → 2.0.4-rc.1

package/src/workflow/driver.ts

@@ -0,0 +1,194 @@
+import type { RunContext } from "@/actor/contexts/run";
+import type { AnyActorInstance } from "@/actor/instance/mod";
+import { makeWorkflowKey, workflowStoragePrefix } from "@/actor/instance/keys";
+import type {
+	EngineDriver,
+	KVEntry,
+	KVWrite,
+	Message,
+	WorkflowMessageDriver,
+} from "@rivetkit/workflow-engine";
+
+const WORKFLOW_STORAGE_PREFIX = workflowStoragePrefix();
+
+function stripWorkflowKey(prefixed: Uint8Array): Uint8Array {
+	return prefixed.slice(WORKFLOW_STORAGE_PREFIX.length);
+}
+
+class ActorWorkflowMessageDriver implements WorkflowMessageDriver {
+	#actor: AnyActorInstance;
+	#runCtx: RunContext<any, any, any, any, any, any, any, any>;
+
+	constructor(
+		actor: AnyActorInstance,
+		runCtx: RunContext<any, any, any, any, any, any, any, any>,
+	) {
+		this.#actor = actor;
+		this.#runCtx = runCtx;
+	}
+
+	async addMessage(message: Message): Promise<void> {
+		await this.#runCtx.keepAwake(
+			this.#actor.queueManager.enqueue(message.name, message.data),
+		);
+	}
+
+	async receiveMessages(opts: {
+		names?: readonly string[];
+		count: number;
+		completable: boolean;
+	}): Promise<Message[]> {
+		const messages = await this.#runCtx.keepAwake(
+			this.#actor.queueManager.receive(
+				opts.names && opts.names.length > 0 ? [...opts.names] : undefined,
+				opts.count,
+				0,
+				undefined,
+				opts.completable,
+			),
+		);
+		return messages.map((message) => ({
+			id: message.id.toString(),
+			name: message.name,
+			data: message.body,
+			sentAt: message.createdAt,
+			...(opts.completable
+				? {
+						complete: async (response?: unknown) => {
+							await this.#runCtx.keepAwake(
+								this.#actor.queueManager.completeMessage(
+									message,
+									response,
+								),
+							);
+						},
+					}
+				: {}),
+		}));
+	}
+
+	async completeMessage(messageId: string, response?: unknown): Promise<void> {
+		let parsedId: bigint;
+		try {
+			parsedId = BigInt(messageId);
+		} catch {
+			return;
+		}
+
+		await this.#runCtx.keepAwake(
+			this.#actor.queueManager.completeMessageById(parsedId, response),
+		);
+	}
+}
+
+export class ActorWorkflowDriver implements EngineDriver {
+	readonly workerPollInterval = 100;
+	readonly messageDriver: WorkflowMessageDriver;
+	#actor: AnyActorInstance;
+	#runCtx: RunContext<any, any, any, any, any, any, any, any>;
+
+	constructor(
+		actor: AnyActorInstance,
+		runCtx: RunContext<any, any, any, any, any, any, any, any>,
+	) {
+		this.#actor = actor;
+		this.#runCtx = runCtx;
+		this.messageDriver = new ActorWorkflowMessageDriver(actor, runCtx);
+	}
+
+	async get(key: Uint8Array): Promise<Uint8Array | null> {
+		const [value] = await this.#runCtx.keepAwake(
+			this.#actor.driver.kvBatchGet(this.#actor.id, [
+				makeWorkflowKey(key),
+			]),
+		);
+		return value ?? null;
+	}
+
+	async set(key: Uint8Array, value: Uint8Array): Promise<void> {
+		await this.#runCtx.keepAwake(
+			this.#actor.driver.kvBatchPut(this.#actor.id, [
+				[makeWorkflowKey(key), value],
+			]),
+		);
+	}
+
+	async delete(key: Uint8Array): Promise<void> {
+		await this.#runCtx.keepAwake(
+			this.#actor.driver.kvBatchDelete(this.#actor.id, [
+				makeWorkflowKey(key),
+			]),
+		);
+	}
+
+	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;
+		}
+		await this.#runCtx.keepAwake(
+			this.#actor.driver.kvBatchDelete(
+				this.#actor.id,
+				entries.map(([key]) => key),
+			),
+		);
+	}
+
+	async list(prefix: Uint8Array): Promise<KVEntry[]> {
+		const entries = await this.#runCtx.keepAwake(
+			this.#actor.driver.kvListPrefix(
+				this.#actor.id,
+				makeWorkflowKey(prefix),
+			),
+		);
+		return entries.map(([key, value]) => ({
+			key: stripWorkflowKey(key),
+			value,
+		}));
+	}
+
+	async batch(writes: KVWrite[]): Promise<void> {
+		if (writes.length === 0) return;
+
+		// Flush actor state together with workflow state to ensure atomicity.
+		// If the server crashes after workflow flush, actor state must also be persisted.
+		await this.#runCtx.keepAwake(
+			Promise.all([
+				this.#actor.driver.kvBatchPut(
+					this.#actor.id,
+					writes.map(({ key, value }) => [makeWorkflowKey(key), value]),
+				),
+				this.#actor.stateManager.saveState({
+					immediate: true,
+					allowStoppingState: true,
+				}),
+			]),
+		);
+	}
+
+	async setAlarm(_workflowId: string, wakeAt: number): Promise<void> {
+		await this.#runCtx.keepAwake(
+			this.#actor.driver.setAlarm(this.#actor, wakeAt),
+		);
+	}
+
+	async clearAlarm(_workflowId: string): Promise<void> {
+		// No dedicated clear alarm support in actor drivers.
+		return;
+	}
+
+	waitForMessages(
+		messageNames: string[],
+		abortSignal: AbortSignal,
+	): Promise<void> {
+		return this.#actor.queueManager.waitForNames(
+			messageNames.length > 0 ? messageNames : undefined,
+			abortSignal,
+		);
+	}
+}

package/src/workflow/inspector.ts

@@ -0,0 +1,268 @@
+import * as cbor from "cbor-x";
+import { createNanoEvents } from "nanoevents";
+import type {
+	BranchStatus,
+	BranchStatusType,
+	EntryKind,
+	EntryStatus,
+	Location,
+	SleepState,
+	WorkflowHistoryEntry,
+	WorkflowHistorySnapshot,
+	WorkflowEntryMetadataSnapshot,
+} from "@rivetkit/workflow-engine";
+import { encodeWorkflowHistoryTransport } from "@/inspector/transport";
+import type * as inspectorSchema from "@/schemas/actor-inspector/mod";
+import * as transport from "@/schemas/transport/mod";
+import { assertUnreachable, bufferToArrayBuffer } from "@/utils";
+
+export interface WorkflowInspectorAdapter {
+	getHistory: () => inspectorSchema.WorkflowHistory | null;
+	onHistoryUpdated: (
+		listener: (history: inspectorSchema.WorkflowHistory) => void,
+	) => () => void;
+}
+
+export function createWorkflowInspectorAdapter(): {
+	adapter: WorkflowInspectorAdapter;
+	update: (snapshot: WorkflowHistorySnapshot) => void;
+} {
+	const emitter = createNanoEvents<{
+		updated: (history: inspectorSchema.WorkflowHistory) => void;
+	}>();
+	let history: inspectorSchema.WorkflowHistory | null = null;
+
+	const adapter: WorkflowInspectorAdapter = {
+		getHistory: () => history,
+		onHistoryUpdated: (listener) => emitter.on("updated", listener),
+	};
+
+	const update = (snapshot: WorkflowHistorySnapshot) => {
+		const transportHistory = toWorkflowHistory(snapshot);
+		const next = encodeWorkflowHistoryTransport(transportHistory);
+		history = next;
+		emitter.emit("updated", next);
+	};
+
+	return { adapter, update };
+}
+
+function encodeCbor(value: unknown): ArrayBuffer {
+	return bufferToArrayBuffer(cbor.encode(value));
+}
+
+function encodeOptionalCbor(value: unknown): ArrayBuffer | null {
+	if (value === undefined) {
+		return null;
+	}
+	return encodeCbor(value);
+}
+
+function toU64(value: number): bigint {
+	return BigInt(Math.max(0, Math.floor(value)));
+}
+
+function toWorkflowLocation(
+	location: Location,
+): transport.WorkflowLocation {
+	return location.map((segment) => {
+		if (typeof segment === "number") {
+			return { tag: "WorkflowNameIndex", val: segment };
+		}
+		return {
+			tag: "WorkflowLoopIterationMarker",
+			val: {
+				loop: segment.loop,
+				iteration: segment.iteration,
+			},
+		};
+	});
+}
+
+function toWorkflowEntryKind(
+	kind: EntryKind,
+): transport.WorkflowEntryKind {
+	switch (kind.type) {
+		case "step":
+			return {
+				tag: "WorkflowStepEntry",
+				val: {
+					output: encodeOptionalCbor(kind.data.output),
+					error: kind.data.error ?? null,
+				},
+			};
+		case "loop":
+			return {
+				tag: "WorkflowLoopEntry",
+				val: {
+					state: encodeCbor(kind.data.state),
+					iteration: kind.data.iteration,
+					output: encodeOptionalCbor(kind.data.output),
+				},
+			};
+		case "sleep":
+			return {
+				tag: "WorkflowSleepEntry",
+				val: {
+					deadline: toU64(kind.data.deadline),
+					state: toWorkflowSleepState(kind.data.state),
+				},
+			};
+		case "message":
+			return {
+				tag: "WorkflowMessageEntry",
+				val: {
+					name: kind.data.name,
+					messageData: encodeCbor(kind.data.data),
+				},
+			};
+		case "rollback_checkpoint":
+			return {
+				tag: "WorkflowRollbackCheckpointEntry",
+				val: { name: kind.data.name },
+			};
+		case "join":
+			return {
+				tag: "WorkflowJoinEntry",
+				val: { branches: toWorkflowBranchStatusMap(kind.data.branches) },
+			};
+		case "race":
+			return {
+				tag: "WorkflowRaceEntry",
+				val: {
+					winner: kind.data.winner ?? null,
+					branches: toWorkflowBranchStatusMap(kind.data.branches),
+				},
+			};
+		case "removed":
+			return {
+				tag: "WorkflowRemovedEntry",
+				val: {
+					originalType: kind.data.originalType,
+					originalName: kind.data.originalName ?? null,
+				},
+			};
+		default:
+			assertUnreachable(kind as never);
+	}
+}
+
+function toWorkflowEntry(
+	entry: WorkflowHistoryEntry,
+): transport.WorkflowEntry {
+	return {
+		id: entry.id,
+		location: toWorkflowLocation(entry.location),
+		kind: toWorkflowEntryKind(entry.kind),
+	};
+}
+
+function toWorkflowEntryStatus(
+	status: EntryStatus,
+): transport.WorkflowEntryStatus {
+	switch (status) {
+		case "pending":
+			return transport.WorkflowEntryStatus.PENDING;
+		case "running":
+			return transport.WorkflowEntryStatus.RUNNING;
+		case "completed":
+			return transport.WorkflowEntryStatus.COMPLETED;
+		case "failed":
+			return transport.WorkflowEntryStatus.FAILED;
+		case "exhausted":
+			return transport.WorkflowEntryStatus.EXHAUSTED;
+		default:
+			assertUnreachable(status as never);
+	}
+}
+
+function toWorkflowSleepState(
+	state: SleepState,
+): transport.WorkflowSleepState {
+	switch (state) {
+		case "pending":
+			return transport.WorkflowSleepState.PENDING;
+		case "completed":
+			return transport.WorkflowSleepState.COMPLETED;
+		case "interrupted":
+			return transport.WorkflowSleepState.INTERRUPTED;
+		default:
+			assertUnreachable(state as never);
+	}
+}
+
+function toWorkflowBranchStatusType(
+	status: BranchStatusType,
+): transport.WorkflowBranchStatusType {
+	switch (status) {
+		case "pending":
+			return transport.WorkflowBranchStatusType.PENDING;
+		case "running":
+			return transport.WorkflowBranchStatusType.RUNNING;
+		case "completed":
+			return transport.WorkflowBranchStatusType.COMPLETED;
+		case "failed":
+			return transport.WorkflowBranchStatusType.FAILED;
+		case "cancelled":
+			return transport.WorkflowBranchStatusType.CANCELLED;
+		default:
+			assertUnreachable(status as never);
+	}
+}
+
+function toWorkflowBranchStatus(
+	status: BranchStatus,
+): transport.WorkflowBranchStatus {
+	return {
+		status: toWorkflowBranchStatusType(status.status),
+		output: encodeOptionalCbor(status.output),
+		error: status.error ?? null,
+	};
+}
+
+function toWorkflowBranchStatusMap(
+	branches: Record<string, BranchStatus>,
+): ReadonlyMap<string, transport.WorkflowBranchStatus> {
+	return new Map(
+		Object.entries(branches).map(([name, status]) => [
+			name,
+			toWorkflowBranchStatus(status),
+		]),
+	);
+}
+
+function toWorkflowEntryMetadata(
+	metadata: WorkflowEntryMetadataSnapshot,
+): transport.WorkflowEntryMetadata {
+	return {
+		status: toWorkflowEntryStatus(metadata.status),
+		error: metadata.error ?? null,
+		attempts: metadata.attempts,
+		lastAttemptAt: toU64(metadata.lastAttemptAt),
+		createdAt: toU64(metadata.createdAt),
+		completedAt:
+			metadata.completedAt === undefined
+				? null
+				: toU64(metadata.completedAt),
+		rollbackCompletedAt:
+			metadata.rollbackCompletedAt === undefined
+				? null
+				: toU64(metadata.rollbackCompletedAt),
+		rollbackError: metadata.rollbackError ?? null,
+	};
+}
+
+function toWorkflowHistory(
+	snapshot: WorkflowHistorySnapshot,
+): transport.WorkflowHistory {
+	const entryMetadata = new Map<string, transport.WorkflowEntryMetadata>();
+	for (const [id, metadata] of snapshot.entryMetadata) {
+		entryMetadata.set(id, toWorkflowEntryMetadata(metadata));
+	}
+
+	return {
+		nameRegistry: snapshot.nameRegistry,
+		entries: snapshot.entries.map((entry) => toWorkflowEntry(entry)),
+		entryMetadata,
+	};
+}

package/src/workflow/mod.ts

@@ -0,0 +1,128 @@
+import { ACTOR_CONTEXT_INTERNAL_SYMBOL } from "@/actor/contexts/base/actor";
+import type { RunContext } from "@/actor/contexts/run";
+import type { AnyDatabaseProvider } from "@/actor/database";
+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 invariant from "invariant";
+import { ActorWorkflowContext } from "./context";
+import { ActorWorkflowDriver } from "./driver";
+import { createWorkflowInspectorAdapter } from "./inspector";
+
+export { Loop } from "@rivetkit/workflow-engine";
+export {
+	ActorWorkflowContext,
+	type WorkflowBranchContextOf,
+	type WorkflowContextOf,
+	type WorkflowLoopContextOf,
+	type WorkflowStepContextOf,
+} from "./context";
+
+export function workflow<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+>(
+	fn: (
+		ctx: ActorWorkflowContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	) => Promise<unknown>,
+): (
+	c: RunContext<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>,
+) => Promise<void> {
+	const workflowInspector = createWorkflowInspectorAdapter();
+
+	async function run(
+		runCtx: RunContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	): Promise<void> {
+		const actor = (
+			runCtx as unknown as {
+				[ACTOR_CONTEXT_INTERNAL_SYMBOL]?: AnyActorInstance;
+			}
+		)[ACTOR_CONTEXT_INTERNAL_SYMBOL];
+		invariant(actor, "workflow() requires an actor instance");
+
+		const driver = new ActorWorkflowDriver(actor, runCtx);
+
+		const handle = runWorkflow(
+			actor.id,
+			async (ctx) => await fn(new ActorWorkflowContext(ctx, runCtx)),
+			undefined,
+			driver,
+			{
+				mode: "live",
+				logger: runCtx.log,
+				onHistoryUpdated: workflowInspector.update,
+			},
+		);
+
+		const onAbort = () => {
+			handle.evict();
+		};
+		if (runCtx.abortSignal.aborted) {
+			onAbort();
+		} else {
+			runCtx.abortSignal.addEventListener("abort", onAbort, {
+				once: true,
+			});
+		}
+
+		try {
+			await handle.result;
+		} catch (error) {
+			runCtx.log.error({
+				msg: "workflow run failed",
+				error: stringifyError(error),
+			});
+			throw error;
+		} finally {
+			runCtx.abortSignal.removeEventListener("abort", onAbort);
+		}
+	}
+
+	const runWithConfig = run as typeof run & {
+		[RUN_FUNCTION_CONFIG_SYMBOL]?: {
+			icon?: string;
+			inspector?: { workflow: typeof workflowInspector.adapter };
+		};
+	};
+	runWithConfig[RUN_FUNCTION_CONFIG_SYMBOL] = {
+		icon: "diagram-project",
+		inspector: { workflow: workflowInspector.adapter },
+	};
+
+	return runWithConfig;
+}