rivetkit 2.0.2 → 2.0.4

package/src/remote-manager-driver/mod.ts

@@ -0,0 +1,274 @@
+import * as cbor from "cbor-x";
+import type { Hono, Context as HonoContext } from "hono";
+import invariant from "invariant";
+import { ActorAlreadyExists } from "@/actor/errors";
+import { deserializeActorKey, serializeActorKey } from "@/actor/keys";
+import type { ClientConfig } from "@/client/client";
+import { noopNext } from "@/common/utils";
+import type {
+	ActorOutput,
+	CreateInput,
+	GetForIdInput,
+	GetOrCreateWithKeyInput,
+	GetWithKeyInput,
+	ManagerDisplayInformation,
+	ManagerDriver,
+} from "@/driver-helpers/mod";
+import type { ManagerInspector } from "@/inspector/manager";
+import type { Encoding, RegistryConfig, UniversalWebSocket } from "@/mod";
+import type { RunConfig } from "@/registry/run-config";
+import { sendHttpRequestToActor } from "./actor-http-client";
+import {
+	buildGuardHeadersForWebSocket,
+	openWebSocketToActor,
+} from "./actor-websocket-client";
+import {
+	createActor,
+	destroyActor,
+	getActor,
+	getActorById,
+	getOrCreateActorById,
+} from "./api-endpoints";
+import { EngineApiError, getEndpoint } from "./api-utils";
+import { logger } from "./log";
+import { createWebSocketProxy } from "./ws-proxy";
+
+// TODO:
+// // Lazily import the dynamic imports so we don't have to turn `createClient` in to an async fn
+// const dynamicImports = (async () => {
+// 	// Import dynamic dependencies
+// 	const [WebSocket, EventSource] = await Promise.all([
+// 		importWebSocket(),
+// 		importEventSource(),
+// 	]);
+// 	return {
+// 		WebSocket,
+// 		EventSource,
+// 	};
+// })();
+
+export class RemoteManagerDriver implements ManagerDriver {
+	#config: ClientConfig;
+
+	constructor(runConfig: ClientConfig) {
+		this.#config = runConfig;
+	}
+
+	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({
+					msg: "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({ msg: "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({
+				msg: "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({
+			msg: "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({
+			msg: "getOrCreateWithKey: actor ready",
+			actorId,
+			name,
+			key,
+			created: response.created,
+		});
+
+		return {
+			actorId,
+			name,
+			key,
+		};
+	}
+
+	async createActor({
+		c,
+		name,
+		key,
+		input,
+	}: CreateInput): Promise<ActorOutput> {
+		logger().info({ msg: "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({ msg: "actor created", actorId, name, key });
+
+		return {
+			actorId,
+			name,
+			key,
+		};
+	}
+
+	async destroyActor(actorId: string): Promise<void> {
+		logger().info({ msg: "destroying actor via engine api", actorId });
+
+		await destroyActor(this.#config, actorId);
+
+		logger().info({ msg: "actor destroyed", actorId });
+	}
+
+	async sendRequest(actorId: string, actorRequest: Request): Promise<Response> {
+		return await sendHttpRequestToActor(this.#config, actorId, actorRequest);
+	}
+
+	async openWebSocket(
+		path: string,
+		actorId: string,
+		encoding: Encoding,
+		params: unknown,
+	): Promise<UniversalWebSocket> {
+		return await openWebSocketToActor(
+			this.#config,
+			path,
+			actorId,
+			encoding,
+			params,
+		);
+	}
+
+	async proxyRequest(
+		_c: HonoContext,
+		actorRequest: Request,
+		actorId: string,
+	): Promise<Response> {
+		return await sendHttpRequestToActor(this.#config, actorId, actorRequest);
+	}
+
+	async proxyWebSocket(
+		c: HonoContext,
+		path: string,
+		actorId: string,
+		encoding: Encoding,
+		params: unknown,
+		authData: unknown,
+	): Promise<Response> {
+		const upgradeWebSocket = this.#config.getUpgradeWebSocket?.();
+		invariant(upgradeWebSocket, "missing getUpgradeWebSocket");
+
+		const endpoint = getEndpoint(this.#config);
+		const guardUrl = `${endpoint}${path}`;
+		const wsGuardUrl = guardUrl.replace("http://", "ws://");
+
+		logger().debug({
+			msg: "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());
+	}
+
+	displayInformation(): ManagerDisplayInformation {
+		return { name: "Remote", properties: {} };
+	}
+}

package/src/remote-manager-driver/ws-proxy.ts

@@ -0,0 +1,180 @@
+import type { Context as HonoContext } from "hono";
+import type { WSContext } from "hono/ws";
+import { stringifyError } from "@/common/utils";
+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({ msg: "client websocket connected", targetUrl });
+
+			if (clientWs.readyState !== 1) {
+				logger().warn({
+					msg: "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({ msg: "target websocket connected", targetUrl });
+
+					if (clientWs.readyState !== 1) {
+						logger().warn({
+							msg: "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({
+						msg: "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({
+					msg: "target websocket closed",
+					targetUrl,
+					code: event.code,
+					reason: event.reason,
+				});
+				closeWebSocketIfOpen(clientWs, event.code, event.reason);
+			});
+
+			state.targetWs.addEventListener("error", (error) => {
+				logger().error({
+					msg: "target websocket error",
+					targetUrl,
+					error: stringifyError(error),
+				});
+				closeWebSocketIfOpen(clientWs, 1011, "Target WebSocket error");
+			});
+		},
+
+		onMessage: async (event: any, clientWs: WSContext) => {
+			if (!state.targetWs || !state.connectPromise) {
+				logger().error({ msg: "websocket state not initialized", targetUrl });
+				return;
+			}
+
+			try {
+				await state.connectPromise;
+				if (state.targetWs.readyState === WebSocket.OPEN) {
+					state.targetWs.send(event.data);
+				} else {
+					logger().warn({
+						msg: "target websocket not open",
+						targetUrl,
+						readyState: state.targetWs.readyState,
+					});
+				}
+			} catch (error) {
+				logger().error({
+					msg: "failed to connect to target websocket",
+					targetUrl,
+					error,
+				});
+				closeWebSocketIfOpen(clientWs, 1011, "Failed to connect to target");
+			}
+		},
+
+		onClose: (event: any, clientWs: WSContext) => {
+			logger().debug({
+				msg: "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({ msg: "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/schemas/actor-persist/mod.ts

@@ -0,0 +1 @@
+export * from "../../../dist/schemas/actor-persist/v1";

package/src/schemas/actor-persist/versioned.ts

@@ -0,0 +1,25 @@
+import {
+	createVersionedDataHandler,
+	type MigrationFn,
+} from "@/common/versioned-data";
+import * as v1 from "../../../dist/schemas/actor-persist/v1";
+
+export const CURRENT_VERSION = 1;
+
+export type CurrentPersistedActor = v1.PersistedActor;
+export type CurrentPersistedConnection = v1.PersistedConnection;
+export type CurrentPersistedSubscription = v1.PersistedSubscription;
+export type CurrentGenericPersistedScheduleEvent =
+	v1.GenericPersistedScheduleEvent;
+export type CurrentPersistedScheduleEventKind = v1.PersistedScheduleEventKind;
+export type CurrentPersistedScheduleEvent = v1.PersistedScheduleEvent;
+
+const migrations = new Map<number, MigrationFn<any, any>>();
+
+export const PERSISTED_ACTOR_VERSIONED =
+	createVersionedDataHandler<CurrentPersistedActor>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (data) => v1.encodePersistedActor(data),
+		deserializeVersion: (bytes) => v1.decodePersistedActor(bytes),
+	});

package/src/schemas/client-protocol/mod.ts

@@ -0,0 +1 @@
+export * from "../../../dist/schemas/client-protocol/v1";

package/src/schemas/client-protocol/versioned.ts

@@ -0,0 +1,63 @@
+import {
+	createVersionedDataHandler,
+	type MigrationFn,
+} from "@/common/versioned-data";
+import * as v1 from "../../../dist/schemas/client-protocol/v1";
+
+export const CURRENT_VERSION = 1;
+
+const migrations = new Map<number, MigrationFn<any, any>>();
+
+export const TO_SERVER_VERSIONED = createVersionedDataHandler<v1.ToServer>({
+	currentVersion: CURRENT_VERSION,
+	migrations,
+	serializeVersion: (data) => v1.encodeToServer(data),
+	deserializeVersion: (bytes) => v1.decodeToServer(bytes),
+});
+
+export const TO_CLIENT_VERSIONED = createVersionedDataHandler<v1.ToClient>({
+	currentVersion: CURRENT_VERSION,
+	migrations,
+	serializeVersion: (data) => v1.encodeToClient(data),
+	deserializeVersion: (bytes) => v1.decodeToClient(bytes),
+});
+
+export const HTTP_ACTION_REQUEST_VERSIONED =
+	createVersionedDataHandler<v1.HttpActionRequest>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (data) => v1.encodeHttpActionRequest(data),
+		deserializeVersion: (bytes) => v1.decodeHttpActionRequest(bytes),
+	});
+
+export const HTTP_ACTION_RESPONSE_VERSIONED =
+	createVersionedDataHandler<v1.HttpActionResponse>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (data) => v1.encodeHttpActionResponse(data),
+		deserializeVersion: (bytes) => v1.decodeHttpActionResponse(bytes),
+	});
+
+export const HTTP_RESPONSE_ERROR_VERSIONED =
+	createVersionedDataHandler<v1.HttpResponseError>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (data) => v1.encodeHttpResponseError(data),
+		deserializeVersion: (bytes) => v1.decodeHttpResponseError(bytes),
+	});
+
+export const HTTP_RESOLVE_REQUEST_VERSIONED =
+	createVersionedDataHandler<v1.HttpResolveRequest>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (_) => new Uint8Array(),
+		deserializeVersion: (bytes) => null,
+	});
+
+export const HTTP_RESOLVE_RESPONSE_VERSIONED =
+	createVersionedDataHandler<v1.HttpResolveResponse>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (data) => v1.encodeHttpResolveResponse(data),
+		deserializeVersion: (bytes) => v1.decodeHttpResolveResponse(bytes),
+	});

package/src/schemas/file-system-driver/mod.ts

@@ -0,0 +1 @@
+export * from "../../../dist/schemas/file-system-driver/v1";

package/src/schemas/file-system-driver/versioned.ts

@@ -0,0 +1,28 @@
+import {
+	createVersionedDataHandler,
+	type MigrationFn,
+} from "@/common/versioned-data";
+import * as v1 from "../../../dist/schemas/file-system-driver/v1";
+
+export const CURRENT_VERSION = 1;
+
+export type CurrentActorState = v1.ActorState;
+export type CurrentActorAlarm = v1.ActorAlarm;
+
+const migrations = new Map<number, MigrationFn<any, any>>();
+
+export const ACTOR_STATE_VERSIONED =
+	createVersionedDataHandler<CurrentActorState>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (data) => v1.encodeActorState(data),
+		deserializeVersion: (bytes) => v1.decodeActorState(bytes),
+	});
+
+export const ACTOR_ALARM_VERSIONED =
+	createVersionedDataHandler<CurrentActorAlarm>({
+		currentVersion: CURRENT_VERSION,
+		migrations,
+		serializeVersion: (data) => v1.encodeActorAlarm(data),
+		deserializeVersion: (bytes) => v1.decodeActorAlarm(bytes),
+	});

package/src/serde.ts

@@ -0,0 +1,90 @@
+import * as cbor from "cbor-x";
+import invariant from "invariant";
+import { assertUnreachable } from "@/common/utils";
+import type { VersionedDataHandler } from "@/common/versioned-data";
+import type { Encoding } from "@/mod";
+import { jsonStringifyCompat } from "./actor/protocol/serde";
+
+export function encodingIsBinary(encoding: Encoding): boolean {
+	if (encoding === "json") {
+		return false;
+	} else if (encoding === "cbor" || encoding === "bare") {
+		return true;
+	} else {
+		assertUnreachable(encoding);
+	}
+}
+
+export function contentTypeForEncoding(encoding: Encoding): string {
+	if (encoding === "json") {
+		return "application/json";
+	} else if (encoding === "cbor" || encoding === "bare") {
+		return "application/octet-stream";
+	} else {
+		assertUnreachable(encoding);
+	}
+}
+
+export function wsBinaryTypeForEncoding(
+	encoding: Encoding,
+): "arraybuffer" | "blob" {
+	if (encoding === "json") {
+		return "blob";
+	} else if (encoding === "cbor" || encoding === "bare") {
+		return "arraybuffer";
+	} else {
+		assertUnreachable(encoding);
+	}
+}
+
+export function serializeWithEncoding<T>(
+	encoding: Encoding,
+	value: T,
+	versionedDataHandler: VersionedDataHandler<T> | undefined,
+): Uint8Array | string {
+	if (encoding === "json") {
+		return jsonStringifyCompat(value);
+	} else if (encoding === "cbor") {
+		return cbor.encode(value);
+	} else if (encoding === "bare") {
+		if (!versionedDataHandler) {
+			throw new Error("VersionedDataHandler is required for 'bare' encoding");
+		}
+		return versionedDataHandler.serializeWithEmbeddedVersion(value);
+	} else {
+		assertUnreachable(encoding);
+	}
+}
+
+export function deserializeWithEncoding<T>(
+	encoding: Encoding,
+	buffer: Uint8Array | string,
+	versionedDataHandler: VersionedDataHandler<T> | undefined,
+): T {
+	if (encoding === "json") {
+		if (typeof buffer === "string") {
+			return JSON.parse(buffer);
+		} else {
+			const decoder = new TextDecoder("utf-8");
+			const jsonString = decoder.decode(buffer);
+			return JSON.parse(jsonString);
+		}
+	} else if (encoding === "cbor") {
+		invariant(
+			typeof buffer !== "string",
+			"buffer cannot be string for cbor encoding",
+		);
+		return cbor.decode(buffer);
+	} else if (encoding === "bare") {
+		invariant(
+			typeof buffer !== "string",
+			"buffer cannot be string for bare encoding",
+		);
+		if (!versionedDataHandler) {
+			throw new Error("VersionedDataHandler is required for 'bare' encoding");
+		}
+		return versionedDataHandler.deserializeWithEmbeddedVersion(buffer);
+	} else {
+		assertUnreachable(encoding);
+	}
+}

package/src/test/config.ts

@@ -0,0 +1,16 @@
+import { z } from "zod";
+import { RunConfigSchema } from "@/registry/run-config";
+
+export const ConfigSchema = RunConfigSchema.removeDefault()
+	.extend({
+		hostname: z
+			.string()
+			.optional()
+			.default(process.env.HOSTNAME ?? "127.0.0.1"),
+		port: z
+			.number()
+			.optional()
+			.default(Number.parseInt(process.env.PORT ?? "8080")),
+	})
+	.default({});
+export type InputConfig = z.input<typeof ConfigSchema>;

package/src/test/log.ts

@@ -0,0 +1,5 @@
+import { getLogger } from "@/common/log";
+
+export function logger() {
+	return getLogger("test");
+}