rivetkit 2.3.0-rc.8 → 2.3.0

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";
@@ -40,13 +40,25 @@ export interface ServerlessHandler {
 	fetch: FetchHandler;
 }
 
-export interface RegistryDiagnostics {
-	mode: string;
-	envoyActiveActorCount?: number | null;
+export interface RegistryRoutes {
+	health(): Promise<Response>;
+	metadata(): Promise<Response>;
+	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> {
 		return this.#config;
@@ -65,20 +77,30 @@ 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(),
+			prometheusMetrics: (request?: Request) =>
+				this.#prometheusMetricsRoute(request),
+		};
 	}
 
-	#ensureServerlessPoolConfigured(config: RegistryConfig): Promise<void> | undefined {
+	#ensureServerlessPoolConfigured(
+		config: RegistryConfig,
+	): Promise<void> | undefined {
 		if (!config.configurePool) return undefined;
 
 		if (!this.#configureServerlessPoolPromise) {
-			this.#configureServerlessPoolPromise = configureServerlessPool(config).catch(
-				(error) => {
-					this.#configureServerlessPoolPromise = undefined;
-					throw error;
-				},
-			);
+			this.#configureServerlessPoolPromise = configureServerlessPool(
+				config,
+			).catch((error) => {
+				this.#configureServerlessPoolPromise = undefined;
+				throw error;
+			});
 			this.#configureServerlessPoolPromise.catch(() => {});
 		}
 
@@ -100,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 } =
@@ -114,12 +137,13 @@ export class Registry<A extends RegistryActors> {
 			serveConfig.serverlessBasePath ?? "/api/rivet",
 		);
 		const isEngineMetadataRequest =
-			request.headers.get("user-agent")?.startsWith("RivetEngine/") ?? false;
+			request.headers.get("user-agent")?.startsWith("RivetEngine/") ??
+			false;
 
 		if (isStartRequest) {
 			try {
 				await this.#ensureServerlessPoolConfigured(config);
-			} catch (error) {
+			} catch (_error) {
 				return new Response(
 					JSON.stringify({
 						group: "guard",
@@ -261,7 +285,7 @@ export class Registry<A extends RegistryActors> {
 		if (isMetadataRequest && !isEngineMetadataRequest) {
 			try {
 				await this.#ensureServerlessPoolConfigured(config);
-			} catch (error) {
+			} catch (_error) {
 				return new Response(
 					JSON.stringify({
 						group: "guard",
@@ -327,21 +351,104 @@ export class Registry<A extends RegistryActors> {
 		await crossPlatformServe(config, port, app, runtime);
 	}
 
-	public async diagnostics(): Promise<RegistryDiagnostics> {
+	/**
+	 * Returns a health response suitable for mounting in a user-owned router.
+	 */
+	async #healthRoute(): Promise<Response> {
+		const configured = await this.#activeConfiguredRegistry();
+		if (!configured) {
+			return jsonRouteResponse(503, {
+				status: "not_started",
+				runtime: "rivetkit",
+				version: VERSION,
+			});
+		}
+
+		const { runtime, registry } = configured;
+		if (!runtime.registryHealth) {
+			return jsonRouteResponse(501, {
+				status: "unsupported",
+				runtime: "rivetkit",
+				version: VERSION,
+			});
+		}
+
+		const response = await runtime.registryHealth(registry);
+		return new Response(new Uint8Array(response.body), {
+			status: response.status,
+			headers: response.headers,
+		});
+	}
+
+	/**
+	 * Returns serverless metadata suitable for mounting in a user-owned router.
+	 */
+	async #metadataRoute(): Promise<Response> {
+		const configured = await this.#activeConfiguredRegistry();
+		if (!configured) {
+			return new Response("registry not started\n", {
+				status: 503,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const { runtime, registry } = configured;
+		if (!runtime.registryMetadata) {
+			return new Response("metadata is not supported by this runtime\n", {
+				status: 501,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const response = await runtime.registryMetadata(registry);
+		return new Response(new Uint8Array(response.body), {
+			status: response.status,
+			headers: response.headers,
+		});
+	}
+
+	/**
+	 * Returns a Prometheus metrics response suitable for mounting in a user-owned router.
+	 */
+	async #prometheusMetricsRoute(_request?: Request): Promise<Response> {
+		const configured = await this.#activeConfiguredRegistry();
+		if (!configured) {
+			return new Response("registry not started\n", {
+				status: 503,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const { runtime, registry } = configured;
+		if (!runtime.registryMetrics) {
+			return new Response("metrics are not supported by this runtime\n", {
+				status: 501,
+				headers: { "content-type": "text/plain; charset=utf-8" },
+			});
+		}
+
+		const response = await runtime.registryMetrics(registry);
+		return new Response(new Uint8Array(response.body), {
+			status: response.status,
+			headers: response.headers,
+		});
+	}
+
+	async #activeConfiguredRegistry(): Promise<
+		Awaited<ReturnType<typeof buildConfiguredRegistry>> | undefined
+	> {
 		const candidates = [
 			this.#runtimeServerlessPromise,
 			this.#runtimeServeConfiguredPromise,
-		].filter((candidate): candidate is ReturnType<typeof buildConfiguredRegistry> =>
-			candidate !== undefined
+		].filter(
+			(
+				candidate,
+			): candidate is ReturnType<typeof buildConfiguredRegistry> =>
+				candidate !== undefined,
 		);
 
-		for (const candidate of candidates) {
-			const { runtime, registry } = await candidate;
-			const diagnostics = await runtime.registryDiagnostics?.(registry);
-			if (diagnostics) return diagnostics;
-		}
-
-		return { mode: "not_started", envoyActiveActorCount: null };
+		if (candidates.length === 0) return undefined;
+		return await candidates[0]!;
 	}
 
 	/**
@@ -349,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 }) => {
@@ -367,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
@@ -392,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);
 		};
@@ -405,37 +505,67 @@ 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.
-			// PID 1 must exit directly because re-raised default signals can be
+			// 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();
 			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> {
+	async #drain(config: RegistryConfig): Promise<void> {
+		const modeAPromise = this.#runtimeServeConfiguredPromise;
+		const modeBPromise = this.#runtimeServerlessPromise;
+
 		const gracePeriodMs =
 			config.shutdown?.gracePeriodMs ??
-			(await this.#actorStopThresholdMs(configuredRegistryPromise)) ??
+			(await this.#actorStopThresholdMs(modeAPromise ?? modeBPromise)) ??
 			30 * 60 * 1000;
 		// Race the entire drain sequence (both modes + serve promise) against
 		// a single grace ceiling. By default, this uses the engine-provided
@@ -443,29 +573,29 @@ export class Registry<A extends RegistryActors> {
 		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(
@@ -492,13 +622,14 @@ export class Registry<A extends RegistryActors> {
 				setTimeout(resolve, gracePeriodMs).unref?.(),
 			),
 		]);
-		this.#removeSignalHandlers();
-		finishShutdownSignal(signal);
 	}
 
 	async #actorStopThresholdMs(
-		configuredRegistryPromise: ReturnType<typeof buildConfiguredRegistry>,
+		configuredRegistryPromise:
+			| ReturnType<typeof buildConfiguredRegistry>
+			| undefined,
 	): Promise<number | undefined> {
+		if (configuredRegistryPromise === undefined) return undefined;
 		try {
 			const { runtime, registry } = await configuredRegistryPromise;
 			const thresholdMs =
@@ -569,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})`);
 		}
 
@@ -590,7 +723,10 @@ function isServerlessStartRequest(request: Request, basePath: string): boolean {
 	return parsed.pathname === `${normalizedBase}/start`;
 }
 
-function isServerlessMetadataRequest(request: Request, basePath: string): boolean {
+function isServerlessMetadataRequest(
+	request: Request,
+	basePath: string,
+): boolean {
 	if (request.method !== "GET") return false;
 	const parsed = new URL(request.url);
 	const normalizedBase =
@@ -598,6 +734,13 @@ function isServerlessMetadataRequest(request: Request, basePath: string): boolea
 	return parsed.pathname === `${normalizedBase}/metadata`;
 }
 
+function jsonRouteResponse(status: number, body: unknown): Response {
+	return new Response(JSON.stringify(body), {
+		status,
+		headers: { "content-type": "application/json" },
+	});
+}
+
 export function setup<A extends RegistryActors>(
 	input: RegistryConfigInput<A>,
 ): Registry<A> {

package/src/registry/napi-runtime.ts

@@ -23,8 +23,8 @@ import type {
 	RuntimeQueueNextBatchOptions,
 	RuntimeQueueTryNextBatchOptions,
 	RuntimeQueueWaitOptions,
+	RuntimeRegistryRouteResponse,
 	RuntimeRequestSaveOpts,
-	RuntimeRegistryDiagnostics,
 	RuntimeServeConfig,
 	RuntimeServerlessRequest,
 	RuntimeServerlessResponseHead,
@@ -202,16 +202,6 @@ export class NapiCoreRuntime implements CoreRuntime {
 		await asNativeRegistry(registry).shutdown();
 	}
 
-	async registryDiagnostics(
-		registry: RegistryHandle,
-	): Promise<RuntimeRegistryDiagnostics> {
-		const diagnostics = await asNativeRegistry(registry).diagnostics();
-		return {
-			mode: diagnostics.mode,
-			envoyActiveActorCount: diagnostics.envoyActiveActorCount,
-		};
-	}
-
 	async registryActorStopThresholdMs(
 		registry: RegistryHandle,
 	): Promise<number | undefined> {
@@ -221,6 +211,39 @@ export class NapiCoreRuntime implements CoreRuntime {
 		);
 	}
 
+	async registryHealth(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse> {
+		const response = await asNativeRegistry(registry).health();
+		return {
+			status: response.status,
+			headers: response.headers,
+			body: response.body,
+		};
+	}
+
+	async registryMetadata(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse> {
+		const response = asNativeRegistry(registry).metadata();
+		return {
+			status: response.status,
+			headers: response.headers,
+			body: response.body,
+		};
+	}
+
+	async registryMetrics(
+		registry: RegistryHandle,
+	): Promise<RuntimeRegistryRouteResponse> {
+		const response = asNativeRegistry(registry).metrics();
+		return {
+			status: response.status,
+			headers: response.headers,
+			body: response.body,
+		};
+	}
+
 	async handleServerlessRequest(
 		registry: RegistryHandle,
 		req: RuntimeServerlessRequest,
@@ -429,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);
 	}
@@ -654,10 +683,15 @@ export class NapiCoreRuntime implements CoreRuntime {
 		ctx: ActorContextHandle,
 		names: string[],
 		options?: RuntimeQueueWaitOptions | undefined | null,
+		signal?: CancellationTokenHandle | undefined | null,
 	): Promise<void> {
 		await asNativeActorContext(ctx)
 			.queue()
-			.waitForNamesAvailable(names, options);
+			.waitForNamesAvailable(
+				names,
+				options,
+				signal ? asNativeCancellationToken(signal) : signal,
+			);
 	}
 
 	async actorQueueEnqueueAndWait(

package/src/registry/native-validation.ts

@@ -1,12 +1,12 @@
 import { RivetError } from "@/actor/errors";
 import {
+	type EventSchemaConfig,
 	isEventSchemaDefinition,
 	isQueueSchemaDefinition,
 	isStandardSchema,
-	validateSchemaSync,
-	type EventSchemaConfig,
 	type PrimitiveSchema,
 	type QueueSchemaConfig,
+	validateSchemaSync,
 } from "@/actor/schema";
 
 const CONN_PARAMS_KEY = "__conn_params__";
@@ -126,19 +126,17 @@ export function validateQueueComplete(
 		response,
 	);
 	if (!result.success) {
-		throw validationError(`queue \`${name}\` completion response`, result.issues);
+		throw validationError(
+			`queue \`${name}\` completion response`,
+			result.issues,
+		);
 	}
 	return result.data;
 }
 
 function validationError(target: string, issues: unknown[]): RivetError {
-	return new RivetError(
-		"actor",
-		"validation_error",
-		`Invalid ${target}`,
-		{
-			public: true,
-			metadata: { issues },
-		},
-	);
+	return new RivetError("actor", "validation_error", `Invalid ${target}`, {
+		public: true,
+		metadata: { issues },
+	});
 }