@rivetkit/workflow-engine 0.0.0-feat-workflows-public-host-apis.1550fe4 → 0.0.0-feat-sqlite-sync.18dce02

package/src/context.ts

@@ -1,5 +1,5 @@
 import type { Logger } from "pino";
-import type { EngineDriver } from "./driver.js";
+import type { EngineDriver, WorkflowSpan } from "./driver.js";
 import {
 	extractErrorInfo,
 	getErrorEventTag,
@@ -345,6 +345,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 	private abortController: AbortController;
 	private currentLocation: Location;
 	private visitedKeys: Set<string>;
+	private runSpan?: WorkflowSpan;
 	private mode: "forward" | "rollback";
 	private rollbackActions?: RollbackAction[];
 	private rollbackCheckpointSet: boolean;
@@ -369,6 +370,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		onError?: WorkflowErrorHandler,
 		logger?: Logger,
 		visitedKeys?: Set<string>,
+		runSpan?: WorkflowSpan,
 	) {
 		this.currentLocation = location;
 		this.abortController = abortController ?? new AbortController();
@@ -379,6 +381,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		this.onError = onError;
 		this.logger = logger;
 		this.visitedKeys = visitedKeys ?? new Set();
+		this.runSpan = runSpan;
 	}
 
 	get abortSignal(): AbortSignal {
@@ -435,6 +438,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			this.onError,
 			this.logger,
 			this.visitedKeys,
+			this.runSpan,
 		);
 	}
 
@@ -906,10 +910,15 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 		// Get timeout configuration
 		const timeout = config.timeout ?? DEFAULT_STEP_TIMEOUT;
 
+		const stepSpan = this.runSpan?.startStep(
+			config.name,
+			metadata.attempts,
+		);
+
 		try {
 			// Execute with timeout
 			const output = await this.executeWithTimeout(
-				config.run(),
+				stepSpan ? stepSpan.run(() => config.run()) : config.run(),
 				timeout,
 				config.name,
 			);
@@ -936,6 +945,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 				await this.flushStorage();
 			}
 
+			stepSpan?.finish("ok");
 			this.log("debug", {
 				msg: "step completed",
 				step: config.name,
@@ -951,6 +961,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			// Timeout errors are treated as critical by default. Steps opt
 			// into retrying on timeout with retryOnTimeout: true.
 			if (error instanceof StepTimeoutError && !config.retryOnTimeout) {
+				stepSpan?.finish("failed", error);
 				metadata.status = "exhausted";
 				metadata.error = String(error);
 				await this.notifyStepError(config, metadata.attempts, error, {
@@ -970,6 +981,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 				error instanceof CriticalError ||
 				error instanceof RollbackError
 			) {
+				stepSpan?.finish("failed", error);
 				metadata.status = "exhausted";
 				metadata.error = String(error);
 				await this.notifyStepError(config, metadata.attempts, error, {
@@ -989,6 +1001,7 @@ export class WorkflowContextImpl implements WorkflowContextInterface {
 			}
 
 			const willRetry = metadata.attempts <= maxRetries;
+			stepSpan?.finish(willRetry ? "retry" : "failed", error);
 			metadata.status = willRetry ? "failed" : "exhausted";
 			metadata.error = String(error);
 

package/src/driver.ts

@@ -16,6 +16,42 @@ export interface KVWrite {
 	value: Uint8Array;
 }
 
+/** How one run of the workflow function ended. */
+export type WorkflowOutcome =
+	| "completed"
+	| "sleeping"
+	| "evicted"
+	| "failed"
+	| "cancelled";
+
+/** How one step attempt ended. `retry` will be tried again, `failed` will not. */
+export type WorkflowStepOutcome = "ok" | "retry" | "failed";
+
+/** One open run of the workflow function. */
+export interface WorkflowSpan {
+	/** Runs `body` attributed to this run. */
+	run<T>(body: () => Promise<T>): Promise<T>;
+	/** Returns nothing when the attempt is not being traced. */
+	startStep(name: string, attempt: number): WorkflowStepSpan | undefined;
+	finish(outcome: WorkflowOutcome): Promise<void>;
+}
+
+/** One open attempt at one step. */
+export interface WorkflowStepSpan {
+	/** Runs `body` attributed to this attempt. */
+	run<T>(body: () => Promise<T>): Promise<T>;
+	/** `error` is what the step callback threw. */
+	finish(outcome: WorkflowStepOutcome, error?: unknown): void;
+}
+
+/**
+ * Receives the start and end of each workflow run and step attempt. A
+ * completed step replayed from history is not an attempt and is not reported.
+ */
+export interface WorkflowTelemetryDriver {
+	startSpan(): Promise<WorkflowSpan>;
+}
+
 /**
  * The engine driver provides the KV and scheduling interface.
  * Implementations must provide these methods to integrate with different backends.
@@ -54,6 +90,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.
 	 */
@@ -111,4 +153,7 @@ export interface EngineDriver {
 		messageNames: string[],
 		abortSignal: AbortSignal,
 	): Promise<void>;
+
+	/** Present when the host traces workflows. */
+	readonly telemetry?: WorkflowTelemetryDriver;
 }

package/src/index.ts

@@ -12,7 +12,16 @@ export {
 	WorkflowContextImpl,
 } from "./context.js";
 // Driver
-export type { EngineDriver, KVEntry, KVWrite } from "./driver.js";
+export type {
+	EngineDriver,
+	KVEntry,
+	KVWrite,
+	WorkflowOutcome,
+	WorkflowSpan,
+	WorkflowStepOutcome,
+	WorkflowStepSpan,
+	WorkflowTelemetryDriver,
+} from "./driver.js";
 export { extractErrorInfo } from "./error-utils.js";
 // Errors
 export {
@@ -143,7 +152,7 @@ import {
 } from "../schemas/serde.js";
 import { type RollbackAction, WorkflowContextImpl } from "./context.js";
 // Main workflow runner
-import type { EngineDriver } from "./driver.js";
+import type { EngineDriver, WorkflowOutcome, WorkflowSpan } from "./driver.js";
 import {
 	extractErrorInfo,
 	getErrorEventTag,
@@ -288,6 +297,7 @@ async function executeRollback<TInput, TOutput>(
 	historyNotifier?: HistoryNotifier,
 	onError?: RunWorkflowOptions["onError"],
 	logger?: Logger,
+	runSpan?: WorkflowSpan,
 ): Promise<void> {
 	const rollbackActions: RollbackAction[] = [];
 	const ctx = new WorkflowContextImpl(
@@ -303,6 +313,8 @@ async function executeRollback<TInput, TOutput>(
 		historyNotifier,
 		onError,
 		logger,
+		undefined,
+		runSpan,
 	);
 
 	try {
@@ -511,16 +523,19 @@ async function executeLiveWorkflow<TInput, TOutput>(
 	let lastResult: WorkflowResult<TOutput> | undefined;
 
 	while (true) {
-		const result = await executeWorkflow(
-			workflowId,
-			workflowFn,
-			input,
-			driver,
-			messageDriver,
-			abortController,
-			onHistoryUpdated,
-			onError,
-			logger,
+		const result = await traceRun(driver, (runSpan) =>
+			executeWorkflow(
+				workflowId,
+				workflowFn,
+				input,
+				driver,
+				messageDriver,
+				abortController,
+				onHistoryUpdated,
+				onError,
+				logger,
+				runSpan,
+			),
 		);
 		lastResult = result;
 
@@ -642,16 +657,19 @@ export function runWorkflow<TInput, TOutput>(
 					options.onError,
 					logger,
 				)
-			: executeWorkflow(
-					workflowId,
-					workflowFn,
-					input,
-					driver,
-					messageDriver,
-					abortController,
-					options.onHistoryUpdated,
-					options.onError,
-					logger,
+			: traceRun(driver, (runSpan) =>
+					executeWorkflow(
+						workflowId,
+						workflowFn,
+						input,
+						driver,
+						messageDriver,
+						abortController,
+						options.onHistoryUpdated,
+						options.onError,
+						logger,
+						runSpan,
+					),
 				);
 
 	return {
@@ -886,6 +904,49 @@ function findReplayBoundaryEntry(
 	return boundary;
 }
 
+function runOutcomeFromState(state: WorkflowState): WorkflowOutcome {
+	switch (state) {
+		case "completed":
+			return "completed";
+		case "sleeping":
+			return "sleeping";
+		case "failed":
+			return "failed";
+		case "cancelled":
+			return "cancelled";
+		case "pending":
+		case "running":
+		case "rolling_back":
+			return "evicted";
+	}
+}
+
+/** Internal: Report one workflow run to the driver's telemetry, when it has any. */
+async function traceRun<TOutput>(
+	driver: EngineDriver,
+	run: (span?: WorkflowSpan) => Promise<WorkflowResult<TOutput>>,
+): Promise<WorkflowResult<TOutput>> {
+	if (!driver.telemetry) {
+		return await run();
+	}
+
+	const span = await driver.telemetry.startSpan();
+	let outcome: WorkflowOutcome = "failed";
+	try {
+		const result = await span.run(() => run(span));
+		outcome = runOutcomeFromState(result.state);
+		return result;
+	} catch (error) {
+		// A run throws EvictedError only when the workflow was already cancelled.
+		if (error instanceof EvictedError) {
+			outcome = "cancelled";
+		}
+		throw error;
+	} finally {
+		await span.finish(outcome);
+	}
+}
+
 /**
  * Internal: Execute the workflow and return the result.
  */
@@ -899,6 +960,7 @@ async function executeWorkflow<TInput, TOutput>(
 	onHistoryUpdated?: (history: WorkflowHistorySnapshot) => void,
 	onError?: RunWorkflowOptions["onError"],
 	logger?: Logger,
+	runSpan?: WorkflowSpan,
 ): Promise<WorkflowResult<TOutput>> {
 	const storage = await loadStorage(driver);
 	const historyNotifier: HistoryNotifier = onHistoryUpdated
@@ -953,6 +1015,7 @@ async function executeWorkflow<TInput, TOutput>(
 				historyNotifier,
 				onError,
 				logger,
+				runSpan,
 			);
 		} catch (error) {
 			if (error instanceof EvictedError) {
@@ -986,6 +1049,8 @@ async function executeWorkflow<TInput, TOutput>(
 		historyNotifier,
 		onError,
 		logger,
+		undefined,
+		runSpan,
 	);
 
 	storage.state = "running";
@@ -1064,6 +1129,7 @@ async function executeWorkflow<TInput, TOutput>(
 				historyNotifier,
 				onError,
 				logger,
+				runSpan,
 			);
 		} catch (rollbackError) {
 			if (rollbackError instanceof EvictedError) {

package/src/keys.ts

@@ -14,6 +14,7 @@ export const KEY_PREFIX = {
 	HISTORY: 2, // History entries: [2, ...locationSegments]
 	WORKFLOW: 3, // Workflow metadata: [3, field]
 	ENTRY_METADATA: 4, // Entry metadata: [4, entryId]
+	// 5 is reserved for RivetKit's workflow trace context.
 } as const;
 
 // Workflow metadata field identifiers

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.
  */
@@ -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,46 @@ 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.
@@ -410,8 +443,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) {