fix(workflows): clear pruned loop history via batched deletes

packages/workflows/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.
  */
@@ -330,18 +333,8 @@ export async function flush(
 	// 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 @@ function splitBatchWrites(writes: KVWrite[]): KVWrite[][] {
 	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 @@ 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();

packages/workflows/tests/storage.test.ts

@@ -1,6 +1,8 @@
 import { beforeEach, describe, expect, it } from "vitest";
 import type { KVWrite } from "../src/driver.js";
 import {
+	deleteEntriesWithPrefix,
+	MAX_CONCURRENT_DELETES,
 	MAX_KV_BATCH_ENTRIES,
 	MAX_KV_BATCH_PAYLOAD_BYTES,
 } from "../src/storage.js";
@@ -201,3 +203,98 @@ describe("Workflow Engine Storage flush", () => {
 		expect(metadata.dirty).toBe(false);
 	});
 });
+
+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);
+	});
+});