rivetkit 2.3.11-rc.8 → 2.3.11

package/src/common/database/mod.ts

@@ -1,23 +1,78 @@
 import { getLogger } from "@/common/log";
 import type {
 	DatabaseProvider,
+	NativeDatabaseProvider,
 	RawAccess,
 	SqliteDatabase,
+	SqliteProfilingOptions,
 	SqliteTransactionDatabase,
+	SqliteTransactionOptions,
 } from "./config";
 import {
 	isManualTransactionControl,
 	isSqliteBindingObject,
 	MIGRATION_TRANSACTION_TIMEOUT_MS,
 	toSqliteBindings,
+	validateTransactionName,
 	validateTransactionTimeout,
 } from "./shared";
 
 export type { RawAccess } from "./config";
 
-interface DatabaseFactoryConfig {
+export interface DatabaseFactoryConfig {
 	onMigrate?: (db: RawAccess) => Promise<void> | void;
 	warnOnManualTransactions?: boolean;
+	/**
+	 * SQLite profiling configuration.
+	 *
+	 * @experimental This entire configuration surface is experimental and
+	 * subject to change without notice.
+	 */
+	profiling?: SqliteProfilingOptions;
+}
+const nativeStateTransactionOpeners = new WeakMap<
+	NativeDatabaseProvider,
+	(
+		timeoutMs?: number,
+		context?: unknown,
+	) => Promise<SqliteTransactionDatabase>
+>();
+type NativeStateTransactionContext = {
+	enter(): Promise<unknown>;
+	exit(scope: unknown): void;
+};
+const nativeStateTransactionClientBinders = new WeakMap<
+	object,
+	(context: NativeStateTransactionContext) => object
+>();
+
+/** @internal */
+export function registerNativeStateTransactionOpener<
+	T extends NativeDatabaseProvider,
+>(
+	provider: T,
+	opener: (
+		timeoutMs?: number,
+		context?: unknown,
+	) => Promise<SqliteTransactionDatabase>,
+): T {
+	nativeStateTransactionOpeners.set(provider, opener);
+	return provider;
+}
+
+/** @internal */
+export function bindNativeStateTransactionContext<T>(
+	client: T,
+	context: NativeStateTransactionContext,
+): T {
+	if (
+		(typeof client !== "object" || client === null) &&
+		typeof client !== "function"
+	) {
+		return client;
+	}
+	const bind = nativeStateTransactionClientBinders.get(client as object);
+	return (bind?.(context) ?? client) as T;
 }
 
 function hasMultipleStatements(query: string): boolean {
@@ -28,8 +83,10 @@ function hasMultipleStatements(query: string): boolean {
 export function db({
 	onMigrate,
 	warnOnManualTransactions = true,
+	profiling,
 }: DatabaseFactoryConfig = {}): DatabaseProvider<RawAccess> {
-	return {
+	const provider: DatabaseProvider<RawAccess> = {
+		sqliteProfiling: profiling,
 		createClient: async (ctx) => {
 			const nativeDatabaseProvider = ctx.nativeDatabaseProvider;
 			if (!nativeDatabaseProvider) {
@@ -52,110 +109,161 @@ export function db({
 			const createClient = (
 				target: SqliteDatabase | SqliteTransactionDatabase,
 				transactionScoped = false,
-			): RawAccess => ({
-				execute: async <
-					TRow extends Record<string, unknown> = Record<
-						string,
-						unknown
-					>,
-				>(
-					query: string,
-					...args: unknown[]
-				): Promise<TRow[]> => {
-					ensureOpen();
-					if (
-						!transactionScoped &&
-						warnOnManualTransactions &&
-						!manualTransactionWarned &&
-						!hasMultipleStatements(query) &&
-						isManualTransactionControl(query)
-					) {
-						manualTransactionWarned = true;
-						getLogger("database").warn(
-							{ actorId: ctx.actorId },
-							"Manual cross-call SQLite transactions can interleave with other actor work. Use db.transaction() for coordinated transactions. Set warnOnManualTransactions: false in your db(...) configuration to disable this warning.",
-						);
-					}
-
-					const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
-					const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
-					const start = performance.now();
-
-					try {
-						if (args.length > 0) {
-							const bindings =
-								args.length === 1 &&
-								isSqliteBindingObject(args[0])
-									? toSqliteBindings(args[0])
-									: toSqliteBindings(args);
-							const { rows, columns } = await target.execute(
-								query,
-								bindings,
-							);
-							return rows.map((row) =>
-								rowToObject<TRow>(row, columns),
+				stateTransactionContext?: NativeStateTransactionContext,
+			): RawAccess => {
+				const client: RawAccess = {
+					execute: async <
+						TRow extends Record<string, unknown> = Record<
+							string,
+							unknown
+						>,
+					>(
+						query: string,
+						...args: unknown[]
+					): Promise<TRow[]> => {
+						ensureOpen();
+						if (
+							!transactionScoped &&
+							warnOnManualTransactions &&
+							!manualTransactionWarned &&
+							!hasMultipleStatements(query) &&
+							isManualTransactionControl(query)
+						) {
+							manualTransactionWarned = true;
+							getLogger("database").warn(
+								{ actorId: ctx.actorId },
+								"Manual cross-call SQLite transactions can interleave with other actor work. Use db.transaction() for coordinated transactions. Set warnOnManualTransactions: false in your db(...) configuration to disable this warning.",
 							);
 						}
 
-						if (!hasMultipleStatements(query)) {
-							const { rows, columns } = await target.execute(
+						const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
+						const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
+						const start = performance.now();
+
+						try {
+							if (args.length > 0) {
+								const bindings =
+									args.length === 1 &&
+									isSqliteBindingObject(args[0])
+										? toSqliteBindings(args[0])
+										: toSqliteBindings(args);
+								const { rows, columns } = await target.execute(
+									query,
+									bindings,
+								);
+								return rows.map((row) =>
+									rowToObject<TRow>(row, columns),
+								);
+							}
+
+							if (!hasMultipleStatements(query)) {
+								const { rows, columns } = await target.execute(
+									query,
+									undefined,
+								);
+								return rows.map((row) =>
+									rowToObject<TRow>(row, columns),
+								);
+							}
+
+							return await execMultiStatement<TRow>(
+								target,
 								query,
-								undefined,
-							);
-							return rows.map((row) =>
-								rowToObject<TRow>(row, columns),
 							);
+						} finally {
+							const durationMs = performance.now() - start;
+							ctx.metrics?.trackSql(query, durationMs);
+							if (ctx.metrics) {
+								const kvReads =
+									ctx.metrics.totalKvReads - kvReadsBefore;
+								const kvWrites =
+									ctx.metrics.totalKvWrites - kvWritesBefore;
+								ctx.log?.debug({
+									msg: "sql query",
+									query: query.slice(0, 120),
+									durationMs,
+									kvReads,
+									kvWrites,
+								});
+							}
 						}
-
-						return await execMultiStatement<TRow>(target, query);
-					} finally {
-						const durationMs = performance.now() - start;
-						ctx.metrics?.trackSql(query, durationMs);
-						if (ctx.metrics) {
-							const kvReads =
-								ctx.metrics.totalKvReads - kvReadsBefore;
-							const kvWrites =
-								ctx.metrics.totalKvWrites - kvWritesBefore;
-							ctx.log?.debug({
-								msg: "sql query",
-								query: query.slice(0, 120),
-								durationMs,
-								kvReads,
-								kvWrites,
-							});
+					},
+					transaction: async <T>(
+						callback: (tx: RawAccess) => Promise<T> | T,
+						options?: SqliteTransactionOptions,
+					): Promise<T> => {
+						validateTransactionTimeout(options?.timeout);
+						validateTransactionName(options?.name);
+						if (
+							transactionScoped &&
+							options?.experimental?.includeState
+						) {
+							throw new Error(
+								"experimental.includeState is not supported for nested transactions",
+							);
 						}
-					}
-				},
-				transaction: async <T>(
-					callback: (tx: RawAccess) => Promise<T> | T,
-					options?: { timeout?: number },
-				): Promise<T> => {
-					validateTransactionTimeout(options?.timeout);
-					const transaction = await db.beginTransaction(
-						options?.timeout,
-					);
-					const tx = createClient(transaction, true);
-					try {
-						const result = await callback(tx);
-						await transaction.commit();
-						return result;
-					} catch (error) {
+						const includeState =
+							options?.experimental?.includeState === true;
+						const stateScope =
+							includeState && stateTransactionContext
+								? await stateTransactionContext.enter()
+								: undefined;
 						try {
-							await transaction.rollback();
-						} catch {
-							// Preserve the callback or commit error after expiry cleanup.
+							const transaction = includeState
+								? await (() => {
+										const beginStateTransaction =
+											ctx.nativeDatabaseProvider &&
+											nativeStateTransactionOpeners.get(
+												ctx.nativeDatabaseProvider,
+											);
+										if (!beginStateTransaction) {
+											throw new Error(
+												"experimental.includeState is only supported by RivetKit's embedded database provider",
+											);
+										}
+										return beginStateTransaction(
+											options?.timeout,
+											stateScope,
+										);
+									})()
+								: await db.beginTransaction(
+										options?.timeout,
+										options?.name,
+									);
+							const tx = createClient(transaction, true);
+							try {
+								const result = await callback(tx);
+								await transaction.commit();
+								return result;
+							} catch (error) {
+								try {
+									await transaction.rollback();
+								} catch {
+									// Preserve the callback or commit error after expiry cleanup.
+								}
+								throw error;
+							}
+						} finally {
+							if (stateScope !== undefined) {
+								stateTransactionContext?.exit(stateScope);
+							}
 						}
-						throw error;
-					}
-				},
-				close: async () => {
-					if (!closed) {
-						closed = true;
-						await db.close();
-					}
-				},
-				nativeMetrics: () => db.nativeMetrics?.() ?? null,
-			});
+					},
+					close: async () => {
+						if (!closed) {
+							closed = true;
+							await db.close();
+						}
+					},
+					nativeMetrics: () => db.nativeMetrics?.() ?? null,
+				};
+				if (!transactionScoped) {
+					nativeStateTransactionClientBinders.set(client, (context) =>
+						createClient(target, false, context),
+					);
+				}
+				return client;
+			};
 			const client = createClient(db);
 			return client;
 		},
@@ -167,6 +275,7 @@ export function db({
 			}
 		},
 	};
+	return provider;
 }
 
 function rowToObject<TRow extends Record<string, unknown>>(
@@ -221,6 +330,9 @@ async function withMigrationSavepoint<T>(
 				throw error;
 			}
 		},
-		{ timeout: MIGRATION_TRANSACTION_TIMEOUT_MS },
+		{
+			name: "rivetkit-migration",
+			timeout: MIGRATION_TRANSACTION_TIMEOUT_MS,
+		},
 	);
 }

package/src/common/database/native-database.ts

@@ -81,7 +81,14 @@ export interface JsNativeDatabaseLike {
 	executeBatch?(
 		statements: NativeBatchStatement[],
 	): Promise<NativeExecuteResult[]>;
-	beginTransaction(timeoutMs?: number): Promise<JsNativeTransactionLike>;
+	beginStateTransaction?(
+		timeoutMs?: number,
+		context?: unknown,
+	): Promise<JsNativeTransactionLike>;
+	beginTransaction(
+		timeoutMs?: number,
+		name?: string,
+	): Promise<JsNativeTransactionLike>;
 	query(
 		sql: string,
 		params?: NativeBindParam[] | null,
@@ -94,6 +101,12 @@ export interface JsNativeDatabaseLike {
 	takeLastKvError?(): string | null;
 	close(): Promise<void>;
 }
+export type StateAwareSqliteDatabase = SqliteDatabase & {
+	beginStateTransaction(
+		timeoutMs?: number,
+		context?: unknown,
+	): Promise<SqliteTransactionDatabase>;
+};
 
 export interface JsNativeTransactionLike {
 	exec(sql: string): Promise<NativeExecResult>;
@@ -312,7 +325,7 @@ class NativeCloseGate {
 
 export function wrapJsNativeDatabase(
 	database: JsNativeDatabaseLike,
-): SqliteDatabase {
+): StateAwareSqliteDatabase {
 	const gate = new NativeCloseGate();
 	let closePromise: Promise<void> | undefined;
 	let lastInsertRowId: number | null = null;
@@ -419,11 +432,37 @@ export function wrapJsNativeDatabase(
 		},
 		async beginTransaction(
 			timeoutMs?: number,
+			name?: string,
 		): Promise<SqliteTransactionDatabase> {
 			const release = gate.enter();
 			let transaction: JsNativeTransactionLike;
 			try {
-				transaction = await database.beginTransaction(timeoutMs);
+				transaction = await database.beginTransaction(timeoutMs, name);
+			} catch (error) {
+				enrichNativeDatabaseError(database, error);
+			} finally {
+				release();
+			}
+			return wrapTransaction(database, transaction, gate, (result) => {
+				if (result.lastInsertRowId !== undefined) {
+					lastInsertRowId = result.lastInsertRowId;
+				}
+			});
+		},
+		async beginStateTransaction(
+			timeoutMs?: number,
+			context?: unknown,
+		): Promise<SqliteTransactionDatabase> {
+			if (!database.beginStateTransaction) {
+				throw new Error("actor state transactions are not configured");
+			}
+			const release = gate.enter();
+			let transaction: JsNativeTransactionLike;
+			try {
+				transaction = await database.beginStateTransaction(
+					timeoutMs,
+					context,
+				);
 			} catch (error) {
 				enrichNativeDatabaseError(database, error);
 			} finally {

package/src/common/database/shared.ts

@@ -15,6 +15,15 @@ export function validateTransactionTimeout(timeout: number | undefined): void {
 	}
 }
 
+export function validateTransactionName(name: string | undefined): void {
+	if (name === undefined) {
+		return;
+	}
+	if (name.length === 0) {
+		throw new Error("db.transaction() name must not be empty");
+	}
+}
+
 type SqliteBindingObject = Record<string, unknown>;
 
 function isSqliteBindingValue(value: unknown): boolean {

package/src/db/drizzle.test.ts

@@ -12,6 +12,7 @@ import { db } from "./drizzle";
 class FakeSqliteDatabase implements SqliteDatabase {
 	executeCalls: Array<{ sql: string; params?: SqliteBindings }> = [];
 	transactionTimeouts: Array<number | undefined> = [];
+	transactionNames: Array<string | undefined> = [];
 
 	async exec(): Promise<void> {}
 
@@ -44,8 +45,10 @@ class FakeSqliteDatabase implements SqliteDatabase {
 
 	async beginTransaction(
 		timeoutMs?: number,
+		name?: string,
 	): Promise<SqliteTransactionDatabase> {
 		this.transactionTimeouts.push(timeoutMs);
+		this.transactionNames.push(name);
 		this.executeCalls.push({ sql: "BEGIN" });
 		return {
 			exec: async () => {},
@@ -114,6 +117,9 @@ describe("Drizzle database transactions", () => {
 		await provider.onMigrate(client);
 
 		expect(nativeDb.transactionTimeouts).toEqual([300_000]);
+		expect(nativeDb.transactionNames).toEqual([
+			"rivetkit-drizzle-migration",
+		]);
 		expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([
 			"BEGIN",
 			"SAVEPOINT __rivet_on_migrate",
@@ -133,10 +139,11 @@ describe("Drizzle database transactions", () => {
 					"inside",
 				);
 			},
-			{ timeout: 120_000 },
+			{ name: "drizzle-insert", timeout: 120_000 },
 		);
 
 		expect(nativeDb.transactionTimeouts).toEqual([120_000]);
+		expect(nativeDb.transactionNames).toEqual(["drizzle-insert"]);
 		expect(nativeDb.executeCalls.map(({ sql }) => sql)).toEqual([
 			"BEGIN",
 			"INSERT INTO items(value) VALUES (?)",

package/src/db/drizzle.ts

@@ -7,6 +7,8 @@ import type {
 	DatabaseProvider,
 	DatabaseProviderContext,
 	RawAccess,
+	SqliteProfilingOptions,
+	SqliteTransactionOptions,
 	SqliteDatabase,
 	SqliteTransactionDatabase,
 } from "@/common/database/config";
@@ -15,6 +17,7 @@ import {
 	isManualTransactionControl,
 	MIGRATION_TRANSACTION_TIMEOUT_MS,
 	toSqliteBindings,
+	validateTransactionName,
 	validateTransactionTimeout,
 } from "@/common/database/shared";
 import { sha256Hex } from "@/utils/crypto";
@@ -42,7 +45,7 @@ type DrizzleDatabase<TSchema extends DrizzleSchema> = Omit<
 	Omit<RawAccess, "transaction"> & {
 		transaction: <T>(
 			callback: (tx: DrizzleDatabase<TSchema>) => Promise<T> | T,
-			options?: { timeout?: number },
+			options?: SqliteTransactionOptions,
 		) => Promise<T>;
 	};
 
@@ -58,11 +61,12 @@ interface DrizzleMigrations {
 	migrations: Record<string, string>;
 }
 
-interface DrizzleDatabaseFactoryConfig<TSchema extends DrizzleSchema> {
+export interface DrizzleDatabaseFactoryConfig<TSchema extends DrizzleSchema> {
 	schema?: TSchema;
 	migrations?: DrizzleMigrations;
 	onMigrate?: (db: DrizzleDatabase<TSchema>) => Promise<void> | void;
 	warnOnManualTransactions?: boolean;
+	profiling?: SqliteProfilingOptions;
 }
 
 interface DrizzleKitConfig {
@@ -86,10 +90,12 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 	migrations,
 	onMigrate,
 	warnOnManualTransactions = true,
+	profiling,
 }: DrizzleDatabaseFactoryConfig<TSchema> = {}): DatabaseProvider<
 	DrizzleDatabase<TSchema>
 > {
 	return {
+		sqliteProfiling: profiling,
 		createClient: async (ctx) => {
 			const override = ctx.overrideDrizzleDatabaseClient
 				? await ctx.overrideDrizzleDatabaseClient()
@@ -194,11 +200,13 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 					transactionCallback: (
 						tx: DrizzleDatabase<TSchema>,
 					) => Promise<T> | T,
-					options?: { timeout?: number },
+					options?: SqliteTransactionOptions,
 				): Promise<T> => {
 					validateTransactionTimeout(options?.timeout);
+					validateTransactionName(options?.name);
 					const transaction = await nativeDb.beginTransaction(
 						options?.timeout,
+						options?.name,
 					);
 					const tx = createDrizzleClient(transaction, true);
 					try {
@@ -286,7 +294,10 @@ async function withMigrationSavepoint<TSchema extends DrizzleSchema, T>(
 				throw error;
 			}
 		},
-		{ timeout: MIGRATION_TRANSACTION_TIMEOUT_MS },
+		{
+			name: "rivetkit-drizzle-migration",
+			timeout: MIGRATION_TRANSACTION_TIMEOUT_MS,
+		},
 	);
 }
 

package/src/db/mod.ts

@@ -2,12 +2,15 @@ export type {
 	AnyDatabaseProvider,
 	DatabaseProvider,
 	DatabaseProviderContext,
+	InferDatabaseClient,
 	NativeDatabaseProvider,
 	RawAccess,
 	RawDatabaseClient,
 	SqliteBindings,
 	SqliteDatabase,
 	SqliteNativeMetrics,
+	SqliteProfilingOptions,
 	SqliteQueryResult,
+	SqliteTransactionOptions,
 } from "@/common/database/config";
 export { db } from "@/common/database/mod";

package/src/inspector/workflow.ts

@@ -0,0 +1,61 @@
+import type { JsonCompatValue } from "@/common/encoding";
+import type { WorkflowHistoryBytes } from "@/common/inspector-transport";
+import { encodeCborCompat } from "@/serde";
+import { bufferToArrayBuffer } from "@/utils";
+
+export type {
+	WorkflowBranchStatus,
+	WorkflowCbor,
+	WorkflowEntry,
+	WorkflowEntryKind,
+	WorkflowEntryMetadata,
+	WorkflowHistory,
+	WorkflowJoinEntry,
+	WorkflowLocation,
+	WorkflowLoopEntry,
+	WorkflowLoopIterationMarker,
+	WorkflowMessageEntry,
+	WorkflowNameIndex,
+	WorkflowPathSegment,
+	WorkflowRaceEntry,
+	WorkflowRemovedEntry,
+	WorkflowRollbackCheckpointEntry,
+	WorkflowSleepEntry,
+	WorkflowStepEntry,
+	WorkflowVersionCheckEntry,
+} from "@/common/bare/transport/v1";
+export {
+	WorkflowBranchStatusType,
+	WorkflowEntryStatus,
+	WorkflowSleepState,
+} from "@/common/bare/transport/v1";
+export type { WorkflowHistoryBytes } from "@/common/inspector-transport";
+export {
+	decodeWorkflowHistoryTransport,
+	encodeWorkflowHistoryTransport,
+} from "@/common/inspector-transport";
+
+/** @experimental State exposed by a durable workflow run handler to the Inspector. */
+export type WorkflowState =
+	| "pending"
+	| "running"
+	| "sleeping"
+	| "failed"
+	| "completed"
+	| "cancelled"
+	| "rolling_back";
+
+/** @experimental The raw workflow Inspector adapter consumed by RivetKit's transport. */
+export interface WorkflowInspectorAdapter {
+	getHistory: () => WorkflowHistoryBytes | null;
+	getState: () => Promise<WorkflowState | null>;
+	onHistoryUpdated: (
+		listener: (history: WorkflowHistoryBytes) => void,
+	) => () => void;
+	replayFromStep: (entryId?: string) => Promise<WorkflowHistoryBytes | null>;
+}
+
+/** @experimental Encodes a workflow Inspector value with RivetKit's CBOR-compatible codec. */
+export function encodeWorkflowInspectorValue(value: unknown): ArrayBuffer {
+	return bufferToArrayBuffer(encodeCborCompat(value as JsonCompatValue));
+}