rivetkit 2.3.3 → 2.3.5-rc.1

package/src/db/drizzle.ts

@@ -8,8 +8,15 @@ import type {
 	DatabaseProviderContext,
 	RawAccess,
 	SqliteDatabase,
+	SqliteTransactionDatabase,
 } from "@/common/database/config";
-import { toSqliteBindings } from "@/common/database/shared";
+import { getLogger } from "@/common/log";
+import {
+	isManualTransactionControl,
+	MIGRATION_TRANSACTION_TIMEOUT_MS,
+	toSqliteBindings,
+	validateTransactionTimeout,
+} from "@/common/database/shared";
 import { sha256Hex } from "@/utils/crypto";
 
 export type { SQLiteTable } from "drizzle-orm/sqlite-core";
@@ -28,8 +35,16 @@ export {
 } from "drizzle-orm/sqlite-core";
 
 type DrizzleSchema = Record<string, unknown>;
-type DrizzleDatabase<TSchema extends DrizzleSchema> =
-	SqliteRemoteDatabase<TSchema> & RawAccess;
+type DrizzleDatabase<TSchema extends DrizzleSchema> = Omit<
+	SqliteRemoteDatabase<TSchema>,
+	"transaction"
+> &
+	Omit<RawAccess, "transaction"> & {
+		transaction: <T>(
+			callback: (tx: DrizzleDatabase<TSchema>) => Promise<T> | T,
+			options?: { timeout?: number },
+		) => Promise<T>;
+	};
 
 interface DrizzleMigrationJournalEntry {
 	idx: number;
@@ -47,6 +62,7 @@ interface DrizzleDatabaseFactoryConfig<TSchema extends DrizzleSchema> {
 	schema?: TSchema;
 	migrations?: DrizzleMigrations;
 	onMigrate?: (db: DrizzleDatabase<TSchema>) => Promise<void> | void;
+	warnOnManualTransactions?: boolean;
 }
 
 interface DrizzleKitConfig {
@@ -69,6 +85,7 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 	schema,
 	migrations,
 	onMigrate,
+	warnOnManualTransactions = true,
 }: DrizzleDatabaseFactoryConfig<TSchema> = {}): DatabaseProvider<
 	DrizzleDatabase<TSchema>
 > {
@@ -90,6 +107,7 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 
 			const nativeDb = await nativeDatabaseProvider.open(ctx.actorId);
 			let closed = false;
+			let manualTransactionWarned = false;
 			const ensureOpen = () => {
 				if (closed) {
 					throw new Error(
@@ -98,107 +116,178 @@ export function db<TSchema extends DrizzleSchema = Record<string, never>>({
 				}
 			};
 
-			const runSql = async (
-				query: string,
-				params: unknown[],
-				method: "run" | "all" | "values" | "get",
-			) => {
-				ensureOpen();
+			const createDrizzleClient = (
+				target: SqliteDatabase | SqliteTransactionDatabase,
+				transactionScoped = false,
+			): DrizzleDatabase<TSchema> => {
+				const runSql = async (
+					query: string,
+					params: unknown[],
+					method: "run" | "all" | "values" | "get",
+				) => {
+					ensureOpen();
+					warnForManualTransaction(query, transactionScoped);
 
-				const start = performance.now();
-				const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
-				const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
-				try {
-					const { rows } = await nativeDb.execute(
+					const start = performance.now();
+					const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
+					const kvWritesBefore = ctx.metrics?.totalKvWrites ?? 0;
+					try {
+						const { rows } = await target.execute(
+							query,
+							toSqliteBindings(params),
+						);
+						if (method === "run") {
+							return { rows: [] };
+						}
+						if (method === "get") {
+							return { rows: rows[0] };
+						}
+						return { rows };
+					} finally {
+						const durationMs = performance.now() - start;
+						ctx.metrics?.trackSql(query, durationMs);
+						if (ctx.metrics) {
+							ctx.log?.debug({
+								msg: "sql query",
+								query: query.slice(0, 120),
+								durationMs,
+								kvReads:
+									ctx.metrics.totalKvReads - kvReadsBefore,
+								kvWrites:
+									ctx.metrics.totalKvWrites - kvWritesBefore,
+							});
+						}
+					}
+				};
+
+				const callback: RemoteCallback = async (
+					query,
+					params,
+					method,
+				) => {
+					return await runSql(query, params, method);
+				};
+
+				const drizzleDb = drizzle(callback, {
+					schema,
+				}) as unknown as DrizzleDatabase<TSchema>;
+				drizzleDb.execute = async <
+					TRow extends Record<string, unknown> = Record<
+						string,
+						unknown
+					>,
+				>(
+					query: string,
+					...args: unknown[]
+				): Promise<TRow[]> => {
+					return await executeRaw<TRow>(
+						target,
+						ctx,
+						ensureOpen,
 						query,
-						toSqliteBindings(params),
+						args,
+						() =>
+							warnForManualTransaction(query, transactionScoped),
 					);
-					if (method === "run") {
-						return { rows: [] };
-					}
-					if (method === "get") {
-						return { rows: rows[0] };
+				};
+				drizzleDb.transaction = async <T>(
+					transactionCallback: (
+						tx: DrizzleDatabase<TSchema>,
+					) => Promise<T> | T,
+					options?: { timeout?: number },
+				): Promise<T> => {
+					validateTransactionTimeout(options?.timeout);
+					const transaction = await nativeDb.beginTransaction(
+						options?.timeout,
+					);
+					const tx = createDrizzleClient(transaction, true);
+					try {
+						const result = await transactionCallback(tx);
+						await transaction.commit();
+						return result;
+					} catch (error) {
+						try {
+							await transaction.rollback();
+						} catch {
+							// Preserve the callback or commit error after expiry cleanup.
+						}
+						throw error;
 					}
-					return { rows };
-				} finally {
-					const durationMs = performance.now() - start;
-					ctx.metrics?.trackSql(query, durationMs);
-					if (ctx.metrics) {
-						ctx.log?.debug({
-							msg: "sql query",
-							query: query.slice(0, 120),
-							durationMs,
-							kvReads: ctx.metrics.totalKvReads - kvReadsBefore,
-							kvWrites:
-								ctx.metrics.totalKvWrites - kvWritesBefore,
-						});
+				};
+				drizzleDb.close = async () => {
+					if (!closed) {
+						closed = true;
+						await nativeDb.close();
 					}
-				}
-			};
+				};
 
-			const callback: RemoteCallback = async (query, params, method) => {
-				return await runSql(query, params, method);
+				return drizzleDb;
 			};
 
-			const drizzleDb = drizzle(callback, {
-				schema,
-			}) as DrizzleDatabase<TSchema>;
-			drizzleDb.execute = async <
-				TRow extends Record<string, unknown> = Record<string, unknown>,
-			>(
+			const warnForManualTransaction = (
 				query: string,
-				...args: unknown[]
-			): Promise<TRow[]> => {
-				return await executeRaw<TRow>(
-					nativeDb,
-					ctx,
-					ensureOpen,
-					query,
-					args,
-				);
-			};
-			drizzleDb.close = async () => {
-				if (!closed) {
-					closed = true;
-					await nativeDb.close();
+				transactionScoped: boolean,
+			) => {
+				if (
+					transactionScoped ||
+					!warnOnManualTransactions ||
+					manualTransactionWarned ||
+					hasMultipleStatements(query) ||
+					!isManualTransactionControl(query)
+				) {
+					return;
 				}
+				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.",
+				);
 			};
 
-			return drizzleDb;
+			return createDrizzleClient(nativeDb);
 		},
 		onMigrate: async (client) => {
 			if (!migrations && !onMigrate) {
 				return;
 			}
-			await withMigrationSavepoint(client, async () => {
+			await withMigrationSavepoint(client, async (leased) => {
 				if (migrations) {
-					await runMigrations(client, migrations);
+					await runMigrations(leased, migrations);
 				}
 				if (onMigrate) {
-					await onMigrate(client);
+					await onMigrate(leased);
 				}
 			});
 		},
 	};
 }
 
-async function withMigrationSavepoint<T>(
-	client: RawAccess,
-	callback: () => Promise<T> | T,
+async function withMigrationSavepoint<TSchema extends DrizzleSchema, T>(
+	client: DrizzleDatabase<TSchema>,
+	callback: (leased: DrizzleDatabase<TSchema>) => Promise<T> | T,
 ): Promise<T> {
-	await client.execute("SAVEPOINT __rivet_on_migrate");
-	try {
-		const result = await callback();
-		await client.execute("RELEASE SAVEPOINT __rivet_on_migrate");
-		return result;
-	} catch (error) {
-		try {
-			await client.execute("ROLLBACK TO SAVEPOINT __rivet_on_migrate");
-		} finally {
-			await client.execute("RELEASE SAVEPOINT __rivet_on_migrate");
-		}
-		throw error;
-	}
+	return await client.transaction(
+		async (leased) => {
+			await leased.execute("SAVEPOINT __rivet_on_migrate");
+			try {
+				const result = await callback(leased);
+				await leased.execute("RELEASE SAVEPOINT __rivet_on_migrate");
+				return result;
+			} catch (error) {
+				try {
+					await leased.execute(
+						"ROLLBACK TO SAVEPOINT __rivet_on_migrate",
+					);
+				} finally {
+					await leased.execute(
+						"RELEASE SAVEPOINT __rivet_on_migrate",
+					);
+				}
+				throw error;
+			}
+		},
+		{ timeout: MIGRATION_TRANSACTION_TIMEOUT_MS },
+	);
 }
 
 async function runMigrations<TSchema extends DrizzleSchema>(
@@ -281,13 +370,15 @@ function rowToObject<TRow extends Record<string, unknown>>(
 }
 
 async function executeRaw<TRow extends Record<string, unknown>>(
-	db: SqliteDatabase,
+	db: SqliteDatabase | SqliteTransactionDatabase,
 	ctx: DatabaseProviderContext,
 	ensureOpen: () => void,
 	query: string,
 	args: unknown[],
+	warnForManualTransaction: () => void,
 ): Promise<TRow[]> {
 	ensureOpen();
+	warnForManualTransaction();
 
 	const start = performance.now();
 	const kvReadsBefore = ctx.metrics?.totalKvReads ?? 0;
@@ -302,7 +393,7 @@ async function executeRaw<TRow extends Record<string, unknown>>(
 		}
 
 		if (!hasMultipleStatements(query)) {
-			const { rows, columns } = await db.execute(query);
+			const { rows, columns } = await db.execute(query, undefined);
 			return rows.map((row) => rowToObject<TRow>(row, columns));
 		}
 

package/src/dynamic/isolate-runtime.ts

@@ -1,4 +1,3 @@
-import type { ActorDriver } from "@/actor/driver";
 import type { Client } from "@/client/client";
 import type { RegistryConfig } from "@/registry/config";
 import type { DynamicActorLoader } from "./internal";
@@ -10,7 +9,6 @@ export interface DynamicActorIsolateRuntimeConfig {
 	input: unknown;
 	region: string;
 	loader: DynamicActorLoader;
-	actorDriver: ActorDriver;
 	inlineClient: Client<any>;
 	test: RegistryConfig["test"];
 }

package/src/engine-client/driver.ts

@@ -61,21 +61,6 @@ export interface EngineControlClient {
 	 * Test-only helper that simulates an abrupt actor crash.
 	 */
 	hardCrashActor?(actorId: string): Promise<void>;
-	kvGet(actorId: string, key: Uint8Array): Promise<string | null>;
-	kvBatchGet(
-		actorId: string,
-		keys: Uint8Array[],
-	): Promise<(Uint8Array | null)[]>;
-	kvBatchPut(
-		actorId: string,
-		entries: [Uint8Array, Uint8Array][],
-	): Promise<void>;
-	kvBatchDelete(actorId: string, keys: Uint8Array[]): Promise<void>;
-	kvDeleteRange(
-		actorId: string,
-		start: Uint8Array,
-		end: Uint8Array,
-	): Promise<void>;
 }
 
 export interface RuntimeDisplayInformation {

package/src/engine-client/mod.ts

@@ -40,7 +40,6 @@ import {
 	getActor,
 	getActorByKey,
 	getOrCreateActor,
-	kvGet,
 	listActorsByName,
 } from "./api-endpoints";
 import { EngineApiError, getEndpoint } from "./api-utils";
@@ -372,46 +371,6 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		return await upgradeWebSocket(() => args)(c, noopNext());
 	}
 
-	async kvGet(actorId: string, key: Uint8Array): Promise<string | null> {
-		await this.#metadataPromise;
-
-		logger().debug({ msg: "getting kv value via engine api", key });
-
-		const response = await kvGet(
-			this.#config,
-			actorId,
-			new TextDecoder("utf8").decode(key),
-		);
-
-		return response.value;
-	}
-
-	async kvBatchGet(
-		_actorId: string,
-		_keys: Uint8Array[],
-	): Promise<(Uint8Array | null)[]> {
-		throw new Error("kvBatchGet not supported on remote engine client");
-	}
-
-	async kvBatchPut(
-		_actorId: string,
-		_entries: [Uint8Array, Uint8Array][],
-	): Promise<void> {
-		throw new Error("kvBatchPut not supported on remote engine client");
-	}
-
-	async kvBatchDelete(_actorId: string, _keys: Uint8Array[]): Promise<void> {
-		throw new Error("kvBatchDelete not supported on remote engine client");
-	}
-
-	async kvDeleteRange(
-		_actorId: string,
-		_start: Uint8Array,
-		_end: Uint8Array,
-	): Promise<void> {
-		throw new Error("kvDeleteRange not supported on remote engine client");
-	}
-
 	displayInformation(): RuntimeDisplayInformation {
 		return { properties: {} };
 	}

package/src/inspector/client.browser.ts

@@ -4,22 +4,24 @@ import * as v2 from "@/common/bare/generated/inspector/v2";
 import * as v3 from "@/common/bare/generated/inspector/v3";
 import * as v4 from "@/common/bare/generated/inspector/v4";
 import * as v5 from "@/common/bare/generated/inspector/v5";
+import * as v6 from "@/common/bare/generated/inspector/v6";
 import {
 	decodeWorkflowHistoryTransport,
 	encodeWorkflowHistoryTransport,
 } from "@/common/inspector-transport";
 
-export * from "@/common/bare/generated/inspector/v5";
+export * from "@/common/bare/generated/inspector/v6";
 export { decodeWorkflowHistoryTransport, encodeWorkflowHistoryTransport };
 export type { WorkflowHistory as TransportWorkflowHistory } from "@/common/bare/transport/v1";
 
-export const CURRENT_VERSION = 5;
+export const CURRENT_VERSION = 6;
 
 const EVENTS_DROPPED_ERROR = "inspector.events_dropped";
 const WORKFLOW_HISTORY_DROPPED_ERROR = "inspector.workflow_history_dropped";
 const QUEUE_DROPPED_ERROR = "inspector.queue_dropped";
 const TRACE_DROPPED_ERROR = "inspector.trace_dropped";
 const DATABASE_DROPPED_ERROR = "inspector.database_dropped";
+const SCHEDULES_DROPPED_ERROR = "inspector.schedules_dropped";
 
 const v1ToClientToV2 = (v1Data: v1.ToClient): v2.ToClient => {
 	if (v1Data.body.tag === "Init") {
@@ -187,12 +189,59 @@ const v5ToClientToV4 = (v5Data: v5.ToClient): v4.ToClient => {
 	return v5Data as unknown as v4.ToClient;
 };
 
+const v5ToClientToV6 = (v5Data: v5.ToClient): v6.ToClient => {
+	if (v5Data.body.tag === "Init") {
+		return {
+			body: {
+				tag: "Init",
+				val: { ...v5Data.body.val, schedules: [] },
+			},
+		};
+	}
+	return v5Data as unknown as v6.ToClient;
+};
+
+const v6ToClientToV5 = (v6Data: v6.ToClient): v5.ToClient => {
+	if (v6Data.body.tag === "Init") {
+		const { schedules: _schedules, ...rest } = v6Data.body.val;
+		return { body: { tag: "Init", val: rest } };
+	}
+	if (
+		v6Data.body.tag === "SchedulesUpdated" ||
+		v6Data.body.tag === "SchedulesResponse" ||
+		v6Data.body.tag === "ScheduleHistoryResponse" ||
+		v6Data.body.tag === "ScheduleDeleteResponse"
+	) {
+		return {
+			body: {
+				tag: "Error",
+				val: { message: SCHEDULES_DROPPED_ERROR },
+			},
+		};
+	}
+	return v6Data as unknown as v5.ToClient;
+};
+
 const v4ToServerToV5 = (v4Data: v4.ToServer): v5.ToServer =>
 	v4Data as unknown as v5.ToServer;
 
 const v5ToServerToV4 = (v5Data: v5.ToServer): v4.ToServer =>
 	v5Data as unknown as v4.ToServer;
 
+const v5ToServerToV6 = (v5Data: v5.ToServer): v6.ToServer =>
+	v5Data as unknown as v6.ToServer;
+
+const v6ToServerToV5 = (v6Data: v6.ToServer): v5.ToServer => {
+	if (
+		v6Data.body.tag === "SchedulesRequest" ||
+		v6Data.body.tag === "ScheduleHistoryRequest" ||
+		v6Data.body.tag === "ScheduleDeleteRequest"
+	) {
+		throw new Error("Cannot convert v6-only schedule requests to v5");
+	}
+	return v6Data as unknown as v5.ToServer;
+};
+
 const v1ToServerToV2 = (v1Data: v1.ToServer): v2.ToServer => {
 	if (
 		v1Data.body.tag === "EventsRequest" ||
@@ -243,7 +292,7 @@ const v4ToServerToV3 = (v4Data: v4.ToServer): v3.ToServer => {
 	return v4Data as unknown as v3.ToServer;
 };
 
-export const TO_SERVER_VERSIONED = createVersionedDataHandler<v5.ToServer>({
+export const TO_SERVER_VERSIONED = createVersionedDataHandler<v6.ToServer>({
 	serializeVersion: (data, version) => {
 		switch (version) {
 			case 1:
@@ -255,7 +304,9 @@ export const TO_SERVER_VERSIONED = createVersionedDataHandler<v5.ToServer>({
 			case 4:
 				return v4.encodeToServer(data as unknown as v4.ToServer);
 			case 5:
-				return v5.encodeToServer(data);
+				return v5.encodeToServer(data as unknown as v5.ToServer);
+			case 6:
+				return v6.encodeToServer(data);
 			default:
 				throw new Error(`Unknown version ${version}`);
 		}
@@ -271,7 +322,9 @@ export const TO_SERVER_VERSIONED = createVersionedDataHandler<v5.ToServer>({
 			case 4:
 				return v4.decodeToServer(bytes) as unknown as v5.ToServer;
 			case 5:
-				return v5.decodeToServer(bytes);
+				return v5.decodeToServer(bytes) as unknown as v6.ToServer;
+			case 6:
+				return v6.decodeToServer(bytes);
 			default:
 				throw new Error(`Unknown version ${version}`);
 		}
@@ -281,8 +334,10 @@ export const TO_SERVER_VERSIONED = createVersionedDataHandler<v5.ToServer>({
 		v2ToServerToV3,
 		v3ToServerToV4,
 		v4ToServerToV5,
+		v5ToServerToV6,
 	],
 	serializeConverters: () => [
+		v6ToServerToV5,
 		v5ToServerToV4,
 		v4ToServerToV3,
 		v3ToServerToV2,
@@ -290,7 +345,7 @@ export const TO_SERVER_VERSIONED = createVersionedDataHandler<v5.ToServer>({
 	],
 });
 
-export const TO_CLIENT_VERSIONED = createVersionedDataHandler<v5.ToClient>({
+export const TO_CLIENT_VERSIONED = createVersionedDataHandler<v6.ToClient>({
 	serializeVersion: (data, version) => {
 		switch (version) {
 			case 1:
@@ -302,7 +357,9 @@ export const TO_CLIENT_VERSIONED = createVersionedDataHandler<v5.ToClient>({
 			case 4:
 				return v4.encodeToClient(data as unknown as v4.ToClient);
 			case 5:
-				return v5.encodeToClient(data);
+				return v5.encodeToClient(data as unknown as v5.ToClient);
+			case 6:
+				return v6.encodeToClient(data);
 			default:
 				throw new Error(`Unknown version ${version}`);
 		}
@@ -318,7 +375,9 @@ export const TO_CLIENT_VERSIONED = createVersionedDataHandler<v5.ToClient>({
 			case 4:
 				return v4.decodeToClient(bytes) as unknown as v5.ToClient;
 			case 5:
-				return v5.decodeToClient(bytes);
+				return v5.decodeToClient(bytes) as unknown as v6.ToClient;
+			case 6:
+				return v6.decodeToClient(bytes);
 			default:
 				throw new Error(`Unknown version ${version}`);
 		}
@@ -328,8 +387,10 @@ export const TO_CLIENT_VERSIONED = createVersionedDataHandler<v5.ToClient>({
 		v2ToClientToV3,
 		v3ToClientToV4,
 		v4ToClientToV5,
+		v5ToClientToV6,
 	],
 	serializeConverters: () => [
+		v6ToClientToV5,
 		v5ToClientToV4,
 		v4ToClientToV3,
 		v3ToClientToV2,

package/src/inspector-tab/mod.ts

@@ -69,6 +69,7 @@ export type BuiltinInspectorTabId =
 	| "database"
 	| "state"
 	| "queue"
+	| "schedules"
 	| "connections"
 	| "console";
 

package/src/registry/config/index.ts

@@ -1,15 +1,6 @@
 import { z } from "zod";
 import { getRunMetadata } from "@/actor/config";
-import type {
-	AnyActorDefinition,
-	BaseActorDefinition,
-} from "@/actor/definition";
-import {
-	KEYS,
-	queueMessagesPrefix,
-	queueMetadataKey,
-	workflowStoragePrefix,
-} from "@/actor/keys";
+import type { AnyActorDefinition } from "@/actor/definition";
 import { buildEngineEndpoint, ENGINE_HOST, ENGINE_PORT } from "@/common/engine";
 import { type Logger, LogLevelSchema } from "@/common/log";
 import { VERSION } from "@/utils";
@@ -440,31 +431,6 @@ export function buildActorNames(
 			// Actor options take precedence over run metadata
 			metadata.icon = options.icon ?? runMeta.icon;
 			metadata.name = options.name ?? runMeta.name;
-			metadata.preload = {
-				keys: [
-					Array.from(KEYS.PERSIST_DATA),
-					Array.from(KEYS.INSPECTOR_TOKEN),
-					Array.from(queueMetadataKey()),
-					Array.from(KEYS.LAST_PUSHED_ALARM),
-				],
-				prefixes: [
-					{
-						prefix: Array.from(workflowStoragePrefix()),
-						maxBytes: options.preloadMaxWorkflowBytes ?? 131_072,
-						partial: false,
-					},
-					{
-						prefix: Array.from(KEYS.CONN_PREFIX),
-						maxBytes: options.preloadMaxConnectionsBytes ?? 65_536,
-						partial: false,
-					},
-					{
-						prefix: Array.from(queueMessagesPrefix()),
-						maxBytes: 65_536,
-						partial: false,
-					},
-				],
-			};
 			// Remove undefined values
 			if (!metadata.icon) delete metadata.icon;
 			if (!metadata.name) delete metadata.name;