rivetkit 2.3.4 → 2.3.5-rc.1

package/src/registry/runtime.ts

@@ -1,5 +1,5 @@
-import { stringifyError } from "@/common/utils";
 import type { SqliteNativeMetrics } from "@/common/database/config";
+import { stringifyError } from "@/common/utils";
 import type { RegistryConfig } from "./config";
 import { logger } from "./log";
 
@@ -15,6 +15,7 @@ export type ActorContextHandle = OpaqueHandle<"actorContext">;
 export type ConnHandle = OpaqueHandle<"conn">;
 export type WebSocketHandle = OpaqueHandle<"webSocket">;
 export type CancellationTokenHandle = OpaqueHandle<"cancellationToken">;
+export type SqliteTransactionHandle = OpaqueHandle<"sqliteTransaction">;
 
 export type RuntimeBytes = Uint8Array;
 
@@ -46,6 +47,11 @@ export interface RuntimeStateDeltaPayload {
 	connHibernationRemoved: string[];
 }
 
+export interface RuntimeWorkflowKvWrite {
+	key: RuntimeBytes;
+	value: RuntimeBytes;
+}
+
 export interface RuntimeRequestSaveOpts {
 	immediate?: boolean;
 	maxWaitMs?: number;
@@ -80,6 +86,50 @@ export interface RuntimeQueueSendResult {
 	response?: RuntimeBytes;
 }
 
+export interface RuntimeScheduledEventInfo {
+	id: string;
+	action: string;
+	args: RuntimeBytes;
+	runAt: number;
+}
+
+export interface RuntimeCronJobInfo {
+	name: string;
+	kind: "cron" | "every";
+	action: string;
+	args: RuntimeBytes;
+	nextRunAt: number;
+	lastRunAt?: number;
+	expression?: string;
+	timezone?: string;
+	intervalMs?: number;
+	maxHistory: number;
+}
+
+export interface RuntimeScheduleErrorInfo {
+	group: string;
+	code: string;
+	message: string;
+	metadata?: unknown;
+}
+
+export interface RuntimeCronFire {
+	action: string;
+	scheduledAt: number;
+	firedAt: number;
+	finishedAt?: number;
+	result: "running" | "ok" | "error" | "skipped";
+	error?: RuntimeScheduleErrorInfo;
+}
+
+export interface RuntimeScheduledFireInfo {
+	kind: "at" | "cron" | "every";
+	id: string;
+	name?: string;
+	scheduledAt: number;
+	firedAt: number;
+}
+
 export interface RuntimeQueueNextBatchOptions {
 	names?: string[];
 	count?: number;
@@ -152,6 +202,10 @@ export type RuntimeSqlBindParam =
 
 export type RuntimeSqlBindParams = RuntimeSqlBindParam[] | null;
 
+export interface RuntimeActorRuntimeSocketEndpointInfo {
+	path: string;
+}
+
 export interface RuntimeSqlQueryResult {
 	columns: string[];
 	rows: unknown[][];
@@ -164,6 +218,11 @@ export interface RuntimeSqlExecuteResult extends RuntimeSqlQueryResult {
 	lastInsertRowId?: number | null;
 }
 
+export interface RuntimeSqlBatchStatement {
+	sql: string;
+	params?: RuntimeSqlBindParams;
+}
+
 export function normalizeRuntimeSqlExecuteResult(
 	result: RuntimeSqlQueryResult & {
 		changes: number;
@@ -183,6 +242,9 @@ export interface RuntimeSqlDatabase {
 		sql: string,
 		params?: RuntimeSqlBindParams,
 	): Promise<RuntimeSqlExecuteResult>;
+	executeBatch(
+		statements: RuntimeSqlBatchStatement[],
+	): Promise<RuntimeSqlExecuteResult[]>;
 	query(
 		sql: string,
 		params?: RuntimeSqlBindParams,
@@ -191,16 +253,30 @@ export interface RuntimeSqlDatabase {
 		sql: string,
 		params?: RuntimeSqlBindParams,
 	): Promise<RuntimeSqlRunResult>;
+	beginTransaction(
+		timeoutMs?: number,
+	): Promise<RuntimeSqlTransactionDatabase>;
 	metrics?(): SqliteNativeMetrics | null;
 	takeLastKvError?(): string | null;
 	close(): Promise<void>;
 }
 
+export interface RuntimeSqlTransactionDatabase {
+	exec(sql: string): Promise<RuntimeSqlExecResult>;
+	execute(
+		sql: string,
+		params?: RuntimeSqlBindParams,
+	): Promise<RuntimeSqlExecuteResult>;
+	commit(): Promise<void>;
+	rollback(): Promise<void>;
+}
+
 export interface RuntimeActorConfig {
 	name?: string;
 	icon?: string;
 	hasDatabase?: boolean;
 	remoteSqlite?: boolean;
+	enableActorRuntimeSocket?: boolean;
 	hasState?: boolean;
 	canHibernateWebsocket?: boolean;
 	stateSaveIntervalMs?: number;
@@ -221,6 +297,7 @@ export interface RuntimeActorConfig {
 	connectionLivenessTimeoutMs?: number;
 	connectionLivenessIntervalMs?: number;
 	maxQueueSize?: number;
+	maxSchedules?: number;
 	maxQueueMessageSize?: number;
 	maxIncomingMessageSize?: number;
 	maxOutgoingMessageSize?: number;
@@ -410,6 +487,10 @@ export interface CoreRuntime {
 		ctx: ActorContextHandle,
 		payload: RuntimeStateDeltaPayload,
 	): Promise<void>;
+	actorSaveStateAndWorkflowBatch(
+		ctx: ActorContextHandle,
+		writes: RuntimeWorkflowKvWrite[],
+	): Promise<void>;
 	actorId(ctx: ActorContextHandle): string;
 	actorName(ctx: ActorContextHandle): string;
 	actorKey(ctx: ActorContextHandle): RuntimeActorKeySegment[];
@@ -491,6 +572,29 @@ export interface CoreRuntime {
 		sql: string,
 		params?: RuntimeSqlBindParams,
 	): Promise<RuntimeSqlExecuteResult>;
+	actorSqlExecuteBatch(
+		ctx: ActorContextHandle,
+		statements: RuntimeSqlBatchStatement[],
+	): Promise<RuntimeSqlExecuteResult[]>;
+	actorSqlBeginTransaction(
+		ctx: ActorContextHandle,
+		timeoutMs?: number,
+	): Promise<SqliteTransactionHandle>;
+	actorSqlTransactionExec(
+		transaction: SqliteTransactionHandle,
+		sql: string,
+	): Promise<RuntimeSqlExecResult>;
+	actorSqlTransactionExecute(
+		transaction: SqliteTransactionHandle,
+		sql: string,
+		params?: RuntimeSqlBindParams,
+	): Promise<RuntimeSqlExecuteResult>;
+	actorSqlTransactionCommit(
+		transaction: SqliteTransactionHandle,
+	): Promise<void>;
+	actorSqlTransactionRollback(
+		transaction: SqliteTransactionHandle,
+	): Promise<void>;
 	actorSqlQuery(
 		ctx: ActorContextHandle,
 		sql: string,
@@ -504,6 +608,9 @@ export interface CoreRuntime {
 	actorSqlMetrics(ctx: ActorContextHandle): SqliteNativeMetrics | null;
 	actorSqlTakeLastKvError(ctx: ActorContextHandle): string | null;
 	actorSqlClose(ctx: ActorContextHandle): Promise<void>;
+	actorRuntimeSocketProvision(
+		ctx: ActorContextHandle,
+	): Promise<RuntimeActorRuntimeSocketEndpointInfo>;
 
 	actorQueueSend(
 		ctx: ActorContextHandle,
@@ -549,13 +656,49 @@ export interface CoreRuntime {
 		durationMs: number,
 		actionName: string,
 		args: RuntimeBytes,
-	): void;
+	): Promise<string>;
 	actorScheduleAt(
 		ctx: ActorContextHandle,
 		timestampMs: number,
 		actionName: string,
 		args: RuntimeBytes,
-	): void;
+	): Promise<string>;
+	actorScheduleCancel(ctx: ActorContextHandle, id: string): Promise<boolean>;
+	actorScheduleGet(
+		ctx: ActorContextHandle,
+		id: string,
+	): Promise<RuntimeScheduledEventInfo | undefined>;
+	actorScheduleList(
+		ctx: ActorContextHandle,
+	): Promise<RuntimeScheduledEventInfo[]>;
+	actorCronSet(
+		ctx: ActorContextHandle,
+		name: string,
+		expression: string,
+		timezone: string | undefined,
+		actionName: string,
+		args: RuntimeBytes,
+		maxHistory: number | undefined,
+	): Promise<void>;
+	actorCronEvery(
+		ctx: ActorContextHandle,
+		name: string,
+		intervalMs: number,
+		actionName: string,
+		args: RuntimeBytes,
+		maxHistory: number | undefined,
+	): Promise<void>;
+	actorCronGet(
+		ctx: ActorContextHandle,
+		name: string,
+	): Promise<RuntimeCronJobInfo | undefined>;
+	actorCronList(ctx: ActorContextHandle): Promise<RuntimeCronJobInfo[]>;
+	actorCronDelete(ctx: ActorContextHandle, name: string): Promise<boolean>;
+	actorCronHistory(
+		ctx: ActorContextHandle,
+		name: string,
+		limit: number | undefined,
+	): Promise<RuntimeCronFire[]>;
 
 	connId(conn: ConnHandle): string;
 	connParams(conn: ConnHandle): RuntimeBytes;

package/src/registry/wasm-runtime.ts

@@ -33,12 +33,16 @@ import type {
 	RuntimeQueueNextBatchOptions,
 	RuntimeQueueTryNextBatchOptions,
 	RuntimeQueueWaitOptions,
+	RuntimeCronFire,
+	RuntimeCronJobInfo,
 	RuntimeRegistryRouteResponse,
 	RuntimeRequestSaveOpts,
 	RuntimeServeConfig,
+	RuntimeScheduledEventInfo,
 	RuntimeServerlessRequest,
 	RuntimeServerlessResponseHead,
 	RuntimeServerlessStreamCallback,
+	RuntimeSqlBatchStatement,
 	RuntimeSqlBindParams,
 	RuntimeSqlDatabase,
 	RuntimeSqlExecResult,
@@ -47,6 +51,8 @@ import type {
 	RuntimeSqlRunResult,
 	RuntimeStateDeltaPayload,
 	RuntimeWebSocketEvent,
+	RuntimeWorkflowKvWrite,
+	SqliteTransactionHandle,
 	WebSocketHandle,
 } from "./runtime";
 import { normalizeRuntimeSqlExecuteResult } from "./runtime";
@@ -464,6 +470,17 @@ export class WasmCoreRuntime implements CoreRuntime {
 		await callHandleAsync(asWasmActorContext(ctx), "saveState", payload);
 	}
 
+	async actorSaveStateAndWorkflowBatch(
+		ctx: ActorContextHandle,
+		writes: RuntimeWorkflowKvWrite[],
+	): Promise<void> {
+		await callHandleAsync(
+			asWasmActorContext(ctx),
+			"saveStateAndWorkflowBatch",
+			writes,
+		);
+	}
+
 	actorId(ctx: ActorContextHandle): string {
 		return callHandle(asWasmActorContext(ctx), "actorId");
 	}
@@ -701,6 +718,74 @@ export class WasmCoreRuntime implements CoreRuntime {
 		return normalizeRuntimeSqlExecuteResult(result);
 	}
 
+	async actorSqlExecuteBatch(
+		ctx: ActorContextHandle,
+		statements: RuntimeSqlBatchStatement[],
+	): Promise<RuntimeSqlExecuteResult[]> {
+		const results = await callWasm(() =>
+			this.#actorSql(ctx).executeBatch(statements),
+		);
+		return results.map(normalizeRuntimeSqlExecuteResult);
+	}
+
+	async actorSqlBeginTransaction(
+		ctx: ActorContextHandle,
+		timeoutMs?: number,
+	): Promise<SqliteTransactionHandle> {
+		return (await callWasm(() =>
+			this.#actorSql(ctx).beginTransaction(timeoutMs),
+		)) as unknown as SqliteTransactionHandle;
+	}
+
+	async actorSqlTransactionExec(
+		transaction: SqliteTransactionHandle,
+		sql: string,
+	): Promise<RuntimeSqlExecResult> {
+		return await callWasm(() =>
+			(
+				transaction as unknown as {
+					exec(sql: string): Promise<RuntimeSqlExecResult>;
+				}
+			).exec(sql),
+		);
+	}
+
+	async actorSqlTransactionExecute(
+		transaction: SqliteTransactionHandle,
+		sql: string,
+		params?: RuntimeSqlBindParams,
+	): Promise<RuntimeSqlExecuteResult> {
+		const result = await callWasm(() =>
+			(
+				transaction as unknown as {
+					execute(
+						sql: string,
+						params?: RuntimeSqlBindParams,
+					): Promise<RuntimeSqlExecuteResult>;
+				}
+			).execute(sql, params),
+		);
+		return normalizeRuntimeSqlExecuteResult(result);
+	}
+
+	async actorSqlTransactionCommit(
+		transaction: SqliteTransactionHandle,
+	): Promise<void> {
+		await callWasm(() =>
+			(transaction as unknown as { commit(): Promise<void> }).commit(),
+		);
+	}
+
+	async actorSqlTransactionRollback(
+		transaction: SqliteTransactionHandle,
+	): Promise<void> {
+		await callWasm(() =>
+			(
+				transaction as unknown as { rollback(): Promise<void> }
+			).rollback(),
+		);
+	}
+
 	async actorSqlQuery(
 		ctx: ActorContextHandle,
 		sql: string,
@@ -736,6 +821,20 @@ export class WasmCoreRuntime implements CoreRuntime {
 		await callWasm(() => database.close());
 	}
 
+	async actorRuntimeSocketProvision(
+		_ctx: ActorContextHandle,
+	): Promise<never> {
+		throw new RivetError(
+			"actor_runtime_socket",
+			"unsupported",
+			"Actor Runtime Socket is only available on Unix native runtimes.",
+			{
+				public: true,
+				metadata: { runtime: "wasm" },
+			},
+		);
+	}
+
 	async actorQueueSend(
 		ctx: ActorContextHandle,
 		name: string,
@@ -845,24 +944,136 @@ export class WasmCoreRuntime implements CoreRuntime {
 		await callHandleAsync(queue, "reset");
 	}
 
-	actorScheduleAfter(
+	async actorScheduleAfter(
 		ctx: ActorContextHandle,
 		durationMs: number | bigint,
 		actionName: string,
 		args: RuntimeBytes,
-	): void {
+	): Promise<string> {
 		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
-		callHandle(schedule, "after", wasmNumber(durationMs), actionName, args);
+		return await callHandleAsync<string>(
+			schedule,
+			"after",
+			wasmNumber(durationMs),
+			actionName,
+			args,
+		);
 	}
 
-	actorScheduleAt(
+	async actorScheduleAt(
 		ctx: ActorContextHandle,
 		timestampMs: number | bigint,
 		actionName: string,
 		args: RuntimeBytes,
-	): void {
+	): Promise<string> {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<string>(
+			schedule,
+			"at",
+			wasmNumber(timestampMs),
+			actionName,
+			args,
+		);
+	}
+
+	async actorScheduleCancel(ctx: ActorContextHandle, id: string) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<boolean>(schedule, "cancel", id);
+	}
+
+	async actorScheduleGet(ctx: ActorContextHandle, id: string) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<RuntimeScheduledEventInfo | undefined>(
+			schedule,
+			"get",
+			id,
+		);
+	}
+
+	async actorScheduleList(ctx: ActorContextHandle) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<RuntimeScheduledEventInfo[]>(
+			schedule,
+			"list",
+		);
+	}
+
+	async actorCronSet(
+		ctx: ActorContextHandle,
+		name: string,
+		expression: string,
+		timezone: string | undefined,
+		actionName: string,
+		args: RuntimeBytes,
+		maxHistory: number | undefined,
+	) {
 		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
-		callHandle(schedule, "at", wasmNumber(timestampMs), actionName, args);
+		await callHandleAsync(
+			schedule,
+			"cronSet",
+			name,
+			expression,
+			timezone,
+			actionName,
+			args,
+			maxHistory,
+		);
+	}
+
+	async actorCronEvery(
+		ctx: ActorContextHandle,
+		name: string,
+		intervalMs: number,
+		actionName: string,
+		args: RuntimeBytes,
+		maxHistory: number | undefined,
+	) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		await callHandleAsync(
+			schedule,
+			"cronEvery",
+			name,
+			intervalMs,
+			actionName,
+			args,
+			maxHistory,
+		);
+	}
+
+	async actorCronGet(ctx: ActorContextHandle, name: string) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<RuntimeCronJobInfo | undefined>(
+			schedule,
+			"cronGet",
+			name,
+		);
+	}
+
+	async actorCronList(ctx: ActorContextHandle) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<RuntimeCronJobInfo[]>(
+			schedule,
+			"cronList",
+		);
+	}
+
+	async actorCronDelete(ctx: ActorContextHandle, name: string) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<boolean>(schedule, "cronDelete", name);
+	}
+
+	async actorCronHistory(
+		ctx: ActorContextHandle,
+		name: string,
+		limit: number | undefined,
+	) {
+		const schedule = childHandle(asWasmActorContext(ctx), "schedule");
+		return await callHandleAsync<RuntimeCronFire[]>(
+			schedule,
+			"cronHistory",
+			name,
+			limit,
+		);
 	}
 
 	connId(conn: ConnHandle): string {

package/src/utils.ts

@@ -10,6 +10,9 @@ export { stringifyError };
 /** @experimental */
 export { assertUnreachable };
 
+/** @experimental */
+export { jsonParseCompat, jsonStringifyCompat } from "./common/encoding";
+
 /**
  * Joins multiple abort signals into one.
  *

package/src/workflow/context.ts

@@ -17,8 +17,8 @@ import type {
 	QueueNextBatchOptions,
 	QueueNextOptions,
 	QueueResultMessageForName,
+	RunContext,
 } from "@/actor/config";
-import { RAW_STATE_SYMBOL, type RunContext } from "@/actor/config";
 import type {
 	AnyActorDefinition,
 	BaseActorDefinition,
@@ -469,7 +469,6 @@ export class WorkflowContext<
 		TEvents,
 		TQueues
 	>;
-
 	constructor(
 		inner: WorkflowContextInterface,
 		runCtx: RunContext<
@@ -968,9 +967,9 @@ export class WorkflowContext<
 		return this.#inner.isEvicted();
 	}
 
-	// Runs a user step body inside a fresh step context, snapshotting actor
-	// state/vars so a thrown step rolls back its mutations, and deactivating the
-	// step context once the body settles so it cannot be used after the step.
+	// Runs a user step body inside a fresh step context and deactivates the step
+	// context once the body settles. Workflows are roll-forward only, so state and
+	// vars mutations made before a failure always remain visible.
 	async #runStep<T>(
 		run: WorkflowStepRun<
 			T,
@@ -995,28 +994,8 @@ export class WorkflowContext<
 			TQueues
 		>(this.#runCtx, () => this.#markGuardTriggered());
 
-		let stateSnapshot: { state: TState } | null = null;
-		try {
-			stateSnapshot = { state: this.#runCtx[RAW_STATE_SYMBOL]() };
-		} catch (error) {
-			this.#runCtx.log.debug({
-				msg: "failed to get state, likely due to being stateless workflow",
-				error,
-			});
-		}
-		if (stateSnapshot) {
-			stateSnapshot.state = structuredClone(stateSnapshot.state);
-		}
-		const varsSnapshot = structuredClone(this.#runCtx.vars);
-
 		try {
 			return await run(stepCtx);
-		} catch (error) {
-			if (stateSnapshot) {
-				this.#runCtx.state = stateSnapshot.state;
-			}
-			this.#runCtx.vars = varsSnapshot;
-			throw error;
 		} finally {
 			stepCtx[DEACTIVATE_STEP]();
 		}
@@ -1040,8 +1019,7 @@ export class WorkflowContext<
 		};
 	}
 
-	// Runs a step rollback compensation with an active step context. Rollbacks
-	// intentionally mutate actor state, so their writes are not snapshotted.
+	// Runs a step rollback compensation with an active step context.
 	async #runRollback<T>(
 		rollback: WorkflowStepRollback<
 			T,

package/src/workflow/driver.test.ts

@@ -0,0 +1,62 @@
+import { describe, expect, test } from "vitest";
+import { ActorWorkflowDriver, chunkWorkflowWrites } from "./driver";
+
+function write(keyBytes: number, valueBytes: number) {
+	return {
+		key: new Uint8Array(keyBytes),
+		value: new Uint8Array(valueBytes),
+	};
+}
+
+describe("workflow sqlite batching", () => {
+	test.each([
+		[127, 1],
+		[128, 1],
+		[129, 2],
+	])("chunks %i rows into %i transaction(s)", (rowCount, chunkCount) => {
+		const chunks = chunkWorkflowWrites(
+			Array.from({ length: rowCount }, () => write(1, 1)),
+		);
+		expect(chunks).toHaveLength(chunkCount);
+		expect(chunks.flat()).toHaveLength(rowCount);
+		expect(chunks.every((chunk) => chunk.length <= 128)).toBe(true);
+	});
+
+	test("splits immediately above the 512 KiB payload budget", () => {
+		// Each persisted key also includes the two-byte workflow namespace.
+		const exact = [write(1, 256 * 1024 - 3), write(1, 256 * 1024 - 3)];
+		expect(chunkWorkflowWrites(exact)).toHaveLength(1);
+
+		const over = [write(1, 256 * 1024 - 3), write(1, 256 * 1024 - 2)];
+		expect(chunkWorkflowWrites(over)).toHaveLength(2);
+	});
+
+	test("delegates actor state and the whole workflow batch to one atomic runtime call", async () => {
+		let receivedWrites: Array<{ key: Uint8Array; value: Uint8Array }> = [];
+		const sql = {
+			query: async () => ({ columns: [], rows: [] }),
+			execute: async () => [],
+			run: async () => ({ changes: 0 }),
+			executeBatch: async () => [],
+		};
+		const actor = {
+			stateManager: {
+				saveStateAndWorkflowBatch: async (
+					writes: Array<{ key: Uint8Array; value: Uint8Array }>,
+				) => {
+					receivedWrites = writes;
+				},
+			},
+			queueManager: {},
+		};
+		const runCtx = {
+			sql,
+			internalKeepAwake: async <T>(promise: Promise<T>) => await promise,
+		};
+		const driver = new ActorWorkflowDriver(actor as never, runCtx as never);
+		const writes = [write(1, 1), write(2, 2)];
+		await driver.batch(writes);
+		expect(driver.atomicBatch).toBe(true);
+		expect(receivedWrites).toBe(writes);
+	});
+});