@rivetkit/workflow-engine 0.0.0-main.72cbc7c → 0.0.0-main.78336a1

package/src/storage.ts

@@ -43,6 +43,9 @@ import type {
 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.
  */
@@ -318,26 +321,20 @@ 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);
+			}
 		}
 	}
 
 	// 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;
 		}
 	}
@@ -393,6 +390,44 @@ function splitBatchWrites(writes: KVWrite[]): KVWrite[][] {
 	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.
@@ -406,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 =
@@ -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;
 }