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

package/src/context.ts

@@ -1,5 +1,10 @@
 import type { Logger } from "pino";
 import type { EngineDriver } from "./driver.js";
+import {
+	extractErrorInfo,
+	getErrorEventTag,
+	markErrorReported,
+} from "./error-utils.js";
 import {
 	CancelledError,
 	CriticalError,
@@ -16,10 +21,12 @@ import {
 	StepExhaustedError,
 	StepFailedError,
 } from "./errors.js";
+import { buildLoopIterationRange, buildEntryMetadataKey } from "./keys.js";
 import {
 	appendLoopIteration,
 	appendName,
 	emptyLocation,
+	isLocationPrefix,
 	locationToKey,
 	registerName,
 } from "./location.js";
@@ -30,6 +37,7 @@ import {
 	getOrCreateMetadata,
 	loadMetadata,
 	setEntry,
+	type PendingDeletions,
 } from "./storage.js";
 import type {
 	BranchConfig,
@@ -47,6 +55,8 @@ import type {
 	StepConfig,
 	Storage,
 	WorkflowContextInterface,
+	WorkflowErrorEvent,
+	WorkflowErrorHandler,
 	WorkflowQueue,
 	WorkflowQueueMessage,
 	WorkflowQueueNextBatchOptions,
@@ -62,9 +72,7 @@ import { sleep } from "./utils.js";
 export const DEFAULT_MAX_RETRIES = 3;
 export const DEFAULT_RETRY_BACKOFF_BASE = 100;
 export const DEFAULT_RETRY_BACKOFF_MAX = 30000;
-export const DEFAULT_LOOP_COMMIT_INTERVAL = 20;
-export const DEFAULT_LOOP_HISTORY_EVERY = 20;
-export const DEFAULT_LOOP_HISTORY_KEEP = 20;
+export const DEFAULT_LOOP_HISTORY_PRUNE_INTERVAL = 20;
 export const DEFAULT_STEP_TIMEOUT = 30000; // 30 seconds
 
 const QUEUE_HISTORY_MESSAGE_MARKER = "__rivetWorkflowQueueMessage";
@@ -108,7 +116,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	private entryInProgress = false;
 	private abortController: AbortController;
 	private currentLocation: Location;
-	private visitedKeys = new Set<string>();
+	private visitedKeys: Set<string>;
 	private mode: "forward" | "rollback";
 	private rollbackActions?: RollbackAction[];
 	private rollbackCheckpointSet: boolean;
@@ -116,6 +124,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	private usedNamesInExecution = new Set<string>();
 	private pendingCompletableMessageIds = new Set<string>();
 	private historyNotifier?: () => void;
+	private onError?: WorkflowErrorHandler;
 	private logger?: Logger;
 
 	constructor(
@@ -129,7 +138,9 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		rollbackActions?: RollbackAction[],
 		rollbackCheckpointSet = false,
 		historyNotifier?: () => void,
+		onError?: WorkflowErrorHandler,
 		logger?: Logger,
+		visitedKeys?: Set<string>,
 	) {
 		this.currentLocation = location;
 		this.abortController = abortController ?? new AbortController();
@@ -137,7 +148,9 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		this.rollbackActions = rollbackActions;
 		this.rollbackCheckpointSet = rollbackCheckpointSet;
 		this.historyNotifier = historyNotifier;
+		this.onError = onError;
 		this.logger = logger;
+		this.visitedKeys = visitedKeys ?? new Set();
 	}
 
 	get abortSignal(): AbortSignal {
@@ -147,7 +160,8 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	get queue(): WorkflowQueue {
 		return {
 			next: async (name, opts) => await this.queueNext(name, opts),
-			nextBatch: async (name, opts) => await this.queueNextBatch(name, opts),
+			nextBatch: async (name, opts) =>
+				await this.queueNextBatch(name, opts),
 			send: async (name, body) => await this.queueSend(name, body),
 		};
 	}
@@ -190,18 +204,65 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			this.rollbackActions,
 			this.rollbackCheckpointSet,
 			this.historyNotifier,
+			this.onError,
 			this.logger,
+			this.visitedKeys,
 		);
 	}
 
 	/**
 	 * Log a debug message using the configured logger.
 	 */
-	private log(level: "debug" | "info" | "warn" | "error", data: Record<string, unknown>): void {
+	private log(
+		level: "debug" | "info" | "warn" | "error",
+		data: Record<string, unknown>,
+	): void {
 		if (!this.logger) return;
 		this.logger[level](data);
 	}
 
+	private async notifyError(event: WorkflowErrorEvent): Promise<void> {
+		if (!this.onError) {
+			return;
+		}
+
+		try {
+			await this.onError(event);
+		} catch (error) {
+			this.log("warn", {
+				msg: "workflow error hook failed",
+				hookEventType: getErrorEventTag(event),
+				error: extractErrorInfo(error),
+			});
+		}
+	}
+
+	private async notifyStepError<T>(
+		config: StepConfig<T>,
+		attempt: number,
+		error: unknown,
+		opts: {
+			willRetry: boolean;
+			retryDelay?: number;
+			retryAt?: number;
+		},
+	): Promise<void> {
+		const maxRetries = config.maxRetries ?? DEFAULT_MAX_RETRIES;
+		await this.notifyError({
+			step: {
+				workflowId: this.workflowId,
+				stepName: config.name,
+				attempt,
+				maxRetries,
+				remainingRetries: Math.max(0, maxRetries - (attempt - 1)),
+				willRetry: opts.willRetry,
+				retryDelay: opts.retryDelay,
+				retryAt: opts.retryAt,
+				error: extractErrorInfo(error),
+			},
+		});
+	}
+
 	/**
 	 * Mark a key as visited.
 	 */
@@ -216,12 +277,12 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	private checkDuplicateName(name: string): void {
 		const fullKey =
 			locationToKey(this.storage, this.currentLocation) + "/" + name;
-			if (this.usedNamesInExecution.has(fullKey)) {
-				throw new HistoryDivergedError(
-					`Duplicate entry name "${name}" at location "${locationToKey(this.storage, this.currentLocation)}". ` +
-						`Each step/loop/sleep/queue.next/join/race must have a unique name within its scope.`,
-				);
-			}
+		if (this.usedNamesInExecution.has(fullKey)) {
+			throw new HistoryDivergedError(
+				`Duplicate entry name "${name}" at location "${locationToKey(this.storage, this.currentLocation)}". ` +
+				`Each step/loop/sleep/queue.next/join/race must have a unique name within its scope.`,
+			);
+		}
 		this.usedNamesInExecution.add(fullKey);
 	}
 
@@ -297,7 +358,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 					// This means workflow code may have changed
 					throw new HistoryDivergedError(
 						`Entry "${key}" exists in history but was not visited. ` +
-							`Workflow code may have changed. Use ctx.removed() to handle migrations.`,
+						`Workflow code may have changed. Use ctx.removed() to handle migrations.`,
 					);
 				}
 			}
@@ -391,19 +452,36 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			);
 
 			// Replay successful result (including void steps).
-			if (metadata.status === "completed" || stepData.output !== undefined) {
+			if (
+				metadata.status === "completed" ||
+				stepData.output !== undefined
+			) {
 				return stepData.output as T;
 			}
 
 			// Check if we should retry
 			const maxRetries = config.maxRetries ?? DEFAULT_MAX_RETRIES;
 
-			if (metadata.attempts >= maxRetries) {
+			if (metadata.attempts > maxRetries) {
 				// Prefer step history error, but fall back to metadata since
 				// driver implementations may persist metadata without the history
 				// entry error (e.g. partial writes/crashes between attempts).
 				const lastError = stepData.error ?? metadata.error;
-				throw new StepExhaustedError(config.name, lastError);
+				const exhaustedError = markErrorReported(
+					new StepExhaustedError(config.name, lastError),
+				);
+				if (metadata.status !== "exhausted") {
+					metadata.status = "exhausted";
+					metadata.dirty = true;
+					await this.flushStorage();
+					await this.notifyStepError(
+						config,
+						metadata.attempts,
+						exhaustedError,
+						{ willRetry: false },
+					);
+				}
+				throw exhaustedError;
 			}
 
 			// Calculate backoff and yield to scheduler
@@ -427,7 +505,11 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			existing ?? createEntry(location, { type: "step", data: {} });
 		if (!existing) {
 			// New entry - register name
-			this.log("debug", { msg: "executing new step", step: config.name, key });
+			this.log("debug", {
+				msg: "executing new step",
+				step: config.name,
+				key,
+			});
 			const nameIndex = registerName(this.storage, config.name);
 			entry.location = [...location];
 			entry.location[entry.location.length - 1] = nameIndex;
@@ -437,6 +519,11 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		}
 
 		const metadata = getOrCreateMetadata(this.storage, entry.id);
+		const maxRetries = config.maxRetries ?? DEFAULT_MAX_RETRIES;
+		const retryBackoffBase =
+			config.retryBackoffBase ?? DEFAULT_RETRY_BACKOFF_BASE;
+		const retryBackoffMax =
+			config.retryBackoffMax ?? DEFAULT_RETRY_BACKOFF_MAX;
 		metadata.status = "running";
 		metadata.attempts++;
 		metadata.lastAttemptAt = Date.now();
@@ -453,63 +540,111 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 				config.name,
 			);
 
-				if (entry.kind.type === "step") {
-					entry.kind.data.output = output;
-				}
-				entry.dirty = true;
-				metadata.status = "completed";
-				metadata.error = undefined;
-				metadata.completedAt = Date.now();
+			if (entry.kind.type === "step") {
+				entry.kind.data.output = output;
+			}
+			entry.dirty = true;
+			metadata.status = "completed";
+			metadata.error = undefined;
+			metadata.completedAt = Date.now();
 
-				// Ephemeral steps don't trigger an immediate flush. This avoids the
+			// Ephemeral steps don't trigger an immediate flush. This avoids the
 			// synchronous write overhead for transient operations. Note that the
 			// step's entry is still marked dirty and WILL be persisted on the
 			// next flush from a non-ephemeral operation. The purpose of ephemeral
 			// is to batch writes, not to avoid persistence entirely.
 			if (!config.ephemeral) {
-				this.log("debug", { msg: "flushing step", step: config.name, key });
+				this.log("debug", {
+					msg: "flushing step",
+					step: config.name,
+					key,
+				});
 				await this.flushStorage();
 			}
 
-			this.log("debug", { msg: "step completed", step: config.name, key });
+			this.log("debug", {
+				msg: "step completed",
+				step: config.name,
+				key,
+			});
 			return output;
 		} catch (error) {
-				// Timeout errors are treated as critical (no retry)
-				if (error instanceof StepTimeoutError) {
-					if (entry.kind.type === "step") {
-						entry.kind.data.error = String(error);
-					}
-					entry.dirty = true;
-					metadata.status = "exhausted";
-					metadata.error = String(error);
-					await this.flushStorage();
-					throw new CriticalError(error.message);
+			// Timeout errors are treated as critical (no retry)
+			if (error instanceof StepTimeoutError) {
+				if (entry.kind.type === "step") {
+					entry.kind.data.error = String(error);
 				}
+				entry.dirty = true;
+				metadata.status = "exhausted";
+				metadata.error = String(error);
+				await this.flushStorage();
+				await this.notifyStepError(config, metadata.attempts, error, {
+					willRetry: false,
+				});
+				throw markErrorReported(new CriticalError(error.message));
+			}
 
 			if (
 				error instanceof CriticalError ||
 				error instanceof RollbackError
-				) {
-					if (entry.kind.type === "step") {
-						entry.kind.data.error = String(error);
-					}
-					entry.dirty = true;
-					metadata.status = "exhausted";
-					metadata.error = String(error);
-					await this.flushStorage();
-					throw error;
-				}
-
+			) {
 				if (entry.kind.type === "step") {
 					entry.kind.data.error = String(error);
 				}
 				entry.dirty = true;
-				metadata.status = "failed";
+				metadata.status = "exhausted";
 				metadata.error = String(error);
-
 				await this.flushStorage();
+				await this.notifyStepError(config, metadata.attempts, error, {
+					willRetry: false,
+				});
+				throw markErrorReported(error);
+			}
+
+			if (entry.kind.type === "step") {
+				entry.kind.data.error = String(error);
+			}
+			entry.dirty = true;
+			const willRetry = metadata.attempts <= maxRetries;
+			metadata.status = willRetry ? "failed" : "exhausted";
+			metadata.error = String(error);
 
-				throw new StepFailedError(config.name, error, metadata.attempts);
+			await this.flushStorage();
+			if (willRetry) {
+				const retryDelay = calculateBackoff(
+					metadata.attempts,
+					retryBackoffBase,
+					retryBackoffMax,
+				);
+				const retryAt = metadata.lastAttemptAt + retryDelay;
+				await this.notifyStepError(
+					config,
+					metadata.attempts,
+					error,
+					{
+						willRetry: true,
+						retryDelay,
+						retryAt,
+					},
+				);
+				throw new StepFailedError(
+					config.name,
+					error,
+					metadata.attempts,
+					retryAt,
+				);
+			}
+
+			const exhaustedError = markErrorReported(
+				new StepExhaustedError(config.name, String(error)),
+			);
+			await this.notifyStepError(
+				config,
+				metadata.attempts,
+				error,
+				{ willRetry: false },
+			);
+			throw exhaustedError;
 		}
 	}
 
@@ -644,7 +779,11 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			}
 
 			const loopData = existing.kind.data;
-			metadata = await loadMetadata(this.storage, this.driver, existing.id);
+			metadata = await loadMetadata(
+				this.storage,
+				this.driver,
+				existing.id,
+			);
 
 			if (rollbackMode) {
 				if (loopData.output !== undefined) {
@@ -692,20 +831,26 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			metadata.dirty = true;
 		}
 
-		// TODO: Add validation for commitInterval (must be > 0)
-		const commitInterval =
-			config.commitInterval ?? DEFAULT_LOOP_COMMIT_INTERVAL;
-		const historyEvery =
-			config.historyEvery ??
-			config.commitInterval ??
-			DEFAULT_LOOP_HISTORY_EVERY;
-		const historyKeep =
-			config.historyKeep ??
-			config.commitInterval ??
-			DEFAULT_LOOP_HISTORY_KEEP;
+		const historyPruneInterval =
+			config.historyPruneInterval ?? DEFAULT_LOOP_HISTORY_PRUNE_INTERVAL;
+		const historySize = config.historySize ?? historyPruneInterval;
+
+		// Track the last iteration we pruned up to so we only delete
+		// newly-expired iterations instead of re-scanning from 0.
+		let lastPrunedUpTo = 0;
+
+		// Deferred flush promise from the previous prune cycle. Awaited at the
+		// start of the next iteration so the flush runs in parallel with user code.
+		let deferredFlush: Promise<void> | null = null;
 
 		// Execute loop iterations
 		while (true) {
+			// Await any deferred flush from the previous prune cycle
+			if (deferredFlush) {
+				await deferredFlush;
+				deferredFlush = null;
+			}
+
 			if (rollbackMode && rollbackSingleIteration) {
 				if (rollbackIterationRan) {
 					return rollbackOutput as T;
@@ -752,13 +897,14 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 					metadata.dirty = true;
 				}
 
-				await this.flushStorage();
-				await this.forgetOldIterations(
+				// Collect pruning deletions and flush
+				const deletions = this.collectLoopPruning(
 					location,
 					iteration + 1,
-					historyEvery,
-					historyKeep,
+					historySize,
+					lastPrunedUpTo,
 				);
+				await this.flushStorageWithDeletions(deletions);
 
 				if (rollbackMode && rollbackSingleIteration) {
 					rollbackOutput = result.value;
@@ -775,72 +921,102 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			}
 			iteration++;
 
-			// Periodic commit
-			if (iteration % commitInterval === 0) {
+			if (!rollbackMode) {
 				if (entry.kind.type === "loop") {
 					entry.kind.data.state = state;
 					entry.kind.data.iteration = iteration;
 				}
 				entry.dirty = true;
+			}
 
-				await this.flushStorage();
-				await this.forgetOldIterations(
+			// Periodically defer the flush so the next iteration can overlap
+			// with loop pruning and any pending dirty state writes.
+			if (iteration % historyPruneInterval === 0) {
+				const deletions = this.collectLoopPruning(
 					location,
 					iteration,
-					historyEvery,
-					historyKeep,
+					historySize,
+					lastPrunedUpTo,
 				);
+				lastPrunedUpTo = Math.max(0, iteration - historySize);
+
+				// Defer the flush to run in parallel with the next iteration
+				deferredFlush = this.flushStorageWithDeletions(deletions);
 			}
 		}
 	}
 
 	/**
-	 * Delete old loop iteration entries to save storage space.
+	 * Collect pending deletions for loop history pruning.
 	 *
-	 * Loop locations always end with a NameIndex (number) because loops are
-	 * created via appendName(). Even for nested loops, the innermost loop's
-	 * location ends with its name index:
-	 *
-	 *   ctx.loop("outer") → location: [outerIndex]
-	 *     iteration 0    → location: [{ loop: outerIndex, iteration: 0 }]
-	 *       ctx.loop("inner") → location: [{ loop: outerIndex, iteration: 0 }, innerIndex]
-	 *
-	 * This function removes iterations older than (currentIteration - historyKeep)
-	 * every historyEvery iterations.
+	 * Only deletes iterations in the range [fromIteration, keepFrom) where
+	 * keepFrom = currentIteration - historySize. This avoids re-scanning
+	 * already-deleted iterations.
 	 */
-	private async forgetOldIterations(
+	private collectLoopPruning(
 		loopLocation: Location,
 		currentIteration: number,
-		historyEvery: number,
-		historyKeep: number,
-	): Promise<void> {
-		if (historyEvery <= 0 || historyKeep <= 0) {
-			return;
+		historySize: number,
+		fromIteration: number,
+	): PendingDeletions | undefined {
+		if (currentIteration <= historySize) {
+			return undefined;
 		}
-		if (currentIteration === 0 || currentIteration % historyEvery !== 0) {
-			return;
+
+		const keepFrom = Math.max(0, currentIteration - historySize);
+		if (fromIteration >= keepFrom) {
+			return undefined;
 		}
-		const keepFrom = Math.max(0, currentIteration - historyKeep);
-		// Get the loop name index from the last segment of loopLocation.
-		// This is always a NameIndex (number) because loop entries are created
-		// via appendName(), not appendLoopIteration().
+
 		const loopSegment = loopLocation[loopLocation.length - 1];
 		if (typeof loopSegment !== "number") {
 			throw new Error("Expected loop location to end with a name index");
 		}
 
-		for (let i = 0; i < keepFrom; i++) {
-			const iterationLocation: Location = [
-				...loopLocation,
-				{ loop: loopSegment, iteration: i },
-			];
-			await deleteEntriesWithPrefix(
-				this.storage,
-				this.driver,
-				iterationLocation,
-				this.historyNotifier,
-			);
+		const range = buildLoopIterationRange(
+			loopLocation,
+			loopSegment,
+			fromIteration,
+			keepFrom,
+		);
+		const metadataKeys: Uint8Array[] = [];
+
+		for (const [key, entry] of this.storage.history.entries) {
+			if (!isLocationPrefix(loopLocation, entry.location)) {
+				continue;
+			}
+
+			const iterationSegment = entry.location[loopLocation.length];
+			if (
+				!iterationSegment ||
+				typeof iterationSegment === "number" ||
+				iterationSegment.loop !== loopSegment ||
+				iterationSegment.iteration < fromIteration ||
+				iterationSegment.iteration >= keepFrom
+			) {
+				continue;
+			}
+
+			metadataKeys.push(buildEntryMetadataKey(entry.id));
+			this.storage.entryMetadata.delete(entry.id);
+			this.storage.history.entries.delete(key);
 		}
+
+		return {
+			prefixes: [],
+			keys: metadataKeys,
+			ranges: [range],
+		};
+	}
+
+	/**
+	 * Flush storage with optional pending deletions so pruning
+	 * happens alongside the state write.
+	 */
+	private async flushStorageWithDeletions(
+		deletions?: PendingDeletions,
+	): Promise<void> {
+		await flush(this.storage, this.driver, this.historyNotifier, deletions);
 	}
 
 	// === Sleep ===
@@ -1246,7 +1422,8 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			);
 			const messageKey = locationToKey(this.storage, messageLocation);
 			this.markVisited(messageKey);
-			const existingMessage = this.storage.history.entries.get(messageKey);
+			const existingMessage =
+				this.storage.history.entries.get(messageKey);
 			if (!existingMessage || existingMessage.kind.type !== "message") {
 				throw new HistoryDivergedError(
 					`Expected queue message "${name}:${i}" in history`,
@@ -1267,7 +1444,9 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		return results;
 	}
 
-	private toWorkflowQueueMessage<T>(message: Message): WorkflowQueueMessage<T> {
+	private toWorkflowQueueMessage<T>(
+		message: Message,
+	): WorkflowQueueMessage<T> {
 		return {
 			id: message.id,
 			name: message.name,
@@ -1377,15 +1556,17 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		if (
 			typeof value === "object" &&
 			value !== null &&
-			(value as Record<string, unknown>)[QUEUE_HISTORY_MESSAGE_MARKER] === 1
+			(value as Record<string, unknown>)[QUEUE_HISTORY_MESSAGE_MARKER] ===
+			1
 		) {
 			const serialized = value as Record<string, unknown>;
-			const id =
-				typeof serialized.id === "string" ? serialized.id : "";
+			const id = typeof serialized.id === "string" ? serialized.id : "";
 			const serializedName =
 				typeof serialized.name === "string" ? serialized.name : name;
 			const createdAt =
-				typeof serialized.createdAt === "number" ? serialized.createdAt : 0;
+				typeof serialized.createdAt === "number"
+					? serialized.createdAt
+					: 0;
 			const completed =
 				typeof serialized.completed === "boolean"
 					? serialized.completed
@@ -1734,7 +1915,10 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 					}
 					if (error instanceof MessageWaitError) {
 						// Track message wait errors, prefer sleep errors with deadlines
-						if (!yieldError || !(yieldError instanceof SleepError)) {
+						if (
+							!yieldError ||
+							!(yieldError instanceof SleepError)
+						) {
 							if (!yieldError) {
 								yieldError = error;
 							} else if (yieldError instanceof MessageWaitError) {

package/src/driver.ts

@@ -53,6 +53,11 @@ export interface EngineDriver {
 	 */
 	deletePrefix(prefix: Uint8Array): Promise<void>;
 
+	/**
+	 * Delete all keys in the half-open range [start, end).
+	 */
+	deleteRange(start: Uint8Array, end: Uint8Array): Promise<void>;
+
 	/**
 	 * List all key-value pairs with a given prefix.
 	 *