rivetkit 2.0.8 → 2.0.10

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,48 @@ 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;
+		}
 
 		// Configure logger
 		if (config.logging?.baseLogger) {
@@ -61,14 +107,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,13 +152,21 @@ 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}`);
 			}
 			if (isInspectorEnabled(config, "manager") && managerDriver.inspector) {
-				console.log(`  - Inspector:     ${getInspectorUrl(config)}`);
+				console.log(`  - Inspector:    ${getInspectorUrl(config)}`);
 			}
 			console.log();
 		}
@@ -116,27 +174,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 +209,79 @@ 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
+				: {};
+
+		// Build the request body
+		const requestBody = {
+			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,
+			},
+		};
 
-		// Create router
-		const managerDriver = driver.manager(this.#config, config);
+		// Make the request to configure the serverless runner
+		const configUrl = `${config.endpoint}/runner-configs/${config.runnerName}?namespace=${config.namespace}`;
 
-		// Create client
-		const client = createClientWithDriver<this>(managerDriver, config);
-
-		const driverLog = managerDriver.extraStartupLog?.() ?? {};
-		logger().info({
-			msg: "rivetkit ready",
-			driver: driver.name,
-			definitions: Object.keys(this.#config.use).length,
-			...driverLog,
+		logger().debug({
+			msg: "configuring serverless runner",
+			url: configUrl,
+			config: requestBody.serverless,
 		});
-		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();
-		}
-
-		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);
-		};
-
-		// 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,
-		);
+		const response = await fetch(configUrl, {
+			method: "PUT",
+			headers: {
+				"Content-Type": "application/json",
+				...(config.token ? { Authorization: `Bearer ${config.token}` } : {}),
+			},
+			body: JSON.stringify(requestBody),
+		});
 
-		// Start server
-		if (!config.disableServer) {
-			(async () => {
-				const out = await crossPlatformServe(hono, undefined);
-				upgradeWebSocket = out.upgradeWebSocket;
-			})();
+		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 +295,8 @@ export function setup<A extends RegistryActors>(
 export type {
 	RegistryConfig,
 	RegistryActors,
-	RunConfig,
-	RunConfigInput,
+	RunnerConfig as RunConfig,
+	RunnerConfigInput as RunConfigInput,
 	DriverConfig,
 };
 export { RegistryConfigSchema };

package/src/registry/run-config.ts

@@ -2,10 +2,12 @@ import type { cors } from "hono/cors";
 import type { Logger } from "pino";
 import { z } from "zod";
 import type { ActorDriverBuilder } from "@/actor/driver";
-import { ClientConfigSchema } from "@/client/config";
 import { LogLevelSchema } from "@/common/log";
+import { EngingConfigSchema as EngineConfigSchema } from "@/drivers/engine/config";
 import { InspectorConfigSchema } from "@/inspector/config";
 import type { ManagerDriverBuilder } from "@/manager/driver";
+import type { GetUpgradeWebSocket } from "@/utils";
+import { getEnvUniversal } from "@/utils";
 
 type CorsOptions = NonNullable<Parameters<typeof cors>[0]>;
 
@@ -19,49 +21,116 @@ export const DriverConfigSchema = z.object({
 export type DriverConfig = z.infer<typeof DriverConfigSchema>;
 
 /** Base config used for the actor config across all platforms. */
-export const RunConfigSchema = ClientConfigSchema.extend({
-	driver: DriverConfigSchema.optional(),
-
-	/** CORS configuration for the router. Uses Hono's CORS middleware options. */
-	cors: z.custom<CorsOptions>().optional(),
-
-	maxIncomingMessageSize: z.number().optional().default(65_536),
-
-	inspector: InspectorConfigSchema,
-
-	/** @experimental */
-	disableServer: z.boolean().optional().default(false),
-
-	/** @experimental */
-	disableActorDriver: z.boolean().optional().default(false),
-
-	/**
-	 * @experimental
-	 *
-	 * Base path for the router. This is used to prefix all routes.
-	 * For example, if the base path is `/api`, then the route `/actors` will be
-	 * available at `/api/actors`.
-	 */
-	basePath: z.string().optional().default("/"),
-
-	/**
-	 * @experimental
-	 *
-	 * Disable welcome message.
-	 * */
-	noWelcome: z.boolean().optional().default(false),
-
-	/**
-	 * @experimental
-	 * */
-	logging: z
-		.object({
-			baseLogger: z.custom<Logger>().optional(),
-			level: LogLevelSchema.optional(),
-		})
-		.optional()
-		.default({}),
-}).default({});
-
-export type RunConfig = z.infer<typeof RunConfigSchema>;
-export type RunConfigInput = z.input<typeof RunConfigSchema>;
+export const RunnerConfigSchema = z
+	.object({
+		driver: DriverConfigSchema.optional(),
+
+		/** CORS configuration for the router. Uses Hono's CORS middleware options. */
+		cors: z.custom<CorsOptions>().optional(),
+
+		/** @experimental */
+		maxIncomingMessageSize: z.number().optional().default(65_536),
+
+		/** @experimental */
+		inspector: InspectorConfigSchema,
+
+		/** @experimental */
+		disableDefaultServer: z.boolean().optional().default(false),
+
+		/** @experimental */
+		defaultServerPort: z.number().default(6420),
+
+		/** @experimental */
+		runEngine: z
+			.boolean()
+			.optional()
+			.default(() => getEnvUniversal("RIVET_RUN_ENGINE") === "1"),
+
+		/** @experimental */
+		runEngineVersion: z
+			.string()
+			.optional()
+			.default(() => getEnvUniversal("RIVET_RUN_ENGINE_VERSION") ?? "25.7.3"),
+
+		/** @experimental */
+		overrideServerAddress: z.string().optional(),
+
+		/** @experimental */
+		disableActorDriver: z.boolean().optional().default(false),
+
+		/**
+		 * @experimental
+		 *
+		 * Whether to run runners normally or have them managed
+		 * serverlessly (by the Rivet Engine for example).
+		 */
+		runnerKind: z
+			.enum(["serverless", "normal"])
+			.optional()
+			.default(() =>
+				getEnvUniversal("RIVET_RUNNER_KIND") === "serverless"
+					? "serverless"
+					: "normal",
+			),
+		totalSlots: z.number().optional(),
+
+		/**
+		 * @experimental
+		 *
+		 * Base path for the router. This is used to prefix all routes.
+		 * For example, if the base path is `/api`, then the route `/actors` will be
+		 * available at `/api/actors`.
+		 */
+		basePath: z.string().optional().default("/"),
+
+		/**
+		 * @experimental
+		 *
+		 * Disable welcome message.
+		 * */
+		noWelcome: z.boolean().optional().default(false),
+
+		/**
+		 * @experimental
+		 * */
+		logging: z
+			.object({
+				baseLogger: z.custom<Logger>().optional(),
+				level: LogLevelSchema.optional(),
+			})
+			.optional()
+			.default({}),
+
+		/**
+		 * @experimental
+		 *
+		 * Automatically configure serverless runners in the engine.
+		 * Can only be used when runnerKind is "serverless".
+		 * If true, uses default configuration. Can also provide custom configuration.
+		 */
+		autoConfigureServerless: z
+			.union([
+				z.boolean(),
+				z.object({
+					url: z.string().optional(),
+					headers: z.record(z.string(), z.string()).optional(),
+					maxRunners: z.number().optional(),
+					minRunners: z.number().optional(),
+					requestLifespan: z.number().optional(),
+					runnersMargin: z.number().optional(),
+					slotsPerRunner: z.number().optional(),
+				}),
+			])
+			.optional(),
+
+		// This is a function to allow for lazy configuration of upgradeWebSocket on the
+		// fly. This is required since the dependencies that upgradeWebSocket
+		// (specifically Node.js) can sometimes only be specified after the router is
+		// created or must be imported async using `await import(...)`
+		getUpgradeWebSocket: z.custom<GetUpgradeWebSocket>().optional(),
+	})
+	.merge(EngineConfigSchema.removeDefault())
+	.default({});
+
+export type RunnerConfig = z.infer<typeof RunnerConfigSchema>;
+export type RunnerConfigInput = z.input<typeof RunnerConfigSchema>;

package/src/registry/serve.ts

@@ -1,7 +1,9 @@
 import { Hono } from "hono";
 import { logger } from "./log";
+import type { RunnerConfig } from "./run-config";
 
 export async function crossPlatformServe(
+	runConfig: RunnerConfig,
 	rivetKitRouter: Hono<any>,
 	userRouter: Hono | undefined,
 ) {
@@ -47,7 +49,7 @@ export async function crossPlatformServe(
 	});
 
 	// Start server
-	const port = 6420;
+	const port = runConfig.defaultServerPort;
 	const server = serve({ fetch: app.fetch, port }, () =>
 		logger().info({ msg: "server listening", port }),
 	);

package/src/remote-manager-driver/mod.ts

@@ -15,6 +15,7 @@ import type {
 	ManagerDriver,
 } from "@/driver-helpers/mod";
 import type { Encoding, UniversalWebSocket } from "@/mod";
+import { uint8ArrayToBase64 } from "@/serde";
 import { combineUrlPath } from "@/utils";
 import { sendHttpRequestToActor } from "./actor-http-client";
 import {
@@ -140,7 +141,7 @@ export class RemoteManagerDriver implements ManagerDriver {
 			key: serializeActorKey(key),
 			runner_name_selector: this.#config.runnerName,
 			input: actorInput
-				? cbor.encode(actorInput).toString("base64")
+				? uint8ArrayToBase64(cbor.encode(actorInput))
 				: undefined,
 			crash_policy: "sleep",
 		});
@@ -175,7 +176,7 @@ export class RemoteManagerDriver implements ManagerDriver {
 			name,
 			runner_name_selector: this.#config.runnerName,
 			key: serializeActorKey(key),
-			input: input ? cbor.encode(input).toString("base64") : null,
+			input: input ? uint8ArrayToBase64(cbor.encode(input)) : undefined,
 			crash_policy: "sleep",
 		});
 		const actorId = result.actor.actor_id;

package/src/serde.ts

@@ -3,7 +3,22 @@ import invariant from "invariant";
 import { assertUnreachable } from "@/common/utils";
 import type { VersionedDataHandler } from "@/common/versioned-data";
 import type { Encoding } from "@/mod";
-import { jsonStringifyCompat } from "./actor/protocol/serde";
+import { jsonParseCompat, jsonStringifyCompat } from "./actor/protocol/serde";
+
+export function uint8ArrayToBase64(uint8Array: Uint8Array): string {
+	// Check if Buffer is available (Node.js)
+	if (typeof Buffer !== "undefined") {
+		return Buffer.from(uint8Array).toString("base64");
+	}
+
+	// Browser environment - use btoa
+	let binary = "";
+	const len = uint8Array.byteLength;
+	for (let i = 0; i < len; i++) {
+		binary += String.fromCharCode(uint8Array[i]);
+	}
+	return btoa(binary);
+}
 
 export function encodingIsBinary(encoding: Encoding): boolean {
 	if (encoding === "json") {
@@ -63,11 +78,11 @@ export function deserializeWithEncoding<T>(
 ): T {
 	if (encoding === "json") {
 		if (typeof buffer === "string") {
-			return JSON.parse(buffer);
+			return jsonParseCompat(buffer);
 		} else {
 			const decoder = new TextDecoder("utf-8");
 			const jsonString = decoder.decode(buffer);
-			return JSON.parse(jsonString);
+			return jsonParseCompat(jsonString);
 		}
 	} else if (encoding === "cbor") {
 		invariant(

package/src/test/config.ts

@@ -1,7 +1,7 @@
 import { z } from "zod";
-import { RunConfigSchema } from "@/registry/run-config";
+import { RunnerConfigSchema } from "@/registry/run-config";
 
-export const ConfigSchema = RunConfigSchema.removeDefault()
+export const ConfigSchema = RunnerConfigSchema.removeDefault()
 	.extend({
 		hostname: z
 			.string()

package/src/test/mod.ts

@@ -10,8 +10,9 @@ import {
 	getInspectorUrl,
 } from "@/inspector/utils";
 import { createManagerRouter } from "@/manager/router";
+import { createClientWithDriver } from "@/mod";
 import type { Registry } from "@/registry/mod";
-import { RunConfigSchema } from "@/registry/run-config";
+import { RunnerConfigSchema } from "@/registry/run-config";
 import { ConfigSchema, type InputConfig } from "./config";
 import { logger } from "./log";
 
@@ -27,15 +28,17 @@ function serve(registry: Registry<any>, inputConfig?: InputConfig): ServerType {
 	}
 
 	// Create router
-	const runConfig = RunConfigSchema.parse(inputConfig);
+	const runConfig = RunnerConfigSchema.parse(inputConfig);
 	const driver = inputConfig.driver ?? createFileSystemOrMemoryDriver(false);
 	const managerDriver = driver.manager(registry.config, config);
+	const client = createClientWithDriver(managerDriver);
 	configureInspectorAccessToken(config, managerDriver);
 	const { router } = createManagerRouter(
 		registry.config,
 		runConfig,
 		managerDriver,
-		undefined,
+		driver,
+		client,
 	);
 
 	// Inject WebSocket

package/src/utils.ts

@@ -31,6 +31,8 @@ export type UpgradeWebSocket = (
 	createEvents: (c: HonoContext) => any,
 ) => HonoHandler;
 
+export type GetUpgradeWebSocket = () => UpgradeWebSocket;
+
 export function getEnvUniversal(key: string): string | undefined {
 	if (typeof Deno !== "undefined") {
 		return Deno.env.get(key);