rivetkit 2.0.2 → 2.0.4-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/log.ts

@@ -0,0 +1,5 @@
+import { getLogger } from "@/common/log";
+
+export function logger() {
+	return getLogger("driver-fs");
+}

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

@@ -0,0 +1,300 @@
+import type { Context as HonoContext } from "hono";
+import invariant from "invariant";
+import { ActorStopping } from "@/actor/errors";
+import { type ActorRouter, createActorRouter } from "@/actor/router";
+import { routeWebSocket } from "@/actor/router-websocket-endpoints";
+import { createClientWithDriver } from "@/client/client";
+import { ClientConfigSchema } from "@/client/config";
+import { createInlineWebSocket } from "@/common/inline-websocket-adapter";
+import { noopNext } from "@/common/utils";
+import type {
+	ActorDriver,
+	ActorOutput,
+	CreateInput,
+	GetForIdInput,
+	GetOrCreateWithKeyInput,
+	GetWithKeyInput,
+	ListActorsInput,
+	ManagerDriver,
+} from "@/driver-helpers/mod";
+import type { ManagerDisplayInformation } from "@/manager/driver";
+import type { Encoding, UniversalWebSocket } from "@/mod";
+import type { DriverConfig, RegistryConfig } from "@/registry/config";
+import type * as schema from "@/schemas/file-system-driver/mod";
+import type { GetUpgradeWebSocket } from "@/utils";
+import type { FileSystemGlobalState } from "./global-state";
+import { logger } from "./log";
+import { generateActorId } from "./utils";
+
+export class FileSystemManagerDriver implements ManagerDriver {
+	#config: RegistryConfig;
+	#state: FileSystemGlobalState;
+	#driverConfig: DriverConfig;
+	#getUpgradeWebSocket: GetUpgradeWebSocket | undefined;
+
+	#actorDriver: ActorDriver;
+	#actorRouter: ActorRouter;
+
+	constructor(
+		config: RegistryConfig,
+		state: FileSystemGlobalState,
+		driverConfig: DriverConfig,
+	) {
+		this.#config = config;
+		this.#state = state;
+		this.#driverConfig = driverConfig;
+
+		// Actors run on the same node as the manager, so we create a dummy actor router that we route requests to
+		const inlineClient = createClientWithDriver(this);
+
+		this.#actorDriver = this.#driverConfig.actor(
+			config,
+			this,
+			inlineClient,
+		);
+		this.#actorRouter = createActorRouter(
+			this.#config,
+			this.#actorDriver,
+			undefined,
+			config.test.enabled,
+		);
+	}
+
+	async sendRequest(
+		actorId: string,
+		actorRequest: Request,
+	): Promise<Response> {
+		return await this.#actorRouter.fetch(actorRequest, {
+			actorId,
+		});
+	}
+
+	async openWebSocket(
+		path: string,
+		actorId: string,
+		encoding: Encoding,
+		params: unknown,
+	): Promise<UniversalWebSocket> {
+		// Normalize the path (add leading slash if needed) but preserve query params
+		const normalizedPath = path.startsWith("/") ? path : `/${path}`;
+
+		// Create a fake request with the full URL including query parameters
+		const fakeUrl = `http://inline-actor${normalizedPath}`;
+		const fakeRequest = new Request(fakeUrl, {
+			method: "GET",
+		});
+
+		// Extract just the pathname for routing (without query params)
+		const pathOnly = normalizedPath.split("?")[0];
+		const { gatewayId, requestId } = createHibernatableRequestMetadata();
+
+		const wsHandler = await routeWebSocket(
+			fakeRequest,
+			pathOnly,
+			{},
+			this.#config,
+			this.#actorDriver,
+			actorId,
+			encoding,
+			params,
+			gatewayId,
+			requestId,
+			true,
+			false,
+		);
+		return createInlineWebSocket(wsHandler);
+	}
+
+	async proxyRequest(
+		c: HonoContext,
+		actorRequest: Request,
+		actorId: string,
+	): Promise<Response> {
+		return await this.#actorRouter.fetch(actorRequest, {
+			actorId,
+		});
+	}
+
+	async proxyWebSocket(
+		c: HonoContext,
+		path: string,
+		actorId: string,
+		encoding: Encoding,
+		params: unknown,
+	): Promise<Response> {
+		const upgradeWebSocket = this.#getUpgradeWebSocket?.();
+		invariant(upgradeWebSocket, "missing getUpgradeWebSocket");
+
+		// Handle raw WebSocket paths
+		const pathOnly = path.split("?")[0];
+		const normalizedPath = pathOnly.startsWith("/")
+			? pathOnly
+			: `/${pathOnly}`;
+		const { gatewayId, requestId } = createHibernatableRequestMetadata();
+		const wsHandler = await routeWebSocket(
+			// TODO: Create new request with new path
+			c.req.raw,
+			normalizedPath,
+			c.req.header(),
+			this.#config,
+			this.#actorDriver,
+			actorId,
+			encoding,
+			params,
+			gatewayId,
+			requestId,
+			true,
+			false,
+		);
+		return upgradeWebSocket(() => wsHandler)(c, noopNext());
+	}
+
+	async buildGatewayUrl(actorId: string): Promise<string> {
+		const port = this.#config.managerPort ?? 6420;
+		return `http://127.0.0.1:${port}/gateway/${encodeURIComponent(actorId)}`;
+	}
+
+	async getForId({
+		actorId,
+	}: GetForIdInput): Promise<ActorOutput | undefined> {
+		// Validate the actor exists
+		const actor = await this.#state.loadActor(actorId);
+		if (!actor.state) {
+			return undefined;
+		}
+		if (this.#state.isActorStopping(actorId)) {
+			throw new ActorStopping(actorId);
+		}
+
+		return actorStateToOutput(actor.state);
+	}
+
+	async getWithKey({
+		name,
+		key,
+	}: GetWithKeyInput): Promise<ActorOutput | undefined> {
+		// Generate the deterministic actor ID
+		const actorId = generateActorId(name, key);
+
+		// Check if actor exists
+		const actor = await this.#state.loadActor(actorId);
+		if (actor.state) {
+			return actorStateToOutput(actor.state);
+		}
+
+		return undefined;
+	}
+
+	async getOrCreateWithKey(
+		input: GetOrCreateWithKeyInput,
+	): Promise<ActorOutput> {
+		// Generate the deterministic actor ID
+		const actorId = generateActorId(input.name, input.key);
+
+		// Use the atomic getOrCreateActor method
+		await this.#state.loadOrCreateActor(
+			actorId,
+			input.name,
+			input.key,
+			input.input,
+		);
+
+		// Start the actor immediately so timestamps are set
+		await this.#actorDriver.loadActor(actorId);
+
+		// Reload state to get updated timestamps
+		const state = await this.#state.loadActorStateOrError(actorId);
+		return actorStateToOutput(state);
+	}
+
+	async createActor({ name, key, input }: CreateInput): Promise<ActorOutput> {
+		// Generate the deterministic actor ID
+		const actorId = generateActorId(name, key);
+
+		await this.#state.createActor(actorId, name, key, input);
+
+		// Start the actor immediately so timestamps are set
+		await this.#actorDriver.loadActor(actorId);
+
+		// Reload state to get updated timestamps
+		const state = await this.#state.loadActorStateOrError(actorId);
+		return actorStateToOutput(state);
+	}
+
+	async listActors({ name }: ListActorsInput): Promise<ActorOutput[]> {
+		const actors: ActorOutput[] = [];
+		const itr = this.#state.getActorsIterator({});
+
+		for await (const actor of itr) {
+			if (actor.name === name) {
+				actors.push(actorStateToOutput(actor));
+			}
+		}
+
+		// Sort by create ts desc (most recent first)
+		actors.sort((a, b) => {
+			const aTs = a.createTs ?? 0;
+			const bTs = b.createTs ?? 0;
+			return bTs - aTs;
+		});
+
+		return actors;
+	}
+
+	async kvGet(actorId: string, key: Uint8Array): Promise<string | null> {
+		const response = await this.#state.kvBatchGet(actorId, [key]);
+		return response[0] !== null
+			? new TextDecoder().decode(response[0])
+			: null;
+	}
+
+	displayInformation(): ManagerDisplayInformation {
+		return {
+			properties: {
+				...(this.#state.persist
+					? { Data: this.#state.storagePath }
+					: {}),
+				Instances: this.#state.actorCountOnStartup.toString(),
+			},
+		};
+	}
+
+	extraStartupLog() {
+		return {
+			instances: this.#state.actorCountOnStartup,
+			data: this.#state.storagePath,
+		};
+	}
+
+	setGetUpgradeWebSocket(getUpgradeWebSocket: GetUpgradeWebSocket): void {
+		this.#getUpgradeWebSocket = getUpgradeWebSocket;
+	}
+}
+
+function actorStateToOutput(state: schema.ActorState): ActorOutput {
+	return {
+		actorId: state.actorId,
+		name: state.name,
+		key: state.key as string[],
+		createTs: Number(state.createdAt),
+		startTs: state.startTs !== null ? Number(state.startTs) : null,
+		connectableTs:
+			state.connectableTs !== null ? Number(state.connectableTs) : null,
+		sleepTs: state.sleepTs !== null ? Number(state.sleepTs) : null,
+		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

@@ -0,0 +1,78 @@
+import { z } from "zod";
+import type { DriverConfig } from "@/registry/config";
+import { importNodeDependencies } from "@/utils/node";
+import { FileSystemActorDriver } from "./actor";
+import {
+	type FileSystemDriverOptions,
+	FileSystemGlobalState,
+} from "./global-state";
+import { FileSystemManagerDriver } from "./manager";
+
+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,
+	options?: CreateFileSystemDriverOptionsInput,
+): DriverConfig {
+	importNodeDependencies();
+
+	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),
+		actor: (config, managerDriver, inlineClient) => {
+			const actorDriver = new FileSystemActorDriver(
+				config,
+				managerDriver,
+				inlineClient,
+				state,
+			);
+
+			state.onRunnerStart(config, inlineClient, actorDriver);
+
+			return actorDriver;
+		},
+		autoStartActorDriver: true,
+	};
+	return driverConfig;
+}
+
+export function createFileSystemDriver(
+	opts?: CreateFileSystemDriverOptionsInput,
+): DriverConfig {
+	const validatedOpts = opts
+		? CreateFileSystemDriverOptionsSchema.parse(opts)
+		: undefined;
+	return createFileSystemOrMemoryDriver(true, validatedOpts);
+}
+
+export function createMemoryDriver(): DriverConfig {
+	return createFileSystemOrMemoryDriver(false);
+}

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`);
+}