rivetkit 2.0.9 → 2.0.10

package/src/actor/router-endpoints.ts

@@ -26,7 +26,7 @@ import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
 import { deconstructError, stringifyError } from "@/common/utils";
 import type { UniversalWebSocket } from "@/common/websocket-interface";
 import { HonoWebSocketAdapter } from "@/manager/hono-websocket-adapter";
-import type { RunConfig } from "@/registry/run-config";
+import type { RunnerConfig } from "@/registry/run-config";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import {
 	HTTP_ACTION_REQUEST_VERSIONED,
@@ -106,7 +106,7 @@ export interface WebSocketOpts {
  */
 export async function handleWebSocketConnect(
 	req: Request | undefined,
-	runConfig: RunConfig,
+	runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	actorId: string,
 	encoding: Encoding,
@@ -325,7 +325,7 @@ export async function handleWebSocketConnect(
  */
 export async function handleSseConnect(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	actorId: string,
 ) {
@@ -437,7 +437,7 @@ export async function handleSseConnect(
  */
 export async function handleAction(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	actionName: string,
 	actorId: string,
@@ -505,7 +505,7 @@ export async function handleAction(
  */
 export async function handleConnectionMessage(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	connId: string,
 	connToken: string,
@@ -542,7 +542,7 @@ export async function handleConnectionMessage(
 
 export async function handleConnectionClose(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	connId: string,
 	connToken: string,

package/src/actor/router.ts

@@ -41,7 +41,7 @@ import {
 	createActorInspectorRouter,
 } from "@/inspector/actor";
 import { isInspectorEnabled, secureInspector } from "@/inspector/utils";
-import type { RunConfig } from "@/registry/run-config";
+import type { RunnerConfig } from "@/registry/run-config";
 import { ConnDriverKind } from "./conn-drivers";
 import type { ActorDriver } from "./driver";
 import { InternalError } from "./errors";
@@ -67,7 +67,7 @@ export type ActorRouter = Hono<{ Bindings: ActorRouterBindings }>;
  * Creates a router that runs on the partitioned instance.
  */
 export function createActorRouter(
-	runConfig: RunConfig,
+	runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	isTest: boolean,
 ): ActorRouter {

package/src/client/actor-conn.ts

@@ -315,13 +315,34 @@ enc
 			});
 		});
 		ws.addEventListener("message", async (ev) => {
-			this.#handleOnMessage(ev.data);
+			try {
+				await this.#handleOnMessage(ev.data);
+			} catch (err) {
+				logger().error({
+					msg: "error in websocket message handler",
+					error: stringifyError(err),
+				});
+			}
 		});
 		ws.addEventListener("close", (ev) => {
-			this.#handleOnClose(ev);
+			try {
+				this.#handleOnClose(ev);
+			} catch (err) {
+				logger().error({
+					msg: "error in websocket close handler",
+					error: stringifyError(err),
+				});
+			}
 		});
 		ws.addEventListener("error", (_ev) => {
-			this.#handleOnError();
+			try {
+				this.#handleOnError();
+			} catch (err) {
+				logger().error({
+					msg: "error in websocket error handler",
+					error: stringifyError(err),
+				});
+			}
 		});
 	}
 

package/src/client/config.ts

@@ -1,38 +1,32 @@
 import z from "zod";
 import { TransportSchema } from "@/actor/protocol/old";
 import { EncodingSchema } from "@/actor/protocol/serde";
-import { getEnvUniversal, type UpgradeWebSocket } from "@/utils";
-
-export type GetUpgradeWebSocket = () => UpgradeWebSocket;
+import { type GetUpgradeWebSocket, getEnvUniversal } from "@/utils";
 
 export const ClientConfigSchema = z.object({
-	/** Configure serving the API */
-	api: z
-		.object({
-			host: z.string().default("127.0.0.1"),
-			port: z.number().default(6420),
-		})
-		.default({}),
-
+	/** Endpoint to connect to for Rivet Engine or RivetKit manager API. */
+	endpoint: z
+		.string()
+		.optional()
+		.transform(
+			(x) =>
+				x ??
+				getEnvUniversal("RIVET_ENGINE") ??
+				getEnvUniversal("RIVET_ENDPOINT"),
+		),
+
+	/** Token to use to authenticate with the API. */
 	token: z
 		.string()
 		.optional()
 		.transform((x) => x ?? getEnvUniversal("RIVET_TOKEN")),
 
-	totalSlots: z.number().optional(),
-
-	headers: z.record(z.string()).optional().default({}),
-
-	/** Endpoint to connect to the Rivet engine. Can be configured via RIVET_ENGINE env var. */
-	endpoint: z
-		.string()
-		.nullable()
-		.default(() => getEnvUniversal("RIVET_ENGINE") ?? null),
-
+	/** Namespace to connect to. */
 	namespace: z
 		.string()
 		.default(() => getEnvUniversal("RIVET_NAMESPACE") ?? "default"),
 
+	/** Name of the runner. This is used to group together runners in to different pools. */
 	runnerName: z
 		.string()
 		.default(() => getEnvUniversal("RIVET_RUNNER") ?? "rivetkit"),
@@ -41,10 +35,9 @@ export const ClientConfigSchema = z.object({
 
 	transport: TransportSchema.default("websocket"),
 
-	// 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(...)`
+	headers: z.record(z.string()).optional().default({}),
+
+	// See RunConfig.getUpgradeWebSocket
 	getUpgradeWebSocket: z.custom<GetUpgradeWebSocket>().optional(),
 });
 

package/src/driver-helpers/mod.ts

@@ -32,5 +32,8 @@ export type {
 	ManagerDisplayInformation,
 	ManagerDriver,
 } from "@/manager/driver";
-export { DriverConfigSchema, RunConfigSchema } from "@/registry/run-config";
+export {
+	DriverConfigSchema,
+	RunnerConfigSchema as RunConfigSchema,
+} from "@/registry/run-config";
 export { serializeEmptyPersistData } from "./utils";

package/src/driver-test-suite/mod.ts

@@ -3,11 +3,16 @@ import { createNodeWebSocket, type NodeWebSocket } from "@hono/node-ws";
 import { bundleRequire } from "bundle-require";
 import invariant from "invariant";
 import { describe } from "vitest";
-import type { Transport } from "@/client/mod";
+import type { Encoding, Transport } from "@/client/mod";
 import { configureInspectorAccessToken } from "@/inspector/utils";
 import { createManagerRouter } from "@/manager/router";
-import type { DriverConfig, Registry, RunConfig } from "@/mod";
-import { RunConfigSchema } from "@/registry/run-config";
+import {
+	createClientWithDriver,
+	type DriverConfig,
+	type Registry,
+	type RunConfig,
+} from "@/mod";
+import { RunnerConfigSchema } from "@/registry/run-config";
 import { getPort } from "@/test/mod";
 import { logger } from "./log";
 import { runActionFeaturesTests } from "./tests/action-features";
@@ -55,6 +60,8 @@ export interface DriverTestConfig {
 
 	transport?: Transport;
 
+	encoding?: Encoding;
+
 	clientType: ClientType;
 
 	cleanup?: () => Promise<void>;
@@ -78,68 +85,81 @@ export interface DriverDeployOutput {
 
 /** Runs all Vitest tests against the provided drivers. */
 export function runDriverTests(
-	driverTestConfigPartial: Omit<DriverTestConfig, "clientType" | "transport">,
+	driverTestConfigPartial: Omit<
+		DriverTestConfig,
+		"clientType" | "transport" | "encoding"
+	>,
 ) {
 	const clientTypes: ClientType[] = driverTestConfigPartial.skip?.inline
 		? ["http"]
 		: ["http", "inline"];
 	for (const clientType of clientTypes) {
-		const driverTestConfig: DriverTestConfig = {
-			...driverTestConfigPartial,
-			clientType,
-		};
-
 		describe(`client type (${clientType})`, () => {
-			runActorDriverTests(driverTestConfig);
-			runManagerDriverTests(driverTestConfig);
+			const encodings: Encoding[] = ["bare", "cbor", "json"];
 
-			const transports: Transport[] = driverTestConfig.skip?.sse
-				? ["websocket"]
-				: ["websocket", "sse"];
-			for (const transport of transports) {
-				describe(`transport (${transport})`, () => {
-					runActorConnTests({
-						...driverTestConfig,
-						transport,
-					});
+			for (const encoding of encodings) {
+				describe(`encoding (${encoding})`, () => {
+					const driverTestConfig: DriverTestConfig = {
+						...driverTestConfigPartial,
+						clientType,
+						encoding,
+					};
 
-					runActorConnStateTests({ ...driverTestConfig, transport });
+					runActorDriverTests(driverTestConfig);
+					runManagerDriverTests(driverTestConfig);
 
-					runActorReconnectTests({ ...driverTestConfig, transport });
+					const transports: Transport[] = driverTestConfig.skip?.sse
+						? ["websocket"]
+						: ["websocket", "sse"];
+					for (const transport of transports) {
+						describe(`transport (${transport})`, () => {
+							runActorConnTests({
+								...driverTestConfig,
+								transport,
+							});
 
-					runRequestAccessTests({ ...driverTestConfig, transport });
+							runActorConnStateTests({ ...driverTestConfig, transport });
 
-					runActorDriverTestsWithTransport({ ...driverTestConfig, transport });
-				});
-			}
+							runActorReconnectTests({ ...driverTestConfig, transport });
 
-			runActorHandleTests(driverTestConfig);
+							runRequestAccessTests({ ...driverTestConfig, transport });
 
-			runActionFeaturesTests(driverTestConfig);
+							runActorDriverTestsWithTransport({
+								...driverTestConfig,
+								transport,
+							});
+						});
+					}
 
-			runActorVarsTests(driverTestConfig);
+					runActorHandleTests(driverTestConfig);
 
-			runActorMetadataTests(driverTestConfig);
+					runActionFeaturesTests(driverTestConfig);
 
-			runActorOnStateChangeTests(driverTestConfig);
+					runActorVarsTests(driverTestConfig);
 
-			runActorErrorHandlingTests(driverTestConfig);
+					runActorMetadataTests(driverTestConfig);
 
-			runActorInlineClientTests(driverTestConfig);
+					runActorOnStateChangeTests(driverTestConfig);
 
-			runRawHttpTests(driverTestConfig);
+					runActorErrorHandlingTests(driverTestConfig);
 
-			runRawHttpRequestPropertiesTests(driverTestConfig);
+					runActorInlineClientTests(driverTestConfig);
 
-			runRawWebSocketTests(driverTestConfig);
+					runRawHttpTests(driverTestConfig);
 
-			// TODO: re-expose this once we can have actor queries on the gateway
-			// runRawHttpDirectRegistryTests(driverTestConfig);
+					runRawHttpRequestPropertiesTests(driverTestConfig);
 
-			// TODO: re-expose this once we can have actor queries on the gateway
-			// runRawWebSocketDirectRegistryTests(driverTestConfig);
+					runRawWebSocketTests(driverTestConfig);
 
-			runActorInspectorTests(driverTestConfig);
+					// TODO: re-expose this once we can have actor queries on the gateway
+					// runRawHttpDirectRegistryTests(driverTestConfig);
+
+					// TODO: re-expose this once we can have actor queries on the gateway
+					// runRawWebSocketDirectRegistryTests(driverTestConfig);
+
+					runActorInspectorTests(driverTestConfig);
+				});
+			}
 		});
 	}
 }
@@ -199,7 +219,7 @@ export async function createTestRuntime(
 		// Build driver config
 		// biome-ignore lint/style/useConst: Assigned later
 		let upgradeWebSocket: any;
-		const config: RunConfig = RunConfigSchema.parse({
+		const config: RunConfig = RunnerConfigSchema.parse({
 			driver,
 			getUpgradeWebSocket: () => upgradeWebSocket!,
 			inspector: {
@@ -210,12 +230,14 @@ export async function createTestRuntime(
 
 		// Create router
 		const managerDriver = driver.manager(registry.config, config);
+		const client = createClientWithDriver(managerDriver);
 		configureInspectorAccessToken(config, managerDriver);
 		const { router } = createManagerRouter(
 			registry.config,
 			config,
 			managerDriver,
-			undefined,
+			driver,
+			client,
 		);
 
 		// Inject WebSocket

package/src/driver-test-suite/utils.ts

@@ -36,17 +36,20 @@ export async function setupDriverTest(
 			namespace,
 			runnerName,
 			transport: driverTestConfig.transport,
+			encoding: driverTestConfig.encoding,
 		});
 	} else if (driverTestConfig.clientType === "inline") {
 		// Use inline client from driver
 		const transport = driverTestConfig.transport ?? "websocket";
+		const encoding = driverTestConfig.encoding ?? "bare";
 		const managerDriver = createTestInlineClientDriver(
 			endpoint,
-			"bare",
+			encoding,
 			transport,
 		);
 		const runConfig = RunConfigSchema.parse({
 			transport: transport,
+			encoding: encoding,
 		});
 		client = createClientWithDriver(managerDriver, runConfig);
 	} else {

package/src/drivers/default.ts

@@ -2,15 +2,21 @@ import { UserError } from "@/actor/errors";
 import { loggerWithoutContext } from "@/actor/log";
 import { createEngineDriver } from "@/drivers/engine/mod";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
-import type { DriverConfig, RunConfig } from "@/registry/run-config";
+import type { DriverConfig, RunnerConfig } from "@/registry/run-config";
 
 /**
  * Chooses the appropriate driver based on the run configuration.
  */
-export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
+export function chooseDefaultDriver(runConfig: RunnerConfig): DriverConfig {
 	if (runConfig.endpoint && runConfig.driver) {
 		throw new UserError(
-			"Cannot specify both 'engine' and 'driver' in configuration",
+			"Cannot specify both 'endpoint' and 'driver' in configuration",
+		);
+	}
+
+	if (runConfig.runnerKind === "serverless" && !runConfig.endpoint) {
+		throw new UserError(
+			"Cannot use 'serverless' runnerKind without the 'endpoint' config set.",
 		);
 	}
 
@@ -18,16 +24,12 @@ export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
 		return runConfig.driver;
 	}
 
-	if (runConfig.endpoint) {
+	if (runConfig.endpoint || runConfig.token) {
 		loggerWithoutContext().debug({
 			msg: "using rivet engine driver",
 			endpoint: runConfig.endpoint,
 		});
-		// TODO: Add all properties from config
-		return createEngineDriver({
-			endpoint: runConfig.endpoint,
-			token: runConfig.token,
-		});
+		return createEngineDriver();
 	}
 
 	loggerWithoutContext().debug({ msg: "using default file system driver" });

package/src/drivers/engine/actor-driver.ts

@@ -1,6 +1,6 @@
 import type {
+	RunnerConfig as EngineRunnerConfig,
 	ActorConfig as RunnerActorConfig,
-	RunnerConfig,
 } from "@rivetkit/engine-runner";
 import { Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
@@ -34,13 +34,13 @@ import {
 	serializeEmptyPersistData,
 } from "@/driver-helpers/mod";
 import type { RegistryConfig } from "@/registry/config";
-import type { RunConfig } from "@/registry/run-config";
+import type { RunnerConfig } from "@/registry/run-config";
+import { getEndpoint } from "@/remote-manager-driver/api-utils";
 import {
 	type LongTimeoutHandle,
 	promiseWithResolvers,
 	setLongTimeout,
 } from "@/utils";
-import type { Config } from "./config";
 import { KEYS } from "./kv";
 import { logger } from "./log";
 
@@ -54,10 +54,9 @@ export type DriverContext = {};
 
 export class EngineActorDriver implements ActorDriver {
 	#registryConfig: RegistryConfig;
-	#runConfig: RunConfig;
+	#runConfig: RunnerConfig;
 	#managerDriver: ManagerDriver;
 	#inlineClient: Client<any>;
-	#config: Config;
 	#runner: Runner;
 	#actors: Map<string, ActorHandler> = new Map();
 	#actorRouter: ActorRouter;
@@ -69,16 +68,22 @@ export class EngineActorDriver implements ActorDriver {
 
 	constructor(
 		registryConfig: RegistryConfig,
-		runConfig: RunConfig,
+		runConfig: RunnerConfig,
 		managerDriver: ManagerDriver,
 		inlineClient: Client<any>,
-		config: Config,
 	) {
 		this.#registryConfig = registryConfig;
 		this.#runConfig = runConfig;
 		this.#managerDriver = managerDriver;
 		this.#inlineClient = inlineClient;
-		this.#config = config;
+
+		// HACK: Override inspector token (which are likely to be
+		// removed later on) with token from x-rivet-token header
+		const token = runConfig.token ?? runConfig.token;
+		if (token && runConfig.inspector && runConfig.inspector.enabled) {
+			runConfig.inspector.token = () => token;
+		}
+
 		this.#actorRouter = createActorRouter(
 			runConfig,
 			this,
@@ -87,15 +92,14 @@ export class EngineActorDriver implements ActorDriver {
 
 		// Create runner configuration
 		let hasDisconnected = false;
-		const runnerConfig: RunnerConfig = {
+		const engineRunnerConfig: EngineRunnerConfig = {
 			version: this.#version,
-			endpoint: config.endpoint,
-			token: runConfig.token ?? config.token,
-			pegboardEndpoint: config.pegboardEndpoint,
-			namespace: config.namespace,
-			totalSlots: runConfig.totalSlots ?? config.totalSlots,
-			runnerName: config.runnerName,
-			runnerKey: config.runnerKey,
+			endpoint: getEndpoint(runConfig),
+			token,
+			namespace: runConfig.namespace ?? runConfig.namespace,
+			totalSlots: runConfig.totalSlots ?? runConfig.totalSlots,
+			runnerName: runConfig.runnerName ?? runConfig.runnerName,
+			runnerKey: runConfig.runnerKey,
 			metadata: {
 				inspectorToken: this.#runConfig.inspector.token(),
 			},
@@ -109,14 +113,14 @@ export class EngineActorDriver implements ActorDriver {
 				if (hasDisconnected) {
 					logger().info({
 						msg: "runner reconnected",
-						namespace: this.#config.namespace,
-						runnerName: this.#config.runnerName,
+						namespace: this.#runConfig.namespace,
+						runnerName: this.#runConfig.runnerName,
 					});
 				} else {
 					logger().debug({
 						msg: "runner connected",
-						namespace: this.#config.namespace,
-						runnerName: this.#config.runnerName,
+						namespace: this.#runConfig.namespace,
+						runnerName: this.#runConfig.runnerName,
 					});
 				}
 
@@ -125,8 +129,8 @@ export class EngineActorDriver implements ActorDriver {
 			onDisconnected: () => {
 				logger().warn({
 					msg: "runner disconnected",
-					namespace: this.#config.namespace,
-					runnerName: this.#config.runnerName,
+					namespace: this.#runConfig.namespace,
+					runnerName: this.#runConfig.runnerName,
 				});
 				hasDisconnected = true;
 			},
@@ -141,13 +145,13 @@ export class EngineActorDriver implements ActorDriver {
 		};
 
 		// Create and start runner
-		this.#runner = new Runner(runnerConfig);
+		this.#runner = new Runner(engineRunnerConfig);
 		this.#runner.start();
 		logger().debug({
 			msg: "engine runner started",
-			endpoint: config.endpoint,
-			namespace: config.namespace,
-			runnerName: config.runnerName,
+			endpoint: runConfig.endpoint,
+			namespace: runConfig.namespace,
+			runnerName: runConfig.runnerName,
 		});
 	}
 
@@ -228,20 +232,20 @@ export class EngineActorDriver implements ActorDriver {
 	async #runnerOnActorStart(
 		actorId: string,
 		generation: number,
-		config: RunnerActorConfig,
+		runConfig: RunnerActorConfig,
 	): Promise<void> {
 		logger().debug({
 			msg: "runner actor starting",
 			actorId,
-			name: config.name,
-			key: config.key,
+			name: runConfig.name,
+			key: runConfig.key,
 			generation,
 		});
 
 		// Deserialize input
 		let input: any;
-		if (config.input) {
-			input = cbor.decode(config.input);
+		if (runConfig.input) {
+			input = cbor.decode(runConfig.input);
 		}
 
 		// Get or create handler
@@ -254,15 +258,12 @@ export class EngineActorDriver implements ActorDriver {
 			this.#actors.set(actorId, handler);
 		}
 
-		const name = config.name as string;
-		invariant(config.key, "actor should have a key");
-		const key = deserializeActorKey(config.key);
+		const name = runConfig.name as string;
+		invariant(runConfig.key, "actor should have a key");
+		const key = deserializeActorKey(runConfig.key);
 
 		// Create actor instance
-		const definition = lookupInRegistry(
-			this.#registryConfig,
-			config.name as string, // TODO: Remove cast
-		);
+		const definition = lookupInRegistry(this.#registryConfig, runConfig.name);
 		handler.actor = definition.instantiate();
 
 		// Start actor
@@ -413,7 +414,7 @@ export class EngineActorDriver implements ActorDriver {
 			// Runner id should be set if the runner started
 			const payload = this.#runner.getServerlessInitPacket();
 			invariant(payload, "runnerId not set");
-			stream.writeSSE({ data: payload });
+			await stream.writeSSE({ data: payload });
 
 			return this.#runnerStopped.promise;
 		});

package/src/drivers/engine/config.ts

@@ -1,35 +1,22 @@
-import type { Hono } from "hono";
 import { z } from "zod";
+import { ClientConfigSchema } from "@/client/config";
 import { getEnvUniversal } from "@/utils";
 
-export const ConfigSchema = z
+export const EngingConfigSchema = z
 	.object({
-		app: z.custom<Hono>().optional(),
-		endpoint: z
-			.string()
-			.default(
-				() => getEnvUniversal("RIVET_ENGINE") ?? "http://localhost:6420",
-			),
-		token: z
-			.string()
-			.optional()
-			.transform((val) => val ?? getEnvUniversal("RIVET_TOKEN")),
-		pegboardEndpoint: z.string().optional(),
-		namespace: z
-			.string()
-			.default(() => getEnvUniversal("RIVET_NAMESPACE") ?? "default"),
-		runnerName: z
-			.string()
-			.default(() => getEnvUniversal("RIVET_RUNNER") ?? "rivetkit"),
-		// TODO: Automatically attempt to determine key by common env vars (e.g. k8s pod name)
+		/** Unique key for this runner. Runners connecting a given key will replace any other runner connected with the same key. */
 		runnerKey: z
 			.string()
 			.default(
 				() => getEnvUniversal("RIVET_RUNNER_KEY") ?? crypto.randomUUID(),
 			),
+
+		/** How many actors this runner can run. */
 		totalSlots: z.number().default(100_000),
 	})
+	// We include the client config since this includes the common properties like endpoint, namespace, etc.
+	.merge(ClientConfigSchema)
 	.default({});
 
-export type InputConfig = z.input<typeof ConfigSchema>;
-export type Config = z.infer<typeof ConfigSchema>;
+export type EngineConfig = z.infer<typeof EngingConfigSchema>;
+export type EngineConfigInput = z.input<typeof EngingConfigSchema>;