rivetkit 0.0.0-skip-ready-wait-flatten-client-options.1cae221 → 0.0.0-sqlite-uds.4e59a38

package/src/registry/index.ts

@@ -1,6 +1,12 @@
-import { ENGINE_ENDPOINT } from "@/common/engine";
+import { isLocalEngineEndpoint } from "@/common/engine";
+import { configureBaseLogger } from "@/common/log";
 import { configureServerlessPool } from "@/serverless/configure";
 import { VERSION } from "@/utils";
+import {
+	getRivetkitPublicDir,
+	getRivetkitRuntimeMode,
+	parsePortEnv,
+} from "@/utils/env-vars";
 import {
 	type RegistryActors,
 	type RegistryConfig,
@@ -13,6 +19,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
@@ -22,13 +44,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;
@@ -47,20 +81,39 @@ 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;
+		if (config.logging?.baseLogger) {
+			configureBaseLogger(config.logging.baseLogger);
+		}
+		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 {
+	/**
+	 * Fires `configureServerlessPool` once per process when the registry
+	 * config opts into it. Cached on the instance so repeated calls (from
+	 * `handler()` and `listen()`) only run the upsert once. The retry loop
+	 * inside `configureServerlessPool` tolerates the engine still warming up.
+	 */
+	#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(() => {});
 		}
 
@@ -82,7 +135,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 } =
@@ -96,12 +150,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",
@@ -243,7 +298,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",
@@ -279,21 +334,152 @@ export class Registry<A extends RegistryActors> {
 		};
 	}
 
-	public async diagnostics(): Promise<RegistryDiagnostics> {
+	/**
+	 * Bind an HTTP listener provided by the native (Rust) runtime and serve
+	 * the registry's serverless endpoints over it. Resolves only after the
+	 * registry is shut down (SIGINT/SIGTERM or `nativeRegistry.shutdown()`).
+	 *
+	 * @param opts.port      Port to listen on. Defaults to `process.env.RIVET_PORT`
+	 *                       if set, otherwise 3000.
+	 * @param opts.host      Address to bind. Defaults to `0.0.0.0`.
+	 * @param opts.publicDir If set, serves static files from this directory
+	 *                       as a fallback below the framework routes.
+	 *
+	 * @example
+	 * ```ts
+	 * await registry.listen();
+	 * await registry.listen({ port: 8080, publicDir: "./public" });
+	 * ```
+	 */
+	public async listen(
+		opts: { port?: number; host?: string; publicDir?: string } = {},
+	): Promise<void> {
+		const port = opts.port ?? parsePortEnv(process.env.RIVET_PORT) ?? 3000;
+		const publicDir = opts.publicDir ?? getRivetkitPublicDir();
+		const config = this.parseConfig();
+
+		// Cache on both promise fields so the shutdown drain sees Mode A and B.
+		const configuredRegistryPromise = buildConfiguredRegistry(config);
+		this.#runtimeServeConfiguredPromise = configuredRegistryPromise;
+		this.#runtimeServerlessPromise = configuredRegistryPromise;
+		this.#installSignalHandlers(config);
+
+		this.#printWelcome(config, "serverless", {
+			port,
+			host: opts.host,
+			publicDir,
+		});
+
+		// Background fire; the retry loop tolerates engine warm-up.
+		this.#ensureServerlessPoolConfigured(config);
+
+		const { runtime, registry, serveConfig } =
+			await configuredRegistryPromise;
+		await runtime.serveListener(
+			registry,
+			{ port, host: opts.host, publicDir },
+			serveConfig,
+		);
+	}
+
+	/**
+	 * 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]!;
 	}
 
 	/**
@@ -301,35 +487,33 @@ 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 }) => {
 					await runtime.serveRegistry(registry, serveConfig);
 				})
-				.catch((err) => {
+				.catch((error) => {
 					// Always-attached catch so the stored promise never leaves a
 					// rejection unhandled. Downstream awaits (e.g. #runShutdown's
 					// Promise.race) attach their own catches and still observe
 					// resolution via the race.
-					logger().warn({ err }, "runtime registry serve errored");
+					logger().warn({ error }, "runtime registry serve errored");
 				});
 			// Install signal handlers once an envoy lifecycle has begun. Only
 			// Mode A ever reaches here. Mode B (handler(request)) intentionally
 			// 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
@@ -344,12 +528,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);
 		};
@@ -357,68 +536,101 @@ 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((err) => {
+		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 (err) {
-						logger().warn(
-							{ err },
-							"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 runtimeServerlessPromise;
+							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 modeBPromise;
+							await runtime.shutdownRegistry(registry);
+						} catch (err) {
+							logger().warn(
+								{ error: err },
 								"runtime registry shutdown errored (mode B)",
 							);
 						}
@@ -441,8 +653,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 {
@@ -461,6 +697,10 @@ export class Registry<A extends RegistryActors> {
 	/**
 	 * Starts the actor envoy for standalone server deployments.
 	 *
+	 * Set `RIVETKIT_RUNTIME_MODE=serverless` to instead bind an HTTP listener
+	 * via `listen()` (Mode B). Mode A (envoy) and Mode B (listener) are
+	 * mutually exclusive per registry instance.
+	 *
 	 * @example
 	 * ```ts
 	 * const registry = setup({ use: { counter } });
@@ -468,6 +708,21 @@ export class Registry<A extends RegistryActors> {
 	 * ```
 	 */
 	public start() {
+		if (getRivetkitRuntimeMode() === "serverless") {
+			// start() defaults publicDir to "/public" unless overridden by env.
+			const publicDir = getRivetkitPublicDir() ?? "/public";
+			// Detached listener; bind failures are fatal so exit hard.
+			this.listen({ publicDir }).catch((error) => {
+				logger().error({ error }, "auto-listen failed; exiting");
+				if (
+					typeof process !== "undefined" &&
+					typeof process.exit === "function"
+				) {
+					process.exit(1);
+				}
+			});
+			return;
+		}
 		const config = this.parseConfig();
 		this.#startEnvoy(config, true);
 	}
@@ -475,6 +730,7 @@ export class Registry<A extends RegistryActors> {
 	#printWelcome(
 		config: RegistryConfig,
 		kind: "serverless" | "serverful",
+		listener?: { port: number; host?: string; publicDir?: string },
 	): void {
 		if (config.noWelcome || this.#welcomePrinted) return;
 		this.#welcomePrinted = true;
@@ -495,7 +751,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})`);
 		}
 
@@ -504,6 +762,15 @@ export class Registry<A extends RegistryActors> {
 		}
 
 		logLine("Actors", Object.keys(config.use).length.toString());
+
+		if (listener) {
+			const host = listener.host ?? "0.0.0.0";
+			logLine("Listening", `http://${host}:${listener.port}`);
+			if (listener.publicDir) {
+				logLine("Public Dir", listener.publicDir);
+			}
+		}
+
 		console.log();
 	}
 }
@@ -516,7 +783,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 =
@@ -524,6 +794,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> {