@rivetkit/workflow-engine 2.1.6 → 2.1.8

package/src/storage.ts

@@ -15,6 +15,7 @@ import {
 import type { EngineDriver, KVWrite } from "./driver.js";
 import {
 	buildEntryMetadataKey,
+	buildEntryMetadataPrefix,
 	buildHistoryKey,
 	buildHistoryPrefix,
 	buildHistoryPrefixAll,
@@ -24,6 +25,7 @@ import {
 	buildWorkflowOutputKey,
 	buildWorkflowStateKey,
 	compareKeys,
+	parseEntryMetadataKey,
 	parseNameKey,
 } from "./keys.js";
 import { isLocationPrefix, locationToKey } from "./location.js";
@@ -129,9 +131,7 @@ export function getOrCreateMetadata(
 /**
  * Load storage from the driver.
  */
-export async function loadStorage(
-	driver: EngineDriver,
-): Promise<Storage> {
+export async function loadStorage(driver: EngineDriver): Promise<Storage> {
 	const storage = createStorage();
 
 	// Load name registry
@@ -155,6 +155,16 @@ export async function loadStorage(
 		storage.history.entries.set(key, parsed);
 	}
 
+	// Load entry metadata so observers can reconstruct workflow state after
+	// the actor wakes and rebuilds storage from persisted history.
+	const metadataEntries = await driver.list(buildEntryMetadataPrefix());
+	for (const entry of metadataEntries) {
+		const entryId = parseEntryMetadataKey(entry.key);
+		const metadata = deserializeEntryMetadata(entry.value);
+		metadata.dirty = false;
+		storage.entryMetadata.set(entryId, metadata);
+	}
+
 	// Load workflow state
 	const stateValue = await driver.get(buildWorkflowStateKey());
 	if (stateValue) {
@@ -207,12 +217,25 @@ export async function loadMetadata(
 }
 
 /**
- * Flush all dirty data to the driver.
+ * Pending deletions collected by collectLoopPruning to be included
+ * in the next flush alongside the state write.
+ */
+export interface PendingDeletions {
+	prefixes: Uint8Array[];
+	keys: Uint8Array[];
+	ranges: { start: Uint8Array; end: Uint8Array }[];
+}
+
+/**
+ * Flush all dirty data to the driver. Optionally includes pending
+ * deletions so that history pruning happens alongside the
+ * state write.
  */
 export async function flush(
 	storage: Storage,
 	driver: EngineDriver,
 	onHistoryUpdated?: () => void,
+	pendingDeletions?: PendingDeletions,
 ): Promise<void> {
 	const writes: KVWrite[] = [];
 	let historyUpdated = false;
@@ -293,6 +316,25 @@ export async function flush(
 		await driver.batch(writes);
 	}
 
+	// Apply pending deletions after the batch write. These are collected
+	// by collectLoopPruning so pruning happens alongside the state write.
+	if (pendingDeletions) {
+		const deleteOps: Promise<void>[] = [];
+		for (const prefix of pendingDeletions.prefixes) {
+			deleteOps.push(driver.deletePrefix(prefix));
+		}
+		for (const range of pendingDeletions.ranges) {
+			deleteOps.push(driver.deleteRange(range.start, range.end));
+		}
+		for (const key of pendingDeletions.keys) {
+			deleteOps.push(driver.delete(key));
+		}
+		if (deleteOps.length > 0) {
+			await Promise.all(deleteOps);
+			historyUpdated = true;
+		}
+	}
+
 	// Update flushed tracking after successful write
 	storage.flushedNameCount = storage.nameRegistry.length;
 	storage.flushedState = storage.state;
@@ -314,30 +356,41 @@ export async function deleteEntriesWithPrefix(
 	prefixLocation: Location,
 	onHistoryUpdated?: () => void,
 ): Promise<void> {
-	// Collect entry IDs for metadata cleanup
-	const entryIds: string[] = [];
+	const deletions = collectDeletionsForPrefix(storage, prefixLocation);
+
+	// Apply deletions to driver
+	await driver.deletePrefix(deletions.prefixes[0]!);
+	await Promise.all(deletions.keys.map((key) => driver.delete(key)));
+
+	if (deletions.keys.length > 0 && onHistoryUpdated) {
+		onHistoryUpdated();
+	}
+}
+
+/**
+ * Remove entries matching a location prefix from memory and collect
+ * the driver-level deletion operations. The returned PendingDeletions
+ * can be applied immediately or batched with a flush.
+ */
+export function collectDeletionsForPrefix(
+	storage: Storage,
+	prefixLocation: Location,
+): PendingDeletions {
+	const pending: PendingDeletions = {
+		prefixes: [buildHistoryPrefix(prefixLocation)],
+		keys: [],
+		ranges: [],
+	};
 
-	// Collect entries to delete and their IDs
 	for (const [key, entry] of storage.history.entries) {
-		// Check if the entry's location starts with the prefix location
 		if (isLocationPrefix(prefixLocation, entry.location)) {
-			entryIds.push(entry.id);
+			pending.keys.push(buildEntryMetadataKey(entry.id));
 			storage.entryMetadata.delete(entry.id);
 			storage.history.entries.delete(key);
 		}
 	}
 
-	// Delete entries from driver using binary prefix
-	await driver.deletePrefix(buildHistoryPrefix(prefixLocation));
-
-	// Delete metadata from driver in parallel
-	await Promise.all(
-		entryIds.map((id) => driver.delete(buildEntryMetadataKey(id))),
-	);
-
-	if (entryIds.length > 0 && onHistoryUpdated) {
-		onHistoryUpdated();
-	}
+	return pending;
 }
 
 /**

package/src/testing.ts

@@ -33,7 +33,11 @@ class InMemoryWorkflowMessageDriver implements WorkflowMessageDriver {
 				: undefined;
 		const selected: Array<{ message: Message; index: number }> = [];
 
-		for (let i = 0; i < this.#messages.length && selected.length < limitedCount; i++) {
+		for (
+			let i = 0;
+			i < this.#messages.length && selected.length < limitedCount;
+			i++
+		) {
 			const message = this.#messages[i];
 			if (nameSet && !nameSet.has(message.name)) {
 				continue;
@@ -64,7 +68,9 @@ class InMemoryWorkflowMessageDriver implements WorkflowMessageDriver {
 	}
 
 	async completeMessage(messageId: string): Promise<void> {
-		const index = this.#messages.findIndex((message) => message.id === messageId);
+		const index = this.#messages.findIndex(
+			(message) => message.id === messageId,
+		);
 		if (index !== -1) {
 			this.#messages.splice(index, 1);
 		}
@@ -78,7 +84,8 @@ class InMemoryWorkflowMessageDriver implements WorkflowMessageDriver {
 			throw new EvictedError();
 		}
 
-		const nameSet = messageNames.length > 0 ? new Set(messageNames) : undefined;
+		const nameSet =
+			messageNames.length > 0 ? new Set(messageNames) : undefined;
 		if (
 			this.#messages.some((message) =>
 				nameSet ? nameSet.has(message.name) : true,
@@ -103,7 +110,9 @@ class InMemoryWorkflowMessageDriver implements WorkflowMessageDriver {
 					waiter.reject(new EvictedError());
 				},
 			};
-			abortSignal.addEventListener("abort", waiter.onAbort, { once: true });
+			abortSignal.addEventListener("abort", waiter.onAbort, {
+				once: true,
+			});
 			this.#waiters.add(waiter);
 		});
 	}
@@ -173,6 +182,18 @@ export class InMemoryDriver implements EngineDriver {
 		}
 	}
 
+	async deleteRange(start: Uint8Array, end: Uint8Array): Promise<void> {
+		await sleep(this.latency);
+		for (const [hexKey, entry] of this.kv) {
+			if (
+				compareKeys(entry.key, start) >= 0 &&
+				compareKeys(entry.key, end) < 0
+			) {
+				this.kv.delete(hexKey);
+			}
+		}
+	}
+
 	async list(prefix: Uint8Array): Promise<KVEntry[]> {
 		await sleep(this.latency);
 		const results: KVEntry[] = [];

package/src/types.ts

@@ -225,8 +225,7 @@ export interface WorkflowQueueNextOptions {
  * Options for receiving a batch of queue messages in workflows.
  */
 export interface WorkflowQueueNextBatchOptions
-	extends WorkflowQueueNextOptions
-{
+	extends WorkflowQueueNextOptions {
 	/** Maximum number of messages to receive. Defaults to 1. */
 	count?: number;
 }
@@ -297,6 +296,47 @@ export interface WorkflowError {
 	metadata?: Record<string, unknown>;
 }
 
+/**
+ * Error event emitted while a workflow is running.
+ */
+export interface WorkflowStepErrorEvent {
+	workflowId: string;
+	stepName: string;
+	attempt: number;
+	maxRetries: number;
+	remainingRetries: number;
+	willRetry: boolean;
+	retryDelay?: number;
+	retryAt?: number;
+	error: WorkflowError;
+}
+
+/**
+ * Error event emitted when a rollback handler fails.
+ */
+export interface WorkflowRollbackErrorEvent {
+	workflowId: string;
+	stepName: string;
+	error: WorkflowError;
+}
+
+/**
+ * Error event emitted for workflow-level failures outside individual steps.
+ */
+export interface WorkflowRunErrorEvent {
+	workflowId: string;
+	error: WorkflowError;
+}
+
+export type WorkflowErrorEvent =
+	| { step: WorkflowStepErrorEvent }
+	| { rollback: WorkflowRollbackErrorEvent }
+	| { workflow: WorkflowRunErrorEvent };
+
+export type WorkflowErrorHandler = (
+	event: WorkflowErrorEvent,
+) => void | Promise<void>;
+
 /**
  * Complete storage state for a workflow.
  */
@@ -384,15 +424,11 @@ export type LoopIterationResult<S, T> = Promise<
 export interface LoopConfig<S, T> {
 	name: string;
 	state?: S;
-	run: (
-		ctx: WorkflowContextInterface,
-		state: S,
-	) => LoopIterationResult<S, T>;
-	commitInterval?: number;
-	/** Trim loop history every N iterations. Defaults to commitInterval or 20. */
-	historyEvery?: number;
-	/** Retain the last N iterations of history. Defaults to commitInterval or 20. */
-	historyKeep?: number;
+	run: (ctx: WorkflowContextInterface, state: S) => LoopIterationResult<S, T>;
+	/** Prune old loop iterations every N iterations. Default: 20. */
+	historyPruneInterval?: number;
+	/** Number of past iterations to retain when pruning. Defaults to historyPruneInterval. */
+	historySize?: number;
 }
 
 /**
@@ -458,6 +494,7 @@ export interface RunWorkflowOptions {
 	mode?: WorkflowRunMode;
 	logger?: Logger;
 	onHistoryUpdated?: (history: WorkflowHistorySnapshot) => void;
+	onError?: WorkflowErrorHandler;
 }
 
 export type WorkflowFunction<TInput = unknown, TOutput = unknown> = (