rivetkit 2.0.3 → 2.0.4

package/src/driver-test-suite/tests/request-access.ts

@@ -114,131 +114,117 @@ export function runRequestAccessTests(driverTestConfig: DriverTestConfig) {
 			await connection.dispose();
 		});
 
-		test("should have access to request object in onAuth", async (c) => {
-			const { client, endpoint } = await setupDriverTest(c, driverTestConfig);
-
-			// Only test in HTTP mode as onAuth only runs for public endpoints
-			if (driverTestConfig.clientType === "http") {
-				// For now, skip this test as onAuth might not be properly invoked in test environment
-				// The onAuth hook is designed for public endpoints and might require special setup
-				console.log("Skipping onAuth test - requires public endpoint setup");
-
-				// TODO: Implement proper public endpoint test for onAuth
-				// This would require setting up the actor with public access and making
-				// requests from outside the internal client
-			}
-		});
-
-		test("should have access to request object in onFetch", async (c) => {
-			const { client, endpoint } = await setupDriverTest(c, driverTestConfig);
-
-			// Create actor
-			const handle = client.requestAccessActor.getOrCreate(["test-fetch"]);
-
-			// Make a raw HTTP request to the actor
-			await handle.resolve(); // Ensure actor is created
-
-			const actorQuery = {
-				getOrCreateForKey: {
-					name: "requestAccessActor",
-					key: ["test-fetch"],
-				},
-			};
-
-			const url = `${endpoint}/registry/actors/raw/http/test-path`;
-			const response = await fetch(url, {
-				method: "POST",
-				headers: {
-					"Content-Type": "application/json",
-					"X-Test-Header": "test-value",
-					"X-RivetKit-Query": JSON.stringify(actorQuery),
-				},
-				body: JSON.stringify({ test: "data" }),
-			});
-
-			if (!response.ok) {
-				const errorText = await response.text();
-				console.error(
-					`HTTP request failed: ${response.status} ${response.statusText}`,
-					errorText,
-				);
-			}
-
-			expect(response.ok).toBe(true);
-			const data = await response.json();
-
-			// Verify request info from onFetch
-			expect((data as any).hasRequest).toBe(true);
-			expect((data as any).requestUrl).toContain("/test-path");
-			expect((data as any).requestMethod).toBe("POST");
-			expect((data as any).requestHeaders).toBeDefined();
-			expect((data as any).requestHeaders["content-type"]).toBe(
-				"application/json",
-			);
-			expect((data as any).requestHeaders["x-test-header"]).toBe("test-value");
-		});
-
-		test("should have access to request object in onWebSocket", async (c) => {
-			const { client, endpoint } = await setupDriverTest(c, driverTestConfig);
-
-			// Only test in environments that support WebSocket
-			if (typeof WebSocket !== "undefined") {
-				// Create actor
-				const handle = client.requestAccessActor.getOrCreate([
-					"test-websocket",
-				]);
-				await handle.resolve(); // Ensure actor is created
-
-				const actorQuery = {
-					getOrCreateForKey: {
-						name: "requestAccessActor",
-						key: ["test-websocket"],
-					},
-				};
-
-				// Encode query as WebSocket subprotocol
-				const queryProtocol = `query.${encodeURIComponent(JSON.stringify(actorQuery))}`;
-
-				// Create raw WebSocket connection
-				const wsUrl = endpoint
-					.replace("http://", "ws://")
-					.replace("https://", "wss://");
-				const ws = new WebSocket(
-					`${wsUrl}/registry/actors/raw/websocket/test-path`,
-					[
-						queryProtocol,
-						"rivetkit", // Required protocol
-					],
-				);
-
-				// Wait for connection and first message
-				await new Promise<void>((resolve, reject) => {
-					ws.onopen = () => {
-						// Connection established
-					};
-
-					ws.onmessage = (event) => {
-						try {
-							const data = JSON.parse(event.data);
-
-							// Verify request info from onWebSocket
-							expect(data.hasRequest).toBe(true);
-							expect(data.requestUrl).toContain("/test-path");
-							expect(data.requestMethod).toBe("GET");
-							expect(data.requestHeaders).toBeDefined();
-
-							ws.close();
-							resolve();
-						} catch (error) {
-							reject(error);
-						}
-					};
-
-					ws.onerror = (error) => {
-						reject(error);
-					};
-				});
-			}
-		});
+		// TODO: re-expose this once we can have actor queries on the gateway
+		// test("should have access to request object in onFetch", async (c) => {
+		// 	const { client, endpoint } = await setupDriverTest(c, driverTestConfig);
+		//
+		// 	// Create actor
+		// 	const handle = client.requestAccessActor.getOrCreate(["test-fetch"]);
+		//
+		// 	// Make a raw HTTP request to the actor
+		// 	await handle.resolve(); // Ensure actor is created
+		//
+		// 	const actorQuery = {
+		// 		getOrCreateForKey: {
+		// 			name: "requestAccessActor",
+		// 			key: ["test-fetch"],
+		// 		},
+		// 	};
+		//
+		// 	const url = `${endpoint}/registry/actors/raw/http/test-path`;
+		// 	const response = await fetch(url, {
+		// 		method: "POST",
+		// 		headers: {
+		// 			"Content-Type": "application/json",
+		// 			"X-Test-Header": "test-value",
+		// 			"X-RivetKit-Query": JSON.stringify(actorQuery),
+		// 		},
+		// 		body: JSON.stringify({ test: "data" }),
+		// 	});
+		//
+		// 	if (!response.ok) {
+		// 		const errorText = await response.text();
+		// 		console.error(
+		// 			`HTTP request failed: ${response.status} ${response.statusText}`,
+		// 			errorText,
+		// 		);
+		// 	}
+		//
+		// 	expect(response.ok).toBe(true);
+		// 	const data = await response.json();
+		//
+		// 	// Verify request info from onFetch
+		// 	expect((data as any).hasRequest).toBe(true);
+		// 	expect((data as any).requestUrl).toContain("/test-path");
+		// 	expect((data as any).requestMethod).toBe("POST");
+		// 	expect((data as any).requestHeaders).toBeDefined();
+		// 	expect((data as any).requestHeaders["content-type"]).toBe(
+		// 		"application/json",
+		// 	);
+		// 	expect((data as any).requestHeaders["x-test-header"]).toBe("test-value");
+		// });
+
+		// test("should have access to request object in onWebSocket", async (c) => {
+		// 	const { client, endpoint } = await setupDriverTest(c, driverTestConfig);
+		//
+		// 	// Only test in environments that support WebSocket
+		// 	if (typeof WebSocket !== "undefined") {
+		// 		// Create actor
+		// 		const handle = client.requestAccessActor.getOrCreate([
+		// 			"test-websocket",
+		// 		]);
+		// 		await handle.resolve(); // Ensure actor is created
+		//
+		// 		const actorQuery = {
+		// 			getOrCreateForKey: {
+		// 				name: "requestAccessActor",
+		// 				key: ["test-websocket"],
+		// 			},
+		// 		};
+		//
+		// 		// Encode query as WebSocket subprotocol
+		// 		const queryProtocol = `query.${encodeURIComponent(JSON.stringify(actorQuery))}`;
+		//
+		// 		// Create raw WebSocket connection
+		// 		const wsUrl = endpoint
+		// 			.replace("http://", "ws://")
+		// 			.replace("https://", "wss://");
+		// 		const ws = new WebSocket(
+		// 			`${wsUrl}/registry/actors/raw/websocket/test-path`,
+		// 			[
+		// 				queryProtocol,
+		// 				"rivetkit", // Required protocol
+		// 			],
+		// 		);
+		//
+		// 		// Wait for connection and first message
+		// 		await new Promise<void>((resolve, reject) => {
+		// 			ws.onopen = () => {
+		// 				// Connection established
+		// 			};
+		//
+		// 			ws.onmessage = (event) => {
+		// 				try {
+		// 					const data = JSON.parse(event.data);
+		//
+		// 					// Verify request info from onWebSocket
+		// 					expect(data.hasRequest).toBe(true);
+		// 					expect(data.requestUrl).toContain("/test-path");
+		// 					expect(data.requestMethod).toBe("GET");
+		// 					expect(data.requestHeaders).toBeDefined();
+		//
+		// 					ws.close();
+		// 					resolve();
+		// 				} catch (error) {
+		// 					reject(error);
+		// 				}
+		// 			};
+		//
+		// 			ws.onerror = (error) => {
+		// 				reject(error);
+		// 			};
+		// 		});
+		// 	}
+		// });
 	});
 }

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

@@ -1,11 +1,9 @@
 import { resolve } from "node:path";
 import { type TestContext, vi } from "vitest";
 import { assertUnreachable } from "@/actor/utils";
-import { createClientWithDriver } from "@/client/client";
 import { type Client, createClient } from "@/client/mod";
 import type { registry } from "../../fixtures/driver-test-suite/registry";
 import type { DriverTestConfig } from "./mod";
-import { createTestInlineClientDriver } from "./test-inline-client-driver";
 
 export const FAKE_TIME = new Date("2024-01-01T00:00:00.000Z");
 
@@ -24,22 +22,27 @@ export async function setupDriverTest(
 
 	// Build drivers
 	const projectPath = resolve(__dirname, "../../fixtures/driver-test-suite");
-	const { endpoint, cleanup } = await driverTestConfig.start(projectPath);
+	const { endpoint, namespace, runnerName, cleanup } =
+		await driverTestConfig.start(projectPath);
 	c.onTestFinished(cleanup);
 
 	let client: Client<typeof registry>;
 	if (driverTestConfig.clientType === "http") {
 		// Create client
-		client = createClient<typeof registry>(endpoint, {
+		client = createClient<typeof registry>({
+			endpoint,
+			namespace,
+			runnerName,
 			transport: driverTestConfig.transport,
 		});
 	} else if (driverTestConfig.clientType === "inline") {
-		// Use inline client from driver
-		const clientDriver = createTestInlineClientDriver(
-			endpoint,
-			driverTestConfig.transport ?? "websocket",
-		);
-		client = createClientWithDriver(clientDriver);
+		throw "TODO";
+		// // Use inline client from driver
+		// const clientDriver = createTestInlineClientDriver(
+		// 	endpoint,
+		// 	driverTestConfig.transport ?? "websocket",
+		// );
+		// client = createClientWithDriver(clientDriver);
 	} else {
 		assertUnreachable(driverTestConfig.clientType);
 	}

package/src/drivers/default.ts

@@ -1,5 +1,5 @@
 import { UserError } from "@/actor/errors";
-import { logger } from "@/actor/log";
+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";
@@ -9,7 +9,7 @@ import { getEnvUniversal } from "@/utils";
  * Chooses the appropriate driver based on the run configuration.
  */
 export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
-	const engineEndpoint = runConfig.engine || getEnvUniversal("RIVET_ENGINE");
+	const engineEndpoint = runConfig.endpoint ?? getEnvUniversal("RIVET_ENGINE");
 
 	if (engineEndpoint && runConfig.driver) {
 		throw new UserError(
@@ -22,10 +22,13 @@ export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
 	}
 
 	if (engineEndpoint) {
-		logger().debug("using rivet engine driver", { endpoint: engineEndpoint });
+		loggerWithoutContext().debug({
+			msg: "using rivet engine driver",
+			endpoint: engineEndpoint,
+		});
 		return createEngineDriver({ endpoint: engineEndpoint });
 	}
 
-	logger().debug("using default file system driver");
+	loggerWithoutContext().debug({ msg: "using default file system driver" });
 	return createFileSystemOrMemoryDriver(true);
 }

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

@@ -6,8 +6,10 @@ import { Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
 import { WSContext } from "hono/ws";
 import invariant from "invariant";
+import { deserializeActorKey } from "@/actor/keys";
 import { EncodingSchema } from "@/actor/protocol/serde";
 import type { Client } from "@/client/client";
+import { getLogger } from "@/common/log";
 import {
 	type ActorDriver,
 	type AnyActorInstance,
@@ -35,7 +37,6 @@ import {
 	PATH_RAW_WEBSOCKET_PREFIX,
 } from "@/mod";
 import type { Config } from "./config";
-import { deserializeActorKey } from "./keys";
 import { KEYS } from "./kv";
 import { logger } from "./log";
 
@@ -94,34 +95,40 @@ export class EngineActorDriver implements ActorDriver {
 			),
 			onConnected: () => {
 				if (hasDisconnected) {
-					logger().info("runner reconnected", {
+					logger().info({
+						msg: "runner reconnected",
 						namespace: this.#config.namespace,
 						runnerName: this.#config.runnerName,
 					});
 				} else {
-					logger().debug("runner connected", {
+					logger().debug({
+						msg: "runner connected",
 						namespace: this.#config.namespace,
 						runnerName: this.#config.runnerName,
 					});
 				}
 			},
 			onDisconnected: () => {
-				logger().warn("runner disconnected", {
+				logger().warn({
+					msg: "runner disconnected",
 					namespace: this.#config.namespace,
 					runnerName: this.#config.runnerName,
 				});
 				hasDisconnected = true;
 			},
+			onShutdown: () => {},
 			fetch: this.#runnerFetch.bind(this),
 			websocket: this.#runnerWebSocket.bind(this),
 			onActorStart: this.#runnerOnActorStart.bind(this),
 			onActorStop: this.#runnerOnActorStop.bind(this),
+			logger: getLogger("engine-runner"),
 		};
 
 		// Create and start runner
 		this.#runner = new Runner(runnerConfig);
 		this.#runner.start();
-		logger().debug("engine runner started", {
+		logger().debug({
+			msg: "engine runner started",
 			endpoint: config.endpoint,
 			namespace: config.namespace,
 			runnerName: config.runnerName,
@@ -195,7 +202,8 @@ export class EngineActorDriver implements ActorDriver {
 		generation: number,
 		config: RunnerActorConfig,
 	): Promise<void> {
-		logger().debug("runner actor starting", {
+		logger().debug({
+			msg: "runner actor starting",
 			actorId,
 			name: config.name,
 			key: config.key,
@@ -248,11 +256,11 @@ export class EngineActorDriver implements ActorDriver {
 		handler.actorStartPromise?.resolve();
 		handler.actorStartPromise = undefined;
 
-		logger().debug("runner actor started", { actorId, name, key });
+		logger().debug({ msg: "runner actor started", actorId, name, key });
 	}
 
 	async #runnerOnActorStop(actorId: string, generation: number): Promise<void> {
-		logger().debug("runner actor stopping", { actorId, generation });
+		logger().debug({ msg: "runner actor stopping", actorId, generation });
 
 		const handler = this.#actors.get(actorId);
 		if (handler?.actor) {
@@ -260,11 +268,12 @@ export class EngineActorDriver implements ActorDriver {
 			this.#actors.delete(actorId);
 		}
 
-		logger().debug("runner actor stopped", { actorId });
+		logger().debug({ msg: "runner actor stopped", actorId });
 	}
 
 	async #runnerFetch(actorId: string, request: Request): Promise<Response> {
-		logger().debug("runner fetch", {
+		logger().debug({
+			msg: "runner fetch",
 			actorId,
 			url: request.url,
 			method: request.method,
@@ -279,7 +288,7 @@ export class EngineActorDriver implements ActorDriver {
 	): Promise<void> {
 		const websocket = websocketRaw as UniversalWebSocket;
 
-		logger().debug("runner websocket", { actorId, url: request.url });
+		logger().debug({ msg: "runner websocket", actorId, url: request.url });
 
 		const url = new URL(request.url);
 
@@ -324,7 +333,7 @@ export class EngineActorDriver implements ActorDriver {
 		const wsContext = new WSContext(websocket);
 
 		wsHandlerPromise.catch((err) => {
-			logger().error("building websocket handlers errored", { err });
+			logger().error({ msg: "building websocket handlers errored", err });
 			wsContext.close(1011, `${err}`);
 		});
 
@@ -354,7 +363,7 @@ export class EngineActorDriver implements ActorDriver {
 	}
 
 	async shutdown(immediate: boolean): Promise<void> {
-		logger().info("stopping engine actor driver");
+		logger().info({ msg: "stopping engine actor driver" });
 		await this.#runner.shutdown(immediate);
 	}
 }

package/src/drivers/engine/config.ts

@@ -8,7 +8,7 @@ export const ConfigSchema = z
 		endpoint: z
 			.string()
 			.default(
-				() => getEnvUniversal("RIVET_ENGINE") ?? "http://localhost:7080",
+				() => getEnvUniversal("RIVET_ENGINE") ?? "http://localhost:6420",
 			),
 		pegboardEndpoint: z.string().optional(),
 		namespace: z
@@ -17,21 +17,13 @@ export const ConfigSchema = z
 		runnerName: z
 			.string()
 			.default(() => getEnvUniversal("RIVET_RUNNER") ?? "rivetkit"),
-		// TODO: Automatically attempt ot determine key by common env vars (e.g. k8s pod name)
+		// TODO: Automatically attempt to determine key by common env vars (e.g. k8s pod name)
 		runnerKey: z
 			.string()
 			.default(
 				() => getEnvUniversal("RIVET_RUNNER_KEY") ?? crypto.randomUUID(),
 			),
 		totalSlots: z.number().default(100_000),
-		addresses: z
-			.record(
-				z.object({
-					host: z.string(),
-					port: z.number(),
-				}),
-			)
-			.default({ main: { host: "127.0.0.1", port: 5051 } }),
 	})
 	.default({});
 

package/src/drivers/engine/kv.ts

@@ -1,3 +1,3 @@
 export const KEYS = {
-	PERSIST_DATA: Uint8Array.from([1, 1]),
+	PERSIST_DATA: Uint8Array.from([1]),
 };

package/src/drivers/engine/log.ts

@@ -1,7 +1,5 @@
 import { getLogger } from "@/common/log";
 
-export const LOGGER_NAME = "driver-engine";
-
 export function logger() {
-	return getLogger(LOGGER_NAME);
+	return getLogger("driver-engine");
 }

package/src/drivers/engine/mod.ts

@@ -2,13 +2,12 @@ import type { Client } from "@/client/client";
 import type { ManagerDriver } from "@/manager/driver";
 import type { RegistryConfig } from "@/registry/config";
 import type { DriverConfig, RunConfig } from "@/registry/run-config";
+import { RemoteManagerDriver } from "@/remote-manager-driver/mod";
 import { EngineActorDriver } from "./actor-driver";
 import { ConfigSchema, type InputConfig } from "./config";
-import { EngineManagerDriver } from "./manager-driver";
 
 export { EngineActorDriver } from "./actor-driver";
 export { type Config, ConfigSchema, type InputConfig } from "./config";
-export { EngineManagerDriver } from "./manager-driver";
 
 export function createEngineDriver(inputConfig?: InputConfig): DriverConfig {
 	const config = ConfigSchema.parse(inputConfig);
@@ -16,7 +15,7 @@ export function createEngineDriver(inputConfig?: InputConfig): DriverConfig {
 	return {
 		name: "engine",
 		manager: (_registryConfig, runConfig) => {
-			return new EngineManagerDriver(config, runConfig);
+			return new RemoteManagerDriver(runConfig);
 		},
 		actor: (
 			registryConfig: RegistryConfig,

package/src/drivers/file-system/actor.ts

@@ -1,5 +1,5 @@
 import type { GenericConnGlobalState } from "@/actor/generic-conn-driver";
-import { logger } from "@/actor/log";
+import { loggerWithoutContext } from "@/actor/log";
 import type { AnyClient } from "@/client/client";
 import type {
 	ActorDriver,