rivetkit 2.3.0-rc.11 → 2.3.0-rc.13

package/src/registry/index.ts

@@ -1,5 +1,5 @@
 import { Hono } from "hono";
-import { ENGINE_ENDPOINT } from "@/common/engine";
+import { isLocalEngineEndpoint } from "@/common/engine";
 import { configureServerlessPool } from "@/serverless/configure";
 import { detectRuntime, VERSION } from "@/utils";
 import { crossPlatformServe, loadRuntimeServeStatic } from "@/utils/serve";
@@ -15,6 +15,22 @@ import type { RuntimeServerlessResponseHead } from "./runtime";
 
 type ShutdownSignal = "SIGINT" | "SIGTERM";
 
+function signalExitCode(signal: ShutdownSignal): number {
+	switch (signal) {
+		case "SIGINT":
+			return 130;
+		case "SIGTERM":
+			return 143;
+	}
+}
+
+function finishShutdownSignal(signal: ShutdownSignal): void {
+	if (process.pid === 1) {
+		process.exit(signalExitCode(signal));
+	}
+	process.kill(process.pid, signal);
+}
+
 export type FetchHandler = (
 	request: Request,
 	...args: any
@@ -30,8 +46,18 @@ export interface RegistryRoutes {
 	prometheusMetrics(request?: Request): Promise<Response>;
 }
 
+/**
+ * Injectable dependencies for {@link Registry}. Production code uses the
+ * defaults. Tests override `buildConfiguredRegistry` to drive lifecycle
+ * orchestration against a fake `CoreRuntime` without an engine.
+ */
+export interface RegistryDeps {
+	buildConfiguredRegistry: typeof buildConfiguredRegistry;
+}
+
 export class Registry<A extends RegistryActors> {
 	#config: RegistryConfigInput<A>;
+	#buildConfiguredRegistry: typeof buildConfiguredRegistry;
 	public readonly routes: RegistryRoutes;
 
 	get config(): RegistryConfigInput<A> {
@@ -51,8 +77,10 @@ export class Registry<A extends RegistryActors> {
 	#shutdownInFlight: Promise<void> | null = null;
 	#signalHandlers: Partial<Record<ShutdownSignal, () => void>> = {};
 
-	constructor(config: RegistryConfigInput<A>) {
+	constructor(config: RegistryConfigInput<A>, deps?: Partial<RegistryDeps>) {
 		this.#config = config;
+		this.#buildConfiguredRegistry =
+			deps?.buildConfiguredRegistry ?? buildConfiguredRegistry;
 		this.routes = {
 			health: () => this.#healthRoute(),
 			metadata: () => this.#metadataRoute(),
@@ -94,7 +122,8 @@ export class Registry<A extends RegistryActors> {
 		this.#printWelcome(config, "serverless");
 
 		if (!this.#runtimeServerlessPromise) {
-			this.#runtimeServerlessPromise = buildConfiguredRegistry(config);
+			this.#runtimeServerlessPromise =
+				this.#buildConfiguredRegistry(config);
 		}
 
 		const { runtime, registry, serveConfig } =
@@ -427,7 +456,8 @@ export class Registry<A extends RegistryActors> {
 	 */
 	#startEnvoy(config: RegistryConfig, printWelcome: boolean) {
 		if (!this.#runtimeServePromise) {
-			const configuredRegistryPromise = buildConfiguredRegistry(config);
+			const configuredRegistryPromise =
+				this.#buildConfiguredRegistry(config);
 			this.#runtimeServeConfiguredPromise = configuredRegistryPromise;
 			this.#runtimeServePromise = configuredRegistryPromise
 				.then(async ({ runtime, registry, serveConfig }) => {
@@ -445,17 +475,14 @@ export class Registry<A extends RegistryActors> {
 			// does not install handlers because it runs on Workers/Vercel/Deno
 			// Deploy where `process.on` is absent or forbidden; those platforms
 			// own their own signal policy.
-			this.#installSignalHandlers(config, configuredRegistryPromise);
+			this.#installSignalHandlers(config);
 		}
 		if (printWelcome) {
 			this.#printWelcome(config, "serverful");
 		}
 	}
 
-	#installSignalHandlers(
-		config: RegistryConfig,
-		configuredRegistryPromise: ReturnType<typeof buildConfiguredRegistry>,
-	): void {
+	#installSignalHandlers(config: RegistryConfig): void {
 		if (this.#shutdownInstalled) return;
 		if (config.shutdown?.disableSignalHandlers) return;
 		// Guard against non-Node runtimes (Workers/Edge) where `process` may
@@ -470,12 +497,7 @@ export class Registry<A extends RegistryActors> {
 		this.#shutdownInstalled = true;
 
 		const install = (signal: ShutdownSignal) => {
-			const handler = () =>
-				this.#onShutdownSignal(
-					signal,
-					config,
-					configuredRegistryPromise,
-				);
+			const handler = () => this.#onShutdownSignal(signal, config);
 			this.#signalHandlers[signal] = handler;
 			process.on(signal, handler);
 		};
@@ -483,64 +505,97 @@ export class Registry<A extends RegistryActors> {
 		install("SIGTERM");
 	}
 
-	#onShutdownSignal(
-		signal: ShutdownSignal,
-		config: RegistryConfig,
-		configuredRegistryPromise: ReturnType<typeof buildConfiguredRegistry>,
-	): void {
+	#onShutdownSignal(signal: ShutdownSignal, config: RegistryConfig): void {
 		if (this.#shutdownInFlight !== null) {
-			// Second delivery of the same (or another) shutdown signal.
-			// Remove our handler only (preserving any user-installed listeners)
-			// and re-raise so Node proceeds with its default exit path.
+			// Second delivery of the same (or another) shutdown signal, or a
+			// drain already started by an explicit `shutdown()` call. Remove
+			// our handler only, preserving any user-installed listeners. PID 1
+			// must exit directly because re-raised default signals can be
+			// swallowed by the container signal path.
 			this.#removeSignalHandlers();
-			process.kill(process.pid, signal);
+			finishShutdownSignal(signal);
 			return;
 		}
-		this.#shutdownInFlight = this.#runShutdown(
-			signal,
-			config,
-			configuredRegistryPromise,
-		).catch((error) => {
-			logger().warn({ error }, "shutdown error");
+		this.#shutdownInFlight = this.#drain(config)
+			.catch((err) => {
+				logger().warn({ err }, "shutdown error");
+			})
+			.then(() => {
+				this.#removeSignalHandlers();
+				finishShutdownSignal(signal);
+			});
+	}
+
+	/**
+	 * Gracefully drains all live registries.
+	 *
+	 * Programmatic counterpart to the SIGINT/SIGTERM handlers: tears down
+	 * every live `CoreRegistry` (both `start()` and `handler()` modes) and
+	 * waits for the serve promise to resolve, all bounded by the shutdown
+	 * grace period. Unlike a signal-driven shutdown, this does not re-raise a
+	 * signal or exit the process. The caller owns process lifetime.
+	 *
+	 * Idempotent: concurrent or repeated calls share a single drain. Safe to
+	 * call even if nothing has been started.
+	 *
+	 * @example
+	 * ```ts
+	 * const registry = setup({ use: { counter } });
+	 * registry.start();
+	 * // ...later, on your own shutdown trigger:
+	 * await registry.shutdown();
+	 * ```
+	 */
+	public async shutdown(): Promise<void> {
+		if (this.#shutdownInFlight !== null) return this.#shutdownInFlight;
+		const config = this.parseConfig();
+		// Uninstall our signal handlers so a later SIGINT/SIGTERM does not
+		// re-trigger a drain on already-torn-down registries. Subsequent
+		// signals fall back to Node's default termination behavior.
+		this.#removeSignalHandlers();
+		this.#shutdownInFlight = this.#drain(config).catch((err) => {
+			logger().warn({ err }, "shutdown error");
 		});
+		return this.#shutdownInFlight;
 	}
 
-	async #runShutdown(
-		signal: ShutdownSignal,
-		config: RegistryConfig,
-		configuredRegistryPromise: ReturnType<typeof buildConfiguredRegistry>,
-	): Promise<void> {
-		const gracePeriodMs = config.shutdown?.gracePeriodMs ?? 30_000;
+	async #drain(config: RegistryConfig): Promise<void> {
+		const modeAPromise = this.#runtimeServeConfiguredPromise;
+		const modeBPromise = this.#runtimeServerlessPromise;
+
+		const gracePeriodMs =
+			config.shutdown?.gracePeriodMs ??
+			(await this.#actorStopThresholdMs(modeAPromise ?? modeBPromise)) ??
+			30 * 60 * 1000;
 		// Race the entire drain sequence (both modes + serve promise) against
-		// a single grace ceiling. Without this, each mode's Rust-side drain
-		// (20s) could stack sequentially and blow past gracePeriodMs before
-		// we re-raise the signal.
+		// a single grace ceiling. By default, this uses the engine-provided
+		// actor stop threshold, matching Pegboard's hard cutoff for actors.
 		const drain = async () => {
 			// Shut down every live `CoreRegistry` we know about. Mode A
 			// (`start()`) and Mode B (`handler()`) each build a separate
-			// runtime registry, so one signal handler fans out to both to
-			// honor the spec invariant "single shutdown tears down both modes".
-			const registries: Promise<void>[] = [
-				(async () => {
-					try {
-						const { runtime, registry } =
-							await configuredRegistryPromise;
-						await runtime.shutdownRegistry(registry);
-					} catch (error) {
-						logger().warn(
-							{ error },
-							"runtime registry shutdown errored (mode A)",
-						);
-					}
-				})(),
-			];
-			const runtimeServerlessPromise = this.#runtimeServerlessPromise;
-			if (runtimeServerlessPromise !== undefined) {
+			// runtime registry, so one drain fans out to both to honor the
+			// spec invariant "single shutdown tears down both modes".
+			const registries: Promise<void>[] = [];
+			if (modeAPromise !== undefined) {
+				registries.push(
+					(async () => {
+						try {
+							const { runtime, registry } = await modeAPromise;
+							await runtime.shutdownRegistry(registry);
+						} catch (err) {
+							logger().warn(
+								{ err },
+								"runtime registry shutdown errored (mode A)",
+							);
+						}
+					})(),
+				);
+			}
+			if (modeBPromise !== undefined) {
 				registries.push(
 					(async () => {
 						try {
-							const { runtime, registry } =
-								await runtimeServerlessPromise;
+							const { runtime, registry } = await modeBPromise;
 							await runtime.shutdownRegistry(registry);
 						} catch (err) {
 							logger().warn(
@@ -567,8 +622,32 @@ export class Registry<A extends RegistryActors> {
 				setTimeout(resolve, gracePeriodMs).unref?.(),
 			),
 		]);
-		this.#removeSignalHandlers();
-		process.kill(process.pid, signal);
+	}
+
+	async #actorStopThresholdMs(
+		configuredRegistryPromise:
+			| ReturnType<typeof buildConfiguredRegistry>
+			| undefined,
+	): Promise<number | undefined> {
+		if (configuredRegistryPromise === undefined) return undefined;
+		try {
+			const { runtime, registry } = await configuredRegistryPromise;
+			const thresholdMs =
+				await runtime.registryActorStopThresholdMs?.(registry);
+			if (
+				thresholdMs !== undefined &&
+				Number.isFinite(thresholdMs) &&
+				thresholdMs > 0
+			) {
+				return thresholdMs;
+			}
+		} catch (err) {
+			logger().warn(
+				{ err },
+				"failed to read actor stop threshold for shutdown grace",
+			);
+		}
+		return undefined;
 	}
 
 	#removeSignalHandlers(): void {
@@ -621,7 +700,9 @@ export class Registry<A extends RegistryActors> {
 
 		if (config.endpoint) {
 			const endpointType =
-				config.endpoint === ENGINE_ENDPOINT ? "local native" : "remote";
+				config.startEngine || isLocalEngineEndpoint(config.endpoint)
+					? "local native"
+					: "remote";
 			logLine("Endpoint", `${config.endpoint} (${endpointType})`);
 		}
 

package/src/registry/napi-runtime.ts

@@ -452,6 +452,12 @@ export class NapiCoreRuntime implements CoreRuntime {
 		return await asNativeActorContext(ctx).waitForTrackedShutdownWork();
 	}
 
+	async actorWaitForTrackedShutdownWorkUnbounded(
+		ctx: ActorContextHandle,
+	): Promise<void> {
+		await asNativeActorContext(ctx).waitForTrackedShutdownWorkUnbounded();
+	}
+
 	actorKeepAwake(ctx: ActorContextHandle, promise: Promise<unknown>): void {
 		asNativeActorContext(ctx).keepAwake(promise);
 	}

package/src/registry/native.ts

@@ -49,10 +49,7 @@ import type {
 	RuntimeKind,
 	SqliteBackend,
 } from "@/registry/config";
-import {
-	decodeCborCompat,
-	encodeCborCompat,
-} from "@/serde";
+import { decodeCborCompat, encodeCborCompat } from "@/serde";
 import { getEnvUniversal, VERSION } from "@/utils";
 import { logger } from "./log";
 import { loadNapiRuntime } from "./napi-runtime";
@@ -74,12 +71,15 @@ import type {
 	RuntimeActorConfig,
 	RuntimeBytes,
 	RuntimeHttpResponse,
+	RuntimeInspectorTabEntry,
 	RuntimeQueueMessage,
 	RuntimeServeConfig,
 	RuntimeStateDeltaPayload,
 	WebSocketHandle,
 } from "./runtime";
 import { loadWasmRuntime } from "./wasm-runtime";
+import nodeFs from "node:fs";
+import nodePath from "node:path";
 import { createWriteThroughProxy } from "./write-through-proxy";
 
 const textEncoder = new TextEncoder();
@@ -424,8 +424,33 @@ function clearNativeRuntimeState(
 async function cleanupNativeSleepRuntimeState(
 	runtime: CoreRuntime,
 	ctx: ActorContextHandle,
+	afterTrackedWorkDrained?: () => Promise<void>,
 ): Promise<void> {
-	await runtime.actorWaitForTrackedShutdownWork(ctx);
+	// The bounded wait gives shutdown work one grace-period chance to finish.
+	// Drained means all tracked shutdown work completed before the deadline, so
+	// we can save final state and clear runtime state immediately. If it did not
+	// drain, close database handles now, then defer the final save and clear until
+	// the tracked work finishes without a deadline.
+	const drained = await runtime.actorWaitForTrackedShutdownWork(ctx);
+	if (!drained) {
+		await closeNativeDatabaseClient(runtime, ctx);
+		await closeNativeSqlDatabase(runtime, ctx);
+		void runtime
+			.actorWaitForTrackedShutdownWorkUnbounded(ctx)
+			.then(async () => {
+				await afterTrackedWorkDrained?.();
+				clearNativeRuntimeState(runtime, ctx);
+			})
+			.catch((error) => {
+				logger().warn({
+					msg: "deferred native sleep cleanup failed",
+					error: stringifyError(error),
+				});
+			});
+		return;
+	}
+
+	await afterTrackedWorkDrained?.();
 	await closeNativeDatabaseClient(runtime, ctx);
 	await closeNativeSqlDatabase(runtime, ctx);
 	clearNativeRuntimeState(runtime, ctx);
@@ -605,6 +630,14 @@ function encodeValue(value: unknown): RuntimeBytes {
 	return encodeCborCompat(value as JsonCompatValue);
 }
 
+function normalizeArgs(value: unknown): unknown[] {
+	return Array.isArray(value)
+		? value
+		: value === undefined || value === null
+			? []
+			: [value];
+}
+
 function unwrapTsfnPayload<T>(error: unknown, payload: T): T {
 	if (error !== null && error !== undefined) {
 		throw error;
@@ -1071,11 +1104,7 @@ function wrapNativeCallback<Args extends Array<unknown>, Result>(
 
 function decodeArgs(value?: RuntimeBytes | null): unknown[] {
 	const decoded = decodeValue<unknown>(value);
-	return Array.isArray(decoded)
-		? decoded
-		: decoded === undefined
-			? []
-			: [decoded];
+	return normalizeArgs(decoded);
 }
 
 function buildRequest(init: {
@@ -2876,6 +2905,7 @@ export class ActorContextHandleAdapter {
 	}
 
 	sleep(): void {
+		this.#flushStateChange();
 		callNativeSync(() => this.#runtime.actorSleep(this.#ctx));
 	}
 
@@ -3300,9 +3330,87 @@ function buildActorConfig(
 		actions: Object.keys((config.actions ?? {}) as Record<string, unknown>)
 			.sort()
 			.map((name) => ({ name })),
+		inspectorTabs: buildInspectorTabs(config.inspector),
 	};
 }
 
+function buildInspectorTabs(
+	inspector: unknown,
+): Array<RuntimeInspectorTabEntry> | undefined {
+	if (!inspector || typeof inspector !== "object") return undefined;
+	const tabs = (inspector as { tabs?: unknown }).tabs;
+	if (!Array.isArray(tabs) || tabs.length === 0) return undefined;
+	return tabs.map((raw) => {
+		const entry = raw as {
+			id: string;
+			label?: string;
+			source?: string;
+			icon?: string;
+			hidden?: boolean;
+		};
+		if (entry.hidden === true) {
+			return { id: entry.id, hidden: true };
+		}
+		// Resolve the author's source path against the current working
+		// directory so the Rust runtime gets an absolute path. The author
+		// runs the actor process from their project root by convention.
+		const resolved =
+			entry.source !== undefined
+				? nodePath.resolve(entry.source)
+				: undefined;
+		if (resolved !== undefined) {
+			validateInspectorTabSource(entry.id, resolved);
+		}
+		return {
+			id: entry.id,
+			label: entry.label,
+			icon: entry.icon,
+			source: resolved,
+		};
+	});
+}
+
+function validateInspectorTabSource(tabId: string, resolved: string): void {
+	// Catch obviously dangerous misconfigurations at registry construction
+	// rather than silently exposing the wrong subtree over the unauthenticated
+	// `/inspector/custom-tabs/<id>/*` route. Fail loudly so misconfigured
+	// actors never start.
+	if (resolved === nodePath.parse(resolved).root) {
+		throw new Error(
+			`inspector.tabs[id="${tabId}"].source resolves to the filesystem root (${resolved}). ` +
+				"Point it at the tab's own static-asset directory instead.",
+		);
+	}
+	let stat: import("node:fs").Stats;
+	try {
+		stat = nodeFs.statSync(resolved);
+	} catch (err) {
+		const code = (err as NodeJS.ErrnoException)?.code;
+		if (code === "ENOENT") {
+			throw new Error(
+				`inspector.tabs[id="${tabId}"].source (${resolved}) does not exist.`,
+			);
+		}
+		if (code === "EACCES") {
+			throw new Error(
+				`inspector.tabs[id="${tabId}"].source (${resolved}) is not readable (EACCES).`,
+			);
+		}
+		throw new Error(
+			`inspector.tabs[id="${tabId}"].source (${resolved}) could not be stat'd: ${
+				(err as Error)?.message ?? err
+			}`,
+		);
+	}
+	if (!stat.isDirectory()) {
+		throw new Error(
+			`inspector.tabs[id="${tabId}"].source (${resolved}) must be a directory, got ${
+				stat.isFile() ? "file" : "non-directory"
+			}.`,
+		);
+	}
+}
+
 export function buildNativeFactory(
 	runtime: CoreRuntime,
 	registryConfig: RegistryConfig,
@@ -3733,14 +3841,38 @@ export function buildNativeFactory(
 						404,
 					);
 				}
-				const body = (await jsRequest.json()) as { args?: unknown[] };
+				const body = (await jsRequest.json()) as {
+					args?: unknown;
+					properties?: unknown;
+				};
+				if (body.args !== undefined && body.properties !== undefined) {
+					return jsonResponse(
+						{ error: "use either args or properties, not both" },
+						{ status: 400 },
+					);
+				}
+				if (
+					body.properties !== undefined &&
+					(body.properties === null ||
+						typeof body.properties !== "object" ||
+						Array.isArray(body.properties))
+				) {
+					return jsonResponse(
+						{ error: "properties must be an object" },
+						{ status: 400 },
+					);
+				}
+				const args =
+					body.properties !== undefined
+						? [body.properties]
+						: normalizeArgs(body.args);
 				try {
 					const output = await action(
 						actorCtx,
 						...validateActionArgs(
 							schemaConfig.actionInputSchemas,
 							actionName,
-							body.args ?? [],
+							args,
 						),
 					);
 					return jsonResponse({ output });
@@ -3939,26 +4071,35 @@ export function buildNativeFactory(
 			async (error: unknown, payload: { ctx: ActorContextHandle }) => {
 				const { ctx } = unwrapTsfnPayload(error, payload);
 				const actorCtx = makeActorCtx(ctx);
+				// TODO: Move this save hook into cleanupNativeSleepRuntimeState
+				// so immediate and deferred sleep cleanup share one save-state
+				// path instead of passing a callback through cleanup.
+				const saveActorState = async () => {
+					if (runtime.kind === "wasm") {
+						// Wasm cannot use the native context save helper here because
+						// the runtime owns the serialized state handoff.
+						await runtime.actorSaveState(
+							ctx,
+							actorCtx.serializeForTick("save"),
+						);
+					} else {
+						await actorCtx.saveState({
+							immediate: true,
+						});
+					}
+				};
 				try {
 					if (onSleep) {
-						try {
-							await onSleep(actorCtx);
-						} finally {
-							if (runtime.kind === "wasm") {
-								// Wasm cannot use the native context save helper here because
-								// the runtime owns the serialized state handoff.
-								await runtime.actorSaveState(
-									ctx,
-									actorCtx.serializeForTick("save"),
-								);
-							} else {
-								await actorCtx.saveState({ immediate: true });
-							}
-						}
+						await onSleep(actorCtx);
 					}
+					await saveActorState();
 				} finally {
 					try {
-						await cleanupNativeSleepRuntimeState(runtime, ctx);
+						await cleanupNativeSleepRuntimeState(
+							runtime,
+							ctx,
+							saveActorState,
+						);
 					} finally {
 						await actorCtx.dispose();
 					}
@@ -4618,6 +4759,8 @@ export async function buildServeConfig(
 	if (config.startEngine) {
 		const { getEnginePath } = await loadEngineCli();
 		serveConfig.engineBinaryPath = getEnginePath();
+		serveConfig.engineHost = config.engineHost;
+		serveConfig.enginePort = config.enginePort;
 	}
 	if (config.test?.enabled) {
 		serveConfig.inspectorTestToken =

package/src/registry/process-metrics.ts

@@ -22,7 +22,10 @@ import * as napi from "@rivetkit/rivetkit-napi";
 type OptionalProcessMetricsNapi = typeof napi & {
 	jsObserveGcDuration?: (kind: string, durationSeconds: number) => void;
 	jsSetEventloopHeartbeatTsMs?: (timestampMs: number) => void;
-	jsSetEventloopLagQuantile?: (quantile: string, valueSeconds: number) => void;
+	jsSetEventloopLagQuantile?: (
+		quantile: string,
+		valueSeconds: number,
+	) => void;
 	jsSetEventloopUtilization?: (utilization: number) => void;
 	jsAddProcessCpuSeconds?: (mode: string, valueSeconds: number) => void;
 	jsSetProcessResidentMemoryBytes?: (bytes: number) => void;
@@ -191,7 +194,10 @@ function collectAndPush(): void {
 		state.lastEventLoopUtilization,
 	);
 	state.lastEventLoopUtilization = nextElu;
-	callIfFn(processMetricsNapi.jsSetEventloopUtilization, eluDelta.utilization);
+	callIfFn(
+		processMetricsNapi.jsSetEventloopUtilization,
+		eluDelta.utilization,
+	);
 
 	// CPU usage delta. `process.cpuUsage()` returns microseconds.
 	const nextCpu = process.cpuUsage();
@@ -230,9 +236,15 @@ function collectAndPush(): void {
 		_getActiveRequests?: () => unknown[];
 	};
 	if (typeof proc._getActiveHandles === "function") {
-		callIfFn(processMetricsNapi.jsSetActiveHandles, proc._getActiveHandles().length);
+		callIfFn(
+			processMetricsNapi.jsSetActiveHandles,
+			proc._getActiveHandles().length,
+		);
 	}
 	if (typeof proc._getActiveRequests === "function") {
-		callIfFn(processMetricsNapi.jsSetActiveRequests, proc._getActiveRequests().length);
+		callIfFn(
+			processMetricsNapi.jsSetActiveRequests,
+			proc._getActiveRequests().length,
+		);
 	}
 }