rivetkit 2.0.1 → 2.0.3

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

@@ -0,0 +1,391 @@
+import * as cbor from "cbor-x";
+import type { Context as HonoContext } from "hono";
+import invariant from "invariant";
+import { ActorAlreadyExists } from "@/actor/errors";
+import {
+	HEADER_AUTH_DATA,
+	HEADER_CONN_PARAMS,
+	HEADER_ENCODING,
+	HEADER_EXPOSE_INTERNAL_ERROR,
+} from "@/actor/router-endpoints";
+import { generateRandomString } from "@/actor/utils";
+import { importWebSocket } from "@/common/websocket";
+import type {
+	ActorOutput,
+	CreateInput,
+	GetForIdInput,
+	GetOrCreateWithKeyInput,
+	GetWithKeyInput,
+	ManagerDriver,
+} from "@/driver-helpers/mod";
+import { type Encoding, noopNext, type RunConfig } from "@/mod";
+import {
+	createActor,
+	destroyActor,
+	getActor,
+	getActorById,
+	getOrCreateActorById,
+} from "./api-endpoints";
+import { EngineApiError } from "./api-utils";
+import type { Config } from "./config";
+import { deserializeActorKey, serializeActorKey } from "./keys";
+import { logger } from "./log";
+import { createWebSocketProxy } from "./ws-proxy";
+
+export class EngineManagerDriver implements ManagerDriver {
+	#config: Config;
+	#runConfig: RunConfig;
+	#importWebSocketPromise: Promise<typeof WebSocket>;
+
+	constructor(config: Config, runConfig: RunConfig) {
+		this.#config = config;
+		this.#runConfig = runConfig;
+		if (!this.#runConfig.inspector.token()) {
+			const token = generateRandomString();
+			this.#runConfig.inspector.token = () => token;
+		}
+		this.#importWebSocketPromise = importWebSocket();
+	}
+
+	async sendRequest(actorId: string, actorRequest: Request): Promise<Response> {
+		logger().debug("sending request to actor via guard", {
+			actorId,
+			method: actorRequest.method,
+			url: actorRequest.url,
+		});
+
+		return this.#forwardHttpRequest(actorRequest, actorId);
+	}
+
+	async openWebSocket(
+		path: string,
+		actorId: string,
+		encoding: Encoding,
+		params: unknown,
+	): Promise<WebSocket> {
+		const WebSocket = await this.#importWebSocketPromise;
+
+		// WebSocket connections go through guard
+		const guardUrl = `${this.#config.endpoint}${path}`;
+
+		logger().debug("opening websocket to actor via guard", {
+			actorId,
+			path,
+			guardUrl,
+		});
+
+		// Create WebSocket connection
+		const ws = new WebSocket(guardUrl, {
+			headers: buildGuardHeadersForWebSocket(actorId, encoding, params),
+		});
+
+		logger().debug("websocket connection opened", { actorId });
+
+		return ws;
+	}
+
+	async proxyRequest(
+		_c: HonoContext,
+		actorRequest: Request,
+		actorId: string,
+	): Promise<Response> {
+		logger().debug("forwarding request to actor via guard", {
+			actorId,
+			method: actorRequest.method,
+			url: actorRequest.url,
+			hasBody: !!actorRequest.body,
+		});
+
+		return this.#forwardHttpRequest(actorRequest, actorId);
+	}
+
+	async proxyWebSocket(
+		c: HonoContext,
+		path: string,
+		actorId: string,
+		encoding: Encoding,
+		params: unknown,
+		authData: unknown,
+	): Promise<Response> {
+		const upgradeWebSocket = this.#runConfig.getUpgradeWebSocket?.();
+		invariant(upgradeWebSocket, "missing getUpgradeWebSocket");
+
+		const guardUrl = `${this.#config.endpoint}${path}`;
+		const wsGuardUrl = guardUrl.replace("http://", "ws://");
+
+		logger().debug("forwarding websocket to actor via guard", {
+			actorId,
+			path,
+			guardUrl,
+		});
+
+		// Build headers
+		const headers = buildGuardHeadersForWebSocket(
+			actorId,
+			encoding,
+			params,
+			authData,
+		);
+		const args = await createWebSocketProxy(c, wsGuardUrl, headers);
+
+		return await upgradeWebSocket(() => args)(c, noopNext());
+	}
+
+	extraStartupLog() {
+		return {
+			engine: this.#config.endpoint,
+			namespace: this.#config.namespace,
+			runner: this.#config.runnerName,
+			address: Object.values(this.#config.addresses)
+				.map((v) => `${v.host}:${v.port}`)
+				.join(", "),
+		};
+	}
+
+	async getForId({
+		c,
+		name,
+		actorId,
+	}: GetForIdInput): Promise<ActorOutput | undefined> {
+		// Fetch from API if not in cache
+		try {
+			const response = await getActor(this.#config, actorId);
+
+			// Validate name matches
+			if (response.actor.name !== name) {
+				logger().debug("actor name mismatch from api", {
+					actorId,
+					apiName: response.actor.name,
+					requestedName: name,
+				});
+				return undefined;
+			}
+
+			const keyRaw = response.actor.key;
+			invariant(keyRaw, `actor ${actorId} should have key`);
+			const key = deserializeActorKey(keyRaw);
+
+			return {
+				actorId,
+				name,
+				key,
+			};
+		} catch (error) {
+			if (
+				error instanceof EngineApiError &&
+				(error as EngineApiError).group === "actor" &&
+				(error as EngineApiError).code === "not_found"
+			) {
+				return undefined;
+			}
+			throw error;
+		}
+	}
+
+	async getWithKey({
+		c,
+		name,
+		key,
+	}: GetWithKeyInput): Promise<ActorOutput | undefined> {
+		logger().debug("getWithKey: searching for actor", { name, key });
+
+		// If not in local cache, fetch by key from API
+		try {
+			const response = await getActorById(this.#config, name, key);
+
+			if (!response.actor_id) {
+				return undefined;
+			}
+
+			const actorId = response.actor_id;
+
+			logger().debug("getWithKey: found actor via api", {
+				actorId,
+				name,
+				key,
+			});
+
+			return {
+				actorId,
+				name,
+				key,
+			};
+		} catch (error) {
+			if (
+				error instanceof EngineApiError &&
+				(error as EngineApiError).group === "actor" &&
+				(error as EngineApiError).code === "not_found"
+			) {
+				return undefined;
+			}
+			throw error;
+		}
+	}
+
+	async getOrCreateWithKey(
+		input: GetOrCreateWithKeyInput,
+	): Promise<ActorOutput> {
+		const { c, name, key, input: actorInput, region } = input;
+
+		logger().info(
+			"getOrCreateWithKey: getting or creating actor via engine api",
+			{
+				name,
+				key,
+			},
+		);
+
+		const response = await getOrCreateActorById(this.#config, {
+			name,
+			key: serializeActorKey(key),
+			runner_name_selector: this.#config.runnerName,
+			input: input ? cbor.encode(actorInput).toString("base64") : undefined,
+			crash_policy: "sleep",
+		});
+
+		const actorId = response.actor_id;
+
+		logger().info("getOrCreateWithKey: actor ready", {
+			actorId,
+			name,
+			key,
+			created: response.created,
+		});
+
+		return {
+			actorId,
+			name,
+			key,
+		};
+	}
+
+	async createActor({
+		c,
+		name,
+		key,
+		input,
+	}: CreateInput): Promise<ActorOutput> {
+		// Check if actor with the same name and key already exists
+		const existingActor = await this.getWithKey({ c, name, key });
+		if (existingActor) {
+			throw new ActorAlreadyExists(name, key);
+		}
+
+		logger().info("creating actor via engine api", { name, key });
+
+		// Create actor via engine API
+		const result = await createActor(this.#config, {
+			name,
+			runner_name_selector: this.#config.runnerName,
+			key: serializeActorKey(key),
+			input: input ? cbor.encode(input).toString("base64") : null,
+			crash_policy: "sleep",
+		});
+		const actorId = result.actor.actor_id;
+
+		logger().info("actor created", { actorId, name, key });
+
+		return {
+			actorId,
+			name,
+			key,
+		};
+	}
+
+	async destroyActor(actorId: string): Promise<void> {
+		logger().info("destroying actor via engine api", { actorId });
+
+		await destroyActor(this.#config, actorId);
+
+		logger().info("actor destroyed", { actorId });
+	}
+
+	async #forwardHttpRequest(
+		actorRequest: Request,
+		actorId: string,
+	): Promise<Response> {
+		// Route through guard port
+		const url = new URL(actorRequest.url);
+		const guardUrl = `${this.#config.endpoint}${url.pathname}${url.search}`;
+
+		// Handle body properly based on method and presence
+		let bodyToSend: ArrayBuffer | null = null;
+		const guardHeaders = buildGuardHeadersForHttp(actorRequest, actorId);
+
+		if (
+			actorRequest.body &&
+			actorRequest.method !== "GET" &&
+			actorRequest.method !== "HEAD"
+		) {
+			if (actorRequest.bodyUsed) {
+				throw new Error("Request body has already been consumed");
+			}
+
+			// TODO: This buffers the entire request in memory every time. We
+			// need to properly implement streaming bodies.
+			// Clone and read the body to ensure it can be sent
+			const clonedRequest = actorRequest.clone();
+			bodyToSend = await clonedRequest.arrayBuffer();
+
+			// If this is a streaming request, we need to convert the headers
+			// for the basic array buffer
+			guardHeaders.delete("transfer-encoding");
+			guardHeaders.set(
+				"content-length",
+				String((bodyToSend as ArrayBuffer).byteLength),
+			);
+		}
+
+		const guardRequest = new Request(guardUrl, {
+			method: actorRequest.method,
+			headers: guardHeaders,
+			body: bodyToSend,
+		});
+
+		return mutableResponse(await fetch(guardRequest));
+	}
+}
+
+function mutableResponse(fetchRes: Response): Response {
+	// We cannot return the raw response from `fetch` since the response type is not mutable.
+	//
+	// In order for middleware to be able to mutate the response, we need to build a new Response object that is mutable.
+	return new Response(fetchRes.body, fetchRes);
+}
+
+function buildGuardHeadersForHttp(
+	actorRequest: Request,
+	actorId: string,
+): Headers {
+	const headers = new Headers();
+	// Copy all headers from the original request
+	for (const [key, value] of actorRequest.headers.entries()) {
+		headers.set(key, value);
+	}
+	// Add guard-specific headers
+	headers.set("x-rivet-target", "actor");
+	headers.set("x-rivet-actor", actorId);
+	headers.set("x-rivet-port", "main");
+	return headers;
+}
+
+function buildGuardHeadersForWebSocket(
+	actorId: string,
+	encoding: Encoding,
+	params?: unknown,
+	authData?: unknown,
+): Record<string, string> {
+	const headers: Record<string, string> = {};
+	headers["x-rivet-target"] = "actor";
+	headers["x-rivet-actor"] = actorId;
+	headers["x-rivet-port"] = "main";
+	headers[HEADER_EXPOSE_INTERNAL_ERROR] = "true";
+	headers[HEADER_ENCODING] = encoding;
+	if (params) {
+		headers[HEADER_CONN_PARAMS] = JSON.stringify(params);
+	}
+	if (authData) {
+		headers[HEADER_AUTH_DATA] = JSON.stringify(authData);
+	}
+	return headers;
+}

package/src/drivers/engine/mod.ts

@@ -0,0 +1,36 @@
+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 { 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);
+
+	return {
+		name: "engine",
+		manager: (_registryConfig, runConfig) => {
+			return new EngineManagerDriver(config, runConfig);
+		},
+		actor: (
+			registryConfig: RegistryConfig,
+			runConfig: RunConfig,
+			managerDriver: ManagerDriver,
+			inlineClient: Client<any>,
+		) => {
+			return new EngineActorDriver(
+				registryConfig,
+				runConfig,
+				managerDriver,
+				inlineClient,
+				config,
+			);
+		},
+	};
+}

package/src/drivers/engine/ws-proxy.ts

@@ -0,0 +1,170 @@
+import type { Context as HonoContext } from "hono";
+import type { WSContext } from "hono/ws";
+import invariant from "invariant";
+import type { CloseEvent } from "ws";
+import { importWebSocket } from "@/common/websocket";
+import type { UpgradeWebSocketArgs } from "@/mod";
+import { logger } from "./log";
+
+/**
+ * Returns Hono `upgradeWebSocket` args that will proxy requests from the client to a destination address.
+ */
+export async function createWebSocketProxy(
+	c: HonoContext,
+	targetUrl: string,
+	headers: Record<string, string>,
+): Promise<UpgradeWebSocketArgs> {
+	const WebSocket = await importWebSocket();
+
+	// HACK: Sanitize WebSocket-specific headers. If we don't do this, some WebSocket implementations (i.e. native WebSocket in Node.js) will fail to connect.
+	for (const [k, v] of c.req.raw.headers.entries()) {
+		if (!k.startsWith("sec-") && k !== "connection" && k !== "upgrade") {
+			headers[k] = v;
+		}
+	}
+
+	// WebSocket state
+	interface WsState {
+		targetWs?: WebSocket;
+		connectPromise?: Promise<void>;
+	}
+	const state: WsState = {};
+
+	return {
+		onOpen: async (event: any, clientWs: WSContext) => {
+			logger().debug("client websocket connected", { targetUrl });
+
+			if (clientWs.readyState !== 1) {
+				logger().warn("client websocket not open on connection", {
+					targetUrl,
+					readyState: clientWs.readyState,
+				});
+				return;
+			}
+
+			// Create WebSocket
+			const targetWs = new WebSocket(targetUrl, { headers });
+			state.targetWs = targetWs;
+
+			// Setup connection promise
+			state.connectPromise = new Promise<void>((resolve, reject) => {
+				targetWs.addEventListener("open", () => {
+					logger().debug("target websocket connected", { targetUrl });
+
+					if (clientWs.readyState !== 1) {
+						logger().warn("client websocket closed before target connected", {
+							targetUrl,
+							clientReadyState: clientWs.readyState,
+						});
+						targetWs.close(1001, "Client disconnected");
+						reject(new Error("Client disconnected"));
+						return;
+					}
+					resolve();
+				});
+
+				targetWs.addEventListener("error", (error) => {
+					logger().warn("target websocket error during connection", {
+						targetUrl,
+					});
+					reject(error);
+				});
+			});
+
+			// Setup bidirectional forwarding
+			state.targetWs.addEventListener("message", (event) => {
+				if (
+					typeof event.data === "string" ||
+					event.data instanceof ArrayBuffer
+				) {
+					clientWs.send(event.data);
+				} else if (event.data instanceof Blob) {
+					event.data.arrayBuffer().then((buffer) => {
+						clientWs.send(buffer);
+					});
+				}
+			});
+
+			state.targetWs.addEventListener("close", (event) => {
+				logger().debug("target websocket closed", {
+					targetUrl,
+					code: event.code,
+					reason: event.reason,
+				});
+				closeWebSocketIfOpen(clientWs, event.code, event.reason);
+			});
+
+			state.targetWs.addEventListener("error", (error) => {
+				logger().error("target websocket error", { targetUrl, error });
+				closeWebSocketIfOpen(clientWs, 1011, "Target WebSocket error");
+			});
+		},
+
+		onMessage: async (event: any, clientWs: WSContext) => {
+			if (!state.targetWs || !state.connectPromise) {
+				logger().error("websocket state not initialized", { targetUrl });
+				return;
+			}
+
+			try {
+				await state.connectPromise;
+				if (state.targetWs.readyState === WebSocket.OPEN) {
+					state.targetWs.send(event.data);
+				} else {
+					logger().warn("target websocket not open", {
+						targetUrl,
+						readyState: state.targetWs.readyState,
+					});
+				}
+			} catch (error) {
+				logger().error("failed to connect to target websocket", {
+					targetUrl,
+					error,
+				});
+				closeWebSocketIfOpen(clientWs, 1011, "Failed to connect to target");
+			}
+		},
+
+		onClose: (event: any, clientWs: WSContext) => {
+			logger().debug("client websocket closed", {
+				targetUrl,
+				code: event.code,
+				reason: event.reason,
+				wasClean: event.wasClean,
+			});
+
+			if (state.targetWs) {
+				if (
+					state.targetWs.readyState === WebSocket.OPEN ||
+					state.targetWs.readyState === WebSocket.CONNECTING
+				) {
+					state.targetWs.close(1000, event.reason || "Client disconnected");
+				}
+			}
+		},
+
+		onError: (event: any, clientWs: WSContext) => {
+			logger().error("client websocket error", { targetUrl, event });
+
+			if (state.targetWs) {
+				if (state.targetWs.readyState === WebSocket.OPEN) {
+					state.targetWs.close(1011, "Client WebSocket error");
+				} else if (state.targetWs.readyState === WebSocket.CONNECTING) {
+					state.targetWs.close();
+				}
+			}
+		},
+	};
+}
+
+function closeWebSocketIfOpen(
+	ws: WebSocket | WSContext,
+	code: number,
+	reason: string,
+): void {
+	if (ws.readyState === 1) {
+		ws.close(code, reason);
+	} else if ("close" in ws && (ws as WebSocket).readyState === WebSocket.OPEN) {
+		ws.close(code, reason);
+	}
+}

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

@@ -0,0 +1,91 @@
+import type { GenericConnGlobalState } from "@/actor/generic-conn-driver";
+import { logger } from "@/actor/log";
+import type { AnyClient } from "@/client/client";
+import type {
+	ActorDriver,
+	AnyActorInstance,
+	ManagerDriver,
+} from "@/driver-helpers/mod";
+import type { RegistryConfig, RunConfig } from "@/mod";
+import { bufferToArrayBuffer } from "@/utils";
+import type { FileSystemGlobalState } from "./global-state";
+
+export type ActorDriverContext = Record<never, never>;
+
+/**
+ * File System implementation of the Actor Driver
+ */
+export class FileSystemActorDriver implements ActorDriver {
+	#registryConfig: RegistryConfig;
+	#runConfig: RunConfig;
+	#managerDriver: ManagerDriver;
+	#inlineClient: AnyClient;
+	#state: FileSystemGlobalState;
+
+	constructor(
+		registryConfig: RegistryConfig,
+		runConfig: RunConfig,
+		managerDriver: ManagerDriver,
+		inlineClient: AnyClient,
+		state: FileSystemGlobalState,
+	) {
+		this.#registryConfig = registryConfig;
+		this.#runConfig = runConfig;
+		this.#managerDriver = managerDriver;
+		this.#inlineClient = inlineClient;
+		this.#state = state;
+	}
+
+	async loadActor(actorId: string): Promise<AnyActorInstance> {
+		return this.#state.startActor(
+			this.#registryConfig,
+			this.#runConfig,
+			this.#inlineClient,
+			this,
+			actorId,
+		);
+	}
+
+	getGenericConnGlobalState(actorId: string): GenericConnGlobalState {
+		return this.#state.getActorOrError(actorId).genericConnGlobalState;
+	}
+
+	/**
+	 * Get the current storage directory path
+	 */
+	get storagePath(): string {
+		return this.#state.storagePath;
+	}
+
+	getContext(_actorId: string): ActorDriverContext {
+		return {};
+	}
+
+	async readPersistedData(actorId: string): Promise<Uint8Array | undefined> {
+		return new Uint8Array(
+			(await this.#state.loadActorStateOrError(actorId)).persistedData,
+		);
+	}
+
+	async writePersistedData(actorId: string, data: Uint8Array): Promise<void> {
+		const state = await this.#state.loadActorStateOrError(actorId);
+
+		// Save state to disk
+		await this.#state.writeActor(actorId, {
+			...state,
+			persistedData: bufferToArrayBuffer(data),
+		});
+	}
+
+	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
+		await this.#state.setActorAlarm(actor.id, timestamp);
+	}
+
+	getDatabase(actorId: string): Promise<unknown | undefined> {
+		return this.#state.createDatabase(actorId);
+	}
+
+	sleep(actorId: string): Promise<void> {
+		return this.#state.sleepActor(actorId);
+	}
+}