@rivetkit/workflow-engine 0.0.0-main.72cbc7c → 0.0.0-main.79e97f1

package/src/context.ts

@@ -21,7 +21,7 @@ import {
 	StepExhaustedError,
 	StepFailedError,
 } from "./errors.js";
-import { buildLoopIterationRange, buildEntryMetadataKey } from "./keys.js";
+import { buildEntryMetadataKey, buildLoopIterationRange } from "./keys.js";
 import {
 	appendLoopIteration,
 	appendName,
@@ -36,13 +36,12 @@ import {
 	flush,
 	getOrCreateMetadata,
 	loadMetadata,
-	setEntry,
 	type PendingDeletions,
+	setEntry,
 } from "./storage.js";
 import type {
 	BranchConfig,
 	BranchOutput,
-	BranchStatus,
 	Entry,
 	EntryKindType,
 	EntryMetadata,
@@ -66,13 +65,12 @@ import type {
 	WorkflowError,
 	WorkflowErrorEvent,
 	WorkflowErrorHandler,
+	WorkflowMessageDriver,
 	WorkflowQueue,
 	WorkflowQueueMessage,
 	WorkflowQueueNextBatchOptions,
 	WorkflowQueueNextOptions,
-	WorkflowMessageDriver,
 } from "./types.js";
-import { sleep } from "./utils.js";
 
 /**
  * Default values for step configuration.
@@ -505,8 +503,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	 * Throws HistoryDivergedError if duplicate detected.
 	 */
 	private checkDuplicateName(name: string): void {
-		const fullKey =
-			locationToKey(this.storage, this.currentLocation) + "/" + name;
+		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)}". ` +
@@ -580,7 +577,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			const isUnderPrefix =
 				prefix === ""
 					? true // Root: all keys are children
-					: key.startsWith(prefix + "/") || key === prefix;
+					: key.startsWith(`${prefix}/`) || key === prefix;
 
 			if (isUnderPrefix) {
 				if (!this.visitedKeys.has(key)) {
@@ -603,25 +600,33 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	}
 
 	/**
-	 * Wait for eviction message.
-	 *
-	 * The event listener uses { once: true } to auto-remove after firing,
-	 * preventing memory leaks if this method is called multiple times.
+	 * Wait for `ms`, rejecting early with EvictedError if the workflow is
+	 * evicted. Both the timer and the abort listener are torn down on either
+	 * outcome, so a completed sleep never leaves a dangling listener on the
+	 * long-lived run abort signal.
 	 */
-	waitForEviction(): Promise<never> {
-		return new Promise((_, reject) => {
-			if (this.abortSignal.aborted) {
-				reject(new EvictedError());
-				return;
+	private async sleepOrEvict(ms: number): Promise<void> {
+		if (this.abortSignal.aborted) {
+			throw new EvictedError();
+		}
+		let timer: ReturnType<typeof setTimeout> | undefined;
+		let onAbort: (() => void) | undefined;
+		try {
+			await new Promise<void>((resolve, reject) => {
+				timer = setTimeout(resolve, ms);
+				onAbort = () => reject(new EvictedError());
+				this.abortSignal.addEventListener("abort", onAbort, {
+					once: true,
+				});
+			});
+		} finally {
+			if (timer !== undefined) {
+				clearTimeout(timer);
 			}
-			this.abortSignal.addEventListener(
-				"abort",
-				() => {
-					reject(new EvictedError());
-				},
-				{ once: true },
-			);
-		});
+			if (onAbort) {
+				this.abortSignal.removeEventListener("abort", onAbort);
+			}
+		}
 	}
 
 	// === Step ===
@@ -825,10 +830,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			const maxRetries = config.maxRetries ?? DEFAULT_MAX_RETRIES;
 
 			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;
+				const lastError = metadata.error;
 				const exhaustedError = new StepExhaustedError(
 					config.name,
 					lastError,
@@ -941,15 +943,16 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			});
 			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;
+			if (entry.kind.type === "step") {
+				entry.kind.data.error = String(error);
+			}
+			entry.dirty = true;
+
+			// Timeout errors are treated as critical by default. Steps opt
+			// into retrying on timeout with retryOnTimeout: true.
+			if (error instanceof StepTimeoutError && !config.retryOnTimeout) {
 				metadata.status = "exhausted";
 				metadata.error = String(error);
-				await this.flushStorage();
 				await this.notifyStepError(config, metadata.attempts, error, {
 					willRetry: false,
 				});
@@ -967,13 +970,8 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 				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();
 				await this.notifyStepError(config, metadata.attempts, error, {
 					willRetry: false,
 				});
@@ -990,15 +988,10 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 				);
 			}
 
-			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);
 
-			await this.flushStorage();
 			if (willRetry) {
 				const retryDelay = calculateBackoff(
 					metadata.attempts,
@@ -1023,7 +1016,10 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 				attachTryStepFailure(
 					new StepExhaustedError(config.name, String(error)),
 					{
-						kind: "exhausted",
+						kind:
+							error instanceof StepTimeoutError
+								? "timeout"
+								: "exhausted",
 						stepName: config.name,
 						attempts: metadata.attempts,
 						error: extractErrorInfo(error),
@@ -1042,7 +1038,9 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	 *
 	 * Note: This does NOT cancel the underlying operation. JavaScript Promises
 	 * cannot be cancelled once started. When a timeout occurs:
-	 * - The step is marked as failed with StepTimeoutError
+	 * - The step is rejected with StepTimeoutError. By default this is treated
+	 *   as a critical failure with no retry. Set retryOnTimeout: true on the
+	 *   step config to retry timeouts like any other error.
 	 * - The underlying async operation continues running in the background
 	 * - Any side effects from the operation may still occur
 	 *
@@ -1493,7 +1491,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 
 		// Short sleep: wait in memory
 		if (remaining < this.driver.workerPollInterval) {
-			await Promise.race([sleep(remaining), this.waitForEviction()]);
+			await this.sleepOrEvict(remaining);
 
 			this.checkEvicted();
 
@@ -2583,4 +2581,97 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		setEntry(this.storage, location, entry);
 		await this.flushStorage();
 	}
+
+	// === Version ===
+
+	async getVersion(name: string, latest: number): Promise<number> {
+		this.assertNotInProgress();
+		this.checkEvicted();
+
+		this.entryInProgress = true;
+		try {
+			return await this.executeGetVersion(name, latest);
+		} finally {
+			this.entryInProgress = false;
+		}
+	}
+
+	private async executeGetVersion(
+		name: string,
+		latest: number,
+	): Promise<number> {
+		if (!Number.isInteger(latest) || latest < 1) {
+			throw new Error(
+				`getVersion("${name}", ${latest}): latest must be an integer >= 1`,
+			);
+		}
+
+		// Check for duplicate name in current execution
+		this.checkDuplicateName(name);
+
+		const location = appendName(this.storage, this.currentLocation, name);
+		const key = locationToKey(this.storage, location);
+		const existing = this.storage.history.entries.get(key);
+
+		// Mark this entry as visited for validateComplete
+		this.markVisited(key);
+
+		this.stopRollbackIfMissing(existing);
+
+		if (existing) {
+			if (existing.kind.type !== "version_check") {
+				throw new HistoryDivergedError(
+					`Expected version_check at ${key}, found ${existing.kind.type}`,
+				);
+			}
+			// Pure replay: this instance is already pinned at this location.
+			return existing.kind.data.resolved;
+		}
+
+		// No recorded version at this location. Decide whether this instance
+		// already executed past this point under older code (old in-flight) or
+		// is reaching it fresh at the live frontier.
+		//
+		// The discriminator is "is there any unvisited history entry under the
+		// current scope?". Entries created earlier in this same run are already
+		// marked visited, so the only unvisited entries under the scope are
+		// leftovers from a prior run, which proves this scope already executed
+		// under code that predates this gate. Such instances resolve to the
+		// implicit floor version 1 (old branch); fresh instances resolve to
+		// `latest`.
+		const resolved = this.hasUnvisitedUnderCurrentScope(key) ? 1 : latest;
+
+		const entry = createEntry(location, {
+			type: "version_check",
+			data: { resolved, latest },
+		});
+		setEntry(this.storage, location, entry);
+		await this.flushStorage();
+
+		return resolved;
+	}
+
+	/**
+	 * Returns true if any history entry under the current location scope
+	 * (other than `excludeKey`) has not yet been visited this run. Used by
+	 * getVersion to detect an old in-flight instance whose scope was already
+	 * executed by code that predates a version gate.
+	 */
+	private hasUnvisitedUnderCurrentScope(excludeKey: string): boolean {
+		const prefix = locationToKey(this.storage, this.currentLocation);
+
+		for (const key of this.storage.history.entries.keys()) {
+			if (key === excludeKey) {
+				continue;
+			}
+			const isUnderPrefix =
+				prefix === ""
+					? true
+					: key.startsWith(`${prefix}/`) || key === prefix;
+			if (isUnderPrefix && !this.visitedKeys.has(key)) {
+				return true;
+			}
+		}
+		return false;
+	}
 }

package/src/driver.ts

@@ -30,6 +30,12 @@ export interface KVWrite {
  * See architecture.md "Isolation Model" for details.
  */
 export interface EngineDriver {
+	/**
+	 * Requires each logical storage flush to reach `batch` as one indivisible
+	 * unit. The driver must reject an oversized unit instead of splitting it.
+	 */
+	readonly atomicBatch?: boolean;
+
 	// === KV Operations ===
 
 	/**

package/src/index.ts

@@ -13,6 +13,7 @@ export {
 } from "./context.js";
 // Driver
 export type { EngineDriver, KVEntry, KVWrite } from "./driver.js";
+export { extractErrorInfo } from "./error-utils.js";
 // Errors
 export {
 	CancelledError,
@@ -29,7 +30,6 @@ export {
 	StepExhaustedError,
 	StepFailedError,
 } from "./errors.js";
-export { extractErrorInfo } from "./error-utils.js";
 
 // Location utilities
 export {
@@ -82,9 +82,6 @@ export type {
 	PathSegment,
 	RaceEntry,
 	RemovedEntry,
-	WorkflowEntryMetadataSnapshot,
-	WorkflowHistoryEntry,
-	WorkflowHistorySnapshot,
 	RollbackCheckpointEntry,
 	RollbackContextInterface,
 	RunWorkflowOptions,
@@ -102,20 +99,23 @@ export type {
 	TryStepFailure,
 	TryStepResult,
 	WorkflowContextInterface,
+	WorkflowEntryMetadataSnapshot,
 	WorkflowError,
 	WorkflowErrorEvent,
 	WorkflowErrorHandler,
 	WorkflowFunction,
 	WorkflowHandle,
-	WorkflowRollbackErrorEvent,
+	WorkflowHistoryEntry,
+	WorkflowHistorySnapshot,
+	WorkflowMessageDriver,
 	WorkflowQueue,
 	WorkflowQueueMessage,
 	WorkflowQueueNextBatchOptions,
 	WorkflowQueueNextOptions,
-	WorkflowMessageDriver,
 	WorkflowResult,
-	WorkflowRunMode,
+	WorkflowRollbackErrorEvent,
 	WorkflowRunErrorEvent,
+	WorkflowRunMode,
 	WorkflowState,
 	WorkflowStepErrorEvent,
 } from "./types.js";
@@ -185,9 +185,9 @@ import type {
 	RunWorkflowOptions,
 	Storage,
 	WorkflowErrorEvent,
-	WorkflowHistorySnapshot,
 	WorkflowFunction,
 	WorkflowHandle,
+	WorkflowHistorySnapshot,
 	WorkflowMessageDriver,
 	WorkflowResult,
 	WorkflowRunMode,
@@ -532,26 +532,46 @@ async function executeLiveWorkflow<TInput, TOutput>(
 		const hasDeadline = result.sleepUntil !== undefined;
 
 		if (hasMessages && hasDeadline) {
-			// Wait for EITHER a message OR the deadline (for queue.next timeout)
+			// Wait for EITHER a message OR the deadline (for queue.next timeout).
+			// Scope both waiters to a per-iteration controller chained to the run
+			// signal so that whichever loses the race is cancelled immediately.
+			// Otherwise the loser (a message waiter or an armed sleep timer) stays
+			// registered on the long-lived run signal and leaks once per cycle.
+			const iterationAbort = new AbortController();
+			const onRunAbort = () => iterationAbort.abort();
+			if (abortController.signal.aborted) {
+				iterationAbort.abort();
+			} else {
+				abortController.signal.addEventListener("abort", onRunAbort, {
+					once: true,
+				});
+			}
+			const messagePromise = awaitWithEviction(
+				driver.waitForMessages(
+					result.waitingForMessages!,
+					iterationAbort.signal,
+				),
+				iterationAbort.signal,
+			);
+			const sleepPromise = waitForSleep(
+				runtime,
+				result.sleepUntil!,
+				iterationAbort.signal,
+			);
+			// Swallow the loser's rejection once it is aborted below so it does
+			// not surface as an unhandled rejection.
+			messagePromise.catch(() => {});
+			sleepPromise.catch(() => {});
 			try {
-				const messagePromise = awaitWithEviction(
-					driver.waitForMessages(
-						result.waitingForMessages!,
-						abortController.signal,
-					),
-					abortController.signal,
-				);
-				const sleepPromise = waitForSleep(
-					runtime,
-					result.sleepUntil!,
-					abortController.signal,
-				);
 				await Promise.race([messagePromise, sleepPromise]);
 			} catch (error) {
 				if (error instanceof EvictedError) {
 					return lastResult;
 				}
 				throw error;
+			} finally {
+				iterationAbort.abort();
+				abortController.signal.removeEventListener("abort", onRunAbort);
 			}
 			continue;
 		}

package/src/location.ts

@@ -159,7 +159,7 @@ export function getChildEntries(
 		const isChild =
 			parentKey === ""
 				? true
-				: key.startsWith(parentKey + "/") || key === parentKey;
+				: key.startsWith(`${parentKey}/`) || key === parentKey;
 
 		if (isChild) {
 			// Return the actual entry's location, not the parent location

package/src/storage.ts

@@ -318,8 +318,12 @@ export async function flush(
 	}
 
 	if (writes.length > 0) {
-		for (const chunk of splitBatchWrites(writes)) {
-			await driver.batch(chunk);
+		if (driver.atomicBatch) {
+			await driver.batch(writes);
+		} else {
+			for (const chunk of splitBatchWrites(writes)) {
+				await driver.batch(chunk);
+			}
 		}
 	}
 

package/src/types.ts

@@ -135,6 +135,17 @@ export interface RemovedEntry {
 	originalName?: string;
 }
 
+/**
+ * Version check entry data - records which version of a code path this
+ * workflow instance is pinned to at a given location.
+ */
+export interface VersionCheckEntry {
+	/** The version this instance resolved to at this location. */
+	resolved: number;
+	/** The `latest` value seen when this entry was first resolved (diagnostics). */
+	latest: number;
+}
+
 /**
  * All possible entry kind types.
  */
@@ -146,7 +157,8 @@ export type EntryKindType =
 	| "rollback_checkpoint"
 	| "join"
 	| "race"
-	| "removed";
+	| "removed"
+	| "version_check";
 
 /**
  * Type-specific entry data.
@@ -159,7 +171,8 @@ export type EntryKind =
 	| { type: "rollback_checkpoint"; data: RollbackCheckpointEntry }
 	| { type: "join"; data: JoinEntry }
 	| { type: "race"; data: RaceEntry }
-	| { type: "removed"; data: RemovedEntry };
+	| { type: "removed"; data: RemovedEntry }
+	| { type: "version_check"; data: VersionCheckEntry };
 
 /**
  * An entry in the workflow history.
@@ -399,6 +412,8 @@ export interface StepConfig<T> {
 	retryBackoffMax?: number;
 	/** Timeout in ms for step execution (default: 30000). Set to 0 to disable. */
 	timeout?: number;
+	/** If true, step timeouts retry like any other error instead of failing immediately as critical. Default: false. */
+	retryOnTimeout?: boolean;
 }
 
 export type TryStepCatchKind =
@@ -455,7 +470,7 @@ export type LoopResult<S, T> =
  * `Loop.continue(undefined)`.
  */
 export type LoopIterationResult<S, T> = Promise<
-	LoopResult<S, T> | (S extends undefined ? void : never)
+	LoopResult<S, T> | (S extends undefined ? undefined : never)
 >;
 
 /**
@@ -537,6 +552,8 @@ export interface WorkflowContextInterface {
 
 	removed(name: string, originalType: EntryKindType): Promise<void>;
 
+	getVersion(name: string, latest: number): Promise<number>;
+
 	isEvicted(): boolean;
 }