rivetkit 2.0.3 → 2.0.4

package/src/client/actor-conn.ts

@@ -1,17 +1,26 @@
 import * as cbor from "cbor-x";
 import invariant from "invariant";
 import pRetry from "p-retry";
-import type { CloseEvent, WebSocket } from "ws";
+import type { CloseEvent } from "ws";
 import type { AnyActorDefinition } from "@/actor/definition";
 import { inputDataToBuffer } from "@/actor/protocol/old";
 import { type Encoding, jsonStringifyCompat } from "@/actor/protocol/serde";
+import { importEventSource } from "@/common/eventsource";
 import type {
 	UniversalErrorEvent,
 	UniversalEventSource,
 	UniversalMessageEvent,
 } from "@/common/eventsource-interface";
 import { assertUnreachable, stringifyError } from "@/common/utils";
+import {
+	HEADER_CONN_ID,
+	HEADER_CONN_PARAMS,
+	HEADER_CONN_TOKEN,
+	HEADER_ENCODING,
+	type ManagerDriver,
+} from "@/driver-helpers/mod";
 import type { ActorQuery } from "@/manager/protocol/query";
+import { PATH_CONNECT_WEBSOCKET, type UniversalWebSocket } from "@/mod";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import {
 	TO_CLIENT_VERSIONED,
@@ -22,17 +31,17 @@ import {
 	encodingIsBinary,
 	serializeWithEncoding,
 } from "@/serde";
-import { bufferToArrayBuffer, getEnvUniversal } from "@/utils";
+import { bufferToArrayBuffer, getEnvUniversal, httpUserAgent } from "@/utils";
 import type { ActorDefinitionActions } from "./actor-common";
-import {
-	ACTOR_CONNS_SYMBOL,
-	type ClientDriver,
-	type ClientRaw,
-	TRANSPORT_SYMBOL,
-} from "./client";
+import { queryActor } from "./actor-query";
+import { ACTOR_CONNS_SYMBOL, type ClientRaw, TRANSPORT_SYMBOL } from "./client";
 import * as errors from "./errors";
 import { logger } from "./log";
-import { type WebSocketMessage as ConnMessage, messageLength } from "./utils";
+import {
+	type WebSocketMessage as ConnMessage,
+	messageLength,
+	sendHttpRequest,
+} from "./utils";
 
 interface ActionInFlight {
 	name: string;
@@ -65,7 +74,7 @@ export interface SendHttpMessageOpts {
 }
 
 export type ConnTransport =
-	| { websocket: WebSocket }
+	| { websocket: UniversalWebSocket }
 	| { sse: UniversalEventSource };
 
 export const CONNECT_SYMBOL = Symbol("connect");
@@ -112,9 +121,9 @@ export class ActorConnRaw {
 	#onOpenPromise?: PromiseWithResolvers<undefined>;
 
 	#client: ClientRaw;
-	#driver: ClientDriver;
+	#driver: ManagerDriver;
 	#params: unknown;
-	#encodingKind: Encoding;
+	#encoding: Encoding;
 	#actorQuery: ActorQuery;
 
 	// TODO: ws message queue
@@ -128,15 +137,15 @@ export class ActorConnRaw {
 	 */
 	public constructor(
 		client: ClientRaw,
-		driver: ClientDriver,
+		driver: ManagerDriver,
 		params: unknown,
-		encodingKind: Encoding,
+		encoding: Encoding,
 		actorQuery: ActorQuery,
 	) {
 		this.#client = client;
 		this.#driver = driver;
 		this.#params = params;
-		this.#encodingKind = encodingKind;
+		this.#encoding = encoding;
 		this.#actorQuery = actorQuery;
 
 		this.#keepNodeAliveInterval = setInterval(() => 60_000);
@@ -160,7 +169,7 @@ export class ActorConnRaw {
 		args: Args;
 		signal?: AbortSignal;
 	}): Promise<Response> {
-		logger().debug("action", { name: opts.name, args: opts.args });
+		logger().debug({ msg: "action", name: opts.name, args: opts.args });
 
 		// If we have an active connection, use the websockactionId
 		const actionId = this.#actionIdCounter;
@@ -214,7 +223,8 @@ enc
 				maxTimeout: 30_000,
 
 				onFailedAttempt: (error) => {
-					logger().warn("failed to reconnect", {
+					logger().warn({
+						msg: "failed to reconnect",
 						attempt: error.attemptNumber,
 						error: stringifyError(error),
 					});
@@ -226,7 +236,7 @@ enc
 		} catch (err) {
 			if ((err as Error).name === "AbortError") {
 				// Ignore abortions
-				logger().info("connection retry aborted");
+				logger().info({ msg: "connection retry aborted" });
 				return;
 			} else {
 				// Unknown error
@@ -260,17 +270,21 @@ enc
 		}
 	}
 
-	async #connectWebSocket({ signal }: { signal?: AbortSignal } = {}) {
-		const ws = await this.#driver.connectWebSocket(
+	async #connectWebSocket() {
+		const { actorId } = await queryActor(
 			undefined,
 			this.#actorQuery,
-			this.#encodingKind,
+			this.#driver,
+		);
+		const ws = await this.#driver.openWebSocket(
+			PATH_CONNECT_WEBSOCKET,
+			actorId,
+			this.#encoding,
 			this.#params,
-			signal ? { signal } : undefined,
 		);
 		this.#transport = { websocket: ws };
 		ws.addEventListener("open", () => {
-			logger().debug("websocket open");
+			logger().debug({ msg: "websocket open" });
 		});
 		ws.addEventListener("message", async (ev) => {
 			this.#handleOnMessage(ev.data);
@@ -283,36 +297,58 @@ enc
 		});
 	}
 
-	async #connectSse({ signal }: { signal?: AbortSignal } = {}) {
-		const eventSource = await this.#driver.connectSse(
+	async #connectSse() {
+		const EventSource = await importEventSource();
+
+		// Get the actor ID
+		const { actorId } = await queryActor(
 			undefined,
 			this.#actorQuery,
-			this.#encodingKind,
-			this.#params,
-			signal ? { signal } : undefined,
+			this.#driver,
 		);
+		logger().debug({ msg: "found actor for sse connection", actorId });
+		invariant(actorId, "Missing actor ID");
+
+		logger().debug({
+			msg: "opening sse connection",
+			actorId,
+			encoding: this.#encoding,
+		});
+
+		const eventSource = new EventSource("http://actor/connect/sse", {
+			fetch: (input, init) => {
+				return this.#driver.sendRequest(
+					actorId,
+					new Request(input, {
+						...init,
+						headers: {
+							...init?.headers,
+							"User-Agent": httpUserAgent(),
+							[HEADER_ENCODING]: this.#encoding,
+							...(this.#params !== undefined
+								? { [HEADER_CONN_PARAMS]: JSON.stringify(this.#params) }
+								: {}),
+						},
+					}),
+				);
+			},
+		}) as UniversalEventSource;
+
 		this.#transport = { sse: eventSource };
-		eventSource.onopen = () => {
-			logger().debug("eventsource open");
-			// #handleOnOpen is called on "i" event
-		};
-		eventSource.onmessage = (ev: UniversalMessageEvent) => {
+
+		eventSource.addEventListener("message", (ev: UniversalMessageEvent) => {
 			this.#handleOnMessage(ev.data);
-		};
-		eventSource.onerror = (_ev: UniversalErrorEvent) => {
-			if (eventSource.readyState === eventSource.CLOSED) {
-				// This error indicates a close event
-				this.#handleOnClose(new Event("error"));
-			} else {
-				// Log error since event source is still open
-				this.#handleOnError();
-			}
-		};
+		});
+
+		eventSource.addEventListener("error", (ev: UniversalErrorEvent) => {
+			this.#handleOnError();
+		});
 	}
 
 	/** Called by the onopen event from drivers. */
 	#handleOnOpen() {
-		logger().debug("socket open", {
+		logger().debug({
+			msg: "socket open",
 			messageQueueLength: this.#messageQueue.length,
 		});
 
@@ -320,7 +356,7 @@ enc
 		if (this.#onOpenPromise) {
 			this.#onOpenPromise.resolve(undefined);
 		} else {
-			logger().warn("#onOpenPromise is undefined");
+			logger().warn({ msg: "#onOpenPromise is undefined" });
 		}
 
 		// Resubscribe to all active events
@@ -340,7 +376,8 @@ enc
 
 	/** Called by the onmessage event from drivers. */
 	async #handleOnMessage(data: any) {
-		logger().trace("received message", {
+		logger().trace({
+			msg: "received message",
 			dataType: typeof data,
 			isBlob: data instanceof Blob,
 			isArrayBuffer: data instanceof ArrayBuffer,
@@ -348,12 +385,12 @@ enc
 
 		const response = await this.#parseMessage(data as ConnMessage);
 		logger().trace(
-			"parsed message",
 			getEnvUniversal("_RIVETKIT_LOG_MESSAGE")
 				? {
+						msg: "parsed message",
 						message: jsonStringifyCompat(response).substring(0, 100) + "...",
 					}
-				: {},
+				: { msg: "parsed message" },
 		);
 
 		if (response.body.tag === "Init") {
@@ -361,36 +398,46 @@ enc
 			this.#actorId = response.body.val.actorId;
 			this.#connectionId = response.body.val.connectionId;
 			this.#connectionToken = response.body.val.connectionToken;
-			logger().trace("received init message", {
+			logger().trace({
+				msg: "received init message",
 				actorId: this.#actorId,
 				connectionId: this.#connectionId,
 			});
 			this.#handleOnOpen();
 		} else if (response.body.tag === "Error") {
 			// Connection error
-			const { code, message, metadata, actionId } = response.body.val;
+			const { group, code, message, metadata, actionId } = response.body.val;
 
 			if (actionId) {
 				const inFlight = this.#takeActionInFlight(Number(actionId));
 
-				logger().warn("action error", {
+				logger().warn({
+					msg: "action error",
 					actionId: actionId,
 					actionName: inFlight?.name,
+					group,
 					code,
 					message,
 					metadata,
 				});
 
-				inFlight.reject(new errors.ActorError(code, message, metadata));
+				inFlight.reject(new errors.ActorError(group, code, message, metadata));
 			} else {
-				logger().warn("connection error", {
+				logger().warn({
+					msg: "connection error",
+					group,
 					code,
 					message,
 					metadata,
 				});
 
 				// Create a connection error
-				const actorError = new errors.ActorError(code, message, metadata);
+				const actorError = new errors.ActorError(
+					group,
+					code,
+					message,
+					metadata,
+				);
 
 				// If we have an onOpenPromise, reject it with the error
 				if (this.#onOpenPromise) {
@@ -409,18 +456,20 @@ enc
 		} else if (response.body.tag === "ActionResponse") {
 			// Action response OK
 			const { id: actionId } = response.body.val;
-			logger().trace("received action response", {
+			logger().trace({
+				msg: "received action response",
 				actionId,
 			});
 
 			const inFlight = this.#takeActionInFlight(Number(actionId));
-			logger().trace("resolving action promise", {
+			logger().trace({
+				msg: "resolving action promise",
 				actionId,
 				actionName: inFlight?.name,
 			});
 			inFlight.resolve(response.body.val);
 		} else if (response.body.tag === "Event") {
-			logger().trace("received event", { name: response.body.val.name });
+			logger().trace({ msg: "received event", name: response.body.val.name });
 			this.#dispatchEvent(response.body.val);
 		} else {
 			assertUnreachable(response.body);
@@ -442,13 +491,15 @@ enc
 		// These properties will be undefined
 		const closeEvent = event as CloseEvent;
 		if (closeEvent.wasClean) {
-			logger().info("socket closed", {
+			logger().info({
+				msg: "socket closed",
 				code: closeEvent.code,
 				reason: closeEvent.reason,
 				wasClean: closeEvent.wasClean,
 			});
 		} else {
-			logger().warn("socket closed", {
+			logger().warn({
+				msg: "socket closed",
 				code: closeEvent.code,
 				reason: closeEvent.reason,
 				wasClean: closeEvent.wasClean,
@@ -514,7 +565,8 @@ enc
 			try {
 				handler(error);
 			} catch (err) {
-				logger().error("Error in connection error handler", {
+				logger().error({
+					msg: "error in connection error handler",
 					error: stringifyError(err),
 				});
 			}
@@ -559,7 +611,7 @@ enc
 	 * @param {string} eventName - The name of the event to subscribe to.
 	 * @param {(...args: Args) => void} callback - The callback function to execute when the event is triggered.
 	 * @returns {EventUnsubscribe} - A function to unsubscribe from the event.
-	 * @see {@link https://rivet.gg/docs/events|Events Documentation}
+	 * @see {@link https://rivet.dev/docs/events|Events Documentation}
 	 */
 	on<Args extends Array<unknown> = unknown[]>(
 		eventName: string,
@@ -575,7 +627,7 @@ enc
 	 * @param {string} eventName - The name of the event to subscribe to.
 	 * @param {(...args: Args) => void} callback - The callback function to execute when the event is triggered.
 	 * @returns {EventUnsubscribe} - A function to unsubscribe from the event.
-	 * @see {@link https://rivet.gg/docs/events|Events Documentation}
+	 * @see {@link https://rivet.dev/docs/events|Events Documentation}
 	 */
 	once<Args extends Array<unknown> = unknown[]>(
 		eventName: string,
@@ -612,16 +664,18 @@ enc
 			if (this.#transport.websocket.readyState === 1) {
 				try {
 					const messageSerialized = serializeWithEncoding(
-						this.#encodingKind,
+						this.#encoding,
 						message,
 						TO_SERVER_VERSIONED,
 					);
 					this.#transport.websocket.send(messageSerialized);
-					logger().trace("sent websocket message", {
+					logger().trace({
+						msg: "sent websocket message",
 						len: messageLength(messageSerialized),
 					});
 				} catch (error) {
-					logger().warn("failed to send message, added to queue", {
+					logger().warn({
+						msg: "failed to send message, added to queue",
 						error,
 					});
 
@@ -644,7 +698,7 @@ enc
 
 		if (!opts?.ephemeral && queueMessage) {
 			this.#messageQueue.push(message);
-			logger().debug("queued connection message");
+			logger().debug({ msg: "queued connection message" });
 		}
 	}
 
@@ -657,29 +711,40 @@ enc
 				throw new errors.InternalError("Missing connection ID or token.");
 
 			logger().trace(
-				"sent http message",
 				getEnvUniversal("_RIVETKIT_LOG_MESSAGE")
 					? {
-							message: jsonStringifyCompat(message).substring(0, 100) + "...",
+							msg: "sent http message",
+							message: `${jsonStringifyCompat(message).substring(0, 100)}...`,
 						}
-					: {},
+					: { msg: "sent http message" },
 			);
 
-			await this.#driver.sendHttpMessage(
-				undefined,
-				this.#actorId,
-				this.#encodingKind,
-				this.#connectionId,
-				this.#connectionToken,
-				message,
-				opts?.signal ? { signal: opts.signal } : undefined,
-			);
+			logger().debug({
+				msg: "sending http message",
+				actorId: this.#actorId,
+				connectionId: this.#connectionId,
+			});
+
+			// Send an HTTP request to the connections endpoint
+			await sendHttpRequest({
+				url: "http://actor/connections/message",
+				method: "POST",
+				headers: {
+					[HEADER_ENCODING]: this.#encoding,
+					[HEADER_CONN_ID]: this.#connectionId,
+					[HEADER_CONN_TOKEN]: this.#connectionToken,
+				},
+				body: message,
+				encoding: this.#encoding,
+				skipParseResponse: true,
+				customFetch: this.#driver.sendRequest.bind(this.#driver, this.#actorId),
+				requestVersionedDataHandler: TO_SERVER_VERSIONED,
+				responseVersionedDataHandler: TO_CLIENT_VERSIONED,
+			});
 		} catch (error) {
 			// TODO: This will not automatically trigger a re-broadcast of HTTP events since SSE is separate from the HTTP action
 
-			logger().warn("failed to send message, added to queue", {
-				error,
-			});
+			logger().warn({ msg: "failed to send message, added to queue", error });
 
 			// Assuming the socket is disconnected and will be reconnected soon
 			//
@@ -694,7 +759,7 @@ enc
 		invariant(this.#transport, "transport must be defined");
 
 		// Decode base64 since SSE sends raw strings
-		if (encodingIsBinary(this.#encodingKind) && "sse" in this.#transport) {
+		if (encodingIsBinary(this.#encoding) && "sse" in this.#transport) {
 			if (typeof data === "string") {
 				const binaryString = atob(data);
 				data = new Uint8Array(
@@ -709,11 +774,7 @@ enc
 
 		const buffer = await inputDataToBuffer(data);
 
-		return deserializeWithEncoding(
-			this.#encodingKind,
-			buffer,
-			TO_CLIENT_VERSIONED,
-		);
+		return deserializeWithEncoding(this.#encoding, buffer, TO_CLIENT_VERSIONED);
 	}
 
 	/**
@@ -725,12 +786,12 @@ enc
 		// Internally, this "disposes" the connection
 
 		if (this.#disposed) {
-			logger().warn("connection already disconnected");
+			logger().warn({ msg: "connection already disconnected" });
 			return;
 		}
 		this.#disposed = true;
 
-		logger().debug("disposing actor conn");
+		logger().debug({ msg: "disposing actor conn" });
 
 		// Clear interval so NodeJS process can exit
 		clearInterval(this.#keepNodeAliveInterval);
@@ -751,11 +812,11 @@ enc
 				ws.readyState === 2 /* CLOSING */ ||
 				ws.readyState === 3 /* CLOSED */
 			) {
-				logger().debug("ws already closed or closing");
+				logger().debug({ msg: "ws already closed or closing" });
 			} else {
 				const { promise, resolve } = Promise.withResolvers();
 				ws.addEventListener("close", () => {
-					logger().debug("ws closed");
+					logger().debug({ msg: "ws closed" });
 					resolve(undefined);
 				});
 				ws.close();

package/src/client/actor-handle.ts

@@ -1,18 +1,30 @@
+import * as cbor from "cbor-x";
 import invariant from "invariant";
 import type { AnyActorDefinition } from "@/actor/definition";
 import type { Encoding } from "@/actor/protocol/serde";
 import { assertUnreachable } from "@/actor/utils";
+import { deconstructError } from "@/common/utils";
 import { importWebSocket } from "@/common/websocket";
+import {
+	HEADER_CONN_PARAMS,
+	HEADER_ENCODING,
+	type ManagerDriver,
+} from "@/driver-helpers/mod";
 import type { ActorQuery } from "@/manager/protocol/query";
+import type * as protocol from "@/schemas/client-protocol/mod";
+import {
+	HTTP_ACTION_REQUEST_VERSIONED,
+	HTTP_ACTION_RESPONSE_VERSIONED,
+} from "@/schemas/client-protocol/versioned";
+import { bufferToArrayBuffer } from "@/utils";
 import type { ActorDefinitionActions } from "./actor-common";
 import { type ActorConn, ActorConnRaw } from "./actor-conn";
-import {
-	type ClientDriver,
-	type ClientRaw,
-	CREATE_ACTOR_CONN_PROXY,
-} from "./client";
+import { queryActor } from "./actor-query";
+import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client";
+import { ActorError } from "./errors";
 import { logger } from "./log";
 import { rawHttpFetch, rawWebSocket } from "./raw-utils";
+import { sendHttpRequest } from "./utils";
 
 /**
  * Provides underlying functions for stateless {@link ActorHandle} for action calls.
@@ -22,8 +34,8 @@ import { rawHttpFetch, rawWebSocket } from "./raw-utils";
  */
 export class ActorHandleRaw {
 	#client: ClientRaw;
-	#driver: ClientDriver;
-	#encodingKind: Encoding;
+	#driver: ManagerDriver;
+	#encoding: Encoding;
 	#actorQuery: ActorQuery;
 	#params: unknown;
 
@@ -36,14 +48,14 @@ export class ActorHandleRaw {
 	 */
 	public constructor(
 		client: any,
-		driver: ClientDriver,
+		driver: ManagerDriver,
 		params: unknown,
-		encodingKind: Encoding,
+		encoding: Encoding,
 		actorQuery: ActorQuery,
 	) {
 		this.#client = client;
 		this.#driver = driver;
-		this.#encodingKind = encodingKind;
+		this.#encoding = encoding;
 		this.#actorQuery = actorQuery;
 		this.#params = params;
 	}
@@ -63,15 +75,64 @@ export class ActorHandleRaw {
 		args: Args;
 		signal?: AbortSignal;
 	}): Promise<Response> {
-		return await this.#driver.action<Args, Response>(
-			undefined,
-			this.#actorQuery,
-			this.#encodingKind,
-			this.#params,
-			opts.name,
-			opts.args,
-			{ signal: opts.signal },
-		);
+		// return await this.#driver.action<Args, Response>(
+		// 	undefined,
+		// 	this.#actorQuery,
+		// 	this.#encodingKind,
+		// 	this.#params,
+		// 	opts.name,
+		// 	opts.args,
+		// 	{ signal: opts.signal },
+		// );
+		try {
+			// Get the actor ID
+			const { actorId } = await queryActor(
+				undefined,
+				this.#actorQuery,
+				this.#driver,
+			);
+			logger().debug({ msg: "found actor for action", actorId });
+			invariant(actorId, "Missing actor ID");
+
+			// Invoke the action
+			logger().debug({
+				msg: "handling action",
+				name: opts.name,
+				encoding: this.#encoding,
+			});
+			const responseData = await sendHttpRequest<
+				protocol.HttpActionRequest,
+				protocol.HttpActionResponse
+			>({
+				url: `http://actor/action/${encodeURIComponent(opts.name)}`,
+				method: "POST",
+				headers: {
+					[HEADER_ENCODING]: this.#encoding,
+					...(this.#params !== undefined
+						? { [HEADER_CONN_PARAMS]: JSON.stringify(this.#params) }
+						: {}),
+				},
+				body: {
+					args: bufferToArrayBuffer(cbor.encode(opts.args)),
+				} satisfies protocol.HttpActionRequest,
+				encoding: this.#encoding,
+				customFetch: this.#driver.sendRequest.bind(this.#driver, actorId),
+				signal: opts?.signal,
+				requestVersionedDataHandler: HTTP_ACTION_REQUEST_VERSIONED,
+				responseVersionedDataHandler: HTTP_ACTION_RESPONSE_VERSIONED,
+			});
+
+			return cbor.decode(new Uint8Array(responseData.output));
+		} catch (err) {
+			// Standardize to ClientActorError instead of the native backend error
+			const { group, code, message, metadata } = deconstructError(
+				err,
+				logger(),
+				{},
+				true,
+			);
+			throw new ActorError(group, code, message, metadata);
+		}
 	}
 
 	/**
@@ -81,7 +142,8 @@ export class ActorHandleRaw {
 	 * @returns {ActorConn<AD>} A connection to the actor.
 	 */
 	connect(): ActorConn<AnyActorDefinition> {
-		logger().debug("establishing connection from handle", {
+		logger().debug({
+			msg: "establishing connection from handle",
 			query: this.#actorQuery,
 		});
 
@@ -89,7 +151,7 @@ export class ActorHandleRaw {
 			this.#client,
 			this.#driver,
 			this.#params,
-			this.#encodingKind,
+			this.#encoding,
 			this.#actorQuery,
 		);
 
@@ -158,12 +220,10 @@ export class ActorHandleRaw {
 				assertUnreachable(this.#actorQuery);
 			}
 
-			const actorId = await this.#driver.resolveActorId(
+			const { actorId } = await queryActor(
 				undefined,
 				this.#actorQuery,
-				this.#encodingKind,
-				this.#params,
-				signal ? { signal } : undefined,
+				this.#driver,
 			);
 
 			this.#actorQuery = { getForId: { actorId, name } };