rivetkit 2.0.42 → 2.1.0-rc.1

package/src/actor/instance/mod.ts

@@ -1,3 +1,12 @@
+import type { OtlpExportTraceServiceRequestJson } from "@rivetkit/traces";
+import {
+	createNoopTraces,
+	createTraces,
+	type SpanHandle,
+	type SpanStatusInput,
+	type Traces,
+} from "@rivetkit/traces";
+import type { SqliteVfs } from "@rivetkit/sqlite-vfs";
 import invariant from "invariant";
 import type { ActorKey } from "@/actor/mod";
 import type { Client } from "@/client/client";
@@ -11,7 +20,11 @@ import {
 	CONN_VERSIONED,
 } from "@/schemas/actor-persist/versioned";
 import { EXTRA_ERROR_LOG } from "@/utils";
-import type { ActorConfig } from "../config";
+import { getRivetExperimentalOtel } from "@/utils/env-vars";
+import {
+	type ActorConfig,
+	getRunFunction,
+} from "../config";
 import type { ConnDriver } from "../conn/driver";
 import { createHttpDriver } from "../conn/drivers/http";
 import {
@@ -25,17 +38,24 @@ import {
 	type PersistedConn,
 } from "../conn/persisted";
 import {
-	type ActionContext,
+	ActionContext,
 	ActorContext,
 	RequestContext,
 	WebSocketContext,
 } from "../contexts";
+
 import type { AnyDatabaseProvider, InferDatabaseClient } from "../database";
 import type { ActorDriver } from "../driver";
 import * as errors from "../errors";
 import { serializeActorKey } from "../keys";
 import { processMessage } from "../protocol/old";
 import { Schedule } from "../schedule";
+import {
+	type EventSchemaConfig,
+	getEventCanSubscribe,
+	getQueueCanPublish,
+	type QueueSchemaConfig,
+} from "../schema";
 import {
 	assertUnreachable,
 	DeadlineError,
@@ -49,8 +69,10 @@ import {
 	convertActorFromBarePersisted,
 	type PersistedActor,
 } from "./persisted";
+import { QueueManager } from "./queue-manager";
 import { ScheduleManager } from "./schedule-manager";
 import { type SaveStateOptions, StateManager } from "./state-manager";
+import { ActorTracesDriver } from "./traces-driver";
 
 export type { SaveStateOptions };
 
@@ -59,47 +81,71 @@ enum CanSleep {
 	NotReady,
 	NotStarted,
 	ActiveConns,
+	ActiveDisconnectCallbacks,
 	ActiveHonoHttpRequests,
+	ActiveKeepAwake,
+	ActiveRun,
 }
 
 /** Actor type alias with all `any` types. Used for `extends` in classes referencing this actor. */
-export type AnyActorInstance = ActorInstance<any, any, any, any, any, any>;
+export type AnyActorInstance = ActorInstance<
+	any,
+	any,
+	any,
+	any,
+	any,
+	any,
+	any,
+	any
+>;
 
 export type ExtractActorState<A extends AnyActorInstance> =
-	A extends ActorInstance<infer State, any, any, any, any, any>
+	A extends ActorInstance<infer State, any, any, any, any, any, any, any>
 		? State
 		: never;
 
 export type ExtractActorConnParams<A extends AnyActorInstance> =
-	A extends ActorInstance<any, infer ConnParams, any, any, any, any>
+	A extends ActorInstance<any, infer ConnParams, any, any, any, any, any, any>
 		? ConnParams
 		: never;
 
 export type ExtractActorConnState<A extends AnyActorInstance> =
-	A extends ActorInstance<any, any, infer ConnState, any, any, any>
+	A extends ActorInstance<any, any, infer ConnState, any, any, any, any, any>
 		? ConnState
 		: never;
 
 // MARK: - Main ActorInstance Class
-export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
+export class ActorInstance<
+	S,
+	CP,
+	CS,
+	V,
+	I,
+	DB extends AnyDatabaseProvider,
+	E extends EventSchemaConfig = Record<never, never>,
+	Q extends QueueSchemaConfig = Record<never, never>,
+> {
 	// MARK: - Core Properties
-	actorContext: ActorContext<S, CP, CS, V, I, DB>;
-	#config: ActorConfig<S, CP, CS, V, I, DB>;
+	actorContext: ActorContext<S, CP, CS, V, I, DB, E, Q>;
+	#config: ActorConfig<S, CP, CS, V, I, DB, E, Q>;
 	driver!: ActorDriver;
 	#inlineClient!: Client<Registry<any>>;
 	#actorId!: string;
 	#name!: string;
 	#key!: ActorKey;
+	#actorKeyString!: string;
 	#region!: string;
 
 	// MARK: - Managers
-	connectionManager!: ConnectionManager<S, CP, CS, V, I, DB>;
+	connectionManager!: ConnectionManager<S, CP, CS, V, I, DB, E, Q>;
+
+	stateManager!: StateManager<S, CP, CS, I, E, Q>;
 
-	stateManager!: StateManager<S, CP, CS, I>;
+	eventManager!: EventManager<S, CP, CS, V, I, DB, E, Q>;
 
-	eventManager!: EventManager<S, CP, CS, V, I, DB>;
+	#scheduleManager!: ScheduleManager<S, CP, CS, V, I, DB, E, Q>;
 
-	#scheduleManager!: ScheduleManager<S, CP, CS, V, I, DB>;
+	queueManager!: QueueManager<S, CP, CS, V, I, DB, E, Q>;
 
 	// MARK: - Logging
 	#log!: Logger;
@@ -126,25 +172,34 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 
 	// MARK: - Variables & Database
 	#vars?: V;
-	#db!: InferDatabaseClient<DB>;
+	#db?: InferDatabaseClient<DB>;
+	#sqliteVfs?: SqliteVfs;
 
 	// MARK: - Background Tasks
 	#backgroundPromises: Promise<void>[] = [];
+	#runPromise?: Promise<void>;
+	#runHandlerActive = false;
+	#activeQueueWaitCount = 0;
 
 	// MARK: - HTTP/WebSocket Tracking
 	#activeHonoHttpRequests = 0;
+	#activeKeepAwakeCount = 0;
 
 	// MARK: - Deprecated (kept for compatibility)
 	#schedule!: Schedule;
 
 	// MARK: - Inspector
 	#inspectorToken?: string;
-	#inspector = new ActorInspector(this);
+	#inspector: ActorInspector;
+
+	// MARK: - Tracing
+	#traces!: Traces<OtlpExportTraceServiceRequestJson>;
 
 	// MARK: - Constructor
-	constructor(config: ActorConfig<S, CP, CS, V, I, DB>) {
+	constructor(config: ActorConfig<S, CP, CS, V, I, DB, E, Q>) {
 		this.#config = config;
 		this.actorContext = new ActorContext(this);
+		this.#inspector = new ActorInspector(this);
 	}
 
 	// MARK: - Public Getters
@@ -186,11 +241,73 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		return this.#inspector;
 	}
 
+	get traces(): Traces<OtlpExportTraceServiceRequestJson> {
+		return this.#traces;
+	}
+
 	get inspectorToken(): string | undefined {
 		return this.#inspectorToken;
 	}
 
-	get conns(): Map<ConnId, Conn<S, CP, CS, V, I, DB>> {
+	// MARK: - Tracing
+	getCurrentTraceSpan(): SpanHandle | null {
+		return this.#traces.getCurrentSpan();
+	}
+
+	startTraceSpan(
+		name: string,
+		attributes?: Record<string, unknown>,
+	): SpanHandle {
+		return this.#traces.startSpan(name, {
+			parent: this.#traces.getCurrentSpan() ?? undefined,
+			attributes: this.#traceAttributes(attributes),
+		});
+	}
+
+	endTraceSpan(handle: SpanHandle, status?: SpanStatusInput): void {
+		this.#traces.endSpan(handle, status ? { status } : undefined);
+	}
+
+	async runInTraceSpan<T>(
+		name: string,
+		attributes: Record<string, unknown> | undefined,
+		fn: () => T | Promise<T>,
+	): Promise<T> {
+		const span = this.startTraceSpan(name, attributes);
+		try {
+			const result = this.#traces.withSpan(span, fn);
+			const resolved = result instanceof Promise ? await result : result;
+			this.#traces.endSpan(span, {
+				status: { code: "OK" },
+			});
+			return resolved;
+		} catch (error) {
+			this.#traces.endSpan(span, {
+				status: {
+					code: "ERROR",
+					message: stringifyError(error),
+				},
+			});
+			throw error;
+		}
+	}
+
+	emitTraceEvent(
+		name: string,
+		attributes?: Record<string, unknown>,
+		handle?: SpanHandle,
+	): void {
+		const span = handle ?? this.#traces.getCurrentSpan();
+		if (!span) {
+			return;
+		}
+		this.#traces.emitEvent(span, name, {
+			attributes: this.#traceAttributes(attributes),
+			timeUnixMs: Date.now(),
+		});
+	}
+
+	get conns(): Map<ConnId, Conn<S, CP, CS, V, I, DB, E, Q>> {
 		return this.connectionManager.connections;
 	}
 
@@ -206,7 +323,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		return Object.keys(this.#config.actions ?? {});
 	}
 
-	get config(): ActorConfig<S, CP, CS, V, I, DB> {
+	get config(): ActorConfig<S, CP, CS, V, I, DB, E, Q> {
 		return this.#config;
 	}
 
@@ -260,8 +377,12 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		this.#actorId = actorId;
 		this.#name = name;
 		this.#key = key;
+		this.#actorKeyString = serializeActorKey(this.#key);
 		this.#region = region;
 
+		// Initialize tracing
+		this.#initializeTraces();
+
 		// Initialize logging
 		this.#initializeLogging();
 
@@ -269,6 +390,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		this.connectionManager = new ConnectionManager(this);
 		this.stateManager = new StateManager(this, actorDriver, this.#config);
 		this.eventManager = new EventManager(this);
+		this.queueManager = new QueueManager(this, actorDriver);
 		this.#scheduleManager = new ScheduleManager(
 			this,
 			actorDriver,
@@ -281,6 +403,8 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		// Load state
 		await this.#loadState();
 
+		await this.queueManager.initialize();
+
 		// Generate or load inspector token
 		await this.#initializeInspectorToken();
 
@@ -318,6 +442,9 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		// timer
 		this.resetSleepTimer();
 
+		// Start run handler in background (does not block startup)
+		this.#startRunHandler();
+
 		// Trigger any pending alarms
 		await this.onAlarm();
 	}
@@ -333,6 +460,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			throw new errors.InternalError("Actor is stopping");
 	}
 
+	async cleanupPersistedConnections(reason?: string): Promise<number> {
+		this.assertReady(true);
+		return await this.connectionManager.cleanupPersistedHibernatableConnections(
+			reason,
+		);
+	}
+
 	// MARK: - Stop
 	async onStop(mode: "sleep" | "destroy") {
 		if (this.#stopCalled) {
@@ -345,46 +479,53 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			mode,
 		});
 
-		// Clear sleep timeout
-		if (this.#sleepTimeout) {
-			clearTimeout(this.#sleepTimeout);
-			this.#sleepTimeout = undefined;
-		}
-
-		// Abort listeners
 		try {
-			this.#abortController.abort();
-		} catch {}
-
-		// Call onStop lifecycle
-		if (mode === "sleep") {
-			await this.#callOnSleep();
-		} else if (mode === "destroy") {
-			await this.#callOnDestroy();
-		} else {
-			assertUnreachable(mode);
-		}
+			// Clear sleep timeout
+			if (this.#sleepTimeout) {
+				clearTimeout(this.#sleepTimeout);
+				this.#sleepTimeout = undefined;
+			}
 
-		// Disconnect non-hibernatable connections
-		await this.#disconnectConnections();
+			// Abort listeners
+			try {
+				this.#abortController.abort();
+			} catch { }
 
-		// Wait for background tasks
-		await this.#waitBackgroundPromises(
-			this.#config.options.waitUntilTimeout,
-		);
+			// Wait for run handler to complete
+			await this.#waitForRunHandler(this.#config.options.runStopTimeout);
 
-		// Clear timeouts and save state
-		this.#rLog.info({ msg: "clearing pending save timeouts" });
-		this.stateManager.clearPendingSaveTimeout();
-		this.#rLog.info({ msg: "saving state immediately" });
-		await this.stateManager.saveState({
-			immediate: true,
-			allowStoppingState: true,
-		});
+			// Call onStop lifecycle
+			if (mode === "sleep") {
+				await this.#callOnSleep();
+			} else if (mode === "destroy") {
+				await this.#callOnDestroy();
+			} else {
+				assertUnreachable(mode);
+			}
+
+			// Disconnect non-hibernatable connections
+			await this.#disconnectConnections();
+
+			// Wait for background tasks
+			await this.#waitBackgroundPromises(
+				this.#config.options.waitUntilTimeout,
+			);
+
+			// Clear timeouts and save state
+			this.#rLog.info({ msg: "clearing pending save timeouts" });
+			this.stateManager.clearPendingSaveTimeout();
+			this.#rLog.info({ msg: "saving state immediately" });
+			await this.stateManager.saveState({
+				immediate: true,
+				allowStoppingState: true,
+			});
 
-		// Wait for write queues
-		await this.stateManager.waitForPendingWrites();
-		await this.#scheduleManager.waitForPendingAlarmWrites();
+			// Wait for write queues
+			await this.stateManager.waitForPendingWrites();
+			await this.#scheduleManager.waitForPendingAlarmWrites();
+		} finally {
+			await this.#cleanupDatabase();
+		}
 	}
 
 	// MARK: - Sleep
@@ -468,49 +609,71 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	async processMessage(
 		message: {
 			body:
-				| {
-						tag: "ActionRequest";
-						val: { id: bigint; name: string; args: unknown };
-				  }
-				| {
-						tag: "SubscriptionRequest";
-						val: { eventName: string; subscribe: boolean };
-				  };
+			| {
+				tag: "ActionRequest";
+				val: { id: bigint; name: string; args: unknown };
+			}
+			| {
+				tag: "SubscriptionRequest";
+				val: { eventName: string; subscribe: boolean };
+			};
 		},
-		conn: Conn<S, CP, CS, V, I, DB>,
+		conn: Conn<S, CP, CS, V, I, DB, E, Q>,
 	) {
 		await processMessage(message, this, conn, {
 			onExecuteAction: async (ctx, name, args) => {
-				this.inspector.emitter.emit("eventFired", {
-					type: "action",
-					name,
-					args,
-					connId: conn.id,
-				});
 				return await this.executeAction(ctx, name, args);
 			},
 			onSubscribe: async (eventName, conn) => {
-				this.inspector.emitter.emit("eventFired", {
-					type: "subscribe",
-					eventName,
-					connId: conn.id,
-				});
 				this.eventManager.addSubscription(eventName, conn, false);
 			},
 			onUnsubscribe: async (eventName, conn) => {
-				this.inspector.emitter.emit("eventFired", {
-					type: "unsubscribe",
-					eventName,
-					connId: conn.id,
-				});
 				this.eventManager.removeSubscription(eventName, conn, false);
 			},
 		});
 	}
 
+	async assertCanSubscribe(
+		ctx: ActionContext<S, CP, CS, V, I, DB, E, Q>,
+		eventName: string,
+	): Promise<void> {
+		const canSubscribe = getEventCanSubscribe(this.#config.events, eventName);
+		if (!canSubscribe) {
+			return;
+		}
+
+		const result = await canSubscribe(ctx);
+		if (typeof result !== "boolean") {
+			throw new errors.InvalidCanSubscribeResponse();
+		}
+		if (!result) {
+			throw new errors.Forbidden();
+		}
+	}
+
+	async assertCanPublish(
+		ctx: ActionContext<S, CP, CS, V, I, DB, E, Q>,
+		queueName: string,
+	): Promise<void> {
+		const canPublish = getQueueCanPublish<
+			ActionContext<S, CP, CS, V, I, DB, E, Q>
+		>(this.#config.queues, queueName);
+		if (!canPublish) {
+			return;
+		}
+
+		const result = await canPublish(ctx);
+		if (typeof result !== "boolean") {
+			throw new errors.InvalidCanPublishResponse();
+		}
+		if (!result) {
+			throw new errors.Forbidden();
+		}
+	}
+
 	// MARK: - Action Execution
 	async executeAction(
-		ctx: ActionContext<S, CP, CS, V, I, DB>,
+		ctx: ActionContext<S, CP, CS, V, I, DB, E, Q>,
 		actionName: string,
 		args: unknown[],
 	): Promise<unknown> {
@@ -532,54 +695,82 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			throw new errors.ActionNotFound(actionName);
 		}
 
-		try {
-			this.#rLog.debug({
-				msg: "executing action",
-				actionName,
-				args,
-			});
+		this.#activeKeepAwakeCount++;
+		this.resetSleepTimer();
+		const actionSpan = this.startTraceSpan(`actor.action.${actionName}`, {
+			"rivet.action.name": actionName,
+		});
+		let spanEnded = false;
 
-			const outputOrPromise = actionFunction.call(
-				undefined,
-				ctx,
-				...args,
-			);
+		try {
+			return await this.#traces.withSpan(actionSpan, async () => {
+				this.#rLog.debug({
+					msg: "executing action",
+					actionName,
+					args,
+				});
 
-			let output: unknown;
-			if (outputOrPromise instanceof Promise) {
-				output = await deadline(
-					outputOrPromise,
-					this.#config.options.actionTimeout,
+				const outputOrPromise = actionFunction.call(
+					undefined,
+					ctx,
+					...args,
 				);
-			} else {
-				output = outputOrPromise;
-			}
 
-			// Process through onBeforeActionResponse if configured
-			if (this.#config.onBeforeActionResponse) {
-				try {
-					const processedOutput = this.#config.onBeforeActionResponse(
-						this.actorContext,
-						actionName,
-						args,
-						output,
+				let output: unknown;
+				const maybeThenable = outputOrPromise as {
+					then?: (
+						onfulfilled?: unknown,
+						onrejected?: unknown,
+					) => unknown;
+				};
+				if (maybeThenable && typeof maybeThenable.then === "function") {
+					output = await deadline(
+						Promise.resolve(outputOrPromise),
+						this.#config.options.actionTimeout,
 					);
-					if (processedOutput instanceof Promise) {
-						output = await processedOutput;
-					} else {
-						output = processedOutput;
+				} else {
+					output = outputOrPromise;
+				}
+
+				// Process through onBeforeActionResponse if configured
+				if (this.#config.onBeforeActionResponse) {
+					try {
+						output = await this.runInTraceSpan(
+							"actor.onBeforeActionResponse",
+							{ "rivet.action.name": actionName },
+							() =>
+								this.#config.onBeforeActionResponse!(
+									this.actorContext,
+									actionName,
+									args,
+									output,
+								),
+						);
+					} catch (error) {
+						this.#rLog.error({
+							msg: "error in `onBeforeActionResponse`",
+							error: stringifyError(error),
+						});
 					}
-				} catch (error) {
-					this.#rLog.error({
-						msg: "error in `onBeforeActionResponse`",
-						error: stringifyError(error),
-					});
 				}
-			}
 
-			return output;
+				return output;
+			});
 		} catch (error) {
-			if (error instanceof DeadlineError) {
+			const isTimeout = error instanceof DeadlineError;
+			const message = isTimeout
+				? "ActionTimedOut"
+				: stringifyError(error);
+			this.#traces.setAttributes(actionSpan, {
+				"error.message": message,
+				"error.type":
+					error instanceof Error ? error.name : typeof error,
+			});
+			this.#traces.endSpan(actionSpan, {
+				status: { code: "ERROR", message },
+			});
+			spanEnded = true;
+			if (isTimeout) {
 				throw new errors.ActionTimedOut();
 			}
 			this.#rLog.error({
@@ -589,13 +780,27 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			});
 			throw error;
 		} finally {
+			if (!spanEnded && actionSpan.isActive()) {
+				this.#traces.endSpan(actionSpan, {
+					status: { code: "OK" },
+				});
+			}
+			this.#activeKeepAwakeCount--;
+			if (this.#activeKeepAwakeCount < 0) {
+				this.#activeKeepAwakeCount = 0;
+				this.#rLog.warn({
+					msg: "active keep awake count went below 0, this is a RivetKit bug",
+					...EXTRA_ERROR_LOG,
+				});
+			}
+			this.resetSleepTimer();
 			this.stateManager.savePersistThrottled();
 		}
 	}
 
 	// MARK: - HTTP/WebSocket Handlers
 	async handleRawRequest(
-		conn: Conn<S, CP, CS, V, I, DB>,
+		conn: Conn<S, CP, CS, V, I, DB, E, Q>,
 		request: Request,
 	): Promise<Response> {
 		this.assertReady();
@@ -603,27 +808,38 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		if (!this.#config.onRequest) {
 			throw new errors.RequestHandlerNotDefined();
 		}
-
-		try {
-			const ctx = new RequestContext(this, conn, request);
-			const response = await this.#config.onRequest(ctx, request);
-			if (!response) {
-				throw new errors.InvalidRequestHandlerResponse();
-			}
-			return response;
-		} catch (error) {
-			this.#rLog.error({
-				msg: "onRequest error",
-				error: stringifyError(error),
-			});
-			throw error;
-		} finally {
-			this.stateManager.savePersistThrottled();
+		const onRequest = this.#config.onRequest;
+
+			return await this.runInTraceSpan(
+				"actor.onRequest",
+				{
+					"http.method": request.method,
+					"http.url": request.url,
+					"rivet.conn.id": conn.id,
+				},
+				async () => {
+					const ctx = new RequestContext(this, conn, request);
+					try {
+						const response = await onRequest(ctx, request);
+						if (!response) {
+							throw new errors.InvalidRequestHandlerResponse();
+						}
+						return response;
+					} catch (error) {
+						this.#rLog.error({
+							msg: "onRequest error",
+							error: stringifyError(error),
+						});
+						throw error;
+					} finally {
+						this.stateManager.savePersistThrottled();
+					}
+				},
+			);
 		}
-	}
 
 	handleRawWebSocket(
-		conn: Conn<S, CP, CS, V, I, DB>,
+		conn: Conn<S, CP, CS, V, I, DB, E, Q>,
 		websocket: UniversalWebSocket,
 		request?: Request,
 	) {
@@ -635,6 +851,12 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			throw new errors.InternalError("onWebSocket handler not defined");
 		}
 
+		const span = this.startTraceSpan("actor.onWebSocket", {
+			"http.url": request?.url,
+			"rivet.conn.id": conn.id,
+		});
+		let spanEnded = false;
+
 		try {
 			// Reset sleep timer when handling WebSocket
 			this.resetSleepTimer();
@@ -643,17 +865,50 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			const ctx = new WebSocketContext(this, conn, request);
 
 			// NOTE: This is async and will run in the background
-			const voidOrPromise = this.#config.onWebSocket(ctx, websocket);
+			const voidOrPromise = this.#traces.withSpan(span, () =>
+				this.#config.onWebSocket!(ctx, websocket),
+			);
 
 			// Save changes from the WebSocket open
 			if (voidOrPromise instanceof Promise) {
-				voidOrPromise.then(() => {
-					this.stateManager.savePersistThrottled();
-				});
+				voidOrPromise
+					.then(() => {
+						if (!spanEnded) {
+							this.endTraceSpan(span, { code: "OK" });
+							spanEnded = true;
+						}
+					})
+					.catch((error) => {
+						if (!spanEnded) {
+							this.endTraceSpan(span, {
+								code: "ERROR",
+								message: stringifyError(error),
+							});
+							spanEnded = true;
+						}
+						this.#rLog.error({
+							msg: "onWebSocket error",
+							error: stringifyError(error),
+						});
+					})
+					.finally(() => {
+						this.stateManager.savePersistThrottled();
+					});
 			} else {
+				if (!spanEnded) {
+					this.endTraceSpan(span, { code: "OK" });
+					spanEnded = true;
+				}
 				this.stateManager.savePersistThrottled();
 			}
 		} catch (error) {
+			if (!spanEnded) {
+				this.endTraceSpan(span, {
+					code: "ERROR",
+					message: stringifyError(error),
+				});
+				spanEnded = true;
+			}
 			this.#rLog.error({
 				msg: "onWebSocket error",
 				error: stringifyError(error),
@@ -693,11 +948,135 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		this.#backgroundPromises.push(nonfailablePromise);
 	}
 
+	/**
+	 * Prevents the actor from sleeping while the given promise is running.
+	 *
+	 * Use this when performing async operations in the `run` handler or other
+	 * background contexts where you need to ensure the actor stays awake.
+	 *
+	 * Returns the resolved value and resets the sleep timer on completion.
+	 * Errors are propagated to the caller.
+	 */
+	async keepAwake<T>(promise: Promise<T>): Promise<T> {
+		this.assertReady();
+
+		this.#activeKeepAwakeCount++;
+		this.resetSleepTimer();
+
+		try {
+			return await promise;
+		} finally {
+			this.#activeKeepAwakeCount--;
+			if (this.#activeKeepAwakeCount < 0) {
+				this.#activeKeepAwakeCount = 0;
+				this.#rLog.warn({
+					msg: "active keep awake count went below 0, this is a RivetKit bug",
+					...EXTRA_ERROR_LOG,
+				});
+			}
+			this.resetSleepTimer();
+		}
+	}
+
+	beginQueueWait() {
+		this.assertReady(true);
+		this.#activeQueueWaitCount++;
+		this.resetSleepTimer();
+	}
+
+	endQueueWait() {
+		this.#activeQueueWaitCount--;
+		if (this.#activeQueueWaitCount < 0) {
+			this.#activeQueueWaitCount = 0;
+			this.#rLog.warn({
+				msg: "active queue wait count went below 0, this is a RivetKit bug",
+				...EXTRA_ERROR_LOG,
+			});
+		}
+		this.resetSleepTimer();
+	}
+
 	// MARK: - Private Helper Methods
+	#initializeTraces() {
+		if (getRivetExperimentalOtel()) {
+			// Experimental mode persists trace data to actor storage so inspector
+			// queries can return OTel payloads.
+			this.#traces = createTraces({
+				driver: new ActorTracesDriver(this.driver, this.#actorId),
+			});
+		} else {
+			// Keep the tracing API calls active while disabling trace persistence
+			// until the experimental flag is enabled.
+			this.#traces = createNoopTraces();
+		}
+	}
+
+	#traceAttributes(
+		attributes?: Record<string, unknown>,
+	): Record<string, unknown> {
+		return {
+			"rivet.actor.id": this.#actorId,
+			"rivet.actor.name": this.#name,
+			"rivet.actor.key": this.#actorKeyString,
+			"rivet.actor.region": this.#region,
+			...(attributes ?? {}),
+		};
+	}
+
+	#patchLoggerForTraces(logger: Logger) {
+		const levels: Array<
+			"trace" | "debug" | "info" | "warn" | "error" | "fatal"
+		> = ["trace", "debug", "info", "warn", "error", "fatal"];
+		for (const level of levels) {
+			const original = logger[level].bind(logger) as (
+				...args: any[]
+			) => unknown;
+			logger[level] = ((...args: unknown[]) => {
+				this.#emitLogEvent(level, args);
+				return original(...(args as any[]));
+			}) as Logger[typeof level];
+		}
+	}
+
+	#emitLogEvent(level: string, args: unknown[]) {
+		const span = this.#traces.getCurrentSpan();
+		if (!span || !span.isActive()) {
+			return;
+		}
+
+		let message: string | undefined;
+		if (args.length >= 2) {
+			message = String(args[1]);
+		} else if (args.length === 1) {
+			const [value] = args;
+			if (typeof value === "string") {
+				message = value;
+			} else if (
+				typeof value === "number" ||
+				typeof value === "boolean"
+			) {
+				message = String(value);
+			} else if (value && typeof value === "object") {
+				const maybeMsg = (value as { msg?: unknown }).msg;
+				if (maybeMsg !== undefined) {
+					message = String(maybeMsg);
+				}
+			}
+		}
+
+		this.#traces.emitEvent(span, "log", {
+			attributes: this.#traceAttributes({
+				"log.level": level,
+				...(message ? { "log.message": message } : {}),
+			}),
+			timeUnixMs: Date.now(),
+		});
+	}
+
 	#initializeLogging() {
 		const logParams = {
 			actor: this.#name,
-			key: serializeActorKey(this.#key),
+			key: this.#actorKeyString,
 			actorId: this.#actorId,
 		};
 
@@ -716,6 +1095,9 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 				logParams,
 			),
 		);
+
+		this.#patchLoggerForTraces(this.#log);
+		this.#patchLoggerForTraces(this.#rLog);
 	}
 
 	async #loadState() {
@@ -753,9 +1135,9 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 
 		// Call onCreate lifecycle
 		if (this.#config.onCreate) {
-			await this.#config.onCreate(
-				this.actorContext as any,
-				persistData.input!,
+			const onCreate = this.#config.onCreate;
+			await this.runInTraceSpan("actor.onCreate", undefined, () =>
+				onCreate(this.actorContext as any, persistData.input!),
 			);
 		}
 	}
@@ -821,18 +1203,24 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	async #initializeVars() {
 		let vars: V | undefined;
 		if ("createVars" in this.#config) {
-			const dataOrPromise = this.#config.createVars(
-				this.actorContext as any,
-				this.driver.getContext(this.#actorId),
+			const createVars = this.#config.createVars;
+			vars = await this.runInTraceSpan(
+				"actor.createVars",
+				undefined,
+				() => {
+					const dataOrPromise = createVars!(
+						this.actorContext as any,
+						this.driver.getContext(this.#actorId),
+					);
+					if (dataOrPromise instanceof Promise) {
+						return deadline(
+							dataOrPromise,
+							this.#config.options.createVarsTimeout,
+						);
+					}
+					return dataOrPromise;
+				},
 			);
-			if (dataOrPromise instanceof Promise) {
-				vars = await deadline(
-					dataOrPromise,
-					this.#config.options.createVarsTimeout,
-				);
-			} else {
-				vars = dataOrPromise;
-			}
 		} else if ("vars" in this.#config) {
 			vars = structuredClone(this.#config.vars);
 		} else {
@@ -846,21 +1234,31 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	async #callOnStart() {
 		this.#rLog.info({ msg: "actor starting" });
 		if (this.#config.onWake) {
-			const result = this.#config.onWake(this.actorContext);
-			if (result instanceof Promise) {
-				await result;
-			}
+			const onWake = this.#config.onWake;
+			await this.runInTraceSpan("actor.onWake", undefined, () =>
+				onWake(this.actorContext),
+			);
 		}
 	}
 
 	async #callOnSleep() {
 		if (this.#config.onSleep) {
+			const onSleep = this.#config.onSleep;
 			try {
 				this.#rLog.debug({ msg: "calling onSleep" });
-				const result = this.#config.onSleep(this.actorContext);
-				if (result instanceof Promise) {
-					await deadline(result, this.#config.options.onSleepTimeout);
-				}
+				await this.runInTraceSpan(
+					"actor.onSleep",
+					undefined,
+					async () => {
+						const result = onSleep(this.actorContext);
+						if (result instanceof Promise) {
+							await deadline(
+								result,
+								this.#config.options.onSleepTimeout,
+							);
+						}
+					},
+				);
 				this.#rLog.debug({ msg: "onSleep completed" });
 			} catch (error) {
 				if (error instanceof DeadlineError) {
@@ -877,15 +1275,22 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 
 	async #callOnDestroy() {
 		if (this.#config.onDestroy) {
+			const onDestroy = this.#config.onDestroy;
 			try {
 				this.#rLog.debug({ msg: "calling onDestroy" });
-				const result = this.#config.onDestroy(this.actorContext);
-				if (result instanceof Promise) {
-					await deadline(
-						result,
-						this.#config.options.onDestroyTimeout,
-					);
-				}
+				await this.runInTraceSpan(
+					"actor.onDestroy",
+					undefined,
+					async () => {
+						const result = onDestroy(this.actorContext);
+						if (result instanceof Promise) {
+							await deadline(
+								result,
+								this.#config.options.onDestroyTimeout,
+							);
+						}
+					},
+				);
 				this.#rLog.debug({ msg: "onDestroy completed" });
 			} catch (error) {
 				if (error instanceof DeadlineError) {
@@ -900,15 +1305,212 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		}
 	}
 
+	#startRunHandler() {
+		const runFn = getRunFunction(this.#config.run);
+		if (!runFn) return;
+
+		this.#rLog.debug({ msg: "starting run handler" });
+		this.#runHandlerActive = true;
+		this.resetSleepTimer();
+
+		const runSpan = this.startTraceSpan("actor.run");
+		const runResult = this.#traces.withSpan(runSpan, () =>
+			runFn(this.actorContext),
+		);
+
+		if (runResult instanceof Promise) {
+			this.#runPromise = runResult
+				.then(() => {
+					if (this.#stopCalled) {
+						if (runSpan.isActive()) {
+							this.endTraceSpan(runSpan, { code: "OK" });
+						}
+						this.#rLog.debug({
+							msg: "run handler exited during actor stop",
+						});
+						return;
+					}
+
+					// Run handler exited normally - this should crash the actor
+					this.emitTraceEvent(
+						"actor.crash",
+						{ "rivet.actor.reason": "run_exited" },
+						runSpan,
+					);
+					this.endTraceSpan(runSpan, {
+						code: "ERROR",
+						message: "run exited unexpectedly",
+					});
+					this.#rLog.warn({
+						msg: "run handler exited unexpectedly, crashing actor to reschedule",
+					});
+					this.startDestroy();
+				})
+				.catch((error) => {
+					if (this.#stopCalled) {
+						if (runSpan.isActive()) {
+							this.endTraceSpan(runSpan, { code: "OK" });
+						}
+						this.#rLog.debug({
+							msg: "run handler threw during actor stop",
+							error: stringifyError(error),
+						});
+						return;
+					}
+
+					// Run handler threw an error - crash the actor
+					this.emitTraceEvent(
+						"actor.crash",
+						{
+							"rivet.actor.reason": "run_error",
+							"error.message": stringifyError(error),
+						},
+						runSpan,
+					);
+					this.endTraceSpan(runSpan, {
+						code: "ERROR",
+						message: stringifyError(error),
+					});
+					this.#rLog.error({
+						msg: "run handler threw error, crashing actor to reschedule",
+						error: stringifyError(error),
+					});
+					this.startDestroy();
+				})
+				.finally(() => {
+					this.#runHandlerActive = false;
+					this.resetSleepTimer();
+				});
+		} else if (runSpan.isActive()) {
+			this.endTraceSpan(runSpan, { code: "OK" });
+			this.#runHandlerActive = false;
+			this.resetSleepTimer();
+		}
+	}
+
+	async #waitForRunHandler(timeoutMs: number) {
+		if (!this.#runPromise) {
+			return;
+		}
+
+		this.#rLog.debug({ msg: "waiting for run handler to complete" });
+
+		const timedOut = await Promise.race([
+			this.#runPromise.then(() => false).catch(() => false),
+			new Promise<true>((resolve) =>
+				setTimeout(() => resolve(true), timeoutMs),
+			),
+		]);
+
+		if (timedOut) {
+			this.#rLog.warn({
+				msg: "run handler did not complete in time, it may have leaked - ensure you use c.aborted (or the abort signal c.abortSignal) to exit gracefully",
+				timeoutMs,
+			});
+		} else {
+			this.#rLog.debug({ msg: "run handler completed" });
+		}
+	}
+
 	async #setupDatabase() {
-		if ("db" in this.#config && this.#config.db) {
-			const client = await this.#config.db.createClient({
-				getDatabase: () => this.driver.getDatabase(this.#actorId),
+		if (!("db" in this.#config) || !this.#config.db) {
+			return;
+		}
+
+		let client: InferDatabaseClient<DB> | undefined;
+		try {
+			// Every actor gets its own SqliteVfs/@rivetkit/sqlite instance. The async
+			// @rivetkit/sqlite build is not re-entrant, and sharing one instance across
+			// actors can cause cross-actor contention and runtime corruption.
+			if (!this.#sqliteVfs && this.driver.createSqliteVfs) {
+				this.#sqliteVfs = await this.driver.createSqliteVfs();
+			}
+
+			client = await this.#config.db.createClient({
+				actorId: this.#actorId,
+				overrideRawDatabaseClient: this.driver.overrideRawDatabaseClient
+					? () => this.driver.overrideRawDatabaseClient!(this.#actorId)
+					: undefined,
+				overrideDrizzleDatabaseClient: this.driver.overrideDrizzleDatabaseClient
+					? () => this.driver.overrideDrizzleDatabaseClient!(this.#actorId)
+					: undefined,
+				kv: {
+					batchPut: (entries) => this.driver.kvBatchPut(this.#actorId, entries),
+					batchGet: (keys) => this.driver.kvBatchGet(this.#actorId, keys),
+					batchDelete: (keys) => this.driver.kvBatchDelete(this.#actorId, keys),
+				},
+				sqliteVfs: this.#sqliteVfs,
 			});
 			this.#rLog.info({ msg: "database migration starting" });
 			await this.#config.db.onMigrate?.(client);
 			this.#rLog.info({ msg: "database migration complete" });
 			this.#db = client;
+		} catch (error) {
+			if (client) {
+				try {
+					await this.#config.db.onDestroy?.(client);
+				} catch (cleanupError) {
+					this.#rLog.error({
+						msg: "database setup cleanup failed",
+						error: stringifyError(cleanupError),
+					});
+				}
+			}
+			if (this.#sqliteVfs) {
+				try {
+					await this.#sqliteVfs.destroy();
+				} catch (cleanupError) {
+					this.#rLog.error({
+						msg: "sqlite vfs teardown after setup failure failed",
+						error: stringifyError(cleanupError),
+					});
+				}
+			}
+			this.#sqliteVfs = undefined;
+			if (error instanceof Error) {
+				this.#rLog.error({
+					msg: "database setup failed",
+					error: stringifyError(error),
+				});
+				throw error;
+			}
+			const wrappedError = new Error(`Database setup failed: ${String(error)}`);
+			this.#rLog.error({
+				msg: "database setup failed with non-Error object",
+				error: String(error),
+				errorType: typeof error,
+			});
+			throw wrappedError;
+		}
+	}
+
+	async #cleanupDatabase() {
+		const client = this.#db;
+		const sqliteVfs = this.#sqliteVfs;
+		const dbConfig = "db" in this.#config ? this.#config.db : undefined;
+		this.#db = undefined;
+		this.#sqliteVfs = undefined;
+
+		if (client && dbConfig) {
+			try {
+				await dbConfig.onDestroy?.(client);
+			} catch (error) {
+				this.#rLog.error({
+					msg: "database cleanup failed",
+					error: stringifyError(error),
+				});
+			}
+		}
+
+		if (sqliteVfs) {
+			try {
+				await sqliteVfs.destroy();
+			} catch (error) {
+				this.#rLog.error({
+					msg: "sqlite vfs cleanup failed",
+					error: stringifyError(error),
+				});
+			}
 		}
 	}
 
@@ -1010,6 +1612,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		if (!this.#started) return CanSleep.NotReady;
 		if (this.#activeHonoHttpRequests > 0)
 			return CanSleep.ActiveHonoHttpRequests;
+		if (this.#activeKeepAwakeCount > 0) return CanSleep.ActiveKeepAwake;
+		if (this.#runHandlerActive && this.#activeQueueWaitCount === 0) {
+			return CanSleep.ActiveRun;
+		}
 
 		for (const _conn of this.connectionManager.connections.values()) {
 			// TODO: Add back
@@ -1018,6 +1624,10 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			// }
 		}
 
+		if (this.connectionManager.pendingDisconnectCount > 0) {
+			return CanSleep.ActiveDisconnectCallbacks;
+		}
+
 		return CanSleep.Yes;
 	}