rivetkit 2.1.4 → 2.1.6-rc.1

package/src/actor/metrics.ts

@@ -0,0 +1,137 @@
+/**
+ * Lightweight in-memory metrics for actor instances.
+ *
+ * Metrics are collected per actor wake cycle and are NOT persisted. They reset
+ * when the actor sleeps and wakes again.
+ */
+
+export interface CounterMetric {
+	type: "counter";
+	help: string;
+	value: number;
+}
+
+export interface GaugeMetric {
+	type: "gauge";
+	help: string;
+	value: number;
+}
+
+export interface LabeledCounterMetric {
+	type: "labeled_counter";
+	help: string;
+	values: Record<string, number>;
+}
+
+export interface LabeledTimingMetric {
+	type: "labeled_timing";
+	help: string;
+	values: Record<string, { calls: number; totalMs: number; keys: number }>;
+}
+
+export type Metric =
+	| CounterMetric
+	| GaugeMetric
+	| LabeledCounterMetric
+	| LabeledTimingMetric;
+
+export type MetricsSnapshot = Record<string, Metric>;
+
+export class ActorMetrics {
+	// KV operations
+	kvGet = { calls: 0, keys: 0, totalMs: 0 };
+	kvGetBatch = { calls: 0, keys: 0, totalMs: 0 };
+	kvPut = { calls: 0, keys: 0, totalMs: 0 };
+	kvPutBatch = { calls: 0, keys: 0, totalMs: 0 };
+	kvDeleteBatch = { calls: 0, keys: 0, totalMs: 0 };
+
+	// SQL statements
+	sqlSelects = 0;
+	sqlInserts = 0;
+	sqlUpdates = 0;
+	sqlDeletes = 0;
+	sqlOther = 0;
+	sqlTotalMs = 0;
+
+	// Actions
+	actionCalls = 0;
+	actionErrors = 0;
+	actionTotalMs = 0;
+
+	// Connections
+	connectionsOpened = 0;
+	connectionsClosed = 0;
+
+	trackSql(query: string, durationMs: number): void {
+		const token = query.trimStart().slice(0, 8).toUpperCase();
+		if (token.startsWith("SELECT") || token.startsWith("PRAGMA") || token.startsWith("WITH")) {
+			this.sqlSelects++;
+		} else if (token.startsWith("INSERT")) {
+			this.sqlInserts++;
+		} else if (token.startsWith("UPDATE")) {
+			this.sqlUpdates++;
+		} else if (token.startsWith("DELETE")) {
+			this.sqlDeletes++;
+		} else {
+			this.sqlOther++;
+		}
+		this.sqlTotalMs += durationMs;
+	}
+
+	snapshot(): MetricsSnapshot {
+		return {
+			kv_operations: {
+				type: "labeled_timing",
+				help: "KV round trips by operation type",
+				values: {
+					get: { ...this.kvGet },
+					getBatch: { ...this.kvGetBatch },
+					put: { ...this.kvPut },
+					putBatch: { ...this.kvPutBatch },
+					deleteBatch: { ...this.kvDeleteBatch },
+				},
+			},
+			sql_statements: {
+				type: "labeled_counter",
+				help: "SQL statements executed by type",
+				values: {
+					select: this.sqlSelects,
+					insert: this.sqlInserts,
+					update: this.sqlUpdates,
+					delete: this.sqlDeletes,
+					other: this.sqlOther,
+				},
+			},
+			sql_duration_ms: {
+				type: "counter",
+				help: "Total SQL execution time in milliseconds",
+				value: this.sqlTotalMs,
+			},
+			action_calls: {
+				type: "counter",
+				help: "Total action invocations",
+				value: this.actionCalls,
+			},
+			action_errors: {
+				type: "counter",
+				help: "Total action errors",
+				value: this.actionErrors,
+			},
+			action_duration_ms: {
+				type: "counter",
+				help: "Total action execution time in milliseconds",
+				value: this.actionTotalMs,
+			},
+			connections_opened: {
+				type: "counter",
+				help: "Total WebSocket connections opened",
+				value: this.connectionsOpened,
+			},
+			connections_closed: {
+				type: "counter",
+				help: "Total WebSocket connections closed",
+				value: this.connectionsClosed,
+			},
+		};
+	}
+}

package/src/actor/protocol/old.ts

@@ -273,18 +273,15 @@ export async function processMessage<
 				subscribe,
 			});
 
-				if (subscribe) {
-					await actor.assertCanSubscribe(
-						new ActionContext<S, CP, CS, V, I, DB, E, Q>(
-							actor,
-							conn,
-						),
-						eventName,
-					);
-					await handler.onSubscribe(eventName, conn);
-				} else {
-					await handler.onUnsubscribe(eventName, conn);
-				}
+			if (subscribe) {
+				await actor.assertCanSubscribe(
+					new ActionContext<S, CP, CS, V, I, DB, E, Q>(actor, conn),
+					eventName,
+				);
+				await handler.onSubscribe(eventName, conn);
+			} else {
+				await handler.onUnsubscribe(eventName, conn);
+			}
 
 			actor.rLog.debug({
 				msg: "subscription request completed",

package/src/actor/router-websocket-endpoints.ts

@@ -86,7 +86,12 @@ export async function routeWebSocket(
 		});
 
 		// Promise used to wait for the websocket close in `disconnect`
-		const closePromiseResolvers = promiseWithResolvers<void>((reason) => loggerWithoutContext().warn({ msg: "unhandled websocket close promise rejection", reason }));
+		const closePromiseResolvers = promiseWithResolvers<void>((reason) =>
+			loggerWithoutContext().warn({
+				msg: "unhandled websocket close promise rejection",
+				reason,
+			}),
+		);
 
 		// Strip query parameters from requestPath for routing purposes.
 		// This handles paths like "/websocket?query=value" which should route
@@ -193,8 +198,8 @@ export async function routeWebSocket(
 			onMessage: (_evt: { data: any }, ws: WSContext) => {
 				ws.close(1011, "actor.not_loaded");
 			},
-			onClose: (_event: any, _ws: WSContext) => { },
-			onError: (_error: unknown) => { },
+			onClose: (_event: any, _ws: WSContext) => {},
+			onError: (_error: unknown) => {},
 		};
 	}
 }
@@ -240,7 +245,8 @@ export async function handleWebSocketConnect(
 				.then(async () => {
 					const message = await parseMessage(value, {
 						encoding: encoding,
-						maxIncomingMessageSize: runConfig.maxIncomingMessageSize,
+						maxIncomingMessageSize:
+							runConfig.maxIncomingMessageSize,
 					});
 					await actor.processMessage(message, conn);
 				})
@@ -339,7 +345,7 @@ export async function handleRawWebSocket(
 		},
 		// Raw websocket messages are handled directly by the actor's event
 		// listeners on the WebSocket object, not through this callback
-		onMessage: (_evt: any, _ws: any) => { },
+		onMessage: (_evt: any, _ws: any) => {},
 		onClose: (evt: any, ws: any) => {
 			// Resolve the close promise
 			closePromiseResolvers.resolve();
@@ -347,7 +353,7 @@ export async function handleRawWebSocket(
 			// Clean up the connection
 			conn.disconnect(evt?.reason);
 		},
-		onError: (error: any, ws: any) => { },
+		onError: (error: any, ws: any) => {},
 	};
 }
 

package/src/actor/router.ts

@@ -43,6 +43,8 @@ export type ActorRouter = Hono<{ Bindings: ActorRouterBindings }>;
 export interface MetadataResponse {
 	runtime: string;
 	version: string;
+	/** "local" for development, "deployed" for production */
+	type: "local" | "deployed";
 }
 
 /**
@@ -90,6 +92,7 @@ export function createActorRouter(
 		return c.json({
 			runtime: "rivetkit",
 			version: VERSION,
+			type: isDev() ? "local" : "deployed",
 		} satisfies MetadataResponse);
 	});
 
@@ -286,6 +289,38 @@ export function createActorRouter(
 			return c.json(result);
 		});
 
+		router.get("/inspector/database/schema", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const schema = await actor.inspector.getDatabaseSchemaJson();
+			return c.json({ schema });
+		});
+
+		router.get("/inspector/database/rows", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			const table = c.req.query("table");
+			if (!table) {
+				return c.json(
+					{ error: "Missing required table query parameter" },
+					400,
+				);
+			}
+
+			const limit = parseInt(c.req.query("limit") ?? "100", 10);
+			const offset = parseInt(c.req.query("offset") ?? "0", 10);
+			const rows = await actor.inspector.getDatabaseTableRowsJson(
+				table,
+				limit,
+				offset,
+			);
+			return c.json({ rows });
+		});
+
 		router.get("/inspector/summary", async (c) => {
 			const authResponse = await inspectorAuth(c);
 			if (authResponse) return authResponse;
@@ -295,6 +330,8 @@ export function createActorRouter(
 			const isStateEnabled = actor.inspector.isStateEnabled();
 			const isDatabaseEnabled = actor.inspector.isDatabaseEnabled();
 			const isWorkflowEnabled = actor.inspector.isWorkflowEnabled();
+			const workflowHistory =
+				actor.inspector.getWorkflowHistoryJson().history;
 
 			const state = isStateEnabled
 				? actor.inspector.getStateJson()
@@ -302,14 +339,6 @@ export function createActorRouter(
 			const connections = actor.inspector.getConnectionsJson();
 			const rpcs = actor.inspector.getRpcs();
 			const queueSize = actor.inspector.getQueueSize();
-			const workflowHistory = actor.inspector.getWorkflowHistory();
-
-			// Convert BigInt values in workflow history to numbers for JSON serialization.
-			const bigIntReplacer = (_key: string, value: unknown) =>
-				typeof value === "bigint" ? Number(value) : value;
-			const safeWorkflowHistory = workflowHistory
-				? JSON.parse(JSON.stringify(workflowHistory, bigIntReplacer))
-				: null;
 
 			return c.json({
 				state,
@@ -319,9 +348,17 @@ export function createActorRouter(
 				isStateEnabled,
 				isDatabaseEnabled,
 				isWorkflowEnabled,
-				workflowHistory: safeWorkflowHistory,
+				workflowHistory,
 			});
 		});
+
+		router.get("/inspector/metrics", async (c) => {
+			const authResponse = await inspectorAuth(c);
+			if (authResponse) return authResponse;
+
+			const actor = await actorDriver.loadActor(c.env.actorId);
+			return c.json(actor.metrics.snapshot());
+		});
 	}
 
 	router.post("/action/:action", async (c) => {

package/src/actor/schema.ts

@@ -11,11 +11,7 @@ export interface EventTypeToken<T, TContext = any> {
 	readonly schema?: PrimitiveSchema;
 }
 
-export interface QueueTypeToken<
-	TMessage,
-	TComplete = never,
-	TContext = any,
-> {
+export interface QueueTypeToken<TMessage, TComplete = never, TContext = any> {
 	readonly _queueMessage?: TMessage;
 	readonly _queueComplete?: TComplete;
 	readonly canPublish?: SchemaHook<TContext>;
@@ -45,11 +41,7 @@ export function event<T, TContext = any>(
 
 export function queue<TMessage, TComplete = never, TContext = any>(
 	options?: QueueOptions<TContext>,
-): QueueTypeToken<
-	TMessage,
-	TComplete,
-	TContext
-> {
+): QueueTypeToken<TMessage, TComplete, TContext> {
 	return (options ?? {}) as QueueTypeToken<TMessage, TComplete, TContext>;
 }
 
@@ -87,16 +79,16 @@ export type SchemaConfig = QueueSchemaConfig;
 
 export type InferSchema<T> =
 	T extends QueueSchemaDefinition<any>
-			? InferSchema<T["message"]>
-			: T extends QueueTypeToken<infer M, unknown, any>
+		? InferSchema<T["message"]>
+		: T extends QueueTypeToken<infer M, unknown, any>
 			? M
 			: T extends EventSchemaDefinition<any>
-			? InferSchema<T["schema"]>
-			: T extends StandardSchemaV1<any, infer O>
-			? O
-			: T extends EventTypeToken<infer R, any>
-				? R
-				: never;
+				? InferSchema<T["schema"]>
+				: T extends StandardSchemaV1<any, infer O>
+					? O
+					: T extends EventTypeToken<infer R, any>
+						? R
+						: never;
 
 export type InferSchemaMap<T extends Record<string, unknown>> = {
 	[K in keyof T]: InferSchema<T[K]>;
@@ -104,10 +96,10 @@ export type InferSchemaMap<T extends Record<string, unknown>> = {
 
 export type InferQueueComplete<T> =
 	T extends QueueTypeToken<unknown, infer C, any>
-			? [C] extends [never]
-				? never
-				: C
-			: T extends QueueSchemaDefinition<any>
+		? [C] extends [never]
+			? never
+			: C
+		: T extends QueueSchemaDefinition<any>
 			? T["complete"] extends PrimitiveSchema
 				? InferSchema<T["complete"]>
 				: never

package/src/client/actor-common.ts

@@ -1,4 +1,16 @@
 import type { ActorDefinition, AnyActorDefinition } from "@/actor/definition";
+import type {
+	EventSchemaConfig,
+	InferEventArgs,
+	InferQueueCompleteMap,
+	InferSchemaMap,
+	QueueSchemaConfig,
+} from "@/actor/schema";
+import type {
+	QueueSendNoWaitOptions,
+	QueueSendResult,
+	QueueSendWaitOptions,
+} from "./queue";
 
 /**
  * Action function returned by Actor connections and handles.
@@ -30,3 +42,56 @@ export type ActorDefinitionActions<AD extends AnyActorDefinition> =
 					: never;
 			}
 		: never;
+
+type ActorQueueSend<TQueues extends QueueSchemaConfig> = {
+	<K extends keyof TQueues & string>(
+		name: K,
+		body: InferSchemaMap<TQueues>[K],
+		options: QueueSendWaitOptions,
+	): Promise<QueueSendResult<InferQueueCompleteMap<TQueues>[K]>>;
+	<K extends keyof TQueues & string>(
+		name: K,
+		body: InferSchemaMap<TQueues>[K],
+		options?: QueueSendNoWaitOptions,
+	): Promise<void>;
+	(
+		name: keyof TQueues extends never ? string : never,
+		body: unknown,
+		options: QueueSendWaitOptions,
+	): Promise<QueueSendResult>;
+	(
+		name: keyof TQueues extends never ? string : never,
+		body: unknown,
+		options?: QueueSendNoWaitOptions,
+	): Promise<void>;
+};
+
+type ActorEventSubscribe<TEvents extends EventSchemaConfig> = {
+	<K extends keyof TEvents & string>(
+		eventName: K,
+		callback: (...args: InferEventArgs<InferSchemaMap<TEvents>[K]>) => void,
+	): () => void;
+	(
+		eventName: keyof TEvents extends never ? string : never,
+		callback: (...args: any[]) => void,
+	): () => void;
+};
+
+export type ActorDefinitionQueueSend<AD extends AnyActorDefinition> =
+	// biome-ignore lint/suspicious/noExplicitAny: safe to use any here
+	AD extends ActorDefinition<any, any, any, any, any, any, any, infer Q, any>
+		? Q extends QueueSchemaConfig
+			? { send: ActorQueueSend<Q> }
+			: never
+		: never;
+
+export type ActorDefinitionEventSubscriptions<AD extends AnyActorDefinition> =
+	// biome-ignore lint/suspicious/noExplicitAny: safe to use any here
+	AD extends ActorDefinition<any, any, any, any, any, any, infer E, any, any>
+		? E extends EventSchemaConfig
+			? {
+					on: ActorEventSubscribe<E>;
+					once: ActorEventSubscribe<E>;
+				}
+			: never
+		: never;

package/src/client/actor-conn.ts

@@ -25,7 +25,11 @@ import {
 import { deserializeWithEncoding, serializeWithEncoding } from "@/serde";
 import { bufferToArrayBuffer, promiseWithResolvers } from "@/utils";
 import { getLogMessage } from "@/utils/env-vars";
-import type { ActorDefinitionActions } from "./actor-common";
+import type {
+	ActorDefinitionActions,
+	ActorDefinitionEventSubscriptions,
+	ActorDefinitionQueueSend,
+} from "./actor-common";
 import { checkForSchedulingError, queryActor } from "./actor-query";
 import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client";
 import * as errors from "./errors";
@@ -160,6 +164,7 @@ export class ActorConnRaw {
 	#client: ClientRaw;
 	#driver: ManagerDriver;
 	#params: unknown;
+	#getParams?: () => Promise<unknown>;
 	#encoding: Encoding;
 	#actorQuery: ActorQuery;
 
@@ -176,12 +181,14 @@ export class ActorConnRaw {
 		client: ClientRaw,
 		driver: ManagerDriver,
 		params: unknown,
+		getParams: (() => Promise<unknown>) | undefined,
 		encoding: Encoding,
 		actorQuery: ActorQuery,
 	) {
 		this.#client = client;
 		this.#driver = driver;
 		this.#params = params;
+		this.#getParams = getParams;
 		this.#encoding = encoding;
 		this.#actorQuery = actorQuery;
 		this.#queueSender = createQueueSender({
@@ -248,7 +255,12 @@ export class ActorConnRaw {
 		const { promise, resolve, reject } = promiseWithResolvers<{
 			id: bigint;
 			output: unknown;
-		}>((reason) => logger().warn({ msg: "unhandled action promise rejection", reason }));
+		}>((reason) =>
+			logger().warn({
+				msg: "unhandled action promise rejection",
+				reason,
+			}),
+		);
 		this.#actionsInFlight.set(actionId, {
 			name: opts.name,
 			resolve,
@@ -379,7 +391,12 @@ export class ActorConnRaw {
 			// Create promise for open
 			if (this.#onOpenPromise)
 				throw new Error("#onOpenPromise already defined");
-			this.#onOpenPromise = promiseWithResolvers((reason) => logger().warn({ msg: "unhandled open promise rejection", reason }));
+			this.#onOpenPromise = promiseWithResolvers((reason) =>
+				logger().warn({
+					msg: "unhandled open promise rejection",
+					reason,
+				}),
+			);
 
 			await this.#connectWebSocket();
 
@@ -390,6 +407,40 @@ export class ActorConnRaw {
 		}
 	}
 
+	#clearQueuedMessages() {
+		if (this.#messageQueue.length === 0) return;
+
+		logger().debug({
+			msg: "clearing queued connection messages",
+			queueLength: this.#messageQueue.length,
+		});
+		this.#messageQueue = [];
+	}
+
+	async #resolveConnectionParams(): Promise<unknown> {
+		if (!this.#getParams) {
+			return this.#params;
+		}
+
+		try {
+			return await this.#getParams();
+		} catch (err) {
+			const errorMessage = stringifyError(err);
+			const error = new errors.ActorError(
+				"client",
+				"get_params_failed",
+				`Failed to resolve connection params: ${errorMessage}`,
+				{ error: errorMessage },
+			);
+
+			this.#clearQueuedMessages();
+			this.#rejectPendingPromises(error, false);
+			this.#dispatchActorError(error);
+
+			throw error;
+		}
+	}
+
 	async #connectWebSocket() {
 		const { actorId } = await queryActor(
 			undefined,
@@ -399,12 +450,13 @@ export class ActorConnRaw {
 
 		// Store actorId early so we can use it for error lookups
 		this.#actorId = actorId;
+		const params = await this.#resolveConnectionParams();
 
 		const ws = await this.#driver.openWebSocket(
 			PATH_CONNECT,
 			actorId,
 			this.#encoding,
-			this.#params,
+			params,
 		);
 		logger().debug({
 			msg: "opened websocket",
@@ -419,7 +471,7 @@ export class ActorConnRaw {
 				connId: this.#connId,
 			});
 		});
-		ws.addEventListener("message", async (ev) => {
+		ws.addEventListener("message", async (ev: { data: ConnMessage }) => {
 			try {
 				await this.#handleOnMessage(ev.data);
 			} catch (err) {
@@ -429,7 +481,7 @@ export class ActorConnRaw {
 				});
 			}
 		});
-		ws.addEventListener("close", async (ev) => {
+		ws.addEventListener("close", async (ev: Event | CloseEvent) => {
 			try {
 				await this.#handleOnClose(ev);
 			} catch (err) {
@@ -439,7 +491,7 @@ export class ActorConnRaw {
 				});
 			}
 		});
-		ws.addEventListener("error", (_ev) => {
+		ws.addEventListener("error", () => {
 			try {
 				this.#handleOnError();
 			} catch (err) {
@@ -1216,7 +1268,12 @@ export class ActorConnRaw {
 				ws.readyState !== 2 /* CLOSING */ &&
 				ws.readyState !== 3 /* CLOSED */
 			) {
-				const { promise, resolve } = promiseWithResolvers((reason) => logger().warn({ msg: "unhandled websocket close promise rejection", reason }));
+				const { promise, resolve } = promiseWithResolvers((reason) =>
+					logger().warn({
+						msg: "unhandled websocket close promise rejection",
+						reason,
+					}),
+				);
 				ws.addEventListener("close", () => resolve(undefined));
 				ws.close(1000, "Disposed");
 				await promise;
@@ -1258,5 +1315,10 @@ export class ActorConnRaw {
  * @template AD The actor class that this connection is for.
  * @see {@link ActorConnRaw}
  */
-export type ActorConn<AD extends AnyActorDefinition> = ActorConnRaw &
+export type ActorConn<AD extends AnyActorDefinition> = Omit<
+	ActorConnRaw,
+	"send" | "on" | "once"
+> &
+	ActorDefinitionQueueSend<AD> &
+	ActorDefinitionEventSubscriptions<AD> &
 	ActorDefinitionActions<AD>;