rivetkit 2.0.42 → 2.1.0-rc.1

package/src/drivers/file-system/kv-limits.ts

@@ -0,0 +1,70 @@
+import type { SqliteRuntimeDatabase } from "./sqlite-runtime";
+
+// Keep these limits in sync with engine/packages/pegboard/src/actor_kv/mod.rs.
+const KV_MAX_KEY_SIZE = 2 * 1024;
+const KV_MAX_VALUE_SIZE = 128 * 1024;
+const KV_MAX_KEYS = 128;
+const KV_MAX_PUT_PAYLOAD_SIZE = 976 * 1024;
+const KV_MAX_STORAGE_SIZE = 10 * 1024 * 1024 * 1024;
+const KV_KEY_WRAPPER_OVERHEAD_SIZE = 2;
+
+export function estimateKvSize(db: SqliteRuntimeDatabase): number {
+	const row = db.get<{ total: number | bigint | null }>(
+		"SELECT COALESCE(SUM(LENGTH(key) + LENGTH(value)), 0) AS total FROM kv",
+	);
+	return row ? Number(row.total ?? 0) : 0;
+}
+
+export function validateKvKey(
+	key: Uint8Array,
+	keyLabel: "key" | "prefix key" = "key",
+): void {
+	if (key.byteLength + KV_KEY_WRAPPER_OVERHEAD_SIZE > KV_MAX_KEY_SIZE) {
+		throw new Error(`${keyLabel} is too long (max 2048 bytes)`);
+	}
+}
+
+export function validateKvKeys(keys: Uint8Array[]): void {
+	if (keys.length > KV_MAX_KEYS) {
+		throw new Error("a maximum of 128 keys is allowed");
+	}
+
+	for (const key of keys) {
+		validateKvKey(key);
+	}
+}
+
+export function validateKvEntries(
+	entries: [Uint8Array, Uint8Array][],
+	totalSize: number,
+): void {
+	if (entries.length > KV_MAX_KEYS) {
+		throw new Error("A maximum of 128 key-value entries is allowed");
+	}
+
+	let payloadSize = 0;
+	for (const [key, value] of entries) {
+		payloadSize +=
+			key.byteLength + KV_KEY_WRAPPER_OVERHEAD_SIZE + value.byteLength;
+	}
+
+	if (payloadSize > KV_MAX_PUT_PAYLOAD_SIZE) {
+		throw new Error("total payload is too large (max 976 KiB)");
+	}
+
+	const storageRemaining = Math.max(0, KV_MAX_STORAGE_SIZE - totalSize);
+	if (payloadSize > storageRemaining) {
+		throw new Error(
+			`not enough space left in storage (${storageRemaining} bytes remaining, current payload is ${payloadSize} bytes)`,
+		);
+	}
+
+	for (const [key, value] of entries) {
+		validateKvKey(key);
+		if (value.byteLength > KV_MAX_VALUE_SIZE) {
+			throw new Error(
+				`value is too large (max ${KV_MAX_VALUE_SIZE / 1024} KiB)`,
+			);
+		}
+	}
+}

package/src/drivers/file-system/manager.ts

@@ -86,6 +86,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 
 		// Extract just the pathname for routing (without query params)
 		const pathOnly = normalizedPath.split("?")[0];
+		const { gatewayId, requestId } = createHibernatableRequestMetadata();
 
 		const wsHandler = await routeWebSocket(
 			fakeRequest,
@@ -96,9 +97,9 @@ export class FileSystemManagerDriver implements ManagerDriver {
 			actorId,
 			encoding,
 			params,
-			undefined,
-			undefined,
-			false,
+			gatewayId,
+			requestId,
+			true,
 			false,
 		);
 		return createInlineWebSocket(wsHandler);
@@ -129,6 +130,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 		const normalizedPath = pathOnly.startsWith("/")
 			? pathOnly
 			: `/${pathOnly}`;
+		const { gatewayId, requestId } = createHibernatableRequestMetadata();
 		const wsHandler = await routeWebSocket(
 			// TODO: Create new request with new path
 			c.req.raw,
@@ -139,9 +141,9 @@ export class FileSystemManagerDriver implements ManagerDriver {
 			actorId,
 			encoding,
 			params,
-			undefined,
-			undefined,
-			false,
+			gatewayId,
+			requestId,
+			true,
 			false,
 		);
 		return upgradeWebSocket(() => wsHandler)(c, noopNext());
@@ -282,3 +284,17 @@ function actorStateToOutput(state: schema.ActorState): ActorOutput {
 		destroyTs: state.destroyTs !== null ? Number(state.destroyTs) : null,
 	};
 }
+
+function createHibernatableRequestMetadata(): {
+	gatewayId: ArrayBuffer;
+	requestId: ArrayBuffer;
+} {
+	const gatewayId = new Uint8Array(4);
+	const requestId = new Uint8Array(4);
+	crypto.getRandomValues(gatewayId);
+	crypto.getRandomValues(requestId);
+	return {
+		gatewayId: gatewayId.buffer.slice(0),
+		requestId: requestId.buffer.slice(0),
+	};
+}

package/src/drivers/file-system/mod.ts

@@ -1,30 +1,52 @@
+import { z } from "zod";
+import type { DriverConfig } from "@/registry/config";
 import { importNodeDependencies } from "@/utils/node";
 import { FileSystemActorDriver } from "./actor";
-import { FileSystemGlobalState } from "./global-state";
+import {
+	type FileSystemDriverOptions,
+	FileSystemGlobalState,
+} from "./global-state";
 import { FileSystemManagerDriver } from "./manager";
-import { DriverConfig } from "@/registry/config";
 
 export { FileSystemActorDriver } from "./actor";
 export { FileSystemGlobalState } from "./global-state";
 export { FileSystemManagerDriver } from "./manager";
 export { getStoragePath } from "./utils";
 
+const CreateFileSystemDriverOptionsSchema = z.object({
+	/** Custom path for storage. */
+	path: z.string().optional(),
+	/** Deprecated: file-system driver KV is now always SQLite-backed. */
+	useNativeSqlite: z.boolean().optional(),
+});
+
+type CreateFileSystemDriverOptionsInput = z.input<
+	typeof CreateFileSystemDriverOptionsSchema
+>;
+
 export function createFileSystemOrMemoryDriver(
 	persist: boolean = true,
-	customPath?: string,
+	options?: CreateFileSystemDriverOptionsInput,
 ): DriverConfig {
 	importNodeDependencies();
 
-	const state = new FileSystemGlobalState(persist, customPath);
+	if (options?.useNativeSqlite === false) {
+		throw new Error(
+			"File-system driver no longer supports non-SQLite KV storage. Remove useNativeSqlite: false.",
+		);
+	}
+
+	const stateOptions: FileSystemDriverOptions = {
+		persist,
+		customPath: options?.path,
+		useNativeSqlite: true,
+	};
+	const state = new FileSystemGlobalState(stateOptions);
 	const driverConfig: DriverConfig = {
 		name: persist ? "file-system" : "memory",
 		displayName: persist ? "File System" : "Memory",
 		manager: (config) =>
-			new FileSystemManagerDriver(
-				config,
-				state,
-				driverConfig,
-			),
+			new FileSystemManagerDriver(config, state, driverConfig),
 		actor: (config, managerDriver, inlineClient) => {
 			const actorDriver = new FileSystemActorDriver(
 				config,
@@ -33,11 +55,7 @@ export function createFileSystemOrMemoryDriver(
 				state,
 			);
 
-			state.onRunnerStart(
-				config,
-				inlineClient,
-				actorDriver,
-			);
+			state.onRunnerStart(config, inlineClient, actorDriver);
 
 			return actorDriver;
 		},
@@ -46,8 +64,13 @@ export function createFileSystemOrMemoryDriver(
 	return driverConfig;
 }
 
-export function createFileSystemDriver(opts?: { path?: string }): DriverConfig {
-	return createFileSystemOrMemoryDriver(true, opts?.path);
+export function createFileSystemDriver(
+	opts?: CreateFileSystemDriverOptionsInput,
+): DriverConfig {
+	const validatedOpts = opts
+		? CreateFileSystemDriverOptionsSchema.parse(opts)
+		: undefined;
+	return createFileSystemOrMemoryDriver(true, validatedOpts);
 }
 
 export function createMemoryDriver(): DriverConfig {

package/src/drivers/file-system/sqlite-runtime.ts

@@ -0,0 +1,210 @@
+import { getRequireFn } from "@/utils/node";
+
+type SqliteRuntimeKind = "bun" | "node" | "better-sqlite3";
+type SqliteDatabaseCtor = new (path: string) => SqliteRawDatabase;
+
+interface SqliteStatement {
+	run(...params: unknown[]): unknown;
+	get<T = Record<string, unknown>>(...params: unknown[]): T | undefined;
+	all<T = Record<string, unknown>>(...params: unknown[]): T[];
+}
+
+interface SqliteRawDatabase {
+	exec(sql: string): unknown;
+	close(): unknown;
+	prepare?(sql: string): SqliteStatement;
+	query?(sql: string): SqliteStatement;
+}
+
+export interface SqliteRuntimeDatabase {
+	exec(sql: string): void;
+	run(sql: string, params?: readonly unknown[]): void;
+	get<T = Record<string, unknown>>(
+		sql: string,
+		params?: readonly unknown[],
+	): T | undefined;
+	all<T = Record<string, unknown>>(
+		sql: string,
+		params?: readonly unknown[],
+	): T[];
+	close(): void;
+}
+
+export interface SqliteRuntime {
+	kind: SqliteRuntimeKind;
+	open(path: string): SqliteRuntimeDatabase;
+}
+
+function normalizeParams(params: readonly unknown[] | undefined): unknown[] {
+	if (!params || params.length === 0) {
+		return [];
+	}
+
+	return params.map((value) => {
+		if (value instanceof Uint8Array) {
+			return Buffer.from(value);
+		}
+		return value;
+	});
+}
+
+function createPreparedDatabaseAdapter(
+	rawDb: SqliteRawDatabase,
+	prepare: (sql: string) => SqliteStatement,
+): SqliteRuntimeDatabase {
+	return {
+		exec: (sql) => {
+			rawDb.exec(sql);
+		},
+		run: (sql, params) => {
+			const stmt = prepare(sql);
+			stmt.run(...normalizeParams(params));
+		},
+		get: <T = Record<string, unknown>>(
+			sql: string,
+			params?: readonly unknown[],
+		) => {
+			const stmt = prepare(sql);
+			return stmt.get<T>(...normalizeParams(params));
+		},
+		all: <T = Record<string, unknown>>(
+			sql: string,
+			params?: readonly unknown[],
+		) => {
+			const stmt = prepare(sql);
+			return stmt.all<T>(...normalizeParams(params));
+		},
+		close: () => {
+			rawDb.close();
+		},
+	};
+}
+
+function configureSqliteRuntimeDatabase(
+	rawDb: SqliteRawDatabase,
+	path: string,
+): void {
+	// Wait briefly when the database file is still being released by another
+	// process during restarts to reduce transient "database is locked" failures.
+	rawDb.exec("PRAGMA busy_timeout = 5000");
+
+	// WAL improves concurrent read/write behavior for file-backed databases.
+	if (path !== ":memory:") {
+		rawDb.exec("PRAGMA journal_mode = WAL");
+	}
+}
+
+export function loadSqliteRuntime(): SqliteRuntime {
+	const requireFn = getRequireFn();
+	const loadErrors: string[] = [];
+
+	try {
+		const bunSqlite = requireFn(/* webpackIgnore: true */ "bun:sqlite") as {
+			Database?: SqliteDatabaseCtor;
+		};
+		const BunDatabase = bunSqlite.Database;
+		if (BunDatabase) {
+			return {
+				kind: "bun",
+				open: (path) => {
+					const rawDb = new BunDatabase(path);
+					configureSqliteRuntimeDatabase(rawDb, path);
+					const query = rawDb.query?.bind(rawDb);
+					if (!query) throw new Error("bun:sqlite database missing query method");
+					return createPreparedDatabaseAdapter(rawDb, query);
+				},
+			};
+		}
+	} catch (error) {
+		loadErrors.push(`bun:sqlite unavailable: ${String(error)}`);
+	}
+
+	try {
+		const nodeSqlite = requireFn(/* webpackIgnore: true */ "node:sqlite") as {
+			DatabaseSync?: SqliteDatabaseCtor;
+		};
+		const NodeDatabaseSync = nodeSqlite.DatabaseSync;
+		if (NodeDatabaseSync) {
+			return {
+				kind: "node",
+				open: (path) => {
+					const rawDb = new NodeDatabaseSync(path);
+					configureSqliteRuntimeDatabase(rawDb, path);
+					const prepare = rawDb.prepare?.bind(rawDb);
+					if (!prepare) {
+						throw new Error("node:sqlite DatabaseSync missing prepare method");
+					}
+					return createPreparedDatabaseAdapter(rawDb, prepare);
+				},
+			};
+		}
+	} catch (error) {
+		loadErrors.push(`node:sqlite unavailable: ${String(error)}`);
+	}
+
+	try {
+		const betterSqlite3Module = requireFn(
+			/* webpackIgnore: true */ "better-sqlite3",
+		) as SqliteDatabaseCtor | { default?: SqliteDatabaseCtor };
+		const BetterSqlite3 =
+			typeof betterSqlite3Module === "function"
+				? betterSqlite3Module
+				: betterSqlite3Module.default;
+		if (BetterSqlite3) {
+			return {
+				kind: "better-sqlite3",
+				open: (path) => {
+					const rawDb = new BetterSqlite3(path);
+					configureSqliteRuntimeDatabase(rawDb, path);
+					const prepare = rawDb.prepare?.bind(rawDb);
+					if (!prepare) {
+						throw new Error("better-sqlite3 database missing prepare method");
+					}
+					return createPreparedDatabaseAdapter(rawDb, prepare);
+				},
+			};
+		}
+	} catch (error) {
+		loadErrors.push(`better-sqlite3 unavailable: ${String(error)}`);
+		throw new Error(
+			`No SQLite runtime available. Tried bun:sqlite, node:sqlite, and better-sqlite3. Install better-sqlite3 (e.g. "pnpm add better-sqlite3") if native runtimes are unavailable.\n${loadErrors.join("\n")}`,
+		);
+	}
+
+	throw new Error(
+		`No SQLite runtime available. Tried bun:sqlite, node:sqlite, and better-sqlite3.\n${loadErrors.join("\n")}`,
+	);
+}
+
+export function computePrefixUpperBound(
+	prefix: Uint8Array,
+): Uint8Array | undefined {
+	if (prefix.length === 0) {
+		return undefined;
+	}
+
+	const upperBound = new Uint8Array(prefix);
+	for (let i = upperBound.length - 1; i >= 0; i--) {
+		if (upperBound[i] !== 0xff) {
+			upperBound[i] += 1;
+			return upperBound.slice(0, i + 1);
+		}
+	}
+	return undefined;
+}
+
+export function ensureUint8Array(
+	value: unknown,
+	fieldName: string,
+): Uint8Array {
+	if (value instanceof Uint8Array) {
+		return value;
+	}
+	if (value instanceof ArrayBuffer) {
+		return new Uint8Array(value);
+	}
+	if (ArrayBuffer.isView(value)) {
+		return new Uint8Array(value.buffer, value.byteOffset, value.byteLength);
+	}
+	throw new Error(`SQLite row field "${fieldName}" is not binary data`);
+}