fix(workflow-engine): clear pruned loop history via batched, bounded deletes

rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts

@@ -116,6 +116,11 @@ class WorkflowStorage {
 		);
 	}
 
+	async batchDelete(keys: Uint8Array[]): Promise<void> {
+		if (keys.length === 0) return;
+		await this.deleteRawKeys(keys.map((key) => makeWorkflowKey(key)));
+	}
+
 	async deletePrefix(prefix: Uint8Array): Promise<void> {
 		const start = makeWorkflowKey(prefix);
 		const end = computeUpperBound(start);
@@ -344,6 +349,10 @@ export class ActorWorkflowDriver implements EngineDriver {
 		await this.#runCtx.internalKeepAwake(this.#storage.delete(key));
 	}
 
+	async batchDelete(keys: Uint8Array[]): Promise<void> {
+		await this.#runCtx.internalKeepAwake(this.#storage.batchDelete(keys));
+	}
+
 	async deletePrefix(prefix: Uint8Array): Promise<void> {
 		await this.#runCtx.internalKeepAwake(
 			this.#storage.deletePrefix(prefix),
@@ -438,6 +447,10 @@ export class ActorWorkflowControlDriver implements EngineDriver {
 		await this.#storage.delete(key);
 	}
 
+	async batchDelete(keys: Uint8Array[]): Promise<void> {
+		await this.#storage.batchDelete(keys);
+	}
+
 	async deletePrefix(prefix: Uint8Array): Promise<void> {
 		await this.#storage.deletePrefix(prefix);
 	}

rivetkit-typescript/packages/workflow-engine/src/driver.ts

@@ -54,6 +54,12 @@ export interface EngineDriver {
 	 */
 	delete(key: Uint8Array): Promise<void>;
 
+	/**
+	 * Batch delete multiple keys in a single operation.
+	 * Should be atomic if possible.
+	 */
+	batchDelete(keys: Uint8Array[]): Promise<void>;
+
 	/**
 	 * Delete all keys with a given prefix.
 	 */

rivetkit-typescript/packages/workflow-engine/src/storage.ts

@@ -1,4 +1,4 @@
 import {
 	deserializeEntry,
 	deserializeEntryMetadata,
 	deserializeName,
@@ -43,6 +43,9 @@
 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.
  */
@@ -330,18 +333,8 @@
 	// 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;
 		}
 	}
@@ -397,6 +390,44 @@
 	return chunks;
 }
 
+/**
+ * Split delete keys into batches within one KV transaction (KV_TX_MAX_ROWS).
+ */
+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.
@@ -410,8 +441,7 @@
 	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();

rivetkit-typescript/packages/workflow-engine/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) {

rivetkit-typescript/packages/workflow-engine/tests/prune-fanout.test.ts

@@ -0,0 +1,111 @@
+import { describe, expect, it } from "vitest";
+import {
+	deleteEntriesWithPrefix,
+	MAX_CONCURRENT_DELETES,
+	MAX_KV_BATCH_ENTRIES,
+} from "../src/storage.js";
+import {
+	appendName,
+	createEntry,
+	createStorage,
+	emptyLocation,
+	InMemoryDriver,
+	setEntry,
+} from "../src/testing.js";
+
+describe("Workflow Engine Storage delete fan-out", () => {
+	// Records batchDelete sizes to assert keys are coalesced, not deleted one-by-one.
+	class BatchDeleteRecordingDriver extends InMemoryDriver {
+		batchSizes: number[] = [];
+		singleDeletes = 0;
+
+		override async batchDelete(keys: Uint8Array[]): Promise<void> {
+			this.batchSizes.push(keys.length);
+			await super.batchDelete(keys);
+		}
+
+		override async delete(key: Uint8Array): Promise<void> {
+			this.singleDeletes++;
+			await super.delete(key);
+		}
+	}
+
+	it("clears a large history prefix in transaction-sized delete batches", async () => {
+		const driver = new BatchDeleteRecordingDriver();
+		driver.latency = 1;
+		const storage = createStorage();
+		const loopLocation = appendName(storage, emptyLocation(), "loop");
+
+		// Span several batches so chunking is exercised.
+		const entryCount = MAX_KV_BATCH_ENTRIES * 3 + 7;
+		for (let i = 0; i < entryCount; i++) {
+			const location = appendName(storage, loopLocation, `iter-${i}`);
+			const entry = createEntry(location, {
+				type: "step",
+				data: { output: i },
+			});
+			setEntry(storage, location, entry);
+		}
+
+		await deleteEntriesWithPrefix(storage, driver, loopLocation);
+
+		// All keys deleted via transaction-sized batches, no per-key fan-out.
+		expect(driver.singleDeletes).toBe(0);
+		expect(driver.batchSizes).toHaveLength(
+			Math.ceil(entryCount / MAX_KV_BATCH_ENTRIES),
+		);
+		for (const size of driver.batchSizes) {
+			expect(size).toBeLessThanOrEqual(MAX_KV_BATCH_ENTRIES);
+		}
+		expect(driver.batchSizes.reduce((a, b) => a + b, 0)).toBe(entryCount);
+		expect(storage.history.entries.size).toBe(0);
+	});
+
+	// Tracks concurrent delete ops so the test can assert the fan-out stays bounded.
+	class ConcurrencyTrackingDriver extends InMemoryDriver {
+		inFlight = 0;
+		peakInFlight = 0;
+
+		async #track<T>(op: Promise<T>): Promise<T> {
+			this.inFlight++;
+			this.peakInFlight = Math.max(this.peakInFlight, this.inFlight);
+			try {
+				return await op;
+			} finally {
+				this.inFlight--;
+			}
+		}
+
+		override batchDelete(keys: Uint8Array[]): Promise<void> {
+			return this.#track(super.batchDelete(keys));
+		}
+
+		override deletePrefix(prefix: Uint8Array): Promise<void> {
+			return this.#track(super.deletePrefix(prefix));
+		}
+	}
+
+	it("bounds concurrent delete ops for a prune larger than the cap", async () => {
+		const driver = new ConcurrencyTrackingDriver();
+		driver.latency = 1;
+		const storage = createStorage();
+		const loopLocation = appendName(storage, emptyLocation(), "loop");
+
+		// Enough keys to yield more batches than MAX_CONCURRENT_DELETES.
+		const entryCount = MAX_CONCURRENT_DELETES * MAX_KV_BATCH_ENTRIES + 1;
+		for (let i = 0; i < entryCount; i++) {
+			const location = appendName(storage, loopLocation, `iter-${i}`);
+			const entry = createEntry(location, {
+				type: "step",
+				data: { output: i },
+			});
+			setEntry(storage, location, entry);
+		}
+
+		await deleteEntriesWithPrefix(storage, driver, loopLocation);
+
+		expect(driver.peakInFlight).toBeLessThanOrEqual(MAX_CONCURRENT_DELETES);
+		expect(driver.peakInFlight).toBe(MAX_CONCURRENT_DELETES);
+		expect(storage.history.entries.size).toBe(0);
+	});
+});