rivetkit 2.1.5 → 2.1.6-rc.1

package/src/db/shared.ts

@@ -1,9 +1,10 @@
 import type { DatabaseProviderContext } from "./config";
-import type { Database } from "@rivetkit/sqlite-vfs";
-import type { KvVfsOptions } from "./sqlite-vfs";
+import type { IDatabase } from "@rivetkit/sqlite-vfs";
+import type { KvVfsOptions } from "@rivetkit/sqlite-vfs";
+import type { ActorMetrics } from "@/actor/metrics";
 
 type ActorKvOperations = DatabaseProviderContext["kv"];
-type SqliteBindings = NonNullable<Parameters<Database["run"]>[1]>;
+type SqliteBindings = NonNullable<Parameters<IDatabase["run"]>[1]>;
 
 function isSqliteBindingValue(value: unknown): boolean {
 	if (
@@ -26,9 +27,7 @@ function isSqliteBindingValue(value: unknown): boolean {
 export function toSqliteBindings(args: unknown[]): SqliteBindings {
 	for (const value of args) {
 		if (!isSqliteBindingValue(value)) {
-			throw new Error(
-				`unsupported sqlite binding type: ${typeof value}`,
-			);
+			throw new Error(`unsupported sqlite binding type: ${typeof value}`);
 		}
 	}
 
@@ -37,24 +36,59 @@ export function toSqliteBindings(args: unknown[]): SqliteBindings {
 
 /**
  * Create a KV store wrapper that uses the actor driver's KV operations.
+ * Tracks per-operation metrics when an ActorMetrics instance is provided.
  */
-export function createActorKvStore(kv: ActorKvOperations): KvVfsOptions {
+export function createActorKvStore(
+	kv: ActorKvOperations,
+	metrics?: ActorMetrics,
+): KvVfsOptions {
 	return {
 		get: async (key: Uint8Array) => {
+			const start = performance.now();
 			const results = await kv.batchGet([key]);
+			if (metrics) {
+				metrics.kvGet.calls++;
+				metrics.kvGet.keys++;
+				metrics.kvGet.totalMs += performance.now() - start;
+			}
 			return results[0] ?? null;
 		},
 		getBatch: async (keys: Uint8Array[]) => {
-			return await kv.batchGet(keys);
+			const start = performance.now();
+			const results = await kv.batchGet(keys);
+			if (metrics) {
+				metrics.kvGetBatch.calls++;
+				metrics.kvGetBatch.keys += keys.length;
+				metrics.kvGetBatch.totalMs += performance.now() - start;
+			}
+			return results;
 		},
 		put: async (key: Uint8Array, value: Uint8Array) => {
+			const start = performance.now();
 			await kv.batchPut([[key, value]]);
+			if (metrics) {
+				metrics.kvPut.calls++;
+				metrics.kvPut.keys++;
+				metrics.kvPut.totalMs += performance.now() - start;
+			}
 		},
 		putBatch: async (entries: [Uint8Array, Uint8Array][]) => {
+			const start = performance.now();
 			await kv.batchPut(entries);
+			if (metrics) {
+				metrics.kvPutBatch.calls++;
+				metrics.kvPutBatch.keys += entries.length;
+				metrics.kvPutBatch.totalMs += performance.now() - start;
+			}
 		},
 		deleteBatch: async (keys: Uint8Array[]) => {
+			const start = performance.now();
 			await kv.batchDelete(keys);
+			if (metrics) {
+				metrics.kvDeleteBatch.calls++;
+				metrics.kvDeleteBatch.keys += keys.length;
+				metrics.kvDeleteBatch.totalMs += performance.now() - start;
+			}
 		},
 	};
 }

package/src/driver-helpers/mod.ts

@@ -30,4 +30,5 @@ export type {
 	ManagerDriver,
 } from "@/manager/driver";
 export { buildManagerRouter } from "@/manager/router";
-export { getInitialActorKvState, importSqliteVfs } from "./utils";
+export { getInitialActorKvState } from "./utils";
+export { SqliteVfsPoolManager } from "./sqlite-pool";

package/src/driver-helpers/sqlite-pool.ts

@@ -0,0 +1,42 @@
+import type { ISqliteVfs } from "@rivetkit/sqlite-vfs";
+import type { RegistryConfig } from "@/registry/config";
+
+/**
+ * Manages a lazily-created SqliteVfsPool shared across actors in a driver.
+ * Uses dynamic import to keep @rivetkit/sqlite-vfs tree-shakeable.
+ */
+export class SqliteVfsPoolManager {
+	#poolPromise:
+		| Promise<{
+				acquire(actorId: string): Promise<ISqliteVfs>;
+				shutdown(): Promise<void>;
+		  }>
+		| undefined;
+	#config: RegistryConfig;
+
+	constructor(config: RegistryConfig) {
+		this.#config = config;
+	}
+
+	async acquire(actorId: string): Promise<ISqliteVfs> {
+		if (!this.#poolPromise) {
+			const poolConfig = this.#config.sqlitePool;
+			this.#poolPromise = import("@rivetkit/sqlite-vfs").then(
+				({ SqliteVfsPool }) =>
+					new SqliteVfsPool({
+						actorsPerInstance: poolConfig.actorsPerInstance,
+						idleDestroyMs: poolConfig.idleDestroyMs,
+					}),
+			);
+		}
+		const pool = await this.#poolPromise;
+		return await pool.acquire(actorId);
+	}
+
+	async shutdown(): Promise<void> {
+		if (this.#poolPromise) {
+			const pool = await this.#poolPromise;
+			await pool.shutdown();
+		}
+	}
+}

package/src/driver-helpers/utils.ts

@@ -6,9 +6,6 @@ import {
 	CURRENT_VERSION,
 } from "@/schemas/actor-persist/versioned";
 import { bufferToArrayBuffer } from "@/utils";
-import type { ActorDriver } from "./mod";
-import type { SqliteVfs } from "@rivetkit/sqlite-vfs";
-
 function serializeEmptyPersistData(input: unknown | undefined): Uint8Array {
 	const persistData: persistSchema.Actor = {
 		input:
@@ -35,20 +32,3 @@ export function getInitialActorKvState(
 	const persistData = serializeEmptyPersistData(input);
 	return [[KEYS.PERSIST_DATA, persistData]];
 }
-
-/**
- * Dynamically import @rivetkit/sqlite-vfs and return a fresh SqliteVfs instance.
- *
- * The module specifier is built with Array.join() so that bundlers (esbuild, tsup,
- * Turbopack) cannot statically analyze or constant-fold the import path. This
- * prevents them from tracing into the WASM dependency tree, which would cause
- * errors in environments that don't support .wasm imports (e.g. Turbopack).
- *
- * Each call returns a new instance so that actors get independent SQLite modules,
- * avoiding cross-actor re-entry on the non-reentrant async build.
- */
-export async function importSqliteVfs(): Promise<SqliteVfs> {
-	const specifier = ["@rivetkit", "sqlite-vfs"].join("/");
-	const { SqliteVfs } = await import(specifier);
-	return new SqliteVfs();
-}

package/src/driver-test-suite/mod.ts

@@ -29,6 +29,7 @@ import { runActorMetadataTests } from "./tests/actor-metadata";
 import { runActorOnStateChangeTests } from "./tests/actor-onstatechange";
 import { runActorQueueTests } from "./tests/actor-queue";
 import { runActorRunTests } from "./tests/actor-run";
+import { runActorSandboxTests } from "./tests/actor-sandbox";
 import { runActorStatelessTests } from "./tests/actor-stateless";
 import { runActorVarsTests } from "./tests/actor-vars";
 import { runActorWorkflowTests } from "./tests/actor-workflow";
@@ -36,6 +37,7 @@ import { runManagerDriverTests } from "./tests/manager-driver";
 import { runRawHttpTests } from "./tests/raw-http";
 import { runRawHttpRequestPropertiesTests } from "./tests/raw-http-request-properties";
 import { runRawWebSocketTests } from "./tests/raw-websocket";
+import { runActorDbKvStatsTests } from "./tests/actor-db-kv-stats";
 import { runRequestAccessTests } from "./tests/request-access";
 
 export interface SkipTests {
@@ -43,6 +45,7 @@ export interface SkipTests {
 	sleep?: boolean;
 	hibernation?: boolean;
 	inline?: boolean;
+	sandbox?: boolean;
 }
 
 export interface DriverTestConfig {
@@ -138,6 +141,8 @@ export function runDriverTests(
 
 						runActorRunTests(driverTestConfig);
 
+						runActorSandboxTests(driverTestConfig);
+
 						runActorInlineClientTests(driverTestConfig);
 
 						runActorKvTests(driverTestConfig);
@@ -159,6 +164,8 @@ export function runDriverTests(
 						// runRawWebSocketDirectRegistryTests(driverTestConfig);
 
 						runActorInspectorTests(driverTestConfig);
+
+						runActorDbKvStatsTests(driverTestConfig);
 					});
 				}
 			});
@@ -269,7 +276,10 @@ export async function createTestRuntime(
 		);
 		nodeWebSocket.injectWebSocket(server);
 		const address = server.address();
-		invariant(address && typeof address !== "string", "missing server address");
+		invariant(
+			address && typeof address !== "string",
+			"missing server address",
+		);
 		const port = address.port;
 		const serverEndpoint = `http://127.0.0.1:${port}`;
 

package/src/driver-test-suite/tests/access-control.ts

@@ -14,7 +14,9 @@ export function runAccessControlTests(driverTestConfig: DriverTestConfig) {
 
 		test("passes connection id into canPublish context", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
-			const handle = client.accessControlActor.getOrCreate(["publish-ctx"]);
+			const handle = client.accessControlActor.getOrCreate([
+				"publish-ctx",
+			]);
 
 			await handle.send("allowedQueue", { value: "one" });
 
@@ -104,7 +106,7 @@ export function runAccessControlTests(driverTestConfig: DriverTestConfig) {
 			await conn.dispose();
 
 			const blockedConn = handle.connect();
-			blockedConn.on("blockedEvent", () => { });
+			blockedConn.on("blockedEvent", () => {});
 			await expect(
 				blockedConn.allowedAction("blocked-subscribe-ready"),
 			).rejects.toMatchObject({
@@ -120,16 +122,21 @@ export function runAccessControlTests(driverTestConfig: DriverTestConfig) {
 			]);
 			const conn = handle.connect();
 
-			const eventPromise = new Promise<{ value: string }>((resolve, reject) => {
-				const unsubscribeError = conn.onError((error) => {
-					reject(error);
-				});
-				const unsubscribeEvent = conn.on("undefinedEvent", (payload) => {
-					unsubscribeError();
-					unsubscribeEvent();
-					resolve(payload as { value: string });
-				});
-			});
+			const eventPromise = new Promise<{ value: string }>(
+				(resolve, reject) => {
+					const unsubscribeError = conn.onError((error) => {
+						reject(error);
+					});
+					const unsubscribeEvent = conn.on(
+						"undefinedEvent",
+						(payload) => {
+							unsubscribeError();
+							unsubscribeEvent();
+							resolve(payload as { value: string });
+						},
+					);
+				},
+			);
 
 			await conn.allowedAction("undefined-subscribe-ready");
 			await conn.allowedBroadcastUndefinedEvent("wildcard");

package/src/driver-test-suite/tests/action-features.ts

@@ -147,14 +147,20 @@ export function runActionFeaturesTests(driverTestConfig: DriverTestConfig) {
 				// Each item is roughly 60 bytes, so 800 items ≈ 48KB
 				const items: string[] = [];
 				for (let i = 0; i < 800; i++) {
-					items.push(`Item ${i} with some additional text to increase size`);
+					items.push(
+						`Item ${i} with some additional text to increase size`,
+					);
 				}
 
 				const result = await instance.processLargeRequest({ items });
 
 				expect(result.itemCount).toBe(800);
-				expect(result.firstItem).toBe("Item 0 with some additional text to increase size");
-				expect(result.lastItem).toBe("Item 799 with some additional text to increase size");
+				expect(result.firstItem).toBe(
+					"Item 0 with some additional text to increase size",
+				);
+				expect(result.lastItem).toBe(
+					"Item 799 with some additional text to increase size",
+				);
 			});
 
 			test("should reject request exceeding maxIncomingMessageSize", async (c) => {
@@ -166,11 +172,13 @@ export function runActionFeaturesTests(driverTestConfig: DriverTestConfig) {
 				// Each item is roughly 60 bytes, so 1500 items ≈ 90KB
 				const items: string[] = [];
 				for (let i = 0; i < 1500; i++) {
-					items.push(`Item ${i} with some additional text to increase size`);
+					items.push(
+						`Item ${i} with some additional text to increase size`,
+					);
 				}
 
 				await expect(
-					instance.processLargeRequest({ items })
+					instance.processLargeRequest({ items }),
 				).rejects.toThrow();
 			});
 
@@ -183,8 +191,12 @@ export function runActionFeaturesTests(driverTestConfig: DriverTestConfig) {
 				const result = await instance.getLargeResponse(800);
 
 				expect(result.items).toHaveLength(800);
-				expect(result.items[0]).toBe("Item 0 with some additional text to increase size");
-				expect(result.items[799]).toBe("Item 799 with some additional text to increase size");
+				expect(result.items[0]).toBe(
+					"Item 0 with some additional text to increase size",
+				);
+				expect(result.items[799]).toBe(
+					"Item 799 with some additional text to increase size",
+				);
 			});
 
 			test("should reject response exceeding maxOutgoingMessageSize", async (c) => {
@@ -195,7 +207,7 @@ export function runActionFeaturesTests(driverTestConfig: DriverTestConfig) {
 				// Request a response that exceeds the default 1MB limit
 				// Each item is roughly 60 bytes, so 20000 items ≈ 1.2MB
 				await expect(
-					instance.getLargeResponse(20000)
+					instance.getLargeResponse(20000),
 				).rejects.toThrow();
 			});
 		});

package/src/driver-test-suite/tests/actor-conn.ts

@@ -243,6 +243,80 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 				await conn1.dispose();
 				await conn2.dispose();
 			});
+
+			test("should call getParams for each connection", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				let connectionCount = 0;
+				const handle = client.counterWithParams.getOrCreate(
+					["test-get-params"],
+					{
+						getParams: async () => ({
+							name: `user${++connectionCount}`,
+						}),
+					},
+				);
+
+				const conn1 = handle.connect();
+				await conn1.getInitializers();
+				await conn1.dispose();
+
+				const conn2 = handle.connect();
+				const initializers = await conn2.getInitializers();
+
+				expect(initializers).toEqual(["user1", "user2"]);
+				expect(connectionCount).toBe(2);
+
+				await conn2.dispose();
+			});
+
+			test("should surface getParams errors and retry connection setup", async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+
+				let attempts = 0;
+				const handle = client.counterWithParams.getOrCreate(
+					["test-get-params-retry"],
+					{
+						getParams: async () => {
+							attempts++;
+							if (attempts === 1) {
+								throw new Error("token unavailable");
+							}
+
+							return { name: "user1" };
+						},
+					},
+				);
+
+				const conn = handle.connect();
+				const receivedErrors: Array<{ group: string; code: string }> =
+					[];
+				conn.onError((error) => {
+					receivedErrors.push({
+						group: error.group,
+						code: error.code,
+					});
+				});
+
+				await expect(conn.getInitializers()).rejects.toMatchObject({
+					group: "client",
+					code: "get_params_failed",
+				});
+
+				await vi.waitFor(
+					async () => {
+						expect(await conn.getInitializers()).toEqual(["user1"]);
+					},
+					{ timeout: 10_000 },
+				);
+
+				expect(receivedErrors).toEqual([
+					{ group: "client", code: "get_params_failed" },
+				]);
+				expect(attempts).toBeGreaterThanOrEqual(2);
+
+				await conn.dispose();
+			});
 		});
 
 		describe("Lifecycle Hooks", () => {
@@ -515,14 +589,20 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 				// Each item is roughly 60 bytes, so 800 items ≈ 48KB
 				const items: string[] = [];
 				for (let i = 0; i < 800; i++) {
-					items.push(`Item ${i} with some additional text to increase size`);
+					items.push(
+						`Item ${i} with some additional text to increase size`,
+					);
 				}
 
 				const result = await connection.processLargeRequest({ items });
 
 				expect(result.itemCount).toBe(800);
-				expect(result.firstItem).toBe("Item 0 with some additional text to increase size");
-				expect(result.lastItem).toBe("Item 799 with some additional text to increase size");
+				expect(result.firstItem).toBe(
+					"Item 0 with some additional text to increase size",
+				);
+				expect(result.lastItem).toBe(
+					"Item 799 with some additional text to increase size",
+				);
 
 				// Verify connection state was updated
 				const lastRequestSize = await connection.getLastRequestSize();
@@ -544,11 +624,13 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 				// Each item is roughly 60 bytes, so 1500 items ≈ 90KB
 				const items: string[] = [];
 				for (let i = 0; i < 1500; i++) {
-					items.push(`Item ${i} with some additional text to increase size`);
+					items.push(
+						`Item ${i} with some additional text to increase size`,
+					);
 				}
 
 				await expect(
-					connection.processLargeRequest({ items })
+					connection.processLargeRequest({ items }),
 				).rejects.toThrow();
 
 				// Clean up
@@ -567,8 +649,12 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 				const result = await connection.getLargeResponse(800);
 
 				expect(result.items).toHaveLength(800);
-				expect(result.items[0]).toBe("Item 0 with some additional text to increase size");
-				expect(result.items[799]).toBe("Item 799 with some additional text to increase size");
+				expect(result.items[0]).toBe(
+					"Item 0 with some additional text to increase size",
+				);
+				expect(result.items[799]).toBe(
+					"Item 799 with some additional text to increase size",
+				);
 
 				// Clean up
 				await connection.dispose();
@@ -585,7 +671,7 @@ export function runActorConnTests(driverTestConfig: DriverTestConfig) {
 				// Request a response that exceeds the default 1MB limit
 				// Each item is roughly 60 bytes, so 20000 items ≈ 1.2MB
 				await expect(
-					connection.getLargeResponse(20000)
+					connection.getLargeResponse(20000),
 				).rejects.toThrow();
 
 				// Clean up