rivetkit 2.0.1 → 2.0.3

package/src/mod.ts

@@ -0,0 +1,20 @@
+export { generateConnId, generateConnToken } from "@/actor/connection";
+export * from "@/actor/mod";
+export {
+	type AnyClient,
+	type Client,
+	type ClientDriver,
+	createClientWithDriver,
+} from "@/client/client";
+export { InlineWebSocketAdapter2 } from "@/common/inline-websocket-adapter2";
+export { noopNext } from "@/common/utils";
+export { createEngineDriver } from "@/drivers/engine/mod";
+export {
+	createFileSystemDriver,
+	createMemoryDriver,
+} from "@/drivers/file-system/mod";
+export { createInlineClientDriver } from "@/inline-client-driver/mod";
+// Re-export important protocol types and utilities needed by drivers
+export type { ActorQuery } from "@/manager/protocol/query";
+export * from "@/registry/mod";
+export { toUint8Array } from "@/utils";

package/src/registry/config.ts

@@ -0,0 +1,32 @@
+//! These configs configs hold anything that's not platform-specific about running actors.
+
+import { z } from "zod";
+import type { ActorDefinition, AnyActorDefinition } from "@/actor/definition";
+
+export const ActorsSchema = z.record(
+	z.string(),
+	z.custom<ActorDefinition<any, any, any, any, any, any, any, any>>(),
+);
+export type RegistryActors = z.infer<typeof ActorsSchema>;
+
+export const TestConfigSchema = z.object({ enabled: z.boolean() });
+export type TestConfig = z.infer<typeof TestConfigSchema>;
+
+/** Base config used for the actor config across all platforms. */
+export const RegistryConfigSchema = z.object({
+	use: z.record(z.string(), z.custom<AnyActorDefinition>()),
+
+	// TODO: Find a better way of passing around the test config
+	/**
+	 * Test configuration.
+	 *
+	 * DO NOT MANUALLY ENABLE. THIS IS USED INTERNALLY.
+	 * @internal
+	 **/
+	test: TestConfigSchema.optional().default({ enabled: false }),
+});
+export type RegistryConfig = z.infer<typeof RegistryConfigSchema>;
+export type RegistryConfigInput<A extends RegistryActors> = Omit<
+	z.input<typeof RegistryConfigSchema>,
+	"use"
+> & { use: A };

package/src/registry/log.ts

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

package/src/registry/mod.ts

@@ -0,0 +1,124 @@
+import type { Hono } from "hono";
+import { createActorRouter } from "@/actor/router";
+import { type Client, createClientWithDriver } from "@/client/client";
+import { chooseDefaultDriver } from "@/drivers/default";
+import { createInlineClientDriver } from "@/inline-client-driver/mod";
+import { getInspectorUrl } from "@/inspector/utils";
+import { createManagerRouter } from "@/manager/router";
+import {
+	type RegistryActors,
+	type RegistryConfig,
+	type RegistryConfigInput,
+	RegistryConfigSchema,
+} from "./config";
+import { logger } from "./log";
+import {
+	type DriverConfig,
+	type RunConfig,
+	type RunConfigInput,
+	RunConfigSchema,
+} from "./run-config";
+import { crossPlatformServe } from "./serve";
+
+interface ServerOutput<A extends Registry<any>> {
+	client: Client<A>;
+	hono: Hono;
+	handler: (req: Request) => Promise<Response>;
+	serve: (hono?: Hono) => void;
+}
+
+export class Registry<A extends RegistryActors> {
+	#config: RegistryConfig;
+
+	public get config(): RegistryConfig {
+		return this.#config;
+	}
+
+	constructor(config: RegistryConfig) {
+		this.#config = config;
+	}
+
+	/**
+	 * Runs the registry for a server.
+	 */
+	public createServer(inputConfig?: RunConfigInput): ServerOutput<this> {
+		const config = RunConfigSchema.parse(inputConfig);
+
+		// Choose the driver based on configuration
+		const driver = chooseDefaultDriver(config);
+
+		// Configure getUpgradeWebSocket lazily so we can assign it in crossPlatformServe
+		let upgradeWebSocket: any;
+		if (!config.getUpgradeWebSocket) {
+			config.getUpgradeWebSocket = () => upgradeWebSocket!;
+		}
+
+		// Create router
+		const managerDriver = driver.manager(this.#config, config);
+		const clientDriver = createInlineClientDriver(managerDriver);
+		const { router: hono } = createManagerRouter(
+			this.#config,
+			config,
+			clientDriver,
+			managerDriver,
+			false,
+		);
+
+		// Create client
+		const client = createClientWithDriver<this>(clientDriver);
+
+		const driverLog = managerDriver.extraStartupLog?.() ?? {};
+		logger().info("rivetkit ready", {
+			driver: driver.name,
+			definitions: Object.keys(this.#config.use).length,
+			...driverLog,
+		});
+		if (config.inspector?.enabled) {
+			logger().info("inspector ready", {
+				url: getInspectorUrl(config),
+			});
+		}
+
+		// Create runner
+		if (config.role === "all" || config.role === "runner") {
+			const inlineClient = createClientWithDriver(
+				createInlineClientDriver(managerDriver),
+			);
+			const _actorDriver = driver.actor(
+				this.#config,
+				config,
+				managerDriver,
+				inlineClient,
+			);
+			// TODO: What do we do with the actor driver here?
+		}
+
+		return {
+			client,
+			hono,
+			handler: async (req: Request) => await hono.fetch(req),
+			serve: async (app) => {
+				const out = await crossPlatformServe(hono, app);
+				upgradeWebSocket = out.upgradeWebSocket;
+			},
+		};
+	}
+
+	/**
+	 * Runs the registry as a standalone server.
+	 */
+	public async runServer(inputConfig?: RunConfigInput) {
+		const { serve } = this.createServer(inputConfig);
+		serve();
+	}
+}
+
+export function setup<A extends RegistryActors>(
+	input: RegistryConfigInput<A>,
+): Registry<A> {
+	const config = RegistryConfigSchema.parse(input);
+	return new Registry(config);
+}
+
+export type { RegistryConfig, RegistryActors, RunConfig, DriverConfig };
+export { RegistryConfigSchema };

package/src/registry/run-config.ts

@@ -0,0 +1,54 @@
+import type { cors } from "hono/cors";
+import { z } from "zod";
+import type { ActorDriverBuilder } from "@/actor/driver";
+import { InspectorConfigSchema } from "@/inspector/config";
+import type { ManagerDriverBuilder } from "@/manager/driver";
+import type { UpgradeWebSocket } from "@/utils";
+
+type CorsOptions = NonNullable<Parameters<typeof cors>[0]>;
+
+export type GetUpgradeWebSocket = () => UpgradeWebSocket;
+
+export const DriverConfigSchema = z.object({
+	/** Machine-readable name to identify this driver by. */
+	name: z.string(),
+	manager: z.custom<ManagerDriverBuilder>(),
+	actor: z.custom<ActorDriverBuilder>(),
+});
+
+export type DriverConfig = z.infer<typeof DriverConfigSchema>;
+
+/** Base config used for the actor config across all platforms. */
+export const RunConfigSchema = z
+	.object({
+		driver: DriverConfigSchema.optional(),
+
+		/** Endpoint to connect to the Rivet engine. Can be configured via RIVET_ENGINE env var. */
+		engine: z.string().optional(),
+
+		// This is a function to allow for lazy configuration of upgradeWebSocket on the
+		// fly. This is required since the dependencies that profie upgradeWebSocket
+		// (specifically Node.js) can sometimes only be specified after the router is
+		// created or must be imported async using `await import(...)`
+		getUpgradeWebSocket: z.custom<GetUpgradeWebSocket>().optional(),
+
+		role: z.enum(["all", "server", "runner"]).optional().default("all"),
+
+		/** CORS configuration for the router. Uses Hono's CORS middleware options. */
+		cors: z.custom<CorsOptions>().optional(),
+
+		maxIncomingMessageSize: z.number().optional().default(65_536),
+
+		inspector: InspectorConfigSchema,
+
+		/**
+		 * Base path for the router. This is used to prefix all routes.
+		 * For example, if the base path is `/api`, then the route `/actors` will be
+		 * available at `/api/actors`.
+		 */
+		basePath: z.string().optional().default("/"),
+	})
+	.default({});
+
+export type RunConfig = z.infer<typeof RunConfigSchema>;
+export type RunConfigInput = z.input<typeof RunConfigSchema>;

package/src/registry/serve.ts

@@ -0,0 +1,53 @@
+import { Hono } from "hono";
+import { getEnvUniversal } from "@/utils";
+import { logger } from "./log";
+
+export async function crossPlatformServe(
+	rivetKitRouter: Hono<any>,
+	userRouter: Hono | undefined,
+) {
+	const app = userRouter ?? new Hono();
+
+	// Import @hono/node-server
+	let serve: any;
+	try {
+		const dep = await import("@hono/node-server");
+		serve = dep.serve;
+	} catch (err) {
+		logger().error(
+			"failed to import @hono/node-server. please run 'npm install @hono/node-server @hono/node-ws'",
+		);
+		process.exit(1);
+	}
+
+	// Mount registry
+	app.route("/registry", rivetKitRouter);
+
+	// Import @hono/node-ws
+	let createNodeWebSocket: any;
+	try {
+		const dep = await import("@hono/node-ws");
+		createNodeWebSocket = dep.createNodeWebSocket;
+	} catch (err) {
+		logger().error(
+			"failed to import @hono/node-ws. please run 'npm install @hono/node-server @hono/node-ws'",
+		);
+		process.exit(1);
+	}
+
+	// Inject WS
+	const { injectWebSocket, upgradeWebSocket } = createNodeWebSocket({
+		app,
+	});
+
+	// Start server
+	const port = Number.parseInt(
+		getEnvUniversal("PORT") ?? getEnvUniversal("PORT_HTTP") ?? "8080",
+	);
+	const server = serve({ fetch: app.fetch, port }, () =>
+		logger().info("server listening", { port }),
+	);
+	injectWebSocket(server);
+
+	return { upgradeWebSocket };
+}

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,84 @@
+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>,
+): Uint8Array | string {
+	if (encoding === "json") {
+		return jsonStringifyCompat(value);
+	} else if (encoding === "cbor") {
+		return cbor.encode(value);
+	} else if (encoding === "bare") {
+		return versionedDataHandler.serializeWithEmbeddedVersion(value);
+	} else {
+		assertUnreachable(encoding);
+	}
+}
+
+export function deserializeWithEncoding<T>(
+	encoding: Encoding,
+	buffer: Uint8Array | string,
+	versionedDataHandler: VersionedDataHandler<T>,
+): 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",
+		);
+		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,7 @@
+import { getLogger } from "@/common/log";
+
+export const LOGGER_NAME = "nodejs";
+
+export function logger() {
+	return getLogger(LOGGER_NAME);
+}

package/src/test/mod.ts

@@ -0,0 +1,153 @@
+import { createServer } from "node:net";
+import { serve as honoServe, type ServerType } from "@hono/node-server";
+import { createNodeWebSocket } from "@hono/node-ws";
+import { type TestContext, vi } from "vitest";
+import { type Client, createClient } from "@/client/mod";
+import { chooseDefaultDriver } from "@/drivers/default";
+import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
+import { createInlineClientDriver } from "@/inline-client-driver/mod";
+import { getInspectorUrl } from "@/inspector/utils";
+import { createManagerRouter } from "@/manager/router";
+import type { Registry } from "@/registry/mod";
+import { RunConfigSchema } from "@/registry/run-config";
+import { ConfigSchema, type InputConfig } from "./config";
+import { logger } from "./log";
+
+function serve(registry: Registry<any>, inputConfig?: InputConfig): ServerType {
+	// Configure default configuration
+	inputConfig ??= {};
+
+	const config = ConfigSchema.parse(inputConfig);
+
+	let upgradeWebSocket: any;
+	if (!config.getUpgradeWebSocket) {
+		config.getUpgradeWebSocket = () => upgradeWebSocket!;
+	}
+
+	// Create router
+	const runConfig = RunConfigSchema.parse(inputConfig);
+	const driver = inputConfig.driver ?? createFileSystemOrMemoryDriver(false);
+	const managerDriver = driver.manager(registry.config, config);
+	const inlineClientDriver = createInlineClientDriver(managerDriver);
+	const { router } = createManagerRouter(
+		registry.config,
+		runConfig,
+		inlineClientDriver,
+		managerDriver,
+		false,
+	);
+
+	// Inject WebSocket
+	const nodeWebSocket = createNodeWebSocket({ app: router });
+	upgradeWebSocket = nodeWebSocket.upgradeWebSocket;
+
+	const server = honoServe({
+		fetch: router.fetch,
+		hostname: config.hostname,
+		port: config.port,
+	});
+	nodeWebSocket.injectWebSocket(server);
+
+	logger().info("rivetkit started", {
+		hostname: config.hostname,
+		port: config.port,
+		definitions: Object.keys(registry.config.use).length,
+	});
+
+	return server;
+}
+
+export interface SetupTestResult<A extends Registry<any>> {
+	client: Client<A>;
+	mockDriver: {
+		actorDriver: {
+			setCreateVarsContext: (ctx: any) => void;
+		};
+	};
+}
+
+// Must use `TestContext` since global hooks do not work when running concurrently
+export async function setupTest<A extends Registry<any>>(
+	c: TestContext,
+	registry: A,
+): Promise<SetupTestResult<A>> {
+	vi.useFakeTimers();
+
+	// Set up mock driver for testing createVars context
+	const mockDriverContext: any = {};
+	const setDriverContextFn = (ctx: any) => {
+		mockDriverContext.current = ctx;
+	};
+
+	// We don't need to modify the driver context anymore since we're testing with the actual context
+
+	// Start server with a random port
+	const port = await getPort();
+	const server = serve(registry, { port });
+	c.onTestFinished(
+		async () => await new Promise((resolve) => server.close(() => resolve())),
+	);
+
+	// Create client
+	const client = createClient<A>(`http://127.0.0.1:${port}`);
+	c.onTestFinished(async () => await client.dispose());
+
+	return {
+		client,
+		mockDriver: {
+			actorDriver: {
+				setCreateVarsContext: setDriverContextFn,
+			},
+		},
+	};
+}
+
+export async function getPort(): Promise<number> {
+	// Pick random port between 10000 and 65535 (avoiding well-known and registered ports)
+	const MIN_PORT = 10000;
+	const MAX_PORT = 65535;
+	const getRandomPort = () =>
+		Math.floor(Math.random() * (MAX_PORT - MIN_PORT + 1)) + MIN_PORT;
+
+	let port = getRandomPort();
+	let maxAttempts = 10;
+
+	while (maxAttempts > 0) {
+		try {
+			// Try to create a server on the port to check if it's available
+			const server = await new Promise<any>((resolve, reject) => {
+				const server = createServer();
+
+				server.once("error", (err: Error & { code?: string }) => {
+					if (err.code === "EADDRINUSE") {
+						reject(new Error(`Port ${port} is in use`));
+					} else {
+						reject(err);
+					}
+				});
+
+				server.once("listening", () => {
+					resolve(server);
+				});
+
+				server.listen(port);
+			});
+
+			// Close the server since we're just checking availability
+			await new Promise<void>((resolve) => {
+				server.close(() => resolve());
+			});
+
+			return port;
+		} catch (err) {
+			// If port is in use, try a different one
+			maxAttempts--;
+			if (maxAttempts <= 0) {
+				break;
+			}
+			port = getRandomPort();
+		}
+	}
+
+	throw new Error("Could not find an available port after multiple attempts");
+}