@rivetkit/workflow-engine 0.0.0-0-0-0-preview-guard-stops.9d82529

package/src/storage.ts

@@ -0,0 +1,468 @@
+import {
+	deserializeEntry,
+	deserializeEntryMetadata,
+	deserializeName,
+	deserializeWorkflowError,
+	deserializeWorkflowOutput,
+	deserializeWorkflowState,
+	serializeEntry,
+	serializeEntryMetadata,
+	serializeName,
+	serializeWorkflowError,
+	serializeWorkflowOutput,
+	serializeWorkflowState,
+} from "../schemas/serde.js";
+import type { EngineDriver, KVWrite } from "./driver.js";
+import {
+	buildEntryMetadataKey,
+	buildEntryMetadataPrefix,
+	buildHistoryKey,
+	buildHistoryPrefix,
+	buildHistoryPrefixAll,
+	buildNameKey,
+	buildNamePrefix,
+	buildWorkflowErrorKey,
+	buildWorkflowOutputKey,
+	buildWorkflowStateKey,
+	compareKeys,
+	parseEntryMetadataKey,
+	parseNameKey,
+} from "./keys.js";
+import { isLocationPrefix, locationToKey } from "./location.js";
+import type {
+	Entry,
+	EntryKind,
+	EntryMetadata,
+	Location,
+	Storage,
+	WorkflowEntryMetadataSnapshot,
+	WorkflowHistoryEntry,
+	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.
+ */
+export function createStorage(): Storage {
+	return {
+		nameRegistry: [],
+		flushedNameCount: 0,
+		history: { entries: new Map() },
+		entryMetadata: new Map(),
+		output: undefined,
+		state: "pending",
+		flushedState: undefined,
+		error: undefined,
+		flushedError: undefined,
+		flushedOutput: undefined,
+	};
+}
+
+/**
+ * Create a snapshot of workflow history for observers.
+ */
+export function createHistorySnapshot(
+	storage: Storage,
+): WorkflowHistorySnapshot {
+	const entryMetadata = new Map<string, WorkflowEntryMetadataSnapshot>();
+	for (const [id, metadata] of storage.entryMetadata) {
+		const { dirty, ...rest } = metadata;
+		entryMetadata.set(id, rest);
+	}
+
+	const entries: WorkflowHistoryEntry[] = [];
+	const entryKeys = Array.from(storage.history.entries.keys()).sort();
+	for (const key of entryKeys) {
+		const entry = storage.history.entries.get(key);
+		if (!entry) continue;
+		const { dirty, ...rest } = entry;
+		entries.push(rest);
+	}
+
+	return {
+		nameRegistry: [...storage.nameRegistry],
+		entries,
+		entryMetadata,
+	};
+}
+
+/**
+ * Generate a UUID v4.
+ */
+export function generateId(): string {
+	return crypto.randomUUID();
+}
+
+/**
+ * Create a new entry.
+ */
+export function createEntry(location: Location, kind: EntryKind): Entry {
+	return {
+		id: generateId(),
+		location,
+		kind,
+		dirty: true,
+	};
+}
+
+/**
+ * Create or get metadata for an entry.
+ */
+export function getOrCreateMetadata(
+	storage: Storage,
+	entryId: string,
+): EntryMetadata {
+	let metadata = storage.entryMetadata.get(entryId);
+	if (!metadata) {
+		metadata = {
+			status: "pending",
+			attempts: 0,
+			lastAttemptAt: 0,
+			createdAt: Date.now(),
+			rollbackCompletedAt: undefined,
+			rollbackError: undefined,
+			dirty: true,
+		};
+		storage.entryMetadata.set(entryId, metadata);
+	}
+	return metadata;
+}
+
+/**
+ * Load storage from the driver.
+ */
+export async function loadStorage(driver: EngineDriver): Promise<Storage> {
+	const storage = createStorage();
+
+	// Load name registry
+	const nameEntries = await driver.list(buildNamePrefix());
+	// Sort by index to ensure correct order
+	nameEntries.sort((a, b) => compareKeys(a.key, b.key));
+	for (const entry of nameEntries) {
+		const index = parseNameKey(entry.key);
+		storage.nameRegistry[index] = deserializeName(entry.value);
+	}
+	// Track how many names are already persisted
+	storage.flushedNameCount = storage.nameRegistry.length;
+
+	// Load history entries
+	const historyEntries = await driver.list(buildHistoryPrefixAll());
+	for (const entry of historyEntries) {
+		const parsed = deserializeEntry(entry.value);
+		parsed.dirty = false;
+		// Use locationToKey to match how context.ts looks up entries
+		const key = locationToKey(storage, parsed.location);
+		storage.history.entries.set(key, parsed);
+	}
+
+	// Load entry metadata so observers can reconstruct workflow state after
+	// the actor wakes and rebuilds storage from persisted history.
+	const metadataEntries = await driver.list(buildEntryMetadataPrefix());
+	for (const entry of metadataEntries) {
+		const entryId = parseEntryMetadataKey(entry.key);
+		const metadata = deserializeEntryMetadata(entry.value);
+		metadata.dirty = false;
+		storage.entryMetadata.set(entryId, metadata);
+	}
+
+	// Load workflow state
+	const stateValue = await driver.get(buildWorkflowStateKey());
+	if (stateValue) {
+		storage.state = deserializeWorkflowState(stateValue);
+		storage.flushedState = storage.state;
+	}
+
+	// Load output if present
+	const outputValue = await driver.get(buildWorkflowOutputKey());
+	if (outputValue) {
+		storage.output = deserializeWorkflowOutput(outputValue);
+		storage.flushedOutput = storage.output;
+	}
+
+	// Load error if present
+	const errorValue = await driver.get(buildWorkflowErrorKey());
+	if (errorValue) {
+		storage.error = deserializeWorkflowError(errorValue);
+		storage.flushedError = storage.error;
+	}
+
+	return storage;
+}
+
+/**
+ * Load metadata for an entry (lazy loading).
+ */
+export async function loadMetadata(
+	storage: Storage,
+	driver: EngineDriver,
+	entryId: string,
+): Promise<EntryMetadata> {
+	// Check if already loaded
+	const existing = storage.entryMetadata.get(entryId);
+	if (existing) {
+		return existing;
+	}
+
+	// Load from driver
+	const value = await driver.get(buildEntryMetadataKey(entryId));
+	if (value) {
+		const metadata = deserializeEntryMetadata(value);
+		metadata.dirty = false;
+		storage.entryMetadata.set(entryId, metadata);
+		return metadata;
+	}
+
+	// Create new metadata
+	return getOrCreateMetadata(storage, entryId);
+}
+
+/**
+ * Pending deletions collected by collectLoopPruning to be included
+ * in the next flush alongside the state write.
+ */
+export interface PendingDeletions {
+	prefixes: Uint8Array[];
+	keys: Uint8Array[];
+	ranges: { start: Uint8Array; end: Uint8Array }[];
+}
+
+/**
+ * Flush all dirty data to the driver. Optionally includes pending
+ * deletions so that history pruning happens alongside the
+ * state write.
+ */
+export async function flush(
+	storage: Storage,
+	driver: EngineDriver,
+	onHistoryUpdated?: () => void,
+	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)
+	for (
+		let i = storage.flushedNameCount;
+		i < storage.nameRegistry.length;
+		i++
+	) {
+		const name = storage.nameRegistry[i];
+		if (name !== undefined) {
+			writes.push({
+				key: buildNameKey(i),
+				value: serializeName(name),
+			});
+			historyUpdated = true;
+		}
+	}
+
+	// Flush dirty entries
+	for (const [, entry] of storage.history.entries) {
+		if (entry.dirty) {
+			writes.push({
+				key: buildHistoryKey(entry.location),
+				value: serializeEntry(entry),
+			});
+			dirtyEntries.push(entry);
+			historyUpdated = true;
+		}
+	}
+
+	// Flush dirty metadata
+	for (const [id, metadata] of storage.entryMetadata) {
+		if (metadata.dirty) {
+			writes.push({
+				key: buildEntryMetadataKey(id),
+				value: serializeEntryMetadata(metadata),
+			});
+			dirtyMetadata.push(metadata);
+			historyUpdated = true;
+		}
+	}
+
+	// Flush workflow state if changed
+	if (storage.state !== storage.flushedState) {
+		writes.push({
+			key: buildWorkflowStateKey(),
+			value: serializeWorkflowState(storage.state),
+		});
+	}
+
+	// Flush output if changed
+	if (
+		storage.output !== undefined &&
+		storage.output !== storage.flushedOutput
+	) {
+		writes.push({
+			key: buildWorkflowOutputKey(),
+			value: serializeWorkflowOutput(storage.output),
+		});
+	}
+
+	// Flush error if changed (compare by message since objects aren't reference-equal)
+	const errorChanged =
+		storage.error !== undefined &&
+		(storage.flushedError === undefined ||
+			storage.error.name !== storage.flushedError.name ||
+			storage.error.message !== storage.flushedError.message);
+	if (errorChanged) {
+		writes.push({
+			key: buildWorkflowErrorKey(),
+			value: serializeWorkflowError(storage.error!),
+		});
+	}
+
+	if (writes.length > 0) {
+		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);
+			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;
+	storage.flushedError = storage.error;
+
+	if (historyUpdated && onHistoryUpdated) {
+		onHistoryUpdated();
+	}
+}
+
+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.
+ */
+export async function deleteEntriesWithPrefix(
+	storage: Storage,
+	driver: EngineDriver,
+	prefixLocation: Location,
+	onHistoryUpdated?: () => void,
+): Promise<void> {
+	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)));
+
+	if (deletions.keys.length > 0 && onHistoryUpdated) {
+		onHistoryUpdated();
+	}
+}
+
+/**
+ * Remove entries matching a location prefix from memory and collect
+ * the driver-level deletion operations. The returned PendingDeletions
+ * can be applied immediately or batched with a flush.
+ */
+export function collectDeletionsForPrefix(
+	storage: Storage,
+	prefixLocation: Location,
+): PendingDeletions {
+	const pending: PendingDeletions = {
+		prefixes: [buildHistoryPrefix(prefixLocation)],
+		keys: [],
+		ranges: [],
+	};
+
+	for (const [key, entry] of storage.history.entries) {
+		if (isLocationPrefix(prefixLocation, entry.location)) {
+			pending.keys.push(buildEntryMetadataKey(entry.id));
+			storage.entryMetadata.delete(entry.id);
+			storage.history.entries.delete(key);
+		}
+	}
+
+	return pending;
+}
+
+/**
+ * Get an entry by location.
+ */
+export function getEntry(
+	storage: Storage,
+	location: Location,
+): Entry | undefined {
+	const key = locationToKey(storage, location);
+	return storage.history.entries.get(key);
+}
+
+/**
+ * Set an entry by location.
+ */
+export function setEntry(
+	storage: Storage,
+	location: Location,
+	entry: Entry,
+): void {
+	const key = locationToKey(storage, location);
+	storage.history.entries.set(key, entry);
+}