rivetkit 2.0.2 → 2.0.4

package/src/client/client.ts

@@ -0,0 +1,554 @@
+import type { AnyActorDefinition } from "@/actor/definition";
+import type { Transport } from "@/actor/protocol/old";
+import type { Encoding } from "@/actor/protocol/serde";
+import type { ManagerDriver } from "@/driver-helpers/mod";
+import type { ActorQuery } from "@/manager/protocol/query";
+import type { Registry } from "@/mod";
+import type { ActorActionFunction } from "./actor-common";
+import {
+	type ActorConn,
+	type ActorConnRaw,
+	CONNECT_SYMBOL,
+} from "./actor-conn";
+import { type ActorHandle, ActorHandleRaw } from "./actor-handle";
+import { queryActor } from "./actor-query";
+import type { ClientConfig } from "./config";
+import { logger } from "./log";
+
+export type { ClientConfig, ClientConfigInput } from "./config";
+
+/** Extract the actor registry from the registry definition. */
+export type ExtractActorsFromRegistry<A extends Registry<any>> =
+	A extends Registry<infer Actors> ? Actors : never;
+
+/** Extract the registry definition from the client. */
+export type ExtractRegistryFromClient<C extends Client<Registry<{}>>> =
+	C extends Client<infer A> ? A : never;
+
+/**
+ * Represents a actor accessor that provides methods to interact with a specific actor.
+ */
+export interface ActorAccessor<AD extends AnyActorDefinition> {
+	/**
+	 * Gets a stateless handle to a actor by its key, but does not create the actor if it doesn't exist.
+	 * The actor name is automatically injected from the property accessor.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string | string[]} [key=[]] - The key to identify the actor. Can be a single string or an array of strings.
+	 * @param {GetWithIdOptions} [opts] - Options for getting the actor.
+	 * @returns {ActorHandle<AD>} - A handle to the actor.
+	 */
+	get(key?: string | string[], opts?: GetWithIdOptions): ActorHandle<AD>;
+
+	/**
+	 * Gets a stateless handle to a actor by its key, creating it if necessary.
+	 * The actor name is automatically injected from the property accessor.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string | string[]} [key=[]] - The key to identify the actor. Can be a single string or an array of strings.
+	 * @param {GetOptions} [opts] - Options for getting the actor.
+	 * @returns {ActorHandle<AD>} - A handle to the actor.
+	 */
+	getOrCreate(
+		key?: string | string[],
+		opts?: GetOrCreateOptions,
+	): ActorHandle<AD>;
+
+	/**
+	 * Gets a stateless handle to a actor by its ID.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string} actorId - The ID of the actor.
+	 * @param {GetWithIdOptions} [opts] - Options for getting the actor.
+	 * @returns {ActorHandle<AD>} - A handle to the actor.
+	 */
+	getForId(actorId: string, opts?: GetWithIdOptions): ActorHandle<AD>;
+
+	/**
+	 * Creates a new actor with the name automatically injected from the property accessor,
+	 * and returns a stateless handle to it with the actor ID resolved.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string | string[]} key - The key to identify the actor. Can be a single string or an array of strings.
+	 * @param {CreateOptions} [opts] - Options for creating the actor (excluding name and key).
+	 * @returns {Promise<ActorHandle<AD>>} - A promise that resolves to a handle to the actor.
+	 */
+	create(
+		key?: string | string[],
+		opts?: CreateOptions,
+	): Promise<ActorHandle<AD>>;
+}
+
+/**
+ * Options for querying actors.
+ * @typedef {Object} QueryOptions
+ * @property {unknown} [parameters] - Parameters to pass to the connection.
+ */
+export interface QueryOptions {
+	/** Parameters to pass to the connection. */
+	params?: unknown;
+	/** Signal to abort the request. */
+	signal?: AbortSignal;
+}
+
+/**
+ * Options for getting a actor by ID.
+ * @typedef {QueryOptions} GetWithIdOptions
+ */
+export interface GetWithIdOptions extends QueryOptions {}
+
+/**
+ * Options for getting a actor.
+ * @typedef {QueryOptions} GetOptions
+ */
+export interface GetOptions extends QueryOptions {}
+
+/**
+ * Options for getting or creating a actor.
+ * @typedef {QueryOptions} GetOrCreateOptions
+ * @property {string} [createInRegion] - Region to create the actor in if it doesn't exist.
+ */
+export interface GetOrCreateOptions extends QueryOptions {
+	/** Region to create the actor in if it doesn't exist. */
+	createInRegion?: string;
+	/** Input data to pass to the actor. */
+	createWithInput?: unknown;
+}
+
+/**
+ * Options for creating a actor.
+ * @typedef {QueryOptions} CreateOptions
+ * @property {string} [region] - The region to create the actor in.
+ */
+export interface CreateOptions extends QueryOptions {
+	/** The region to create the actor in. */
+	region?: string;
+	/** Input data to pass to the actor. */
+	input?: unknown;
+}
+
+/**
+ * Represents a region to connect to.
+ * @typedef {Object} Region
+ * @property {string} id - The region ID.
+ * @property {string} name - The region name.
+ * @see {@link https://rivet.dev/docs/edge|Edge Networking}
+ * @see {@link https://rivet.dev/docs/regions|Available Regions}
+ */
+export interface Region {
+	/**
+	 * The region slug.
+	 */
+	id: string;
+
+	/**
+	 * The human-friendly region name.
+	 */
+	name: string;
+}
+
+export const ACTOR_CONNS_SYMBOL = Symbol("actorConns");
+export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy");
+export const TRANSPORT_SYMBOL = Symbol("transport");
+
+/**
+ * Client for managing & connecting to actors.
+ *
+ * @template A The actors map type that defines the available actors.
+ * @see {@link https://rivet.dev/docs/manage|Create & Manage Actors}
+ */
+export class ClientRaw {
+	#disposed = false;
+
+	[ACTOR_CONNS_SYMBOL] = new Set<ActorConnRaw>();
+
+	#driver: ManagerDriver;
+	#encodingKind: Encoding;
+	[TRANSPORT_SYMBOL]: Transport;
+
+	/**
+	 * Creates an instance of Client.
+	 *
+	 * @param {string} managerEndpoint - The manager endpoint. See {@link https://rivet.dev/docs/setup|Initial Setup} for instructions on getting the manager endpoint.
+	 * @param {ClientConfig} [opts] - Options for configuring the client.
+	 * @see {@link https://rivet.dev/docs/setup|Initial Setup}
+	 */
+	public constructor(driver: ManagerDriver, opts?: ClientConfig) {
+		this.#driver = driver;
+
+		this.#encodingKind = opts?.encoding ?? "bare";
+		this[TRANSPORT_SYMBOL] = opts?.transport ?? "websocket";
+	}
+
+	/**
+	 * Gets a stateless handle to a actor by its ID.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string} name - The name of the actor.
+	 * @param {string} actorId - The ID of the actor.
+	 * @param {GetWithIdOptions} [opts] - Options for getting the actor.
+	 * @returns {ActorHandle<AD>} - A handle to the actor.
+	 */
+	getForId<AD extends AnyActorDefinition>(
+		name: string,
+		actorId: string,
+		opts?: GetWithIdOptions,
+	): ActorHandle<AD> {
+		logger().debug({
+			msg: "get handle to actor with id",
+			name,
+			actorId,
+			params: opts?.params,
+		});
+
+		const actorQuery: ActorQuery = {
+			getForId: {
+				name,
+				actorId,
+			},
+		};
+
+		const handle = this.#createHandle(opts?.params, actorQuery);
+		return createActorProxy(handle) as ActorHandle<AD>;
+	}
+
+	/**
+	 * Gets a stateless handle to a actor by its key, but does not create the actor if it doesn't exist.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string} name - The name of the actor.
+	 * @param {string | string[]} [key=[]] - The key to identify the actor. Can be a single string or an array of strings.
+	 * @param {GetWithIdOptions} [opts] - Options for getting the actor.
+	 * @returns {ActorHandle<AD>} - A handle to the actor.
+	 */
+	get<AD extends AnyActorDefinition>(
+		name: string,
+		key?: string | string[],
+		opts?: GetWithIdOptions,
+	): ActorHandle<AD> {
+		// Convert string to array of strings
+		const keyArray: string[] = typeof key === "string" ? [key] : key || [];
+
+		logger().debug({
+			msg: "get handle to actor",
+			name,
+			key: keyArray,
+			parameters: opts?.params,
+		});
+
+		const actorQuery: ActorQuery = {
+			getForKey: {
+				name,
+				key: keyArray,
+			},
+		};
+
+		const handle = this.#createHandle(opts?.params, actorQuery);
+		return createActorProxy(handle) as ActorHandle<AD>;
+	}
+
+	/**
+	 * Gets a stateless handle to a actor by its key, creating it if necessary.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string} name - The name of the actor.
+	 * @param {string | string[]} [key=[]] - The key to identify the actor. Can be a single string or an array of strings.
+	 * @param {GetOptions} [opts] - Options for getting the actor.
+	 * @returns {ActorHandle<AD>} - A handle to the actor.
+	 */
+	getOrCreate<AD extends AnyActorDefinition>(
+		name: string,
+		key?: string | string[],
+		opts?: GetOrCreateOptions,
+	): ActorHandle<AD> {
+		// Convert string to array of strings
+		const keyArray: string[] = typeof key === "string" ? [key] : key || [];
+
+		logger().debug({
+			msg: "get or create handle to actor",
+			name,
+			key: keyArray,
+			parameters: opts?.params,
+			createInRegion: opts?.createInRegion,
+		});
+
+		const actorQuery: ActorQuery = {
+			getOrCreateForKey: {
+				name,
+				key: keyArray,
+				input: opts?.createWithInput,
+				region: opts?.createInRegion,
+			},
+		};
+
+		const handle = this.#createHandle(opts?.params, actorQuery);
+		return createActorProxy(handle) as ActorHandle<AD>;
+	}
+
+	/**
+	 * Creates a new actor with the provided key and returns a stateless handle to it.
+	 * Resolves the actor ID and returns a handle with getForId query.
+	 *
+	 * @template AD The actor class that this handle is for.
+	 * @param {string} name - The name of the actor.
+	 * @param {string | string[]} key - The key to identify the actor. Can be a single string or an array of strings.
+	 * @param {CreateOptions} [opts] - Options for creating the actor (excluding name and key).
+	 * @returns {Promise<ActorHandle<AD>>} - A promise that resolves to a handle to the actor.
+	 */
+	async create<AD extends AnyActorDefinition>(
+		name: string,
+		key?: string | string[],
+		opts?: CreateOptions,
+	): Promise<ActorHandle<AD>> {
+		// Convert string to array of strings
+		const keyArray: string[] = typeof key === "string" ? [key] : key || [];
+
+		const createQuery = {
+			create: {
+				...opts,
+				// Do these last to override `opts`
+				name,
+				key: keyArray,
+			},
+		} satisfies ActorQuery;
+
+		logger().debug({
+			msg: "create actor handle",
+			name,
+			key: keyArray,
+			parameters: opts?.params,
+			create: createQuery.create,
+		});
+
+		// Create the actor
+		const { actorId } = await queryActor(undefined, createQuery, this.#driver);
+		logger().debug({
+			msg: "created actor with ID",
+			name,
+			key: keyArray,
+			actorId,
+		});
+
+		// Create handle with actor ID
+		const getForIdQuery = {
+			getForId: {
+				name,
+				actorId,
+			},
+		} satisfies ActorQuery;
+		const handle = this.#createHandle(opts?.params, getForIdQuery);
+
+		const proxy = createActorProxy(handle) as ActorHandle<AD>;
+
+		return proxy;
+	}
+
+	#createHandle(params: unknown, actorQuery: ActorQuery): ActorHandleRaw {
+		return new ActorHandleRaw(
+			this,
+			this.#driver,
+			params,
+			this.#encodingKind,
+			actorQuery,
+		);
+	}
+
+	[CREATE_ACTOR_CONN_PROXY]<AD extends AnyActorDefinition>(
+		conn: ActorConnRaw,
+	): ActorConn<AD> {
+		// Save to connection list
+		this[ACTOR_CONNS_SYMBOL].add(conn);
+
+		// Start connection
+		conn[CONNECT_SYMBOL]();
+
+		return createActorProxy(conn) as ActorConn<AD>;
+	}
+
+	/**
+	 * Disconnects from all actors.
+	 *
+	 * @returns {Promise<void>} A promise that resolves when all connections are closed.
+	 */
+	async dispose(): Promise<void> {
+		if (this.#disposed) {
+			logger().warn({ msg: "client already disconnected" });
+			return;
+		}
+		this.#disposed = true;
+
+		logger().debug({ msg: "disposing client" });
+
+		const disposePromises = [];
+
+		// Dispose all connections
+		for (const conn of this[ACTOR_CONNS_SYMBOL].values()) {
+			disposePromises.push(conn.dispose());
+		}
+
+		await Promise.all(disposePromises);
+	}
+}
+
+/**
+ * Client type with actor accessors.
+ * This adds property accessors for actor names to the ClientRaw base class.
+ *
+ * @template A The actor registry type.
+ */
+export type Client<A extends Registry<any>> = ClientRaw & {
+	[K in keyof ExtractActorsFromRegistry<A>]: ActorAccessor<
+		ExtractActorsFromRegistry<A>[K]
+	>;
+};
+
+export type AnyClient = Client<Registry<any>>;
+
+export function createClientWithDriver<A extends Registry<any>>(
+	driver: ManagerDriver,
+	config?: ClientConfig,
+): Client<A> {
+	const client = new ClientRaw(driver, config);
+
+	// Create proxy for accessing actors by name
+	return new Proxy(client, {
+		get: (target: ClientRaw, prop: string | symbol, receiver: unknown) => {
+			// Get the real property if it exists
+			if (typeof prop === "symbol" || prop in target) {
+				const value = Reflect.get(target, prop, receiver);
+				// Preserve method binding
+				if (typeof value === "function") {
+					return value.bind(target);
+				}
+				return value;
+			}
+
+			// Handle actor accessor for string properties (actor names)
+			if (typeof prop === "string") {
+				// Return actor accessor object with methods
+				return {
+					// Handle methods (stateless action)
+					get: (
+						key?: string | string[],
+						opts?: GetWithIdOptions,
+					): ActorHandle<ExtractActorsFromRegistry<A>[typeof prop]> => {
+						return target.get<ExtractActorsFromRegistry<A>[typeof prop]>(
+							prop,
+							key,
+							opts,
+						);
+					},
+					getOrCreate: (
+						key?: string | string[],
+						opts?: GetOptions,
+					): ActorHandle<ExtractActorsFromRegistry<A>[typeof prop]> => {
+						return target.getOrCreate<
+							ExtractActorsFromRegistry<A>[typeof prop]
+						>(prop, key, opts);
+					},
+					getForId: (
+						actorId: string,
+						opts?: GetWithIdOptions,
+					): ActorHandle<ExtractActorsFromRegistry<A>[typeof prop]> => {
+						return target.getForId<ExtractActorsFromRegistry<A>[typeof prop]>(
+							prop,
+							actorId,
+							opts,
+						);
+					},
+					create: async (
+						key: string | string[],
+						opts: CreateOptions = {},
+					): Promise<
+						ActorHandle<ExtractActorsFromRegistry<A>[typeof prop]>
+					> => {
+						return await target.create<
+							ExtractActorsFromRegistry<A>[typeof prop]
+						>(prop, key, opts);
+					},
+				} as ActorAccessor<ExtractActorsFromRegistry<A>[typeof prop]>;
+			}
+
+			return undefined;
+		},
+	}) as Client<A>;
+}
+
+/**
+ * Creates a proxy for a actor that enables calling actions without explicitly using `.action`.
+ **/
+function createActorProxy<AD extends AnyActorDefinition>(
+	handle: ActorHandleRaw | ActorConnRaw,
+): ActorHandle<AD> | ActorConn<AD> {
+	// Stores returned action functions for faster calls
+	const methodCache = new Map<string, ActorActionFunction>();
+	return new Proxy(handle, {
+		get(target: ActorHandleRaw, prop: string | symbol, receiver: unknown) {
+			// Handle built-in Symbol properties
+			if (typeof prop === "symbol") {
+				return Reflect.get(target, prop, receiver);
+			}
+
+			// Handle built-in Promise methods and existing properties
+			if (prop === "constructor" || prop in target) {
+				const value = Reflect.get(target, prop, receiver);
+				// Preserve method binding
+				if (typeof value === "function") {
+					return value.bind(target);
+				}
+				return value;
+			}
+
+			// Create action function that preserves 'this' context
+			if (typeof prop === "string") {
+				// If JS is attempting to calling this as a promise, ignore it
+				if (prop === "then") return undefined;
+
+				let method = methodCache.get(prop);
+				if (!method) {
+					method = (...args: unknown[]) => target.action({ name: prop, args });
+					methodCache.set(prop, method);
+				}
+				return method;
+			}
+		},
+
+		// Support for 'in' operator
+		has(target: ActorHandleRaw, prop: string | symbol) {
+			// All string properties are potentially action functions
+			if (typeof prop === "string") {
+				return true;
+			}
+			// For symbols, defer to the target's own has behavior
+			return Reflect.has(target, prop);
+		},
+
+		// Support instanceof checks
+		getPrototypeOf(target: ActorHandleRaw) {
+			return Reflect.getPrototypeOf(target);
+		},
+
+		// Prevent property enumeration of non-existent action methods
+		ownKeys(target: ActorHandleRaw) {
+			return Reflect.ownKeys(target);
+		},
+
+		// Support proper property descriptors
+		getOwnPropertyDescriptor(target: ActorHandleRaw, prop: string | symbol) {
+			const targetDescriptor = Reflect.getOwnPropertyDescriptor(target, prop);
+			if (targetDescriptor) {
+				return targetDescriptor;
+			}
+			if (typeof prop === "string") {
+				// Make action methods appear non-enumerable
+				return {
+					configurable: true,
+					enumerable: false,
+					writable: false,
+					value: (...args: unknown[]) => target.action({ name: prop, args }),
+				};
+			}
+			return undefined;
+		},
+	}) as ActorHandle<AD> | ActorConn<AD>;
+}

package/src/client/config.ts

@@ -0,0 +1,44 @@
+import z from "zod";
+import { TransportSchema } from "@/actor/protocol/old";
+import { EncodingSchema } from "@/actor/protocol/serde";
+import { getEnvUniversal, type UpgradeWebSocket } from "@/utils";
+
+export type GetUpgradeWebSocket = () => UpgradeWebSocket;
+
+export const ClientConfigSchema = z.object({
+	/** Configure serving the API */
+	api: z
+		.object({
+			host: z.string().default("127.0.0.1"),
+			port: z.number().default(6420),
+		})
+		.default({}),
+
+	/** Endpoint to connect to the Rivet engine. Can be configured via RIVET_ENGINE env var. */
+	endpoint: z
+		.string()
+		.nullable()
+		.default(() => getEnvUniversal("RIVET_ENGINE") ?? null),
+
+	namespace: z
+		.string()
+		.default(() => getEnvUniversal("RIVET_NAMESPACE") ?? "default"),
+
+	runnerName: z
+		.string()
+		.default(() => getEnvUniversal("RIVET_RUNNER") ?? "rivetkit"),
+
+	encoding: EncodingSchema.default("bare"),
+
+	transport: TransportSchema.default("websocket"),
+
+	// This is a function to allow for lazy configuration of upgradeWebSocket on the
+	// fly. This is required since the dependencies that 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(),
+});
+
+export type ClientConfig = z.infer<typeof ClientConfigSchema>;
+
+export type ClientConfigInput = z.input<typeof ClientConfigSchema>;

package/src/client/errors.ts

@@ -0,0 +1,42 @@
+import { MAX_CONN_PARAMS_SIZE } from "@/common//network";
+
+export class ActorClientError extends Error {}
+
+export class InternalError extends ActorClientError {}
+
+export class ManagerError extends ActorClientError {
+	constructor(error: string, opts?: ErrorOptions) {
+		super(`Manager error: ${error}`, opts);
+	}
+}
+
+export class MalformedResponseMessage extends ActorClientError {
+	constructor(cause?: unknown) {
+		super(`Malformed response message: ${cause}`, { cause });
+	}
+}
+
+export class ActorError extends ActorClientError {
+	__type = "ActorError";
+
+	constructor(
+		public readonly group: string,
+		public readonly code: string,
+		message: string,
+		public readonly metadata?: unknown,
+	) {
+		super(message);
+	}
+}
+
+export class HttpRequestError extends ActorClientError {
+	constructor(message: string, opts?: { cause?: unknown }) {
+		super(`HTTP request error: ${message}`, { cause: opts?.cause });
+	}
+}
+
+export class ActorConnDisposed extends ActorClientError {
+	constructor() {
+		super("Attempting to interact with a disposed actor connection.");
+	}
+}

package/src/client/log.ts

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

package/src/client/mod.ts

@@ -0,0 +1,60 @@
+import type { Registry } from "@/registry/mod";
+import { RemoteManagerDriver } from "@/remote-manager-driver/mod";
+import {
+	type Client,
+	type ClientConfigInput,
+	createClientWithDriver,
+} from "./client";
+import { ClientConfigSchema } from "./config";
+
+export {
+	ActorDefinition,
+	AnyActorDefinition,
+} from "@/actor/definition";
+export type { Transport } from "@/actor/protocol/old";
+export type { Encoding } from "@/actor/protocol/serde";
+export {
+	ActorClientError,
+	ActorError,
+	InternalError,
+	MalformedResponseMessage,
+	ManagerError,
+} from "@/client/errors";
+export type { CreateRequest } from "@/manager/protocol/query";
+export type { ActorActionFunction } from "./actor-common";
+export type { ActorConn, EventUnsubscribe } from "./actor-conn";
+export { ActorConnRaw } from "./actor-conn";
+export type { ActorHandle } from "./actor-handle";
+export { ActorHandleRaw } from "./actor-handle";
+export type {
+	ActorAccessor,
+	Client,
+	ClientRaw,
+	CreateOptions,
+	ExtractActorsFromRegistry,
+	ExtractRegistryFromClient,
+	GetOptions,
+	GetWithIdOptions,
+	QueryOptions,
+	Region,
+} from "./client";
+
+/**
+ * Creates a client with the actor accessor proxy.
+ */
+export function createClient<A extends Registry<any>>(
+	endpointOrConfig?: string | ClientConfigInput,
+): Client<A> {
+	// Parse config
+	const configInput =
+		endpointOrConfig === undefined
+			? {}
+			: typeof endpointOrConfig === "string"
+				? { engine: endpointOrConfig }
+				: endpointOrConfig;
+	const config = ClientConfigSchema.parse(configInput);
+
+	// Create client
+	const driver = new RemoteManagerDriver(config);
+	return createClientWithDriver<A>(driver, config);
+}