rivetkit 2.0.2 → 2.0.4

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

@@ -0,0 +1,230 @@
+import { describe, expect, test } from "vitest";
+import type { DriverTestConfig } from "../mod";
+import { setupDriverTest } from "../utils";
+
+export function runRequestAccessTests(driverTestConfig: DriverTestConfig) {
+	describe("Request Access in Lifecycle Hooks", () => {
+		test("should have access to request object in onBeforeConnect and createConnState", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			// Create actor with request tracking enabled
+			const handle = client.requestAccessActor.getOrCreate(["test-request"], {
+				params: { trackRequest: true },
+			});
+			const connection = handle.connect();
+
+			// Get request info that was captured in onBeforeConnect
+			const requestInfo = await connection.getRequestInfo();
+
+			// Verify request was accessible in HTTP mode, but not in inline mode
+			if (driverTestConfig.clientType === "http") {
+				// Check onBeforeConnect
+				expect(requestInfo.onBeforeConnect.hasRequest).toBe(true);
+				expect(requestInfo.onBeforeConnect.requestUrl).toBeDefined();
+				expect(requestInfo.onBeforeConnect.requestMethod).toBeDefined();
+				expect(requestInfo.onBeforeConnect.requestHeaders).toBeDefined();
+
+				// Check createConnState
+				expect(requestInfo.createConnState.hasRequest).toBe(true);
+				expect(requestInfo.createConnState.requestUrl).toBeDefined();
+				expect(requestInfo.createConnState.requestMethod).toBeDefined();
+				expect(requestInfo.createConnState.requestHeaders).toBeDefined();
+			} else {
+				// Inline client may or may not have request object depending on the driver
+				//
+				// e.g.
+				// - File system does not have a request for inline requests
+				// - Rivet Engine proxies the request so it has access to the request object
+			}
+
+			// Clean up
+			await connection.dispose();
+		});
+
+		test("should not have request when trackRequest is false", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			// Create actor without request tracking
+			const handle = client.requestAccessActor.getOrCreate(
+				["test-no-request"],
+				{
+					params: { trackRequest: false },
+				},
+			);
+			const connection = handle.connect();
+
+			// Get request info
+			const requestInfo = await connection.getRequestInfo();
+
+			// Verify request was not tracked
+			expect(requestInfo.onBeforeConnect.hasRequest).toBe(false);
+			expect(requestInfo.onBeforeConnect.requestUrl).toBeNull();
+			expect(requestInfo.onBeforeConnect.requestMethod).toBeNull();
+			expect(
+				Object.keys(requestInfo.onBeforeConnect.requestHeaders),
+			).toHaveLength(0);
+
+			expect(requestInfo.createConnState.hasRequest).toBe(false);
+			expect(requestInfo.createConnState.requestUrl).toBeNull();
+			expect(requestInfo.createConnState.requestMethod).toBeNull();
+			expect(
+				Object.keys(requestInfo.createConnState.requestHeaders),
+			).toHaveLength(0);
+
+			// Clean up
+			await connection.dispose();
+		});
+
+		test("should capture request headers and method", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			// Create actor and connect with request tracking
+			const handle = client.requestAccessActor.getOrCreate(["test-headers"], {
+				params: { trackRequest: true },
+			});
+			const connection = handle.connect();
+
+			// Get request info
+			const requestInfo = await connection.getRequestInfo();
+
+			if (driverTestConfig.clientType === "http") {
+				// Verify request details were captured in both hooks
+				expect(requestInfo.onBeforeConnect.hasRequest).toBe(true);
+				expect(requestInfo.onBeforeConnect.requestMethod).toBeTruthy();
+				expect(requestInfo.onBeforeConnect.requestUrl).toBeTruthy();
+				expect(requestInfo.onBeforeConnect.requestHeaders).toBeTruthy();
+				expect(typeof requestInfo.onBeforeConnect.requestHeaders).toBe(
+					"object",
+				);
+
+				expect(requestInfo.createConnState.hasRequest).toBe(true);
+				expect(requestInfo.createConnState.requestMethod).toBeTruthy();
+				expect(requestInfo.createConnState.requestUrl).toBeTruthy();
+				expect(requestInfo.createConnState.requestHeaders).toBeTruthy();
+				expect(typeof requestInfo.createConnState.requestHeaders).toBe(
+					"object",
+				);
+			} else {
+				// Inline client may or may not have request object depending on the driver
+				//
+				// See "should have access to request object in onBeforeConnect and createConnState"
+			}
+
+			// Clean up
+			await connection.dispose();
+		});
+
+		// 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

@@ -0,0 +1,71 @@
+import { resolve } from "node:path";
+import { type TestContext, vi } from "vitest";
+import { assertUnreachable } from "@/actor/utils";
+import { type Client, createClient } from "@/client/mod";
+import type { registry } from "../../fixtures/driver-test-suite/registry";
+import type { DriverTestConfig } from "./mod";
+
+export const FAKE_TIME = new Date("2024-01-01T00:00:00.000Z");
+
+// Must use `TestContext` since global hooks do not work when running concurrently
+export async function setupDriverTest(
+	c: TestContext,
+	driverTestConfig: DriverTestConfig,
+): Promise<{
+	client: Client<typeof registry>;
+	endpoint: string;
+}> {
+	if (!driverTestConfig.useRealTimers) {
+		vi.useFakeTimers();
+		vi.setSystemTime(FAKE_TIME);
+	}
+
+	// Build drivers
+	const projectPath = resolve(__dirname, "../../fixtures/driver-test-suite");
+	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,
+			namespace,
+			runnerName,
+			transport: driverTestConfig.transport,
+		});
+	} else if (driverTestConfig.clientType === "inline") {
+		throw "TODO";
+		// // Use inline client from driver
+		// const clientDriver = createTestInlineClientDriver(
+		// 	endpoint,
+		// 	driverTestConfig.transport ?? "websocket",
+		// );
+		// client = createClientWithDriver(clientDriver);
+	} else {
+		assertUnreachable(driverTestConfig.clientType);
+	}
+
+	// Cleanup client
+	if (!driverTestConfig.HACK_skipCleanupNet) {
+		c.onTestFinished(async () => await client.dispose());
+	}
+
+	return {
+		client,
+		endpoint,
+	};
+}
+
+export async function waitFor(
+	driverTestConfig: DriverTestConfig,
+	ms: number,
+): Promise<void> {
+	if (driverTestConfig.useRealTimers) {
+		return new Promise((resolve) => setTimeout(resolve, ms));
+	} else {
+		vi.advanceTimersByTime(ms);
+		return Promise.resolve();
+	}
+}

package/src/drivers/default.ts

@@ -0,0 +1,34 @@
+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 { getEnvUniversal } from "@/utils";
+
+/**
+ * Chooses the appropriate driver based on the run configuration.
+ */
+export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
+	const engineEndpoint = runConfig.endpoint ?? getEnvUniversal("RIVET_ENGINE");
+
+	if (engineEndpoint && runConfig.driver) {
+		throw new UserError(
+			"Cannot specify both 'engine' and 'driver' in configuration",
+		);
+	}
+
+	if (runConfig.driver) {
+		return runConfig.driver;
+	}
+
+	if (engineEndpoint) {
+		loggerWithoutContext().debug({
+			msg: "using rivet engine driver",
+			endpoint: engineEndpoint,
+		});
+		return createEngineDriver({ endpoint: engineEndpoint });
+	}
+
+	loggerWithoutContext().debug({ msg: "using default file system driver" });
+	return createFileSystemOrMemoryDriver(true);
+}

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

@@ -0,0 +1,369 @@
+import type {
+	ActorConfig as RunnerActorConfig,
+	RunnerConfig,
+} from "@rivetkit/engine-runner";
+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,
+	HEADER_AUTH_DATA,
+	HEADER_CONN_PARAMS,
+	HEADER_ENCODING,
+	type ManagerDriver,
+	serializeEmptyPersistData,
+} from "@/driver-helpers/mod";
+import type {
+	ActorRouter,
+	RegistryConfig,
+	RunConfig,
+	UniversalWebSocket,
+	UpgradeWebSocketArgs,
+} from "@/mod";
+import {
+	createActorRouter,
+	createGenericConnDrivers,
+	GenericConnGlobalState,
+	handleRawWebSocketHandler,
+	handleWebSocketConnect,
+	lookupInRegistry,
+	PATH_CONNECT_WEBSOCKET,
+	PATH_RAW_WEBSOCKET_PREFIX,
+} from "@/mod";
+import type { Config } from "./config";
+import { KEYS } from "./kv";
+import { logger } from "./log";
+
+interface ActorHandler {
+	actor?: AnyActorInstance;
+	actorStartPromise?: PromiseWithResolvers<void>;
+	genericConnGlobalState: GenericConnGlobalState;
+	persistedData?: Uint8Array;
+}
+
+export type DriverContext = {};
+
+export class EngineActorDriver implements ActorDriver {
+	#registryConfig: RegistryConfig;
+	#runConfig: RunConfig;
+	#managerDriver: ManagerDriver;
+	#inlineClient: Client<any>;
+	#config: Config;
+	#runner: Runner;
+	#actors: Map<string, ActorHandler> = new Map();
+	#actorRouter: ActorRouter;
+	#version: number = 1; // Version for the runner protocol
+
+	constructor(
+		registryConfig: RegistryConfig,
+		runConfig: RunConfig,
+		managerDriver: ManagerDriver,
+		inlineClient: Client<any>,
+		config: Config,
+	) {
+		this.#registryConfig = registryConfig;
+		this.#runConfig = runConfig;
+		this.#managerDriver = managerDriver;
+		this.#inlineClient = inlineClient;
+		this.#config = config;
+		this.#actorRouter = createActorRouter(runConfig, this);
+
+		// Create runner configuration
+		let hasDisconnected = false;
+		const runnerConfig: RunnerConfig = {
+			version: this.#version,
+			endpoint: config.endpoint,
+			pegboardEndpoint: config.pegboardEndpoint,
+			namespace: config.namespace,
+			totalSlots: config.totalSlots,
+			runnerName: config.runnerName,
+			runnerKey: config.runnerKey,
+			metadata: {
+				inspectorToken: this.#runConfig.inspector.token(),
+			},
+			prepopulateActorNames: Object.fromEntries(
+				Object.keys(this.#registryConfig.use).map((name) => [
+					name,
+					{ metadata: {} },
+				]),
+			),
+			onConnected: () => {
+				if (hasDisconnected) {
+					logger().info({
+						msg: "runner reconnected",
+						namespace: this.#config.namespace,
+						runnerName: this.#config.runnerName,
+					});
+				} else {
+					logger().debug({
+						msg: "runner connected",
+						namespace: this.#config.namespace,
+						runnerName: this.#config.runnerName,
+					});
+				}
+			},
+			onDisconnected: () => {
+				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({
+			msg: "engine runner started",
+			endpoint: config.endpoint,
+			namespace: config.namespace,
+			runnerName: config.runnerName,
+		});
+	}
+
+	async #loadActorHandler(actorId: string): Promise<ActorHandler> {
+		// Check if actor is already loaded
+		const handler = this.#actors.get(actorId);
+		if (!handler) throw new Error(`Actor handler does not exist ${actorId}`);
+		if (handler.actorStartPromise) await handler.actorStartPromise.promise;
+		if (!handler.actor) throw new Error("Actor should be loaded");
+		return handler;
+	}
+
+	async loadActor(actorId: string): Promise<AnyActorInstance> {
+		const handler = await this.#loadActorHandler(actorId);
+		if (!handler.actor) throw new Error(`Actor ${actorId} failed to load`);
+		return handler.actor;
+	}
+
+	getGenericConnGlobalState(actorId: string): GenericConnGlobalState {
+		const handler = this.#actors.get(actorId);
+		if (!handler) {
+			throw new Error(`Actor ${actorId} not loaded`);
+		}
+		return handler.genericConnGlobalState;
+	}
+
+	getContext(actorId: string): DriverContext {
+		return {};
+	}
+
+	async readPersistedData(actorId: string): Promise<Uint8Array | undefined> {
+		const handler = this.#actors.get(actorId);
+		if (!handler) throw new Error(`Actor ${actorId} not loaded`);
+		if (handler.persistedData) return handler.persistedData;
+
+		const [value] = await this.#runner.kvGet(actorId, [KEYS.PERSIST_DATA]);
+
+		if (value !== null) {
+			handler.persistedData = value;
+			return value;
+		} else {
+			return undefined;
+		}
+	}
+
+	async writePersistedData(actorId: string, data: Uint8Array): Promise<void> {
+		const handler = this.#actors.get(actorId);
+		if (!handler) throw new Error(`Actor ${actorId} not loaded`);
+
+		handler.persistedData = data;
+
+		await this.#runner.kvPut(actorId, [[KEYS.PERSIST_DATA, data]]);
+	}
+
+	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
+		// TODO: Set timeout
+		// TODO: Use alarm on sleep
+		// TODO: Send alarm to runner
+	}
+
+	async getDatabase(_actorId: string): Promise<unknown | undefined> {
+		return undefined;
+	}
+
+	// Runner lifecycle callbacks
+	async #runnerOnActorStart(
+		actorId: string,
+		generation: number,
+		config: RunnerActorConfig,
+	): Promise<void> {
+		logger().debug({
+			msg: "runner actor starting",
+			actorId,
+			name: config.name,
+			key: config.key,
+			generation,
+		});
+
+		// Deserialize input
+		let input: any;
+		if (config.input) {
+			input = cbor.decode(config.input);
+		}
+
+		// Get or create handler
+		let handler = this.#actors.get(actorId);
+		if (!handler) {
+			handler = {
+				genericConnGlobalState: new GenericConnGlobalState(),
+				actorStartPromise: Promise.withResolvers(),
+				persistedData: serializeEmptyPersistData(input),
+			};
+			this.#actors.set(actorId, handler);
+		}
+
+		const name = config.name as string;
+		invariant(config.key, "actor should have a key");
+		const key = deserializeActorKey(config.key);
+
+		// Create actor instance
+		const definition = lookupInRegistry(
+			this.#registryConfig,
+			config.name as string, // TODO: Remove cast
+		);
+		handler.actor = definition.instantiate();
+
+		// Start actor
+		const connDrivers = createGenericConnDrivers(
+			handler.genericConnGlobalState,
+		);
+		await handler.actor.start(
+			connDrivers,
+			this,
+			this.#inlineClient,
+			actorId,
+			name,
+			key,
+			"unknown", // TODO: Add regions
+		);
+
+		// Resolve promise if waiting
+		handler.actorStartPromise?.resolve();
+		handler.actorStartPromise = undefined;
+
+		logger().debug({ msg: "runner actor started", actorId, name, key });
+	}
+
+	async #runnerOnActorStop(actorId: string, generation: number): Promise<void> {
+		logger().debug({ msg: "runner actor stopping", actorId, generation });
+
+		const handler = this.#actors.get(actorId);
+		if (handler?.actor) {
+			await handler.actor._stop();
+			this.#actors.delete(actorId);
+		}
+
+		logger().debug({ msg: "runner actor stopped", actorId });
+	}
+
+	async #runnerFetch(actorId: string, request: Request): Promise<Response> {
+		logger().debug({
+			msg: "runner fetch",
+			actorId,
+			url: request.url,
+			method: request.method,
+		});
+		return await this.#actorRouter.fetch(request, { actorId });
+	}
+
+	async #runnerWebSocket(
+		actorId: string,
+		websocketRaw: any,
+		request: Request,
+	): Promise<void> {
+		const websocket = websocketRaw as UniversalWebSocket;
+
+		logger().debug({ msg: "runner websocket", actorId, url: request.url });
+
+		const url = new URL(request.url);
+
+		// Parse headers
+		const encodingRaw = request.headers.get(HEADER_ENCODING);
+		const connParamsRaw = request.headers.get(HEADER_CONN_PARAMS);
+		const authDataRaw = request.headers.get(HEADER_AUTH_DATA);
+
+		const encoding = EncodingSchema.parse(encodingRaw);
+		const connParams = connParamsRaw ? JSON.parse(connParamsRaw) : undefined;
+		const authData = authDataRaw ? JSON.parse(authDataRaw) : undefined;
+
+		// Fetch WS handler
+		//
+		// We store the promise since we need to add WebSocket event listeners immediately that will wait for the promise to resolve
+		let wsHandlerPromise: Promise<UpgradeWebSocketArgs>;
+		if (url.pathname === PATH_CONNECT_WEBSOCKET) {
+			wsHandlerPromise = handleWebSocketConnect(
+				request,
+				this.#runConfig,
+				this,
+				actorId,
+				encoding,
+				connParams,
+				authData,
+			);
+		} else if (url.pathname.startsWith(PATH_RAW_WEBSOCKET_PREFIX)) {
+			wsHandlerPromise = handleRawWebSocketHandler(
+				request,
+				url.pathname + url.search,
+				this,
+				actorId,
+				authData,
+			);
+		} else {
+			throw new Error(`Unreachable path: ${url.pathname}`);
+		}
+
+		// TODO: Add close
+
+		// Connect the Hono WS hook to the adapter
+		const wsContext = new WSContext(websocket);
+
+		wsHandlerPromise.catch((err) => {
+			logger().error({ msg: "building websocket handlers errored", err });
+			wsContext.close(1011, `${err}`);
+		});
+
+		if (websocket.readyState === 1) {
+			wsHandlerPromise.then((x) => x.onOpen?.(new Event("open"), wsContext));
+		} else {
+			websocket.addEventListener("open", (event) => {
+				wsHandlerPromise.then((x) => x.onOpen?.(event, wsContext));
+			});
+		}
+
+		websocket.addEventListener("message", (event) => {
+			wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));
+		});
+
+		websocket.addEventListener("close", (event) => {
+			wsHandlerPromise.then((x) => x.onClose?.(event, wsContext));
+		});
+
+		websocket.addEventListener("error", (event) => {
+			wsHandlerPromise.then((x) => x.onError?.(event, wsContext));
+		});
+	}
+
+	async sleep(actorId: string) {
+		this.#runner.sleepActor(actorId);
+	}
+
+	async shutdown(immediate: boolean): Promise<void> {
+		logger().info({ msg: "stopping engine actor driver" });
+		await this.#runner.shutdown(immediate);
+	}
+}