@rivetkit/workflow-engine 0.0.0-pr.4600.32b0fc8

package/src/storage.ts

@@ -0,0 +1,417 @@
+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";
+
+/**
+ * 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[] = [];
+	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),
+			});
+			entry.dirty = false;
+			historyUpdated = true;
+		}
+	}
+
+	// Flush dirty metadata
+	for (const [id, metadata] of storage.entryMetadata) {
+		if (metadata.dirty) {
+			writes.push({
+				key: buildEntryMetadataKey(id),
+				value: serializeEntryMetadata(metadata),
+			});
+			metadata.dirty = false;
+			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) {
+		await driver.batch(writes);
+	}
+
+	// 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
+	storage.flushedNameCount = storage.nameRegistry.length;
+	storage.flushedState = storage.state;
+	storage.flushedOutput = storage.output;
+	storage.flushedError = storage.error;
+
+	if (historyUpdated && onHistoryUpdated) {
+		onHistoryUpdated();
+	}
+}
+
+/**
+ * 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);
+}

package/src/testing.ts

@@ -0,0 +1,313 @@
+import type { EngineDriver, KVEntry, KVWrite } from "./driver.js";
+import { EvictedError } from "./errors.js";
+import { compareKeys, keyStartsWith, keyToHex } from "./keys.js";
+import type { Message, WorkflowMessageDriver } from "./types.js";
+import { sleep } from "./utils.js";
+
+interface Waiter {
+	nameSet?: Set<string>;
+	resolve: () => void;
+	reject: (error: Error) => void;
+	abortSignal: AbortSignal;
+	onAbort: () => void;
+}
+
+class InMemoryWorkflowMessageDriver implements WorkflowMessageDriver {
+	#messages: Message[] = [];
+	#waiters = new Set<Waiter>();
+
+	async addMessage(message: Message): Promise<void> {
+		this.#messages.push(message);
+		this.#notifyWaiters(message.name);
+	}
+
+	async receiveMessages(opts: {
+		names?: readonly string[];
+		count: number;
+		completable: boolean;
+	}): Promise<Message[]> {
+		const limitedCount = Math.max(1, opts.count);
+		const nameSet =
+			opts.names && opts.names.length > 0
+				? new Set(opts.names)
+				: undefined;
+		const selected: Array<{ message: Message; index: number }> = [];
+
+		for (
+			let i = 0;
+			i < this.#messages.length && selected.length < limitedCount;
+			i++
+		) {
+			const message = this.#messages[i];
+			if (nameSet && !nameSet.has(message.name)) {
+				continue;
+			}
+			selected.push({ message, index: i });
+		}
+
+		if (selected.length === 0) {
+			return [];
+		}
+
+		if (!opts.completable) {
+			for (let i = selected.length - 1; i >= 0; i--) {
+				this.#messages.splice(selected[i].index, 1);
+			}
+			return selected.map((entry) => entry.message);
+		}
+
+		return selected.map((entry) => {
+			const { message } = entry;
+			return {
+				...message,
+				complete: async () => {
+					await this.completeMessage(message.id);
+				},
+			};
+		});
+	}
+
+	async completeMessage(messageId: string): Promise<void> {
+		const index = this.#messages.findIndex(
+			(message) => message.id === messageId,
+		);
+		if (index !== -1) {
+			this.#messages.splice(index, 1);
+		}
+	}
+
+	async waitForMessages(
+		messageNames: string[],
+		abortSignal: AbortSignal,
+	): Promise<void> {
+		if (abortSignal.aborted) {
+			throw new EvictedError();
+		}
+
+		const nameSet =
+			messageNames.length > 0 ? new Set(messageNames) : undefined;
+		if (
+			this.#messages.some((message) =>
+				nameSet ? nameSet.has(message.name) : true,
+			)
+		) {
+			return;
+		}
+
+		await new Promise<void>((resolve, reject) => {
+			const waiter: Waiter = {
+				nameSet,
+				resolve: () => {
+					this.#removeWaiter(waiter);
+					resolve();
+				},
+				reject: (error) => {
+					this.#removeWaiter(waiter);
+					reject(error);
+				},
+				abortSignal,
+				onAbort: () => {
+					waiter.reject(new EvictedError());
+				},
+			};
+			abortSignal.addEventListener("abort", waiter.onAbort, {
+				once: true,
+			});
+			this.#waiters.add(waiter);
+		});
+	}
+
+	#removeWaiter(waiter: Waiter): void {
+		if (this.#waiters.delete(waiter)) {
+			waiter.abortSignal.removeEventListener("abort", waiter.onAbort);
+		}
+	}
+
+	#notifyWaiters(name: string): void {
+		for (const waiter of [...this.#waiters]) {
+			if (waiter.nameSet && !waiter.nameSet.has(name)) {
+				continue;
+			}
+			waiter.resolve();
+		}
+	}
+
+	clear(): void {
+		this.#messages = [];
+		for (const waiter of [...this.#waiters]) {
+			waiter.reject(new Error("cleared"));
+		}
+	}
+}
+
+/**
+ * In-memory implementation of EngineDriver for testing.
+ * Uses binary keys (Uint8Array) with hex encoding for internal Map storage.
+ */
+export class InMemoryDriver implements EngineDriver {
+	// Map from hex-encoded key to { originalKey, value }
+	private kv = new Map<string, { key: Uint8Array; value: Uint8Array }>();
+	private alarms = new Map<string, number>();
+	#inMemoryMessageDriver = new InMemoryWorkflowMessageDriver();
+
+	/** Simulated latency per operation (ms) */
+	latency = 10;
+
+	/** How often the worker polls for work */
+	workerPollInterval = 100;
+	messageDriver: WorkflowMessageDriver = this.#inMemoryMessageDriver;
+
+	async get(key: Uint8Array): Promise<Uint8Array | null> {
+		await sleep(this.latency);
+		const entry = this.kv.get(keyToHex(key));
+		return entry?.value ?? null;
+	}
+
+	async set(key: Uint8Array, value: Uint8Array): Promise<void> {
+		await sleep(this.latency);
+		this.kv.set(keyToHex(key), { key, value });
+	}
+
+	async delete(key: Uint8Array): Promise<void> {
+		await sleep(this.latency);
+		this.kv.delete(keyToHex(key));
+	}
+
+	async deletePrefix(prefix: Uint8Array): Promise<void> {
+		await sleep(this.latency);
+		for (const [hexKey, entry] of this.kv) {
+			if (keyStartsWith(entry.key, prefix)) {
+				this.kv.delete(hexKey);
+			}
+		}
+	}
+
+	async deleteRange(start: Uint8Array, end: Uint8Array): Promise<void> {
+		await sleep(this.latency);
+		for (const [hexKey, entry] of this.kv) {
+			if (
+				compareKeys(entry.key, start) >= 0 &&
+				compareKeys(entry.key, end) < 0
+			) {
+				this.kv.delete(hexKey);
+			}
+		}
+	}
+
+	async list(prefix: Uint8Array): Promise<KVEntry[]> {
+		await sleep(this.latency);
+		const results: KVEntry[] = [];
+		for (const entry of this.kv.values()) {
+			if (keyStartsWith(entry.key, prefix)) {
+				results.push({ key: entry.key, value: entry.value });
+			}
+		}
+		// Sort by key lexicographically
+		return results.sort((a, b) => compareKeys(a.key, b.key));
+	}
+
+	async batch(writes: KVWrite[]): Promise<void> {
+		await sleep(this.latency);
+		for (const { key, value } of writes) {
+			this.kv.set(keyToHex(key), { key, value });
+		}
+	}
+
+	async setAlarm(workflowId: string, wakeAt: number): Promise<void> {
+		await sleep(this.latency);
+		this.alarms.set(workflowId, wakeAt);
+	}
+
+	async clearAlarm(workflowId: string): Promise<void> {
+		await sleep(this.latency);
+		this.alarms.delete(workflowId);
+	}
+
+	async waitForMessages(
+		messageNames: string[],
+		abortSignal: AbortSignal,
+	): Promise<void> {
+		const driver = this.messageDriver as WorkflowMessageDriver & {
+			waitForMessages?: (
+				messageNames: string[],
+				abortSignal: AbortSignal,
+			) => Promise<void>;
+		};
+		if (driver.waitForMessages) {
+			await driver.waitForMessages(messageNames, abortSignal);
+			return;
+		}
+
+		while (true) {
+			if (abortSignal.aborted) {
+				throw new EvictedError();
+			}
+			const messages = await this.messageDriver.receiveMessages({
+				names: messageNames.length > 0 ? messageNames : undefined,
+				count: 1,
+				completable: true,
+			});
+			if (messages.length > 0) {
+				return;
+			}
+			await sleep(Math.max(1, this.latency));
+		}
+	}
+
+	/**
+	 * Get the alarm time for a workflow (for testing).
+	 */
+	getAlarm(workflowId: string): number | undefined {
+		return this.alarms.get(workflowId);
+	}
+
+	/**
+	 * Check if any alarms are due and return their workflow IDs.
+	 */
+	getDueAlarms(): string[] {
+		const now = Date.now();
+		const due: string[] = [];
+		for (const [workflowId, wakeAt] of this.alarms) {
+			if (wakeAt <= now) {
+				due.push(workflowId);
+			}
+		}
+		return due;
+	}
+
+	/**
+	 * Clear all data (for testing).
+	 */
+	clear(): void {
+		this.kv.clear();
+		this.alarms.clear();
+		this.#inMemoryMessageDriver.clear();
+	}
+
+	/**
+	 * Get a snapshot of all data (for testing/debugging).
+	 */
+	snapshot(): {
+		kv: Record<string, Uint8Array>;
+		alarms: Record<string, number>;
+	} {
+		const kvSnapshot: Record<string, Uint8Array> = {};
+		for (const [hexKey, entry] of this.kv) {
+			kvSnapshot[hexKey] = entry.value;
+		}
+		return {
+			kv: kvSnapshot,
+			alarms: Object.fromEntries(this.alarms),
+		};
+	}
+
+	/**
+	 * Get all hex-encoded keys (for testing).
+	 */
+	keys(): string[] {
+		return [...this.kv.keys()];
+	}
+}
+
+// Re-export main exports for convenience
+export * from "./index.js";