rivetkit 2.3.0-rc.9 → 2.3.0

package/src/registry/index.ts

@@ -1,6 +1,8 @@
-import { ENGINE_ENDPOINT } from "@/common/engine";
+import { Hono } from "hono";
+import { isLocalEngineEndpoint } from "@/common/engine";
 import { configureServerlessPool } from "@/serverless/configure";
-import { VERSION } from "@/utils";
+import { detectRuntime, VERSION } from "@/utils";
+import { crossPlatformServe, loadRuntimeServeStatic } from "@/utils/serve";
 import {
 	type RegistryActors,
 	type RegistryConfig,
@@ -44,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> {
@@ -65,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(),
@@ -75,16 +89,18 @@ export class Registry<A extends RegistryActors> {
 		};
 	}
 
-	#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(() => {});
 		}
 
@@ -106,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 } =
@@ -120,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",
@@ -267,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",
@@ -303,6 +321,36 @@ export class Registry<A extends RegistryActors> {
 		};
 	}
 
+	/**
+	 * Starts an HTTP server that dispatches every request through the
+	 * serverless handler. Uses `crossPlatformServe` to pick the right
+	 * runtime (Node, Bun, Deno).
+	 *
+	 * @param opts.port      Port to listen on. Defaults to 3000.
+	 * @param opts.publicDir If set, serves static files from this directory
+	 *                       before falling through to the registry handler.
+	 *
+	 * @example
+	 * ```ts
+	 * await registry.listen();
+	 * await registry.listen({ port: 8080, publicDir: "./public" });
+	 * ```
+	 */
+	public async listen(
+		opts: { port?: number; publicDir?: string } = {},
+	): Promise<void> {
+		const port = opts.port ?? 3000;
+		const config = this.parseConfig();
+		const runtime = detectRuntime();
+		const app = new Hono();
+		if (opts.publicDir) {
+			const serveStatic = await loadRuntimeServeStatic(runtime);
+			app.use("*", serveStatic({ root: opts.publicDir }));
+		}
+		app.all("*", (c) => this.handler(c.req.raw));
+		await crossPlatformServe(config, port, app, runtime);
+	}
+
 	/**
 	 * Returns a health response suitable for mounting in a user-owned router.
 	 */
@@ -392,8 +440,11 @@ export class Registry<A extends RegistryActors> {
 		const candidates = [
 			this.#runtimeServerlessPromise,
 			this.#runtimeServeConfiguredPromise,
-		].filter((candidate): candidate is ReturnType<typeof buildConfiguredRegistry> =>
-			candidate !== undefined
+		].filter(
+			(
+				candidate,
+			): candidate is ReturnType<typeof buildConfiguredRegistry> =>
+				candidate !== undefined,
 		);
 
 		if (candidates.length === 0) return undefined;
@@ -405,35 +456,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
@@ -448,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);
 		};
@@ -461,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((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> {
+	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
@@ -499,33 +573,33 @@ 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 (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)",
 							);
 						}
@@ -548,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 =
@@ -625,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})`);
 		}
 
@@ -646,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 =

package/src/registry/napi-runtime.ts

@@ -23,8 +23,8 @@ import type {
 	RuntimeQueueNextBatchOptions,
 	RuntimeQueueTryNextBatchOptions,
 	RuntimeQueueWaitOptions,
-	RuntimeRequestSaveOpts,
 	RuntimeRegistryRouteResponse,
+	RuntimeRequestSaveOpts,
 	RuntimeServeConfig,
 	RuntimeServerlessRequest,
 	RuntimeServerlessResponseHead,
@@ -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);
 	}
@@ -677,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 },
+	});
 }