rivetkit 2.3.0-rc.8 → 2.3.0-rc.9

package/src/registry/native.ts

@@ -2,7 +2,6 @@ import { VirtualWebSocket } from "@rivetkit/virtual-websocket";
 import {
 	ACTOR_CONTEXT_INTERNAL_SYMBOL,
 	CONN_STATE_MANAGER_SYMBOL,
-	RAW_STATE_SYMBOL,
 	getRunFunction,
 	getRunInspectorConfig,
 	type WorkflowInspectorConfig,
@@ -34,10 +33,6 @@ import { HEADER_CONN_PARAMS } from "@/common/actor-router-consts";
 import type { AnyDatabaseProvider } from "@/common/database/config";
 import { wrapJsNativeDatabase } from "@/common/database/native-database";
 import { decodeWorkflowHistoryTransport } from "@/common/inspector-transport";
-import {
-	assertJsonCompatValue,
-	type JsonCompatValue,
-} from "@/common/encoding";
 import { deconstructError, stringifyError } from "@/common/utils";
 import type {
 	RivetCloseEvent,
@@ -59,7 +54,6 @@ import {
 } from "@/serde";
 import { getEnvUniversal, VERSION } from "@/utils";
 import { logger } from "./log";
-import { createWriteThroughProxy } from "./write-through-proxy";
 import { loadNapiRuntime } from "./napi-runtime";
 import {
 	type NativeValidationConfig,
@@ -249,6 +243,14 @@ type NativePersistActorState = {
 	state: unknown;
 	isInOnStateChange: boolean;
 	connStates: Map<string, NativePersistConnState>;
+	// Memoized deep write-through proxy and the state object it wraps. Rebuilt
+	// only when the underlying state object identity changes.
+	stateProxy?: unknown;
+	stateProxyTarget?: unknown;
+	// Set when a coalesced save and onStateChange flush is pending for the
+	// current event loop tick.
+	saveScheduled?: boolean;
+	pendingSaveHandle?: ReturnType<typeof setImmediate>;
 };
 type NativeDestroyGate = {
 	destroyCompletion?: Promise<void>;
@@ -343,15 +345,6 @@ function databaseNotConfiguredError(): RivetError {
 	);
 }
 
-function databaseClientNotReadyError(): RivetError {
-	return new RivetError(
-		"actor",
-		"database_client_not_ready",
-		"actor database client was not initialized before user code ran. this is an internal lifecycle error; the migration callback should have pre-warmed the client. file an issue if you can reproduce.",
-		{ public: true },
-	);
-}
-
 function stateNotEnabledError(): RivetError {
 	return new RivetError(
 		"actor",
@@ -421,20 +414,7 @@ async function cleanupNativeSleepRuntimeState(
 	runtime: CoreRuntime,
 	ctx: ActorContextHandle,
 ): Promise<void> {
-	const waitStarted = Date.now();
-	const drained = await runtime.actorWaitForTrackedShutdownWork(ctx);
-	const waitMs = Date.now() - waitStarted;
-	if (drained) {
-		logger().debug({
-			msg: "sleep cleanup: tracked shutdown work drained",
-			waitMs,
-		});
-	} else {
-		logger().warn({
-			msg: "sleep cleanup: shutdown deadline reached before tracked work drained; closing DB anyway",
-			waitMs,
-		});
-	}
+	await runtime.actorWaitForTrackedShutdownWork(ctx);
 	await closeNativeDatabaseClient(runtime, ctx);
 	await closeNativeSqlDatabase(runtime, ctx);
 	clearNativeRuntimeState(runtime, ctx);
@@ -611,7 +591,7 @@ function decodeValue<T>(value?: RuntimeBytes | null): T {
 }
 
 function encodeValue(value: unknown): RuntimeBytes {
-	return encodeCborCompat(value as JsonCompatValue);
+	return encodeCborCompat(value);
 }
 
 function unwrapTsfnPayload<T>(error: unknown, payload: T): T {
@@ -1086,6 +1066,54 @@ function decodeArgs(value?: RuntimeBytes | null): unknown[] {
 			: [decoded];
 }
 
+function createWriteThroughProxy<T>(
+	value: T,
+	commit: (next: T) => void,
+	beforeChange?: () => void,
+): T {
+	if (!value || typeof value !== "object") {
+		return value;
+	}
+
+	const proxies = new WeakMap<object, object>();
+	const wrap = (target: object): object => {
+		const cached = proxies.get(target);
+		if (cached) {
+			return cached;
+		}
+
+		const proxy = new Proxy(target, {
+			get(innerTarget, property, receiver) {
+				const result = Reflect.get(innerTarget, property, receiver);
+				return result && typeof result === "object"
+					? wrap(result as object)
+					: result;
+			},
+			set(innerTarget, property, nextValue, receiver) {
+				beforeChange?.();
+				const updated = Reflect.set(
+					innerTarget,
+					property,
+					nextValue,
+					receiver,
+				);
+				commit(value);
+				return updated;
+			},
+			deleteProperty(innerTarget, property) {
+				beforeChange?.();
+				const updated = Reflect.deleteProperty(innerTarget, property);
+				commit(value);
+				return updated;
+			},
+		});
+
+		proxies.set(target, proxy);
+		return proxy;
+	};
+
+	return wrap(value as object) as T;
+}
 
 function buildRequest(init: {
 	method: string;
@@ -1176,21 +1204,14 @@ class NativeConnAdapter {
 		);
 	}
 
-	[RAW_STATE_SYMBOL](): unknown {
-		return this.#readState();
-	}
-
 	get state(): unknown {
 		const nextState = this.#readState();
 		return createWriteThroughProxy(nextState, (nextValue) => {
 			this.#writeState(nextValue, { writeNative: true });
-		}, (newValue) => {
-			assertJsonCompatValue(newValue);
 		});
 	}
 
 	set state(value: unknown) {
-		assertJsonCompatValue(value);
 		this.#writeState(value, { writeNative: true });
 	}
 
@@ -2302,90 +2323,6 @@ class TrackedWebSocketHandleAdapter implements UniversalWebSocket {
 	}
 }
 
-class NativeConnectionMap implements ReadonlyMap<string, NativeConnAdapter> {
-	#runtime: CoreRuntime;
-	#ctx: ActorContextHandle;
-	#schemas: NativeValidationConfig;
-
-	constructor(
-		runtime: CoreRuntime,
-		ctx: ActorContextHandle,
-		schemas: NativeValidationConfig,
-	) {
-		this.#runtime = runtime;
-		this.#ctx = ctx;
-		this.#schemas = schemas;
-	}
-
-	#connToAdapter(conn: ConnHandle): NativeConnAdapter {
-		return new NativeConnAdapter(
-			this.#runtime,
-			conn,
-			this.#schemas,
-			this.#ctx,
-			(connId) =>
-				callNativeSync(() =>
-					this.#runtime.actorQueueHibernationRemoval(
-						this.#ctx,
-						connId,
-					),
-				),
-		);
-	}
-
-	get size(): number {
-		return callNativeSync(() => this.#runtime.actorConns(this.#ctx)).length;
-	}
-
-	get(key: string): NativeConnAdapter | undefined {
-		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
-		const conn = conns.find(
-			(c) => this.#runtime.connId(c) === key,
-		);
-		if (!conn) return undefined;
-		return this.#connToAdapter(conn);
-	}
-
-	has(key: string): boolean {
-		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
-		return conns.some((c) => this.#runtime.connId(c) === key);
-	}
-
-	keys(): MapIterator<string> {
-		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
-		return conns.map((c) => this.#runtime.connId(c))[Symbol.iterator]() satisfies MapIterator<string>;
-	}
-
-	values(): MapIterator<NativeConnAdapter> {
-		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
-		return conns.map((c) => this.#connToAdapter(c))[Symbol.iterator]() satisfies MapIterator<NativeConnAdapter>;
-	}
-
-	entries(): MapIterator<[string, NativeConnAdapter]> {
-		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
-		return conns.map(
-			(c) => [this.#runtime.connId(c), this.#connToAdapter(c)] as [string, NativeConnAdapter],
-		)[Symbol.iterator]() satisfies MapIterator<[string, NativeConnAdapter]>;
-	}
-
-	forEach(
-		callback: (value: NativeConnAdapter, key: string, map: ReadonlyMap<string, NativeConnAdapter>) => void,
-		thisArg?: unknown,
-	): void {
-		const conns = callNativeSync(() => this.#runtime.actorConns(this.#ctx));
-		for (const conn of conns) {
-			const id = this.#runtime.connId(conn);
-			callback.call(thisArg, this.#connToAdapter(conn), id, this);
-		}
-	}
-
-	[Symbol.iterator](): MapIterator<[string, NativeConnAdapter]> {
-		return this.entries();
-	}
-
-	readonly [Symbol.toStringTag] = "NativeConnectionMap";
-}
-
 export class ActorContextHandleAdapter {
 	#runtime: CoreRuntime;
 	#ctx: ActorContextHandle;
@@ -2394,9 +2331,9 @@ export class ActorContextHandleAdapter {
 	#abortSignalCleanup?: () => void;
 	#client?: AnyClient;
 	#clientFactory?: () => AnyClient;
-	#connMap?: NativeConnectionMap;
 	#databaseProvider?: Exclude<AnyDatabaseProvider, undefined>;
 	#db?: unknown;
+	#dbProxy?: unknown;
 	#dispatchCancelToken?: CancellationTokenHandle;
 	#kv?: NativeKvAdapter;
 	#queue?: NativeQueueAdapter;
@@ -2459,42 +2396,60 @@ export class ActorContextHandleAdapter {
 			throw databaseNotConfiguredError();
 		}
 
-		if (this.#db) {
-			return this.#db;
-		}
+		if (!this.#dbProxy) {
+			this.#dbProxy = new Proxy(
+				{},
+				{
+					get: (_target, property) => {
+						if (property === "then") {
+							return undefined;
+						}
 
-		const runtimeState = getNativeRuntimeState(this.#runtime, this.#ctx);
-		const cachedClient = runtimeState.databaseClient;
-		if (cachedClient) {
-			this.#db = cachedClient.client;
-			return this.#db;
+						return async (...args: Array<unknown>) => {
+							const client = await this.ensureDatabaseClient();
+							const value = Reflect.get(
+								client as object,
+								property,
+							);
+							if (typeof value !== "function") {
+								return value;
+							}
+							return await value.apply(client, args);
+						};
+					},
+				},
+			);
 		}
 
-		throw databaseClientNotReadyError();
-	}
-
-	[RAW_STATE_SYMBOL](): unknown {
-		if (!this.#stateEnabled) {
-			throw stateNotEnabledError();
-		}
-		return this.#readState();
+		return this.#dbProxy;
 	}
 
 	get state(): unknown {
 		if (!this.#stateEnabled) {
 			throw stateNotEnabledError();
 		}
+		const actorState = getNativePersistState(this.#runtime, this.#ctx);
 		const nextState = this.#readState();
-		return createWriteThroughProxy(
-			nextState,
-			(nextValue) => {
-				this.#writeState(nextValue, { scheduleSave: true });
-			},
-			(newValue) => {
-				this.#assertCanMutateState();
-				assertJsonCompatValue(newValue);
-			},
-		);
+		// Reading `c.state` rebuilds the deep write-through proxy, which
+		// allocates fresh on-change caches and rewraps the whole tree. Memoize
+		// the proxy keyed on the underlying state object so repeated reads and
+		// deep read cascades reuse a single proxy.
+		if (
+			actorState.stateProxy === undefined ||
+			actorState.stateProxyTarget !== nextState
+		) {
+			actorState.stateProxyTarget = nextState;
+			actorState.stateProxy = createWriteThroughProxy(
+				nextState,
+				(nextValue) => {
+					this.#writeState(nextValue, { scheduleSave: true });
+				},
+				() => {
+					this.#assertCanMutateState();
+				},
+			);
+		}
+		return actorState.stateProxy;
 	}
 
 	set state(value: unknown) {
@@ -2502,7 +2457,6 @@ export class ActorContextHandleAdapter {
 			throw stateNotEnabledError();
 		}
 		this.#assertCanMutateState();
-		assertJsonCompatValue(value);
 		this.#writeState(value, { scheduleSave: true });
 	}
 
@@ -2569,11 +2523,27 @@ export class ActorContextHandleAdapter {
 		return callNativeSync(() => this.#runtime.actorRegion(this.#ctx));
 	}
 
-	get conns(): ReadonlyMap<string, NativeConnAdapter> {
-		if (!this.#connMap) {
-			this.#connMap = new NativeConnectionMap(this.#runtime, this.#ctx, this.#schemas);
-		}
-		return this.#connMap;
+	get conns(): Map<string, NativeConnAdapter> {
+		return new Map(
+			callNativeSync(() => this.#runtime.actorConns(this.#ctx)).map(
+				(conn) => [
+					this.#runtime.connId(conn),
+					new NativeConnAdapter(
+						this.#runtime,
+						conn,
+						this.#schemas,
+						this.#ctx,
+						(connId) =>
+							callNativeSync(() =>
+								this.#runtime.actorQueueHibernationRemoval(
+									this.#ctx,
+									connId,
+								),
+							),
+					),
+				],
+			),
+		);
 	}
 
 	get log() {
@@ -2804,41 +2774,22 @@ export class ActorContextHandleAdapter {
 	}
 
 	keepAwake<T>(promise: Promise<T>): Promise<T> {
-		const startedAt = Date.now();
-		logger().debug({
-			msg: "keepAwake registered",
-			at: startedAt,
-		});
 		const trackedPromise = Promise.resolve(promise)
-			.then(
-				() => {
-					logger().debug({
-						msg: "keepAwake promise resolved",
-						durationMs: Date.now() - startedAt,
-					});
-				},
-				(error) => {
-					logger().warn({
-						msg: "keepAwake promise rejected",
-						durationMs: Date.now() - startedAt,
-						error: stringifyError(error),
-					});
-				},
-			)
+			.catch((error) => {
+				logger().warn({
+					msg: "keepAwake promise rejected",
+					error: stringifyError(error),
+				});
+			})
 			.then(() => null);
 		try {
 			callNativeSync(() =>
 				this.#runtime.actorKeepAwake(this.#ctx, trackedPromise),
 			);
 		} catch (error) {
-			if (isClosedTaskRegistrationError(error)) {
-				logger().warn({
-					msg: "keepAwake registration dropped (teardown already started); promise will not delay grace",
-					error: stringifyError(error),
-				});
-				return promise;
+			if (!isClosedTaskRegistrationError(error)) {
+				throw error;
 			}
-			throw error;
 		}
 		return promise;
 	}
@@ -2934,6 +2885,9 @@ export class ActorContextHandleAdapter {
 	}
 
 	async dispose(): Promise<void> {
+		// Flush any save coalesced for this tick before the context is torn
+		// down so the request-save and onStateChange always run.
+		this.#flushStateChange();
 		this.#abortSignalCleanup?.();
 		this.#sql = undefined;
 	}
@@ -2969,13 +2923,12 @@ export class ActorContextHandleAdapter {
 			scheduleSave: boolean;
 		},
 	): void {
-		encodeValue(value);
 		const actorState = getNativePersistState(this.#runtime, this.#ctx);
 		actorState.state = value;
 		if (!options.scheduleSave) {
 			return;
 		}
-		this.#handleStateChange();
+		this.#scheduleSave();
 	}
 
 	#assertCanMutateState(): void {
@@ -2985,9 +2938,33 @@ export class ActorContextHandleAdapter {
 		}
 	}
 
-	#handleStateChange(): void {
+	// Coalesce the request-save and onStateChange work to once per event loop
+	// tick. A synchronous burst of mutations (for example
+	// `Object.assign(c.state, ...)`) would otherwise cross the NAPI boundary and
+	// run onStateChange once per field, re-serializing the whole state each time
+	// and pinning the event loop on large state.
+	#scheduleSave(): void {
+		const actorState = getNativePersistState(this.#runtime, this.#ctx);
+		if (actorState.saveScheduled) {
+			return;
+		}
+		actorState.saveScheduled = true;
+		actorState.pendingSaveHandle = setImmediate(() => {
+			this.#flushStateChange();
+		});
+	}
+
+	#flushStateChange(): void {
 		const actorState = getNativePersistState(this.#runtime, this.#ctx);
-		encodeValue(actorState.state);
+		if (!actorState.saveScheduled) {
+			return;
+		}
+		actorState.saveScheduled = false;
+		if (actorState.pendingSaveHandle !== undefined) {
+			clearImmediate(actorState.pendingSaveHandle);
+			actorState.pendingSaveHandle = undefined;
+		}
+
 		callNativeSync(() =>
 			this.#runtime.actorRequestSave(this.#ctx, { immediate: false }),
 		);
@@ -3571,10 +3548,6 @@ export function buildNativeFactory(
 							getNativeWorkflowInspector(ctx) !== undefined,
 					});
 				} catch (error) {
-					logger().error({
-						msg: "error replaying workflow history",
-						error,
-					});
 					return errorResponse(error);
 				}
 			}
@@ -3754,10 +3727,6 @@ export function buildNativeFactory(
 					);
 					return jsonResponse({ output });
 				} catch (error) {
-					logger().error({
-						msg: "Error handling inspector action request",
-						error,
-					});
 					return errorResponse(error);
 				}
 			}
@@ -3772,10 +3741,6 @@ export function buildNativeFactory(
 				{ status: 404 },
 			);
 		} catch (error) {
-			logger().error({
-				msg: "Error handling inspector request",
-				error,
-			});
 			return errorResponse(error);
 		} finally {
 			await actorCtx.dispose();
@@ -4686,14 +4651,6 @@ export async function buildConfiguredRegistry(config: RegistryConfig): Promise<{
 	serveConfig: RuntimeServeConfig;
 }> {
 	const runtime = await loadConfiguredRuntime(config);
-	if (runtime.kind === "napi") {
-		// Start Node.js runtime health metrics collection (event loop lag,
-		// GC, heap, CPU, libuv handles). Only available on the native NAPI
-		// runtime; wasm/edge hosts do not expose perf_hooks/v8 the same
-		// way and have no Rust-side prometheus collectors loaded.
-		const { startProcessMetrics } = await import("./process-metrics");
-		startProcessMetrics();
-	}
 	return buildRegistryWithRuntime(
 		normalizeRuntimeConfig(config, runtime),
 		runtime,

package/src/registry/runtime.ts

@@ -257,9 +257,10 @@ export interface RuntimeServerlessResponseHead {
 	headers: Record<string, string>;
 }
 
-export interface RuntimeRegistryDiagnostics {
-	mode: string;
-	envoyActiveActorCount?: number | null;
+export interface RuntimeRegistryRouteResponse {
+	status: number;
+	headers: Record<string, string>;
+	body: RuntimeBytes;
 }
 
 export type RuntimeServerlessStreamEvent =
@@ -319,9 +320,15 @@ export interface CoreRuntime {
 		cancelToken: CancellationTokenHandle,
 		config: RuntimeServeConfig,
 	): Promise<RuntimeServerlessResponseHead>;
-	registryDiagnostics?(
+	registryHealth?(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse>;
+	registryMetadata?(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse>;
+	registryMetrics?(
 		registry: RegistryHandle,
-	): Promise<RuntimeRegistryDiagnostics>;
+	): Promise<RuntimeRegistryRouteResponse>;
 	createActorFactory(
 		callbacks: object,
 		config?: RuntimeActorConfig | undefined | null,

package/src/registry/wasm-runtime.ts

@@ -33,6 +33,7 @@ import type {
 	RuntimeQueueTryNextBatchOptions,
 	RuntimeQueueWaitOptions,
 	RuntimeRequestSaveOpts,
+	RuntimeRegistryRouteResponse,
 	RuntimeServeConfig,
 	RuntimeServerlessRequest,
 	RuntimeServerlessResponseHead,
@@ -270,8 +271,18 @@ export class WasmCoreRuntime implements CoreRuntime {
 		await callWasm(() => asWasmRegistry(registry).shutdown());
 	}
 
-	async registryDiagnostics(): Promise<{ mode: string; envoyActiveActorCount: null }> {
-		return { mode: "wasm", envoyActiveActorCount: null };
+	async registryHealth(): Promise<RuntimeRegistryRouteResponse> {
+		return {
+			status: 200,
+			headers: { "content-type": "application/json" },
+			body: new TextEncoder().encode(
+				JSON.stringify({
+					status: "ok",
+					runtime: "rivetkit",
+					version: "wasm",
+				}),
+			),
+		};
 	}
 
 	async handleServerlessRequest(

package/src/serde.ts

@@ -3,7 +3,7 @@ import invariant from "invariant";
 import type { VersionedDataHandler } from "vbare";
 import type { z } from "zod/v4";
 import { assertUnreachable } from "@/common/utils";
-import type { JsonCompatValue, Encoding } from "@/common/encoding";
+import type { Encoding } from "@/common/encoding";
 import {
 	encodeJsonCompatValue,
 	jsonParseCompat,
@@ -46,7 +46,7 @@ export function contentTypeForEncoding(encoding: Encoding): string {
 	}
 }
 
-export function encodeCborCompat(value: JsonCompatValue): Uint8Array {
+export function encodeCborCompat(value: unknown): Uint8Array {
 	return cbor.encode(encodeJsonCompatValue(value));
 }
 

package/src/workflow/context.ts

@@ -1,5 +1,5 @@
 // @ts-nocheck
-import { RAW_STATE_SYMBOL, type RunContext } from "@/actor/config";
+import type { RunContext } from "@/actor/config";
 import type {
 	QueueFilterName,
 	QueueNextBatchOptions,
@@ -247,14 +247,14 @@ export class ActorWorkflowContext<
 			}
 			return await this.#wrapActive(() =>
 				this.#inner.step(nameOrConfig, () =>
-					this.#withActorAccessAndStateRollback(run),
+					this.#withActorAccess(run),
 				),
 			);
 		}
 		const stepConfig = nameOrConfig as StepConfig<T>;
 		const config: StepConfig<T> = {
 			...stepConfig,
-			run: () => this.#withActorAccessAndStateRollback(stepConfig.run),
+			run: () => this.#withActorAccess(stepConfig.run),
 		};
 		return await this.#wrapActive(() => this.#inner.step(config));
 	}
@@ -271,14 +271,14 @@ export class ActorWorkflowContext<
 			}
 			return await this.#wrapActive(() =>
 				this.#inner.tryStep(nameOrConfig, () =>
-					this.#withActorAccessAndStateRollback(run),
+					this.#withActorAccess(run),
 				),
 			);
 		}
 		const stepConfig = nameOrConfig as TryStepConfig<T>;
 		const config: TryStepConfig<T> = {
 			...stepConfig,
-			run: () => this.#withActorAccessAndStateRollback(stepConfig.run),
+			run: () => this.#withActorAccess(stepConfig.run),
 		};
 		return await this.#wrapActive(() => this.#inner.tryStep(config));
 	}
@@ -612,33 +612,6 @@ export class ActorWorkflowContext<
 		}
 	}
 
-	async #withActorAccessAndStateRollback<T>(
-		run: () => Promise<T>,
-	): Promise<T> {
-		let stateSnapshot: { state: TState } | null = null;
-		try {
-			stateSnapshot = { state: this.#runCtx[RAW_STATE_SYMBOL]() };
-		} catch (error) {
-			this.#runCtx.log.debug({
-				msg: "failed to get state, likely due to being stateless workflow",
-				error,
-			});
-		}
-		if (stateSnapshot) {
-			stateSnapshot.state = structuredClone(stateSnapshot.state);
-		}
-		const varsSnapshot = structuredClone(this.#runCtx.vars);
-		try {
-			return await this.#withActorAccess(run);
-		} catch (error) {
-			if (stateSnapshot) {
-				this.#runCtx.state = stateSnapshot.state;
-			}
-			this.#runCtx.vars = varsSnapshot;
-			throw error;
-		}
-	}
-
 	#ensureActorAccess(feature: string): void {
 		if (!this.#allowActorAccess) {
 			this.#guardViolation = true;