@rivetkit/workflow-engine 0.0.0-main.14140ce → 0.0.0-main.17f6330

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

@@ -40,6 +40,9 @@ import type {
 	WorkflowHistorySnapshot,
 } from "./types.js";
 
+export const MAX_KV_BATCH_ENTRIES = 128;
+export const MAX_KV_BATCH_PAYLOAD_BYTES = 976 * 1024;
+
 /**
  * Create an empty storage instance.
  */
@@ -238,6 +241,8 @@ export async function flush(
 	pendingDeletions?: PendingDeletions,
 ): Promise<void> {
 	const writes: KVWrite[] = [];
+	const dirtyEntries: Entry[] = [];
+	const dirtyMetadata: EntryMetadata[] = [];
 	let historyUpdated = false;
 
 	// Flush only new names (those added since last flush)
@@ -263,7 +268,7 @@ export async function flush(
 				key: buildHistoryKey(entry.location),
 				value: serializeEntry(entry),
 			});
-			entry.dirty = false;
+			dirtyEntries.push(entry);
 			historyUpdated = true;
 		}
 	}
@@ -275,7 +280,7 @@ export async function flush(
 				key: buildEntryMetadataKey(id),
 				value: serializeEntryMetadata(metadata),
 			});
-			metadata.dirty = false;
+			dirtyMetadata.push(metadata);
 			historyUpdated = true;
 		}
 	}
@@ -313,7 +318,9 @@ export async function flush(
 	}
 
 	if (writes.length > 0) {
-		await driver.batch(writes);
+		for (const chunk of splitBatchWrites(writes)) {
+			await driver.batch(chunk);
+		}
 	}
 
 	// Apply pending deletions after the batch write. These are collected
@@ -336,6 +343,12 @@ export async function flush(
 	}
 
 	// Update flushed tracking after successful write
+	for (const entry of dirtyEntries) {
+		entry.dirty = false;
+	}
+	for (const metadata of dirtyMetadata) {
+		metadata.dirty = false;
+	}
 	storage.flushedNameCount = storage.nameRegistry.length;
 	storage.flushedState = storage.state;
 	storage.flushedOutput = storage.output;
@@ -346,6 +359,40 @@ export async function flush(
 	}
 }
 
+function splitBatchWrites(writes: KVWrite[]): KVWrite[][] {
+	const chunks: KVWrite[][] = [];
+	let chunk: KVWrite[] = [];
+	let chunkBytes = 0;
+
+	for (const write of writes) {
+		const writeBytes = write.key.byteLength + write.value.byteLength;
+		if (writeBytes > MAX_KV_BATCH_PAYLOAD_BYTES) {
+			throw new Error(
+				`KV batch write is ${writeBytes} bytes, exceeding the ${MAX_KV_BATCH_PAYLOAD_BYTES} byte limit`,
+			);
+		}
+
+		if (
+			chunk.length >= MAX_KV_BATCH_ENTRIES ||
+			(chunk.length > 0 &&
+				chunkBytes + writeBytes > MAX_KV_BATCH_PAYLOAD_BYTES)
+		) {
+			chunks.push(chunk);
+			chunk = [];
+			chunkBytes = 0;
+		}
+
+		chunk.push(write);
+		chunkBytes += writeBytes;
+	}
+
+	if (chunk.length > 0) {
+		chunks.push(chunk);
+	}
+
+	return chunks;
+}
+
 /**
  * Delete entries with a given location prefix (used for loop forgetting).
  * Also cleans up associated metadata from both memory and driver.

package/src/types.ts

@@ -399,6 +399,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 =
@@ -422,11 +424,7 @@ export interface TryStepConfig<T> extends StepConfig<T> {
 	catch?: readonly TryStepCatchKind[];
 }
 
-export type TryBlockCatchKind =
-	| "step"
-	| "join"
-	| "race"
-	| "rollback";
+export type TryBlockCatchKind = "step" | "join" | "race" | "rollback";
 
 export interface TryBlockFailure {
 	source: "step" | "join" | "race" | "block";
@@ -459,7 +457,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)
 >;
 
 /**
@@ -473,6 +471,12 @@ export interface LoopConfig<S, T> {
 	historyPruneInterval?: number;
 	/** Number of past iterations to retain when pruning. Defaults to historyPruneInterval. */
 	historySize?: number;
+	/** @deprecated Use historyPruneInterval. */
+	commitInterval?: number;
+	/** @deprecated Use historyPruneInterval. */
+	historyEvery?: number;
+	/** @deprecated Use historySize. */
+	historyKeep?: number;
 }
 
 /**