@rivetkit/workflow-engine 2.1.6 → 2.1.8

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,20 +137,32 @@ 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 {
+	buildEntryMetadataKey,
 	buildEntryMetadataPrefix,
+	buildHistoryKey,
 	buildWorkflowErrorKey,
 	buildWorkflowInputKey,
 	buildWorkflowOutputKey,
 	buildWorkflowStateKey,
 } from "./keys.js";
+import { isLocationPrefix } from "./location.js";
 import {
 	createHistorySnapshot,
 	flush,
@@ -154,9 +171,12 @@ import {
 	loadStorage,
 } from "./storage.js";
 import type {
+	Entry,
+	EntryMetadata,
 	RollbackContextInterface,
 	RunWorkflowOptions,
 	Storage,
+	WorkflowErrorEvent,
 	WorkflowHistorySnapshot,
 	WorkflowFunction,
 	WorkflowHandle,
@@ -184,6 +204,12 @@ interface LiveRuntime {
 
 type HistoryNotifier = (() => void) | undefined;
 
+type ReplayEntryRecord = {
+	key: string;
+	entry: Entry;
+	metadata: EntryMetadata;
+};
+
 function createLiveRuntime(): LiveRuntime {
 	return {
 		isSleeping: false,
@@ -252,6 +278,7 @@ async function executeRollback<TInput, TOutput>(
 	abortController: AbortController,
 	storage: Storage,
 	historyNotifier?: HistoryNotifier,
+	onError?: RunWorkflowOptions["onError"],
 	logger?: Logger,
 ): Promise<void> {
 	const rollbackActions: RollbackAction[] = [];
@@ -266,6 +293,7 @@ async function executeRollback<TInput, TOutput>(
 		rollbackActions,
 		false,
 		historyNotifier,
+		onError,
 		logger,
 	);
 
@@ -299,6 +327,7 @@ async function executeRollback<TInput, TOutput>(
 			continue;
 		}
 
+		let rollbackEvent: WorkflowErrorEvent | undefined;
 		try {
 			await awaitWithEviction(
 				action.rollback(rollbackContext, action.output),
@@ -312,14 +341,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 +424,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 +497,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 +511,7 @@ async function executeLiveWorkflow<TInput, TOutput>(
 			messageDriver,
 			abortController,
 			onHistoryUpdated,
+			onError,
 			logger,
 		);
 		lastResult = result;
@@ -549,6 +611,7 @@ export function runWorkflow<TInput, TOutput>(
 					abortController,
 					liveRuntime,
 					options.onHistoryUpdated,
+					options.onError,
 					logger,
 				)
 			: executeWorkflow(
@@ -559,6 +622,7 @@ export function runWorkflow<TInput, TOutput>(
 					messageDriver,
 					abortController,
 					options.onHistoryUpdated,
+					options.onError,
 					logger,
 				);
 
@@ -675,6 +739,125 @@ export function runWorkflow<TInput, TOutput>(
 	};
 }
 
+/**
+ * Remove a step and every later workflow entry, then schedule the workflow to
+ * start again immediately. Omitting `entryId` replays the workflow from the
+ * beginning.
+ */
+export async function replayWorkflowFromStep(
+	workflowId: string,
+	driver: EngineDriver,
+	entryId?: string,
+	options?: {
+		scheduleAlarm?: boolean;
+	},
+): Promise<WorkflowHistorySnapshot> {
+	const storage = await loadStorage(driver);
+	const entries = await Promise.all(
+		Array.from(storage.history.entries.entries()).map(
+			async ([key, entry]) => ({
+				key,
+				entry,
+				metadata: await loadMetadata(storage, driver, entry.id),
+			}),
+		),
+	);
+	const ordered = [...entries].sort((a, b) => {
+		if (a.metadata.createdAt !== b.metadata.createdAt) {
+			return a.metadata.createdAt - b.metadata.createdAt;
+		}
+		return a.key.localeCompare(b.key);
+	});
+
+	let entriesToDelete = ordered;
+	if (entryId !== undefined) {
+		const target = entries.find(({ entry }) => entry.id === entryId);
+		if (!target) {
+			throw new Error(`Workflow step not found: ${entryId}`);
+		}
+		if (target.entry.kind.type !== "step") {
+			throw new Error("Workflow replay target must be a step");
+		}
+
+		const replayBoundary = findReplayBoundaryEntry(entries, target);
+		const targetIndex = ordered.findIndex(
+			({ entry }) => entry.id === replayBoundary.entry.id,
+		);
+		entriesToDelete = ordered.slice(targetIndex);
+	}
+
+	const entryIdsToDelete = new Set(
+		entriesToDelete.map(({ entry }) => entry.id),
+	);
+	if (
+		entries.some(
+			({ entry, metadata }) =>
+				metadata.status === "running" &&
+				!entryIdsToDelete.has(entry.id),
+		)
+	) {
+		throw new Error(
+			"Cannot replay a workflow while a step is currently running",
+		);
+	}
+
+	await Promise.all(
+		entriesToDelete.flatMap(({ entry }) => [
+			driver.delete(buildHistoryKey(entry.location)),
+			driver.delete(buildEntryMetadataKey(entry.id)),
+		]),
+	);
+
+	for (const { key, entry } of entriesToDelete) {
+		storage.history.entries.delete(key);
+		storage.entryMetadata.delete(entry.id);
+	}
+
+	storage.output = undefined;
+	storage.flushedOutput = undefined;
+	storage.error = undefined;
+	storage.flushedError = undefined;
+	storage.state = "sleeping";
+	storage.flushedState = "sleeping";
+
+	await Promise.all([
+		driver.delete(buildWorkflowOutputKey()),
+		driver.delete(buildWorkflowErrorKey()),
+		driver.set(buildWorkflowStateKey(), serializeWorkflowState("sleeping")),
+	]);
+	if (options?.scheduleAlarm ?? true) {
+		await driver.setAlarm(workflowId, Date.now());
+	}
+
+	return createHistorySnapshot(storage);
+}
+
+function findReplayBoundaryEntry(
+	entries: ReplayEntryRecord[],
+	target: ReplayEntryRecord,
+): ReplayEntryRecord {
+	let boundary = target;
+	let boundaryDepth = -1;
+
+	for (const candidate of entries) {
+		if (candidate.entry.kind.type !== "loop") {
+			continue;
+		}
+		if (
+			candidate.entry.location.length >= target.entry.location.length ||
+			!isLocationPrefix(candidate.entry.location, target.entry.location)
+		) {
+			continue;
+		}
+		if (candidate.entry.location.length > boundaryDepth) {
+			boundary = candidate;
+			boundaryDepth = candidate.entry.location.length;
+		}
+	}
+
+	return boundary;
+}
+
 /**
  * Internal: Execute the workflow and return the result.
  */
@@ -686,6 +869,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 +923,7 @@ async function executeWorkflow<TInput, TOutput>(
 				abortController,
 				storage,
 				historyNotifier,
+				onError,
 				logger,
 			);
 		} catch (error) {
@@ -771,6 +956,7 @@ async function executeWorkflow<TInput, TOutput>(
 		undefined,
 		false,
 		historyNotifier,
+		onError,
 		logger,
 	);
 
@@ -815,12 +1001,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 +1034,7 @@ async function executeWorkflow<TInput, TOutput>(
 				abortController,
 				storage,
 				historyNotifier,
+				onError,
 				logger,
 			);
 		} catch (rollbackError) {
@@ -850,59 +1046,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;
 			}