@rivetkit/workflow-engine 0.0.0-main.6151ab0 → 0.0.0-main.78336a1

package/src/storage.ts

@@ -40,6 +40,12 @@ import type {
 	WorkflowHistorySnapshot,
 } from "./types.js";
 
+export const MAX_KV_BATCH_ENTRIES = 128;
+export const MAX_KV_BATCH_PAYLOAD_BYTES = 976 * 1024;
+
+/** Max delete ops (one transaction/permit each) run at once, under the 128-permit cap. */
+export const MAX_CONCURRENT_DELETES = 64;
+
 /**
  * Create an empty storage instance.
  */
@@ -238,6 +244,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 +271,7 @@ export async function flush(
 				key: buildHistoryKey(entry.location),
 				value: serializeEntry(entry),
 			});
-			entry.dirty = false;
+			dirtyEntries.push(entry);
 			historyUpdated = true;
 		}
 	}
@@ -275,7 +283,7 @@ export async function flush(
 				key: buildEntryMetadataKey(id),
 				value: serializeEntryMetadata(metadata),
 			});
-			metadata.dirty = false;
+			dirtyMetadata.push(metadata);
 			historyUpdated = true;
 		}
 	}
@@ -313,29 +321,31 @@ export async function flush(
 	}
 
 	if (writes.length > 0) {
-		await driver.batch(writes);
+		if (driver.atomicBatch) {
+			await driver.batch(writes);
+		} else {
+			for (const chunk of splitBatchWrites(writes)) {
+				await driver.batch(chunk);
+			}
+		}
 	}
 
 	// 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);
+		const didChange = await runDeletes(driver, pendingDeletions);
+		if (didChange) {
 			historyUpdated = true;
 		}
 	}
 
 	// 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 +356,78 @@ 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;
+}
+
+/**
+ * Split delete keys into batches within one KV transaction (MAX_KV_BATCH_ENTRIES).
+ */
+function splitBatchDeletes(keys: Uint8Array[]): Uint8Array[][] {
+	const chunks: Uint8Array[][] = [];
+	for (let i = 0; i < keys.length; i += MAX_KV_BATCH_ENTRIES) {
+		chunks.push(keys.slice(i, i + MAX_KV_BATCH_ENTRIES));
+	}
+	return chunks;
+}
+
+/**
+ * Apply deletions concurrently in bounded rounds; returns whether anything was deleted.
+ */
+async function runDeletes(
+	driver: EngineDriver,
+	deletions: PendingDeletions,
+): Promise<boolean> {
+	const ops = [
+		...deletions.prefixes.map((prefix) => () => driver.deletePrefix(prefix)),
+		...deletions.ranges.map(
+			(range) => () => driver.deleteRange(range.start, range.end),
+		),
+		...splitBatchDeletes(deletions.keys).map(
+			(chunk) => () => driver.batchDelete(chunk),
+		),
+	];
+	if (ops.length === 0) {
+		return false;
+	}
+	for (let i = 0; i < ops.length; i += MAX_CONCURRENT_DELETES) {
+		await Promise.all(
+			ops.slice(i, i + MAX_CONCURRENT_DELETES).map((op) => op()),
+		);
+	}
+	return true;
+}
+
 /**
  * Delete entries with a given location prefix (used for loop forgetting).
  * Also cleans up associated metadata from both memory and driver.
@@ -359,8 +441,7 @@ export async function deleteEntriesWithPrefix(
 	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)));
+	await runDeletes(driver, deletions);
 
 	if (deletions.keys.length > 0 && onHistoryUpdated) {
 		onHistoryUpdated();

package/src/testing.ts

@@ -173,6 +173,13 @@ export class InMemoryDriver implements EngineDriver {
 		this.kv.delete(keyToHex(key));
 	}
 
+	async batchDelete(keys: Uint8Array[]): Promise<void> {
+		await sleep(this.latency);
+		for (const key of keys) {
+			this.kv.delete(keyToHex(key));
+		}
+	}
+
 	async deletePrefix(prefix: Uint8Array): Promise<void> {
 		await sleep(this.latency);
 		for (const [hexKey, entry] of this.kv) {

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 =
@@ -422,11 +437,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 +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)
 >;
 
 /**
@@ -473,6 +484,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;
 }
 
 /**
@@ -535,6 +552,8 @@ export interface WorkflowContextInterface {
 
 	removed(name: string, originalType: EntryKindType): Promise<void>;
 
+	getVersion(name: string, latest: number): Promise<number>;
+
 	isEvicted(): boolean;
 }