chore(core): add x-rivet-endpoint header for serverless runners & cleanup severless runner setup (#1359)

Depends on rivet-dev/rivet#3077

Author: NathanFlurry

packages/rivetkit/scripts/dump-openapi.ts

@@ -3,7 +3,12 @@ import { resolve } from "node:path";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
 import type { ManagerDriver } from "@/manager/driver";
 import { createManagerRouter } from "@/manager/router";
-import { type RegistryConfig, RegistryConfigSchema, setup } from "@/mod";
+import {
+	createClientWithDriver,
+	type RegistryConfig,
+	RegistryConfigSchema,
+	setup,
+} from "@/mod";
 import { type RunnerConfig, RunnerConfigSchema } from "@/registry/run-config";
 import { VERSION } from "@/utils";
 
@@ -34,11 +39,14 @@ function main() {
 		getOrCreateInspectorAccessToken: unimplemented,
 	};
 
+	const client = createClientWithDriver(managerDriver);
+
 	const { openapi } = createManagerRouter(
 		registryConfig,
 		driverConfig,
 		managerDriver,
-		undefined,
+		driverConfig.driver!,
+		client,
 	);
 
 	const openApiDoc = openapi.getOpenAPIDocument({

packages/rivetkit/src/driver-test-suite/mod.ts

@@ -6,7 +6,12 @@ import { describe } from "vitest";
 import type { Transport } from "@/client/mod";
 import { configureInspectorAccessToken } from "@/inspector/utils";
 import { createManagerRouter } from "@/manager/router";
-import type { DriverConfig, Registry, RunConfig } from "@/mod";
+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";
@@ -210,12 +215,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

packages/rivetkit/src/drivers/default.ts

@@ -29,7 +29,7 @@ export function chooseDefaultDriver(runConfig: RunnerConfig): DriverConfig {
 			msg: "using rivet engine driver",
 			endpoint: runConfig.endpoint,
 		});
-		return createEngineDriver(runConfig);
+		return createEngineDriver();
 	}
 
 	loggerWithoutContext().debug({ msg: "using default file system driver" });

packages/rivetkit/src/drivers/engine/actor-driver.ts

@@ -41,7 +41,6 @@ import {
 	promiseWithResolvers,
 	setLongTimeout,
 } from "@/utils";
-import type { EngineConfig } from "./config";
 import { KEYS } from "./kv";
 import { logger } from "./log";
 
@@ -58,7 +57,6 @@ export class EngineActorDriver implements ActorDriver {
 	#runConfig: RunnerConfig;
 	#managerDriver: ManagerDriver;
 	#inlineClient: Client<any>;
-	#config: EngineConfig;
 	#runner: Runner;
 	#actors: Map<string, ActorHandler> = new Map();
 	#actorRouter: ActorRouter;
@@ -73,17 +71,15 @@ export class EngineActorDriver implements ActorDriver {
 		runConfig: RunnerConfig,
 		managerDriver: ManagerDriver,
 		inlineClient: Client<any>,
-		config: EngineConfig,
 	) {
 		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 ?? config.token;
+		const token = runConfig.token ?? runConfig.token;
 		if (token && runConfig.inspector && runConfig.inspector.enabled) {
 			runConfig.inspector.token = () => token;
 		}
@@ -98,12 +94,12 @@ export class EngineActorDriver implements ActorDriver {
 		let hasDisconnected = false;
 		const engineRunnerConfig: EngineRunnerConfig = {
 			version: this.#version,
-			endpoint: getEndpoint(config),
+			endpoint: getEndpoint(runConfig),
 			token,
-			namespace: runConfig.namespace ?? config.namespace,
-			totalSlots: runConfig.totalSlots ?? config.totalSlots,
-			runnerName: runConfig.runnerName ?? config.runnerName,
-			runnerKey: config.runnerKey,
+			namespace: runConfig.namespace ?? runConfig.namespace,
+			totalSlots: runConfig.totalSlots ?? runConfig.totalSlots,
+			runnerName: runConfig.runnerName ?? runConfig.runnerName,
+			runnerKey: runConfig.runnerKey,
 			metadata: {
 				inspectorToken: this.#runConfig.inspector.token(),
 			},
@@ -117,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,
 					});
 				}
 
@@ -133,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;
 			},
@@ -153,9 +149,9 @@ export class EngineActorDriver implements ActorDriver {
 		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,
 		});
 	}
 
@@ -236,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
@@ -262,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

packages/rivetkit/src/drivers/engine/mod.ts

@@ -13,11 +13,7 @@ export {
 	EngingConfigSchema as ConfigSchema,
 } from "./config";
 
-export function createEngineDriver(
-	inputConfig?: EngineConfigInput,
-): DriverConfig {
-	const config = EngingConfigSchema.parse(inputConfig);
-
+export function createEngineDriver(): DriverConfig {
 	return {
 		name: "engine",
 		manager: (_registryConfig, runConfig) => {
@@ -34,7 +30,6 @@ export function createEngineDriver(
 				runConfig,
 				managerDriver,
 				inlineClient,
-				config,
 			);
 		},
 	};

packages/rivetkit/src/manager/router-schema.ts

@@ -0,0 +1,20 @@
+import { z } from "zod";
+
+export const ServerlessStartHeadersSchema = z.object({
+	endpoint: z.string({ required_error: "x-rivet-endpoint header is required" }),
+	token: z
+		.string({ invalid_type_error: "x-rivet-token header must be a string" })
+		.optional(),
+	totalSlots: z.coerce
+		.number({
+			invalid_type_error: "x-rivet-total-slots header must be a number",
+		})
+		.int("x-rivet-total-slots header must be an integer")
+		.gte(1, "x-rivet-total-slots header must be positive"),
+	runnerName: z.string({
+		required_error: "x-rivet-runner-name header is required",
+	}),
+	namespace: z.string({
+		required_error: "x-rivet-namespace-id header is required",
+	}),
+});