rivetkit 2.1.5 → 2.1.6-rc.1

package/src/client/actor-handle.ts

@@ -23,7 +23,10 @@ import {
 	HttpActionResponseSchema,
 } from "@/schemas/client-protocol-zod/mod";
 import { bufferToArrayBuffer } from "@/utils";
-import type { ActorDefinitionActions } from "./actor-common";
+import type {
+	ActorDefinitionActions,
+	ActorDefinitionQueueSend,
+} from "./actor-common";
 import { type ActorConn, ActorConnRaw } from "./actor-conn";
 import { checkForSchedulingError, queryActor } from "./actor-query";
 import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client";
@@ -51,6 +54,7 @@ export class ActorHandleRaw {
 	#encoding: Encoding;
 	#actorQuery: ActorQuery;
 	#params: unknown;
+	#getParams?: () => Promise<unknown>;
 	#queueSender: ReturnType<typeof createQueueSender>;
 
 	/**
@@ -64,6 +68,7 @@ export class ActorHandleRaw {
 		client: any,
 		driver: ManagerDriver,
 		params: unknown,
+		getParams: (() => Promise<unknown>) | undefined,
 		encoding: Encoding,
 		actorQuery: ActorQuery,
 	) {
@@ -72,6 +77,7 @@ export class ActorHandleRaw {
 		this.#encoding = encoding;
 		this.#actorQuery = actorQuery;
 		this.#params = params;
+		this.#getParams = getParams;
 		this.#queueSender = createQueueSender({
 			encoding: this.#encoding,
 			params: this.#params,
@@ -86,6 +92,14 @@ export class ActorHandleRaw {
 		});
 	}
 
+	async #resolveConnectionParams(): Promise<unknown> {
+		if (this.#getParams) {
+			return await this.#getParams();
+		}
+
+		return this.#params;
+	}
+
 	send(
 		name: string,
 		body: unknown,
@@ -227,6 +241,7 @@ export class ActorHandleRaw {
 			this.#client,
 			this.#driver,
 			this.#params,
+			this.#getParams,
 			this.#encoding,
 			this.#actorQuery,
 		);
@@ -253,11 +268,12 @@ export class ActorHandleRaw {
 	/**
 	 * Opens a raw WebSocket connection to this actor.
 	 */
-	webSocket(path?: string, protocols?: string | string[]) {
+	async webSocket(path?: string, protocols?: string | string[]) {
+		const params = await this.#resolveConnectionParams();
 		return rawWebSocket(
 			this.#driver,
 			this.#actorQuery,
-			this.#params,
+			params,
 			path,
 			protocols,
 		);
@@ -335,10 +351,11 @@ export class ActorHandleRaw {
  */
 export type ActorHandle<AD extends AnyActorDefinition> = Omit<
 	ActorHandleRaw,
-	"connect"
+	"connect" | "send"
 > & {
 	// Add typed version of ActorConn (instead of using AnyActorDefinition)
 	connect(): ActorConn<AD>;
 	// Resolve method returns the actor ID
 	resolve(): Promise<string>;
-} & ActorDefinitionActions<AD>;
+} & ActorDefinitionQueueSend<AD> &
+	ActorDefinitionActions<AD>;

package/src/client/client.ts

@@ -86,6 +86,11 @@ export interface ActorAccessor<AD extends AnyActorDefinition> {
 export interface QueryOptions {
 	/** Parameters to pass to the connection. */
 	params?: unknown;
+	/**
+	 * Lazily resolves connection parameters for `.connect()` and `.webSocket()`.
+	 * This is called each time a new connection is opened.
+	 */
+	getParams?: () => Promise<unknown>;
 	/** Signal to abort the request. */
 	signal?: AbortSignal;
 }
@@ -200,7 +205,11 @@ export class ClientRaw {
 			},
 		};
 
-		const handle = this.#createHandle(opts?.params, actorQuery);
+		const handle = this.#createHandle(
+			opts?.params,
+			opts?.getParams,
+			actorQuery,
+		);
 		return createActorProxy(handle) as ActorHandle<AD>;
 	}
 
@@ -235,7 +244,11 @@ export class ClientRaw {
 			},
 		};
 
-		const handle = this.#createHandle(opts?.params, actorQuery);
+		const handle = this.#createHandle(
+			opts?.params,
+			opts?.getParams,
+			actorQuery,
+		);
 		return createActorProxy(handle) as ActorHandle<AD>;
 	}
 
@@ -273,7 +286,11 @@ export class ClientRaw {
 			},
 		};
 
-		const handle = this.#createHandle(opts?.params, actorQuery);
+		const handle = this.#createHandle(
+			opts?.params,
+			opts?.getParams,
+			actorQuery,
+		);
 		return createActorProxy(handle) as ActorHandle<AD>;
 	}
 
@@ -332,18 +349,27 @@ export class ClientRaw {
 				actorId,
 			},
 		} satisfies ActorQuery;
-		const handle = this.#createHandle(opts?.params, getForIdQuery);
+		const handle = this.#createHandle(
+			opts?.params,
+			opts?.getParams,
+			getForIdQuery,
+		);
 
 		const proxy = createActorProxy(handle) as ActorHandle<AD>;
 
 		return proxy;
 	}
 
-	#createHandle(params: unknown, actorQuery: ActorQuery): ActorHandleRaw {
+	#createHandle(
+		params: unknown,
+		getParams: (() => Promise<unknown>) | undefined,
+		actorQuery: ActorQuery,
+	): ActorHandleRaw {
 		return new ActorHandleRaw(
 			this,
 			this.#driver,
 			params,
+			getParams,
 			this.#encodingKind,
 			actorQuery,
 		);
@@ -554,5 +580,5 @@ function createActorProxy<AD extends AnyActorDefinition>(
 			}
 			return undefined;
 		},
-	}) as ActorHandle<AD> | ActorConn<AD>;
+	}) as unknown as ActorHandle<AD> | ActorConn<AD>;
 }

package/src/client/config.ts

@@ -11,19 +11,9 @@ import {
 import type { RegistryConfig } from "@/registry/config";
 import { tryParseEndpoint } from "@/utils/endpoint-parser";
 
-/**
- * Gets the default endpoint for the client.
- *
- * In browser: uses current origin + /api/rivet
- *
- * Server-side: uses 127.0.0.1:6420
- */
-function getDefaultEndpoint(): string {
-	if (typeof window !== "undefined" && window.location?.origin) {
-		return `${window.location.origin}/api/rivet`;
-	}
-	return "http://127.0.0.1:6420";
-}
+const DEFAULT_ENDPOINT = "http://localhost:6420";
+
+let hasWarnedMissingEndpoint = false;
 
 /**
  * Base client config schema without transforms so it can be merged in to other schemas.
@@ -38,18 +28,25 @@ export const ClientConfigSchemaBase = z.object({
 	 *
 	 * Can also be set via RIVET_ENDPOINT environment variables.
 	 *
-	 * Defaults to current origin + /api/rivet in browser, or 127.0.0.1:6420 server-side.
+	 * Defaults to http://localhost:6420.
 	 */
 	endpoint: z
 		.string()
 		.optional()
-		.transform(
-			(val) =>
-				val ??
-				getRivetEngine() ??
-				getRivetEndpoint() ??
-				getDefaultEndpoint(),
-		),
+		.transform((val) => {
+			const resolved =
+				val ?? getRivetEngine() ?? getRivetEndpoint();
+			if (!resolved && !hasWarnedMissingEndpoint) {
+				hasWarnedMissingEndpoint = true;
+				console.warn(
+					`[rivetkit] No endpoint provided to client. Defaulting to ${DEFAULT_ENDPOINT}. ` +
+						`Starting in 2.2.0, an explicit endpoint will be required. ` +
+						`Pass an endpoint to createClient() or createRivetKit(), ` +
+						`or set the RIVET_ENDPOINT environment variable.`,
+				);
+			}
+			return resolved ?? DEFAULT_ENDPOINT;
+		}),
 
 	/** Token to use to authenticate with the API. */
 	token: z

package/src/client/mod.ts

@@ -34,6 +34,7 @@ export { ActorHandleRaw } from "./actor-handle";
 export type {
 	ActorAccessor,
 	Client,
+	ClientConfigInput,
 	ClientRaw,
 	CreateOptions,
 	ExtractActorsFromRegistry,

package/src/client/queue.ts

@@ -43,9 +43,9 @@ export interface QueueSendNoWaitOptions {
 
 export type QueueSendOptions = QueueSendWaitOptions | QueueSendNoWaitOptions;
 
-export interface QueueSendResult {
+export interface QueueSendResult<TResponse = unknown> {
 	status: "completed" | "timedOut";
-	response?: unknown;
+	response?: TResponse;
 }
 
 interface QueueSenderOptions {
@@ -54,7 +54,9 @@ interface QueueSenderOptions {
 	customFetch: (request: Request) => Promise<Response>;
 }
 
-export function createQueueSender(senderOptions: QueueSenderOptions): QueueSender {
+export function createQueueSender(
+	senderOptions: QueueSenderOptions,
+): QueueSender {
 	async function send(
 		name: string,
 		body: unknown,
@@ -100,8 +102,7 @@ export function createQueueSender(senderOptions: QueueSenderOptions): QueueSende
 			requestVersion: CLIENT_PROTOCOL_CURRENT_VERSION,
 			requestVersionedDataHandler: HTTP_QUEUE_SEND_REQUEST_VERSIONED,
 			responseVersion: CLIENT_PROTOCOL_CURRENT_VERSION,
-			responseVersionedDataHandler:
-				HTTP_QUEUE_SEND_RESPONSE_VERSIONED,
+			responseVersionedDataHandler: HTTP_QUEUE_SEND_RESPONSE_VERSIONED,
 			requestZodSchema: HttpQueueSendRequestSchema,
 			responseZodSchema: HttpQueueSendResponseSchema,
 			requestToJson: (value): HttpQueueSendRequestJson => ({
@@ -112,7 +113,8 @@ export function createQueueSender(senderOptions: QueueSenderOptions): QueueSende
 				name: value.name ?? name,
 				body: bufferToArrayBuffer(cbor.encode(value.body)),
 				wait: value.wait ?? false,
-				timeout: value.timeout !== undefined ? BigInt(value.timeout) : null,
+				timeout:
+					value.timeout !== undefined ? BigInt(value.timeout) : null,
 			}),
 			responseFromJson: (json): QueueSendResult => {
 				if (json.response === undefined) {

package/src/common/inline-websocket-adapter.ts

@@ -111,7 +111,10 @@ export class InlineWebSocketAdapter {
 		try {
 			this.#handler.onError(err, this.#wsContext);
 		} catch (handlerErr) {
-			logger().error({ msg: "error in onError handler", error: handlerErr });
+			logger().error({
+				msg: "error in onError handler",
+				error: handlerErr,
+			});
 		}
 
 		// Fire error event to both sides
@@ -129,7 +132,10 @@ export class InlineWebSocketAdapter {
 		this.#readyState = 2; // CLOSING
 
 		try {
-			this.#handler.onClose({ code, reason, wasClean: true }, this.#wsContext);
+			this.#handler.onClose(
+				{ code, reason, wasClean: true },
+				this.#wsContext,
+			);
 		} catch (err) {
 			logger().error({ msg: "error closing websocket", error: err });
 		} finally {

package/src/common/router.ts

@@ -5,10 +5,7 @@ import {
 	getRequestEncoding,
 	getRequestExposeInternalError,
 } from "@/actor/router-endpoints";
-import {
-	buildActorNames,
-	type RegistryConfig,
-} from "@/registry/config";
+import { buildActorNames, type RegistryConfig } from "@/registry/config";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import {
 	CURRENT_VERSION as CLIENT_PROTOCOL_CURRENT_VERSION,

package/src/common/utils.ts

@@ -299,10 +299,7 @@ export function deconstructError(
 
 export function stringifyError(error: unknown): string {
 	if (error instanceof Error) {
-		if (
-			typeof process !== "undefined" &&
-			getLogErrorStack()
-		) {
+		if (typeof process !== "undefined" && getLogErrorStack()) {
 			return `${error.name}: ${error.message}${error.stack ? `\n${error.stack}` : ""}`;
 		} else {
 			return `${error.name}: ${error.message}`;
@@ -335,5 +332,5 @@ function getErrorMessage(err: unknown): string {
 
 /** Generates a `Next` handler to pass to middleware in order to be able to call arbitrary middleware. */
 export function noopNext(): Next {
-	return async () => { };
+	return async () => {};
 }

package/src/db/config.ts

@@ -1,4 +1,5 @@
-import type { SqliteVfs } from "@rivetkit/sqlite-vfs";
+import type { ISqliteVfs } from "@rivetkit/sqlite-vfs";
+import type { ActorMetrics } from "@/actor/metrics";
 
 export type AnyDatabaseProvider = DatabaseProvider<any> | undefined;
 
@@ -35,11 +36,15 @@ export interface DatabaseProviderContext {
 	};
 
 	/**
-	 * SQLite VFS instance for creating KV-backed databases.
-	 * This should be actor-scoped because @rivetkit/sqlite is not re-entrant per
-	 * module instance.
+	 * SQLite VFS handle for creating KV-backed databases.
+	 * May be a standalone VFS or a pooled handle from SqliteVfsPool.
 	 */
-	sqliteVfs?: SqliteVfs;
+	sqliteVfs?: ISqliteVfs;
+
+	/**
+	 * Actor metrics instance. When provided, KV and SQL operations are tracked.
+	 */
+	metrics?: ActorMetrics;
 }
 
 export type DatabaseProvider<DB extends RawAccess> = {

package/src/db/drizzle/mod.ts

@@ -1,4 +1,4 @@
-import type { Database } from "@rivetkit/sqlite-vfs";
+import type { IDatabase } from "@rivetkit/sqlite-vfs";
 import {
 	drizzle as proxyDrizzle,
 	type SqliteRemoteDatabase,
@@ -31,7 +31,7 @@ interface DatabaseFactoryConfig<
  * Create a sqlite-proxy async callback from a @rivetkit/sqlite Database
  */
 function createProxyCallback(
-	waDb: Database,
+	waDb: IDatabase,
 	mutex: AsyncMutex,
 	isClosed: () => boolean,
 ) {
@@ -40,7 +40,7 @@ function createProxyCallback(
 		params: any[],
 		method: "run" | "all" | "values" | "get",
 	): Promise<{ rows: any }> => {
-		return mutex.run(async () => {
+		return await mutex.run(async () => {
 			if (isClosed()) {
 				throw new Error("database is closed");
 			}
@@ -69,30 +69,25 @@ function createProxyCallback(
  * Migrations use the same embedded format as drizzle-orm's durable-sqlite.
  */
 async function runInlineMigrations(
-	waDb: Database,
-	mutex: AsyncMutex,
+	waDb: IDatabase,
 	migrations: any,
 ): Promise<void> {
 	// Create migrations table
-	await mutex.run(() =>
-		waDb.exec(`
+	await waDb.exec(`
 		CREATE TABLE IF NOT EXISTS __drizzle_migrations (
 			id INTEGER PRIMARY KEY AUTOINCREMENT,
 			hash TEXT NOT NULL,
 			created_at INTEGER
 		)
-	`),
-	);
+	`);
 
 	// Get the last applied migration
 	let lastCreatedAt = 0;
-	await mutex.run(() =>
-		waDb.exec(
-			"SELECT id, hash, created_at FROM __drizzle_migrations ORDER BY created_at DESC LIMIT 1",
-			(row) => {
-				lastCreatedAt = Number(row[2]) || 0;
-			},
-		),
+	await waDb.exec(
+		"SELECT id, hash, created_at FROM __drizzle_migrations ORDER BY created_at DESC LIMIT 1",
+		(row) => {
+			lastCreatedAt = Number(row[2]) || 0;
+		},
 	);
 
 	// Apply pending migrations from journal entries
@@ -109,14 +104,12 @@ async function runInlineMigrations(
 		if (!sql) continue;
 
 		// Execute migration SQL
-		await mutex.run(() => waDb.exec(sql));
+		await waDb.exec(sql);
 
 		// Record migration
-		await mutex.run(() =>
-			waDb.run(
-				"INSERT INTO __drizzle_migrations (hash, created_at) VALUES (?, ?)",
-				[entry.tag, entry.when],
-			),
+		await waDb.run(
+			"INSERT INTO __drizzle_migrations (hash, created_at) VALUES (?, ?)",
+			[entry.tag, entry.when],
 		);
 	}
 }
@@ -126,9 +119,13 @@ export function db<
 >(
 	config?: DatabaseFactoryConfig<TSchema>,
 ): DatabaseProvider<SqliteRemoteDatabase<TSchema> & RawAccess> {
-	// Store the @rivetkit/sqlite Database instance alongside the drizzle client
-	let waDbInstance: Database | null = null;
-	const mutex = new AsyncMutex();
+	// Map from drizzle client to the underlying @rivetkit/sqlite Database.
+	// Uses WeakMap so entries are garbage-collected when the client is.
+	// This replaces the previous shared `waDbInstance` variable which caused
+	// a race when multiple actors of the same type created databases
+	// concurrently: the last writer won, and earlier actors' migrations
+	// ran on the wrong database.
+	const clientToRawDb = new WeakMap<object, IDatabase>();
 
 	return {
 		createClient: async (ctx) => {
@@ -139,9 +136,12 @@ export function db<
 				);
 			}
 
-			const kvStore = createActorKvStore(ctx.kv);
+			const kvStore = createActorKvStore(ctx.kv, ctx.metrics);
 			const waDb = await ctx.sqliteVfs.open(ctx.actorId, kvStore);
-			waDbInstance = waDb;
+			// Per-client mutex so actors of the same type do not serialize
+			// against each other. Each actor has its own database handle and
+			// its own closed flag, so there is no shared state to guard.
+			const mutex = new AsyncMutex();
 			let closed = false;
 			const ensureOpen = () => {
 				if (closed) {
@@ -155,14 +155,17 @@ export function db<
 			// Create the drizzle instance using sqlite-proxy
 			const client = proxyDrizzle<TSchema>(callback, config);
 
-			return Object.assign(client, {
+			const result = Object.assign(client, {
 				execute: async <
-					TRow extends Record<string, unknown> = Record<string, unknown>,
+					TRow extends Record<string, unknown> = Record<
+						string,
+						unknown
+					>,
 				>(
 					query: string,
 					...args: unknown[]
 				): Promise<TRow[]> => {
-					return mutex.run(async () => {
+					return await mutex.run(async () => {
 						ensureOpen();
 
 						if (args.length > 0) {
@@ -172,7 +175,11 @@ export function db<
 							);
 							return result.rows.map((row: unknown[]) => {
 								const obj: Record<string, unknown> = {};
-								for (let i = 0; i < result.columns.length; i++) {
+								for (
+									let i = 0;
+									i < result.columns.length;
+									i++
+								) {
 									obj[result.columns[i]] = row[i];
 								}
 								return obj;
@@ -199,22 +206,25 @@ export function db<
 					});
 				},
 				close: async () => {
-					await mutex.run(async () => {
-						if (closed) {
-							return;
-						}
+					const shouldClose = await mutex.run(async () => {
+						if (closed) return false;
 						closed = true;
-						await waDb.close();
-						waDbInstance = null;
+						return true;
 					});
+					if (shouldClose) {
+						await waDb.close();
+					}
 				},
 			} satisfies RawAccess);
+
+			clientToRawDb.set(result, waDb);
+			return result;
 		},
-		onMigrate: async (_client) => {
-			if (config?.migrations && waDbInstance) {
+		onMigrate: async (client) => {
+			const waDb = clientToRawDb.get(client as object);
+			if (config?.migrations && waDb) {
 				await runInlineMigrations(
-					waDbInstance,
-					mutex,
+					waDb,
 					config.migrations,
 				);
 			}

package/src/db/mod.ts

@@ -21,7 +21,10 @@ export function db({
 				// Use the override
 				return {
 					execute: async <
-						TRow extends Record<string, unknown> = Record<string, unknown>,
+						TRow extends Record<string, unknown> = Record<
+							string,
+							unknown
+						>,
 					>(
 						query: string,
 						...args: unknown[]
@@ -36,10 +39,12 @@ export function db({
 
 			// Construct KV-backed client using actor driver's KV operations
 			if (!ctx.sqliteVfs) {
-				throw new Error("SqliteVfs instance not provided in context. The driver must provide a sqliteVfs instance.");
+				throw new Error(
+					"SqliteVfs instance not provided in context. The driver must provide a sqliteVfs instance.",
+				);
 			}
 
-			const kvStore = createActorKvStore(ctx.kv);
+			const kvStore = createActorKvStore(ctx.kv, ctx.metrics);
 			const db = await ctx.sqliteVfs.open(ctx.actorId, kvStore);
 			let closed = false;
 			const mutex = new AsyncMutex();
@@ -51,7 +56,10 @@ export function db({
 
 			return {
 				execute: async <
-					TRow extends Record<string, unknown> = Record<string, unknown>,
+					TRow extends Record<string, unknown> = Record<
+						string,
+						unknown
+					>,
 				>(
 					query: string,
 					...args: unknown[]
@@ -59,57 +67,74 @@ export function db({
 					return await mutex.run(async () => {
 						ensureOpen();
 
+						const start = performance.now();
+
 						// `db.exec` does not support binding `?` placeholders.
 						// Use `db.query` for statements that return rows and `db.run` for
 						// statements that mutate data when parameters are provided.
 						// Keep using `db.exec` for non-parameterized SQL because it
 						// supports multi-statement migrations.
+						let result: TRow[];
 						if (args.length > 0) {
 							const bindings = toSqliteBindings(args);
-							const token = query.trimStart().slice(0, 16).toUpperCase();
+							const token = query
+								.trimStart()
+								.slice(0, 16)
+								.toUpperCase();
 							const returnsRows =
 								token.startsWith("SELECT") ||
 								token.startsWith("PRAGMA") ||
-								token.startsWith("WITH");
+								token.startsWith("WITH") ||
+								/\bRETURNING\b/i.test(query);
 
 							if (returnsRows) {
-								const { rows, columns } = await db.query(query, bindings);
-								return rows.map((row: unknown[]) => {
+								const { rows, columns } = await db.query(
+									query,
+									bindings,
+								);
+								result = rows.map((row: unknown[]) => {
 									const rowObj: Record<string, unknown> = {};
 									for (let i = 0; i < columns.length; i++) {
 										rowObj[columns[i]] = row[i];
 									}
 									return rowObj;
 								}) as TRow[];
+							} else {
+								await db.run(query, bindings);
+								result = [] as TRow[];
 							}
-
-							await db.run(query, bindings);
-							return [] as TRow[];
+						} else {
+							const results: Record<string, unknown>[] = [];
+							let columnNames: string[] | null = null;
+							await db.exec(
+								query,
+								(row: unknown[], columns: string[]) => {
+									if (!columnNames) {
+										columnNames = columns;
+									}
+									const rowObj: Record<string, unknown> = {};
+									for (let i = 0; i < row.length; i++) {
+										rowObj[columnNames[i]] = row[i];
+									}
+									results.push(rowObj);
+								},
+							);
+							result = results as TRow[];
 						}
 
-						const results: Record<string, unknown>[] = [];
-						let columnNames: string[] | null = null;
-						await db.exec(query, (row: unknown[], columns: string[]) => {
-							if (!columnNames) {
-								columnNames = columns;
-							}
-							const rowObj: Record<string, unknown> = {};
-							for (let i = 0; i < row.length; i++) {
-								rowObj[columnNames[i]] = row[i];
-							}
-							results.push(rowObj);
-						});
-						return results as TRow[];
+						ctx.metrics?.trackSql(query, performance.now() - start);
+						return result;
 					});
 				},
 				close: async () => {
-					await mutex.run(async () => {
-						if (closed) {
-							return;
-						}
+					const shouldClose = await mutex.run(async () => {
+						if (closed) return false;
 						closed = true;
-						await db.close();
+						return true;
 					});
+					if (shouldClose) {
+						await db.close();
+					}
 				},
 			} satisfies RawAccess;
 		},