rivetkit 2.0.9 → 2.0.11

package/src/manager/router.ts

@@ -11,9 +11,9 @@ import { createMiddleware } from "hono/factory";
 import { streamSSE } from "hono/streaming";
 import invariant from "invariant";
 import { z } from "zod";
-import { ActorNotFound, Unsupported } from "@/actor/errors";
+import { ActorNotFound, InvalidRequest, Unsupported } from "@/actor/errors";
 import { serializeActorKey } from "@/actor/keys";
-import type { Encoding, Transport } from "@/client/mod";
+import type { Client, Encoding, Transport } from "@/client/mod";
 import {
 	WS_PROTOCOL_ACTOR,
 	WS_PROTOCOL_CONN_ID,
@@ -28,7 +28,12 @@ import {
 	handleRouteNotFound,
 	loggerMiddleware,
 } from "@/common/router";
-import { deconstructError, noopNext, stringifyError } from "@/common/utils";
+import {
+	assertUnreachable,
+	deconstructError,
+	noopNext,
+	stringifyError,
+} from "@/common/utils";
 import { type ActorDriver, HEADER_ACTOR_ID } from "@/driver-helpers/mod";
 import type {
 	TestInlineDriverCallRequest,
@@ -50,11 +55,14 @@ import {
 	type Actor as ApiActor,
 } from "@/manager-api/actors";
 import { RivetIdSchema } from "@/manager-api/common";
+import type { AnyClient } from "@/mod";
 import type { RegistryConfig } from "@/registry/config";
-import type { RunConfig } from "@/registry/run-config";
+import type { DriverConfig, RunnerConfig } from "@/registry/run-config";
+import { VERSION } from "@/utils";
 import type { ActorOutput, ManagerDriver } from "./driver";
 import { actorGateway, createTestWebSocketProxy } from "./gateway";
 import { logger } from "./log";
+import { ServerlessStartHeadersSchema } from "./router-schema";
 
 function buildOpenApiResponses<T>(schema: T) {
 	return {
@@ -77,9 +85,10 @@ function buildOpenApiResponses<T>(schema: T) {
 
 export function createManagerRouter(
 	registryConfig: RegistryConfig,
-	runConfig: RunConfig,
+	runConfig: RunnerConfig,
 	managerDriver: ManagerDriver,
-	serverlessActorDriverBuilder: (() => ActorDriver) | undefined,
+	driverConfig: DriverConfig,
+	client: AnyClient,
 ): { router: Hono; openapi: OpenAPIHono } {
 	const router = new OpenAPIHono({ strict: false }).basePath(
 		runConfig.basePath,
@@ -87,10 +96,38 @@ export function createManagerRouter(
 
 	router.use("*", loggerMiddleware(logger()));
 
-	if (serverlessActorDriverBuilder) {
-		addServerlessRoutes(runConfig, serverlessActorDriverBuilder, router);
-	} else {
+	// HACK: Add Sec-WebSocket-Protocol header to fix KIT-339
+	//
+	// Some Deno WebSocket providers do not auto-set the protocol, which
+	// will cause some WebSocket clients to fail
+	router.use(
+		"*",
+		createMiddleware(async (c, next) => {
+			const upgrade = c.req.header("upgrade");
+			const isWebSocket = upgrade?.toLowerCase() === "websocket";
+			const isGet = c.req.method === "GET";
+
+			if (isGet && isWebSocket) {
+				c.header("Sec-WebSocket-Protocol", "rivet");
+			}
+
+			await next();
+		}),
+	);
+
+	if (runConfig.runnerKind === "serverless") {
+		addServerlessRoutes(
+			driverConfig,
+			registryConfig,
+			runConfig,
+			managerDriver,
+			client,
+			router,
+		);
+	} else if (runConfig.runnerKind === "normal") {
 		addManagerRoutes(registryConfig, runConfig, managerDriver, router);
+	} else {
+		assertUnreachable(runConfig.runnerKind);
 	}
 
 	// Error handling
@@ -101,11 +138,11 @@ export function createManagerRouter(
 }
 
 function addServerlessRoutes(
-	runConfig: RunConfig,
-	serverlessActorDriverBuilder: (
-		token: string | undefined,
-		totalSlots: number | undefined,
-	) => ActorDriver,
+	driverConfig: DriverConfig,
+	registryConfig: RegistryConfig,
+	runConfig: RunnerConfig,
+	managerDriver: ManagerDriver,
+	client: AnyClient,
 	router: OpenAPIHono,
 ) {
 	// Apply CORS
@@ -120,28 +157,68 @@ function addServerlessRoutes(
 
 	// Serverless start endpoint
 	router.get("/start", async (c) => {
-		const token = c.req.header("x-rivet-token");
-		let totalSlots: number | undefined = parseInt(
-			c.req.header("x-rivetkit-total-slots") as any,
-		);
-		if (isNaN(totalSlots)) totalSlots = undefined;
+		// Parse headers
+		const parseResult = ServerlessStartHeadersSchema.safeParse({
+			endpoint: c.req.header("x-rivet-endpoint"),
+			token: c.req.header("x-rivet-token") ?? undefined,
+			totalSlots: c.req.header("x-rivet-total-slots"),
+			runnerName: c.req.header("x-rivet-runner-name"),
+			namespace: c.req.header("x-rivet-namespace-id"),
+		});
+		if (!parseResult.success) {
+			throw new InvalidRequest(
+				parseResult.error.issues[0]?.message ??
+					"invalid serverless start headers",
+			);
+		}
+		const { endpoint, token, totalSlots, runnerName, namespace } =
+			parseResult.data;
+
+		logger().debug({
+			msg: "received serverless runner start request",
+			endpoint,
+			totalSlots,
+			runnerName,
+			namespace,
+		});
 
-		const actorDriver = serverlessActorDriverBuilder(token, totalSlots);
+		// Override config
+		//
+		// We can't do a structuredClone here since this holds functions
+		const newRunConfig = Object.assign({}, runConfig);
+		newRunConfig.endpoint = endpoint;
+		newRunConfig.token = token;
+		newRunConfig.totalSlots = totalSlots;
+		newRunConfig.runnerName = runnerName;
+		newRunConfig.namespace = namespace;
+
+		// Create new actor driver with updated config
+		const actorDriver = driverConfig.actor(
+			registryConfig,
+			newRunConfig,
+			managerDriver,
+			client,
+		);
 		invariant(
 			actorDriver.serverlessHandleStart,
 			"missing serverlessHandleStart on ActorDriver",
 		);
+
 		return await actorDriver.serverlessHandleStart(c);
 	});
 
 	router.get("/health", (c) => {
-		return c.text("ok");
+		return c.json({
+			status: "ok",
+			runtime: "rivetkit",
+			version: VERSION,
+		});
 	});
 }
 
 function addManagerRoutes(
 	registryConfig: RegistryConfig,
-	runConfig: RunConfig,
+	runConfig: RunnerConfig,
 	managerDriver: ManagerDriver,
 	router: OpenAPIHono,
 ) {
@@ -244,7 +321,9 @@ function addManagerRoutes(
 			}
 
 			return c.json<ActorsListResponse>({
-				actors: actors.map(createApiActor),
+				actors: actors.map((actor) =>
+					createApiActor(actor, runConfig.runnerName),
+				),
 			});
 		});
 	}
@@ -278,7 +357,7 @@ function addManagerRoutes(
 
 			if (existingActor) {
 				return c.json<ActorsGetOrCreateResponse>({
-					actor: createApiActor(existingActor),
+					actor: createApiActor(existingActor, runConfig.runnerName),
 					created: false,
 				});
 			}
@@ -295,7 +374,7 @@ function addManagerRoutes(
 			});
 
 			return c.json<ActorsGetOrCreateResponse>({
-				actor: createApiActor(newActor),
+				actor: createApiActor(newActor, runConfig.runnerName),
 				created: true,
 			});
 		});
@@ -333,7 +412,7 @@ function addManagerRoutes(
 			});
 
 			// Transform ActorOutput to match ActorSchema
-			const actor = createApiActor(actorOutput);
+			const actor = createApiActor(actorOutput, runConfig.runnerName);
 
 			return c.json<ActorsCreateResponse>({ actor });
 		});
@@ -564,7 +643,12 @@ function addManagerRoutes(
 	}
 
 	router.get("/health", (c) => {
-		return c.text("ok");
+		return c.json({
+			status: "ok",
+			rivetkit: {
+				version: VERSION,
+			},
+		});
 	});
 
 	managerDriver.modifyManagerRouter?.(
@@ -573,13 +657,16 @@ function addManagerRoutes(
 	);
 }
 
-function createApiActor(actor: ActorOutput): ApiActor {
+function createApiActor(
+	actor: ActorOutput,
+	runnerName: string = "default",
+): ApiActor {
 	return {
 		actor_id: actor.actorId,
 		name: actor.name,
 		key: serializeActorKey(actor.key),
 		namespace_id: "default", // Assert default namespace
-		runner_name_selector: "rivetkit", // Assert rivetkit runner
+		runner_name_selector: runnerName,
 		create_ts: Date.now(),
 		connectable_ts: null,
 		destroy_ts: null,

package/src/registry/mod.ts

@@ -1,7 +1,9 @@
+import invariant from "invariant";
 import { type Client, createClientWithDriver } from "@/client/client";
 import { configureBaseLogger, configureDefaultLogger } from "@/common/log";
 import type { ActorDriver } from "@/driver-helpers/mod";
 import { chooseDefaultDriver } from "@/drivers/default";
+import { ENGINE_ENDPOINT, ensureEngineProcess } from "@/engine-process/mod";
 import {
 	configureInspectorAccessToken,
 	getInspectorUrl,
@@ -18,12 +20,16 @@ import {
 import { logger } from "./log";
 import {
 	type DriverConfig,
-	type RunConfig,
-	type RunConfigInput,
-	RunConfigSchema,
+	type RunnerConfig,
+	type RunnerConfigInput,
+	RunnerConfigSchema,
 } from "./run-config";
 import { crossPlatformServe } from "./serve";
 
+export type ServerlessActorDriverBuilder = (
+	updateConfig: (config: RunnerConfig) => void,
+) => ActorDriver;
+
 interface ServerOutput<A extends Registry<any>> {
 	/** Client to communicate with the actors. */
 	client: Client<A>;
@@ -45,8 +51,49 @@ export class Registry<A extends RegistryActors> {
 	/**
 	 * Runs the registry for a server.
 	 */
-	public start(inputConfig?: RunConfigInput): ServerOutput<this> {
-		const config = RunConfigSchema.parse(inputConfig);
+	public start(inputConfig?: RunnerConfigInput): ServerOutput<this> {
+		const config = RunnerConfigSchema.parse(inputConfig);
+
+		// Validate autoConfigureServerless is only used with serverless runner
+		if (config.autoConfigureServerless && config.runnerKind !== "serverless") {
+			throw new Error(
+				"autoConfigureServerless can only be configured when runnerKind is 'serverless'",
+			);
+		}
+
+		// Promise for any async operations we need to wait to complete
+		const readyPromises = [];
+
+		// Start engine
+		if (config.runEngine) {
+			logger().debug({
+				msg: "run engine requested",
+				version: config.runEngineVersion,
+			});
+
+			// Set config to point to the engine
+			invariant(
+				config.endpoint === undefined,
+				"cannot specify 'endpoint' with 'runEngine'",
+			);
+			config.endpoint = ENGINE_ENDPOINT;
+			config.disableActorDriver = true;
+
+			// Start the engine
+			const engineProcessPromise = ensureEngineProcess({
+				version: config.runEngineVersion,
+			});
+
+			// Chain ready promise
+			readyPromises.push(engineProcessPromise);
+		}
+
+		// Configure for serverless
+		if (config.runnerKind === "serverless") {
+			config.defaultServerPort = 8080;
+			config.overrideServerAddress = config.endpoint;
+			config.disableActorDriver = true;
+		}
 
 		// Configure logger
 		if (config.logging?.baseLogger) {
@@ -61,14 +108,18 @@ export class Registry<A extends RegistryActors> {
 		// Choose the driver based on configuration
 		const driver = chooseDefaultDriver(config);
 
-		// TODO: Find cleaner way of disabling by default
+		// Set defaults based on the driver
 		if (driver.name === "engine") {
 			config.inspector.enabled = { manager: false, actor: true };
-			config.disableServer = true;
+
+			// We need to leave the default server enabled for dev
+			if (config.runnerKind !== "serverless") {
+				config.disableDefaultServer = true;
+			}
 		}
 		if (driver.name === "cloudflare-workers") {
 			config.inspector.enabled = { manager: false, actor: true };
-			config.disableServer = true;
+			config.disableDefaultServer = true;
 			config.disableActorDriver = true;
 			config.noWelcome = true;
 		}
@@ -102,7 +153,15 @@ export class Registry<A extends RegistryActors> {
 			const displayInfo = managerDriver.displayInformation();
 			console.log();
 			console.log(`  RivetKit ${pkg.version} (${displayInfo.name})`);
-			console.log(`  - Endpoint:     http://127.0.0.1:6420`);
+			if (!config.disableDefaultServer) {
+				console.log(`  - Endpoint:     ${config.endpoint}`);
+			} else if (config.overrideServerAddress) {
+				console.log(`  - Endpoint:     ${config.overrideServerAddress}`);
+			}
+			if (config.runEngine) {
+				const padding = " ".repeat(Math.max(0, 13 - "Engine".length));
+				console.log(`  - Engine:${padding}v${config.runEngineVersion}`);
+			}
 			for (const [k, v] of Object.entries(displayInfo.properties)) {
 				const padding = " ".repeat(Math.max(0, 13 - k.length));
 				console.log(`  - ${k}:${padding}${v}`);
@@ -116,27 +175,32 @@ export class Registry<A extends RegistryActors> {
 		// HACK: We need to find a better way to let the driver itself decide when to start the actor driver
 		// Create runner
 		//
-		// Even though we do not use the return value, this is required to start the code that will handle incoming actors
+		// Even though we do not use the returned ActorDriver, this is required to start the code that will handle incoming actors
 		if (!config.disableActorDriver) {
-			const _actorDriver = driver.actor(
-				this.#config,
-				config,
-				managerDriver,
-				client,
-			);
+			Promise.all(readyPromises).then(async () => {
+				driver.actor(this.#config, config, managerDriver, client);
+			});
+		}
+
+		// Configure serverless runner if enabled when actor driver is disabled
+		if (config.runnerKind === "serverless" && config.autoConfigureServerless) {
+			Promise.all(readyPromises).then(async () => {
+				await configureServerlessRunner(config);
+			});
 		}
 
 		const { router: hono } = createManagerRouter(
 			this.#config,
 			config,
 			managerDriver,
-			undefined,
+			driver,
+			client,
 		);
 
 		// Start server
-		if (!config.disableServer) {
+		if (!config.disableDefaultServer) {
 			(async () => {
-				const out = await crossPlatformServe(hono, undefined);
+				const out = await crossPlatformServe(config, hono, undefined);
 				upgradeWebSocket = out.upgradeWebSocket;
 			})();
 		}
@@ -146,116 +210,107 @@ export class Registry<A extends RegistryActors> {
 			fetch: hono.fetch.bind(hono),
 		};
 	}
+}
 
-	public startServerless(inputConfig?: RunConfigInput): ServerOutput<this> {
-		const config = RunConfigSchema.parse(inputConfig);
-
-		// Configure logger
-		if (config.logging?.baseLogger) {
-			// Use provided base logger
-			configureBaseLogger(config.logging.baseLogger);
-		} else {
-			// Configure default logger with log level from config
-			// getPinoLevel will handle env variable priority
-			configureDefaultLogger(config.logging?.level);
+async function configureServerlessRunner(config: RunnerConfig): Promise<void> {
+	try {
+		// Ensure we have required config values
+		if (!config.runnerName) {
+			throw new Error("runnerName is required for serverless configuration");
 		}
-
-		// Choose the driver based on configuration
-		const driver = chooseDefaultDriver(config);
-
-		// TODO: Find cleaner way of disabling by default
-		if (driver.name === "engine") {
-			config.inspector.enabled = false;
-			config.disableServer = true;
-			config.disableActorDriver = true;
+		if (!config.namespace) {
+			throw new Error("namespace is required for serverless configuration");
 		}
-		if (driver.name === "cloudflare-workers") {
-			config.inspector.enabled = false;
-			config.disableServer = true;
-			config.disableActorDriver = true;
-			config.noWelcome = true;
+		if (!config.endpoint) {
+			throw new Error("endpoint is required for serverless configuration");
 		}
 
-		// Configure getUpgradeWebSocket lazily so we can assign it in crossPlatformServe
-		let upgradeWebSocket: any;
-		if (!config.getUpgradeWebSocket) {
-			config.getUpgradeWebSocket = () => upgradeWebSocket!;
-		}
+		// Prepare the configuration
+		const customConfig =
+			typeof config.autoConfigureServerless === "object"
+				? config.autoConfigureServerless
+				: {};
 
-		// Create router
-		const managerDriver = driver.manager(this.#config, config);
+		// Make the request to fetch all datacenters
+		const dcsUrl = `${config.endpoint}/datacenters`;
 
-		// Create client
-		const client = createClientWithDriver<this>(managerDriver, config);
+		logger().debug({
+			msg: "fetching datacenters",
+			url: dcsUrl,
+		});
 
-		const driverLog = managerDriver.extraStartupLog?.() ?? {};
-		logger().info({
-			msg: "rivetkit ready",
-			driver: driver.name,
-			definitions: Object.keys(this.#config.use).length,
-			...driverLog,
+		const dcsResponse = await fetch(dcsUrl, {
+			headers: {
+				...(config.token ? { Authorization: `Bearer ${config.token}` } : {}),
+			},
 		});
-		if (config.inspector?.enabled && managerDriver.inspector) {
-			logger().info({ msg: "inspector ready", url: getInspectorUrl(config) });
-		}
 
-		// Print welcome information
-		if (!config.noWelcome) {
-			const displayInfo = managerDriver.displayInformation();
-			console.log();
-			console.log(`  RivetKit ${pkg.version} (${displayInfo.name})`);
-			console.log(`  - Endpoint:     http://127.0.0.1:6420`);
-			for (const [k, v] of Object.entries(displayInfo.properties)) {
-				const padding = " ".repeat(Math.max(0, 13 - k.length));
-				console.log(`  - ${k}:${padding}${v}`);
-			}
-			if (config.inspector?.enabled && managerDriver.inspector) {
-				console.log(`  - Inspector:    ${getInspectorUrl(config)}`);
-			}
-			console.log();
+		if (!dcsResponse.ok) {
+			const errorText = await dcsResponse.text();
+			throw new Error(
+				`failed to configure serverless runner: ${dcsResponse.status} ${dcsResponse.statusText} - ${errorText}`,
+			);
 		}
 
-		let serverlessActorDriverBuilder:
-			| ((token?: string, totalSlots?: number) => ActorDriver)
-			| undefined = (
-			token: string | undefined,
-			totalSlots: number | undefined,
-		) => {
-			// Override config
-			if (token) config.token = token;
-			if (totalSlots) config.totalSlots = totalSlots;
-
-			return driver.actor(this.#config, config, managerDriver, client);
+		const dcsRes = (await dcsResponse.json()) as {
+			datacenters: { name: string }[];
 		};
 
-		// HACK: We need to find a better way to let the driver itself decide when to start the actor driver
-		// Create runner
-		//
-		// Even though we do not use the return value, this is required to start the code that will handle incoming actors
-		if (!config.disableActorDriver) {
-			const _actorDriver = serverlessActorDriverBuilder();
-			serverlessActorDriverBuilder = undefined;
-		}
-
-		const { router: hono } = createManagerRouter(
-			this.#config,
-			config,
-			managerDriver,
-			serverlessActorDriverBuilder,
+		// Build the request body
+		const serverlessConfig = {
+			serverless: {
+				url:
+					customConfig.url ||
+					`http://localhost:${config.defaultServerPort}/start`,
+				headers: customConfig.headers || {},
+				max_runners: customConfig.maxRunners ?? 100,
+				min_runners: customConfig.minRunners ?? 0,
+				request_lifespan: customConfig.requestLifespan ?? 15 * 60,
+				runners_margin: customConfig.runnersMargin ?? 0,
+				slots_per_runner:
+					customConfig.slotsPerRunner ?? config.totalSlots ?? 1000,
+			},
+		};
+		const requestBody = Object.fromEntries(
+			dcsRes.datacenters.map((dc) => [dc.name, serverlessConfig]),
 		);
 
-		// Start server
-		if (!config.disableServer) {
-			(async () => {
-				const out = await crossPlatformServe(hono, undefined);
-				upgradeWebSocket = out.upgradeWebSocket;
-			})();
+		// Make the request to configure the serverless runner
+		const configUrl = `${config.endpoint}/runner-configs/${config.runnerName}?namespace=${config.namespace}`;
+
+		logger().debug({
+			msg: "configuring serverless runner",
+			url: configUrl,
+			config: serverlessConfig.serverless,
+		});
+
+		const response = await fetch(configUrl, {
+			method: "PUT",
+			headers: {
+				"Content-Type": "application/json",
+				...(config.token ? { Authorization: `Bearer ${config.token}` } : {}),
+			},
+			body: JSON.stringify(requestBody),
+		});
+
+		if (!response.ok) {
+			const errorText = await response.text();
+			throw new Error(
+				`failed to configure serverless runner: ${response.status} ${response.statusText} - ${errorText}`,
+			);
 		}
 
-		return {
-			client,
-			fetch: hono.fetch.bind(hono),
-		};
+		logger().info({
+			msg: "serverless runner configured successfully",
+			runnerName: config.runnerName,
+			namespace: config.namespace,
+		});
+	} catch (error) {
+		logger().error({
+			msg: "failed to configure serverless runner",
+			error,
+		});
+		throw error;
 	}
 }
 
@@ -269,8 +324,8 @@ export function setup<A extends RegistryActors>(
 export type {
 	RegistryConfig,
 	RegistryActors,
-	RunConfig,
-	RunConfigInput,
+	RunnerConfig as RunConfig,
+	RunnerConfigInput as RunConfigInput,
 	DriverConfig,
 };
 export { RegistryConfigSchema };