@rivetkit/workflow-engine 2.1.5 → 2.1.6-rc.1

package/src/error-utils.ts

@@ -0,0 +1,87 @@
+import type { WorkflowError, WorkflowErrorEvent } from "./types.js";
+
+const WORKFLOW_ERROR_REPORTED_SYMBOL = Symbol("workflow.error.reported");
+
+/**
+ * Extract structured error information from an error.
+ */
+export function extractErrorInfo(error: unknown): WorkflowError {
+	if (error instanceof Error) {
+		const result: WorkflowError = {
+			name: error.name,
+			message: error.message,
+			stack: error.stack,
+		};
+
+		const metadata: Record<string, unknown> = {};
+		for (const key of Object.keys(error)) {
+			if (key !== "name" && key !== "message" && key !== "stack") {
+				const value = (error as unknown as Record<string, unknown>)[
+					key
+				];
+				if (
+					typeof value === "string" ||
+					typeof value === "number" ||
+					typeof value === "boolean" ||
+					value === null
+				) {
+					metadata[key] = value;
+				}
+			}
+		}
+		if (Object.keys(metadata).length > 0) {
+			result.metadata = metadata;
+		}
+
+		return result;
+	}
+
+	return {
+		name: "Error",
+		message: String(error),
+	};
+}
+
+/**
+ * Mark an error after it has been reported to the error hook.
+ */
+export function markErrorReported<T extends Error>(error: T): T {
+	(
+		error as T & {
+			[WORKFLOW_ERROR_REPORTED_SYMBOL]?: boolean;
+		}
+	)[WORKFLOW_ERROR_REPORTED_SYMBOL] = true;
+	return error;
+}
+
+/**
+ * Check if an error was already reported to the error hook.
+ */
+export function isErrorReported(error: unknown): boolean {
+	if (!(error instanceof Error)) {
+		return false;
+	}
+
+	return Boolean(
+		(
+			error as Error & {
+				[WORKFLOW_ERROR_REPORTED_SYMBOL]?: boolean;
+			}
+		)[WORKFLOW_ERROR_REPORTED_SYMBOL],
+	);
+}
+
+/**
+ * Return the outer tag name for a workflow error event.
+ */
+export function getErrorEventTag(
+	event: WorkflowErrorEvent,
+): "step" | "rollback" | "workflow" {
+	if ("step" in event) {
+		return "step";
+	}
+	if ("rollback" in event) {
+		return "rollback";
+	}
+	return "workflow";
+}

package/src/errors.ts

@@ -115,6 +115,7 @@ export class StepFailedError extends Error {
 		public readonly stepName: string,
 		public readonly originalError: unknown,
 		public readonly attempts: number,
+		public readonly retryAt: number,
 	) {
 		super(`Step "${stepName}" failed (attempt ${attempts})`);
 		this.name = "StepFailedError";

package/src/index.ts

@@ -4,9 +4,7 @@ import type { Logger } from "pino";
 
 // Context
 export {
-	DEFAULT_LOOP_COMMIT_INTERVAL,
-	DEFAULT_LOOP_HISTORY_EVERY,
-	DEFAULT_LOOP_HISTORY_KEEP,
+	DEFAULT_LOOP_HISTORY_PRUNE_INTERVAL,
 	DEFAULT_MAX_RETRIES,
 	DEFAULT_RETRY_BACKOFF_BASE,
 	DEFAULT_RETRY_BACKOFF_MAX,
@@ -31,6 +29,7 @@ export {
 	StepExhaustedError,
 	StepFailedError,
 } from "./errors.js";
+export { extractErrorInfo } from "./error-utils.js";
 
 // Location utilities
 export {
@@ -95,8 +94,12 @@ export type {
 	StepEntry,
 	Storage,
 	WorkflowContextInterface,
+	WorkflowError,
+	WorkflowErrorEvent,
+	WorkflowErrorHandler,
 	WorkflowFunction,
 	WorkflowHandle,
+	WorkflowRollbackErrorEvent,
 	WorkflowQueue,
 	WorkflowQueueMessage,
 	WorkflowQueueNextBatchOptions,
@@ -104,7 +107,9 @@ export type {
 	WorkflowMessageDriver,
 	WorkflowResult,
 	WorkflowRunMode,
+	WorkflowRunErrorEvent,
 	WorkflowState,
+	WorkflowStepErrorEvent,
 } from "./types.js";
 
 // Loop result helpers
@@ -132,11 +137,20 @@ import { type RollbackAction, WorkflowContextImpl } from "./context.js";
 // Main workflow runner
 import type { EngineDriver } from "./driver.js";
 import {
+	extractErrorInfo,
+	getErrorEventTag,
+	isErrorReported,
+	markErrorReported,
+} from "./error-utils.js";
+import {
+	CriticalError,
 	EvictedError,
 	MessageWaitError,
 	RollbackCheckpointError,
+	RollbackError,
 	RollbackStopError,
 	SleepError,
+	StepExhaustedError,
 	StepFailedError,
 } from "./errors.js";
 import {
@@ -157,6 +171,7 @@ import type {
 	RollbackContextInterface,
 	RunWorkflowOptions,
 	Storage,
+	WorkflowErrorEvent,
 	WorkflowHistorySnapshot,
 	WorkflowFunction,
 	WorkflowHandle,
@@ -252,6 +267,7 @@ async function executeRollback<TInput, TOutput>(
 	abortController: AbortController,
 	storage: Storage,
 	historyNotifier?: HistoryNotifier,
+	onError?: RunWorkflowOptions["onError"],
 	logger?: Logger,
 ): Promise<void> {
 	const rollbackActions: RollbackAction[] = [];
@@ -266,6 +282,7 @@ async function executeRollback<TInput, TOutput>(
 		rollbackActions,
 		false,
 		historyNotifier,
+		onError,
 		logger,
 	);
 
@@ -299,6 +316,7 @@ async function executeRollback<TInput, TOutput>(
 			continue;
 		}
 
+		let rollbackEvent: WorkflowErrorEvent | undefined;
 		try {
 			await awaitWithEviction(
 				action.rollback(rollbackContext, action.output),
@@ -312,14 +330,45 @@ async function executeRollback<TInput, TOutput>(
 			}
 			metadata.rollbackError =
 				error instanceof Error ? error.message : String(error);
+			if (onError) {
+					rollbackEvent = {
+						rollback: {
+							workflowId,
+							stepName: action.name,
+							error: extractErrorInfo(error),
+						},
+					};
+				}
+			if (error instanceof Error) {
+				markErrorReported(error);
+			}
 			throw error;
 		} finally {
 			metadata.dirty = true;
 			await flush(storage, driver, historyNotifier);
+			if (rollbackEvent && onError) {
+				await notifyError(onError, logger, rollbackEvent);
+			}
 		}
 	}
 }
 
+async function notifyError(
+	onError: NonNullable<RunWorkflowOptions["onError"]>,
+	logger: Logger | undefined,
+	event: WorkflowErrorEvent,
+): Promise<void> {
+	try {
+		await onError(event);
+	} catch (error) {
+		logger?.warn({
+			msg: "workflow error hook failed",
+			hookEventType: getErrorEventTag(event),
+			error: extractErrorInfo(error),
+		});
+	}
+}
+
 async function setSleepState<TOutput>(
 	storage: Storage,
 	driver: EngineDriver,
@@ -364,12 +413,12 @@ async function setRetryState<TOutput>(
 	storage: Storage,
 	driver: EngineDriver,
 	workflowId: string,
+	retryAt: number,
 	historyNotifier?: HistoryNotifier,
 ): Promise<WorkflowResult<TOutput>> {
 	storage.state = "sleeping";
 	await flush(storage, driver, historyNotifier);
 
-	const retryAt = Date.now() + 100;
 	await driver.setAlarm(workflowId, retryAt);
 
 	return { state: "sleeping", sleepUntil: retryAt };
@@ -437,6 +486,7 @@ async function executeLiveWorkflow<TInput, TOutput>(
 	abortController: AbortController,
 	runtime: LiveRuntime,
 	onHistoryUpdated?: (history: WorkflowHistorySnapshot) => void,
+	onError?: RunWorkflowOptions["onError"],
 	logger?: Logger,
 ): Promise<WorkflowResult<TOutput>> {
 	let lastResult: WorkflowResult<TOutput> | undefined;
@@ -450,6 +500,7 @@ async function executeLiveWorkflow<TInput, TOutput>(
 			messageDriver,
 			abortController,
 			onHistoryUpdated,
+			onError,
 			logger,
 		);
 		lastResult = result;
@@ -549,6 +600,7 @@ export function runWorkflow<TInput, TOutput>(
 					abortController,
 					liveRuntime,
 					options.onHistoryUpdated,
+					options.onError,
 					logger,
 				)
 			: executeWorkflow(
@@ -559,6 +611,7 @@ export function runWorkflow<TInput, TOutput>(
 					messageDriver,
 					abortController,
 					options.onHistoryUpdated,
+					options.onError,
 					logger,
 				);
 
@@ -686,6 +739,7 @@ async function executeWorkflow<TInput, TOutput>(
 	messageDriver: WorkflowMessageDriver,
 	abortController: AbortController,
 	onHistoryUpdated?: (history: WorkflowHistorySnapshot) => void,
+	onError?: RunWorkflowOptions["onError"],
 	logger?: Logger,
 ): Promise<WorkflowResult<TOutput>> {
 	const storage = await loadStorage(driver);
@@ -739,6 +793,7 @@ async function executeWorkflow<TInput, TOutput>(
 				abortController,
 				storage,
 				historyNotifier,
+				onError,
 				logger,
 			);
 		} catch (error) {
@@ -771,6 +826,7 @@ async function executeWorkflow<TInput, TOutput>(
 		undefined,
 		false,
 		historyNotifier,
+		onError,
 		logger,
 	);
 
@@ -815,12 +871,21 @@ async function executeWorkflow<TInput, TOutput>(
 				storage,
 				driver,
 				workflowId,
+				error.retryAt,
 				historyNotifier,
 			);
 		}
 
 		if (error instanceof RollbackCheckpointError) {
 			await setFailedState(storage, driver, error, historyNotifier);
+			if (onError && !isErrorReported(error)) {
+				await notifyError(onError, logger, {
+					workflow: {
+						workflowId,
+						error: extractErrorInfo(error),
+					},
+				});
+			}
 			throw error;
 		}
 
@@ -839,6 +904,7 @@ async function executeWorkflow<TInput, TOutput>(
 				abortController,
 				storage,
 				historyNotifier,
+				onError,
 				logger,
 			);
 		} catch (rollbackError) {
@@ -850,59 +916,22 @@ async function executeWorkflow<TInput, TOutput>(
 
 		storage.state = "failed";
 		await flush(storage, driver, historyNotifier);
-
-		throw error;
-	}
-}
-
-/**
- * Extract structured error information from an error.
- */
-function extractErrorInfo(error: unknown): {
-	name: string;
-	message: string;
-	stack?: string;
-	metadata?: Record<string, unknown>;
-} {
-	if (error instanceof Error) {
-		const result: {
-			name: string;
-			message: string;
-			stack?: string;
-			metadata?: Record<string, unknown>;
-		} = {
-			name: error.name,
-			message: error.message,
-			stack: error.stack,
-		};
-
-		// Extract custom properties from error
-		const metadata: Record<string, unknown> = {};
-		for (const key of Object.keys(error)) {
-			if (key !== "name" && key !== "message" && key !== "stack") {
-				const value = (error as unknown as Record<string, unknown>)[
-					key
-				];
-				// Only include serializable values
-				if (
-					typeof value === "string" ||
-					typeof value === "number" ||
-					typeof value === "boolean" ||
-					value === null
-				) {
-					metadata[key] = value;
-				}
+		if (onError && !isErrorReported(error)) {
+			await notifyError(onError, logger, {
+				workflow: {
+					workflowId,
+					error: extractErrorInfo(error),
+				},
+			});
+			if (
+				error instanceof CriticalError ||
+				error instanceof RollbackError ||
+				error instanceof StepExhaustedError
+			) {
+				markErrorReported(error);
 			}
 		}
-		if (Object.keys(metadata).length > 0) {
-			result.metadata = metadata;
-		}
 
-		return result;
+		throw error;
 	}
-
-	return {
-		name: "Error",
-		message: String(error),
-	};
 }

package/src/keys.ts

@@ -142,6 +142,32 @@ export function buildHistoryPrefix(location: Location): Uint8Array {
 	return pack([KEY_PREFIX.HISTORY, ...locationToTupleElements(location)]);
 }
 
+/**
+ * Build a key range for loop iteration history entries.
+ * Range: [2, ...locationSegments, [loopSegment, fromIteration]] to
+ *        [2, ...locationSegments, [loopSegment, toIteration]])
+ */
+export function buildLoopIterationRange(
+	loopLocation: Location,
+	loopSegment: number,
+	fromIteration: number,
+	toIteration: number,
+): { start: Uint8Array; end: Uint8Array } {
+	const loopLocationSegments = locationToTupleElements(loopLocation);
+	return {
+		start: pack([
+			KEY_PREFIX.HISTORY,
+			...loopLocationSegments,
+			[loopSegment, fromIteration],
+		]),
+		end: pack([
+			KEY_PREFIX.HISTORY,
+			...loopLocationSegments,
+			[loopSegment, toIteration],
+		]),
+	};
+}
+
 /**
  * Build a prefix for listing all history entries.
  * Prefix: [2]

package/src/location.ts

@@ -107,7 +107,10 @@ export function isLocationPrefix(
 		const prefixSegment = prefix[i];
 		const locationSegment = location[i];
 
-		if (typeof prefixSegment === "number" && typeof locationSegment === "number") {
+		if (
+			typeof prefixSegment === "number" &&
+			typeof locationSegment === "number"
+		) {
 			if (prefixSegment !== locationSegment) {
 				return false;
 			}

package/src/storage.ts

@@ -129,9 +129,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
@@ -207,12 +205,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 +304,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 +344,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;
 }
 
 /**
@@ -361,4 +402,4 @@ export function setEntry(
 ): void {
 	const key = locationToKey(storage, location);
 	storage.history.entries.set(key, entry);
-}
+}

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[] = [];