rivetkit 2.0.24-rc.1 → 2.0.24

package/src/client/actor-conn.ts

@@ -6,11 +6,9 @@ import type { AnyActorDefinition } from "@/actor/definition";
 import { inputDataToBuffer } from "@/actor/protocol/old";
 import { type Encoding, jsonStringifyCompat } from "@/actor/protocol/serde";
 import {
-	HEADER_CONN_ID,
 	HEADER_CONN_PARAMS,
-	HEADER_CONN_TOKEN,
 	HEADER_ENCODING,
-	PATH_CONNECT_WEBSOCKET,
+	PATH_CONNECT,
 } from "@/common/actor-router-consts";
 import { importEventSource } from "@/common/eventsource";
 import type {
@@ -27,6 +25,12 @@ import {
 	TO_CLIENT_VERSIONED,
 	TO_SERVER_VERSIONED,
 } from "@/schemas/client-protocol/versioned";
+import {
+	type ToClient as ToClientJson,
+	ToClientSchema,
+	type ToServer as ToServerJson,
+	ToServerSchema,
+} from "@/schemas/client-protocol-zod/mod";
 import {
 	deserializeWithEncoding,
 	encodingIsBinary,
@@ -40,7 +44,7 @@ import {
 } from "@/utils";
 import type { ActorDefinitionActions } from "./actor-common";
 import { queryActor } from "./actor-query";
-import { ACTOR_CONNS_SYMBOL, type ClientRaw, TRANSPORT_SYMBOL } from "./client";
+import { ACTOR_CONNS_SYMBOL, type ClientRaw } from "./client";
 import * as errors from "./errors";
 import { logger } from "./log";
 import {
@@ -51,7 +55,7 @@ import {
 
 interface ActionInFlight {
 	name: string;
-	resolve: (response: protocol.ActionResponse) => void;
+	resolve: (response: { id: bigint; output: unknown }) => void;
 	reject: (error: Error) => void;
 }
 
@@ -79,10 +83,6 @@ export interface SendHttpMessageOpts {
 	signal?: AbortSignal;
 }
 
-export type ConnTransport =
-	| { websocket: UniversalWebSocket }
-	| { sse: UniversalEventSource };
-
 export const CONNECT_SYMBOL = Symbol("connect");
 
 /**
@@ -96,17 +96,22 @@ export class ActorConnRaw {
 	/* Will be aborted on dispose. */
 	#abortController = new AbortController();
 
-	/** If attempting to connect. Helpful for knowing if in a retry loop when reconnecting. */
 	#connecting = false;
 
-	// Connection info, used for reconnection and HTTP requests
 	#actorId?: string;
 	#connectionId?: string;
-	#connectionToken?: string;
-
-	#transport?: ConnTransport;
 
-	#messageQueue: protocol.ToServer[] = [];
+	#messageQueue: Array<{
+		body:
+			| {
+					tag: "ActionRequest";
+					val: { id: bigint; name: string; args: unknown };
+			  }
+			| {
+					tag: "SubscriptionRequest";
+					val: { eventName: string; subscribe: boolean };
+			  };
+	}> = [];
 	#actionsInFlight = new Map<number, ActionInFlight>();
 
 	// biome-ignore lint/suspicious/noExplicitAny: Unknown subscription type
@@ -126,6 +131,8 @@ export class ActorConnRaw {
 	/** Promise used to indicate the socket has connected successfully. This will be rejected if the connection fails. */
 	#onOpenPromise?: ReturnType<typeof promiseWithResolvers<undefined>>;
 
+	#websocket?: UniversalWebSocket;
+
 	#client: ClientRaw;
 	#driver: ManagerDriver;
 	#params: unknown;
@@ -181,13 +188,21 @@ export class ActorConnRaw {
 		const actionId = this.#actionIdCounter;
 		this.#actionIdCounter += 1;
 
-		const { promise, resolve, reject } =
-			promiseWithResolvers<protocol.ActionResponse>();
+		const { promise, resolve, reject } = promiseWithResolvers<{
+			id: bigint;
+			output: unknown;
+		}>();
 		this.#actionsInFlight.set(actionId, {
 			name: opts.name,
 			resolve,
 			reject,
 		});
+		logger().debug({
+			msg: "added action to in-flight map",
+			actionId,
+			actionName: opts.name,
+			inFlightCount: this.#actionsInFlight.size,
+		});
 
 		this.#sendMessage({
 			body: {
@@ -195,10 +210,10 @@ export class ActorConnRaw {
 				val: {
 					id: BigInt(actionId),
 					name: opts.name,
-					args: bufferToArrayBuffer(cbor.encode(opts.args)),
+					args: opts.args,
 				},
 			},
-		} satisfies protocol.ToServer);
+		});
 
 		// TODO: Throw error if disconnect is called
 
@@ -208,7 +223,7 @@ export class ActorConnRaw {
 				`Request ID ${actionId} does not match response ID ${responseId}`,
 			);
 
-		return cbor.decode(new Uint8Array(output)) as Response;
+		return output as Response;
 	}
 
 	/**
@@ -264,14 +279,7 @@ enc
 				throw new Error("#onOpenPromise already defined");
 			this.#onOpenPromise = promiseWithResolvers();
 
-			// Connect transport
-			if (this.#client[TRANSPORT_SYMBOL] === "websocket") {
-				await this.#connectWebSocket();
-			} else if (this.#client[TRANSPORT_SYMBOL] === "sse") {
-				await this.#connectSse();
-			} else {
-				assertUnreachable(this.#client[TRANSPORT_SYMBOL]);
-			}
+			await this.#connectWebSocket();
 
 			// Wait for result
 			await this.#onOpenPromise.promise;
@@ -287,31 +295,19 @@ enc
 			this.#driver,
 		);
 
-		// Check if we have connection info for reconnection
-		const isReconnection = this.#connectionId && this.#connectionToken;
-		if (isReconnection) {
-			logger().debug({
-				msg: "attempting websocket reconnection",
-				connectionId: this.#connectionId,
-			});
-		}
-
 		const ws = await this.#driver.openWebSocket(
-			PATH_CONNECT_WEBSOCKET,
+			PATH_CONNECT,
 			actorId,
 			this.#encoding,
 			this.#params,
-			// Pass connection ID and token for reconnection if available
-			isReconnection ? this.#connectionId : undefined,
-			isReconnection ? this.#connectionToken : undefined,
 		);
 		logger().debug({
-			msg: "transport set to new websocket",
+			msg: "opened websocket",
 			connectionId: this.#connectionId,
 			readyState: ws.readyState,
 			messageQueueLength: this.#messageQueue.length,
 		});
-		this.#transport = { websocket: ws };
+		this.#websocket = ws;
 		ws.addEventListener("open", () => {
 			logger().debug({
 				msg: "client websocket open",
@@ -350,70 +346,6 @@ enc
 		});
 	}
 
-	async #connectSse() {
-		const EventSource = await importEventSource();
-
-		// Get the actor ID
-		const { actorId } = await queryActor(
-			undefined,
-			this.#actorQuery,
-			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 isReconnection = this.#connectionId && this.#connectionToken;
-
-		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,
-										),
-									}
-								: {}),
-							...(isReconnection
-								? {
-										[HEADER_CONN_ID]: this.#connectionId,
-										[HEADER_CONN_TOKEN]:
-											this.#connectionToken,
-									}
-								: {}),
-						},
-					}),
-				);
-			},
-		}) as UniversalEventSource;
-
-		this.#transport = { sse: eventSource };
-
-		eventSource.addEventListener("message", (ev: UniversalMessageEvent) => {
-			// Ignore pings
-			if (ev.type === "ping") return;
-
-			this.#handleOnMessage(ev.data);
-		});
-
-		eventSource.addEventListener("error", (ev: UniversalErrorEvent) => {
-			this.#handleOnError();
-		});
-	}
-
 	/** Called by the onopen event from drivers. */
 	#handleOnOpen() {
 		logger().debug({
@@ -470,10 +402,9 @@ enc
 		);
 
 		if (response.body.tag === "Init") {
-			// Store connection info for reconnection
+			// Store connection info
 			this.#actorId = response.body.val.actorId;
 			this.#connectionId = response.body.val.connectionId;
-			this.#connectionToken = response.body.val.connectionToken;
 			logger().trace({
 				msg: "received init message",
 				actorId: this.#actorId,
@@ -535,9 +466,11 @@ enc
 		} else if (response.body.tag === "ActionResponse") {
 			// Action response OK
 			const { id: actionId } = response.body.val;
-			logger().trace({
+			logger().debug({
 				msg: "received action response",
-				actionId,
+				actionId: Number(actionId),
+				inFlightCount: this.#actionsInFlight.size,
+				inFlightIds: Array.from(this.#actionsInFlight.keys()),
 			});
 
 			const inFlight = this.#takeActionInFlight(Number(actionId));
@@ -598,7 +531,7 @@ enc
 			});
 
 			const disconnectError = new Error(
-				wasClean ? "Connection closed" : "Connection lost",
+				`${wasClean ? "Connection closed" : "Connection lost"} (code: ${closeEvent.code}, reason: ${closeEvent.reason})`,
 			);
 
 			for (const actionInfo of this.#actionsInFlight.values()) {
@@ -607,7 +540,7 @@ enc
 			this.#actionsInFlight.clear();
 		}
 
-		this.#transport = undefined;
+		this.#websocket = undefined;
 
 		// Automatically reconnect. Skip if already attempting to connect.
 		if (!this.#disposed && !this.#connecting) {
@@ -636,22 +569,39 @@ enc
 	#takeActionInFlight(id: number): ActionInFlight {
 		const inFlight = this.#actionsInFlight.get(id);
 		if (!inFlight) {
+			logger().error({
+				msg: "action not found in in-flight map",
+				lookupId: id,
+				inFlightCount: this.#actionsInFlight.size,
+				inFlightIds: Array.from(this.#actionsInFlight.keys()),
+				inFlightActions: Array.from(
+					this.#actionsInFlight.entries(),
+				).map(([id, action]) => ({
+					id,
+					name: action.name,
+				})),
+			});
 			throw new errors.InternalError(`No in flight response for ${id}`);
 		}
 		this.#actionsInFlight.delete(id);
+		logger().debug({
+			msg: "removed action from in-flight map",
+			actionId: id,
+			actionName: inFlight.name,
+			inFlightCount: this.#actionsInFlight.size,
+		});
 		return inFlight;
 	}
 
-	#dispatchEvent(event: protocol.Event) {
-		const { name, args: argsRaw } = event;
-		const args = cbor.decode(new Uint8Array(argsRaw));
+	#dispatchEvent(event: { name: string; args: unknown }) {
+		const { name, args } = event;
 
 		const listeners = this.#eventSubscriptions.get(name);
 		if (!listeners) return;
 
 		// Create a new array to avoid issues with listeners being removed during iteration
 		for (const listener of [...listeners]) {
-			listener.callback(...args);
+			listener.callback(...(args as unknown[]));
 
 			// Remove if this was a one-time listener
 			if (listener.once) {
@@ -757,18 +707,27 @@ enc
 		};
 	}
 
-	#sendMessage(message: protocol.ToServer, opts?: SendHttpMessageOpts) {
+	#sendMessage(
+		message: {
+			body:
+				| {
+						tag: "ActionRequest";
+						val: { id: bigint; name: string; args: unknown };
+				  }
+				| {
+						tag: "SubscriptionRequest";
+						val: { eventName: string; subscribe: boolean };
+				  };
+		},
+		opts?: SendHttpMessageOpts,
+	) {
 		if (this.#disposed) {
 			throw new errors.ActorConnDisposed();
 		}
 
 		let queueMessage = false;
-		if (!this.#transport) {
-			// No transport connected yet
-			logger().debug({ msg: "no transport, queueing message" });
-			queueMessage = true;
-		} else if ("websocket" in this.#transport) {
-			const readyState = this.#transport.websocket.readyState;
+		if (this.#websocket) {
+			const readyState = this.#websocket.readyState;
 			logger().debug({
 				msg: "websocket send attempt",
 				readyState,
@@ -790,8 +749,30 @@ enc
 						this.#encoding,
 						message,
 						TO_SERVER_VERSIONED,
+						ToServerSchema,
+						// JSON: args is the raw value
+						(msg): ToServerJson => msg as ToServerJson,
+						// BARE: args needs to be CBOR-encoded to ArrayBuffer
+						(msg): protocol.ToServer => {
+							if (msg.body.tag === "ActionRequest") {
+								return {
+									body: {
+										tag: "ActionRequest",
+										val: {
+											id: msg.body.val.id,
+											name: msg.body.val.name,
+											args: bufferToArrayBuffer(
+												cbor.encode(msg.body.val.args),
+											),
+										},
+									},
+								};
+							} else {
+								return msg as protocol.ToServer;
+							}
+						},
 					);
-					this.#transport.websocket.send(messageSerialized);
+					this.#websocket.send(messageSerialized);
 					logger().trace({
 						msg: "sent websocket message",
 						len: messageLength(messageSerialized),
@@ -813,15 +794,10 @@ enc
 				});
 				queueMessage = true;
 			}
-		} else if ("sse" in this.#transport) {
-			if (this.#transport.sse.readyState === 1) {
-				// Spawn in background since #sendMessage cannot be async
-				this.#sendHttpMessage(message, opts);
-			} else {
-				queueMessage = true;
-			}
 		} else {
-			assertUnreachable(this.#transport);
+			// No websocket connected yet
+			logger().debug({ msg: "no websocket, queueing message" });
+			queueMessage = true;
 		}
 
 		if (!opts?.ephemeral && queueMessage) {
@@ -836,83 +812,23 @@ enc
 		}
 	}
 
-	async #sendHttpMessage(
-		message: protocol.ToServer,
-		opts?: SendHttpMessageOpts,
-	) {
-		try {
-			if (!this.#actorId || !this.#connectionId || !this.#connectionToken)
-				throw new errors.InternalError(
-					"Missing connection ID or token.",
-				);
-
-			logger().trace(
-				getEnvUniversal("_RIVETKIT_LOG_MESSAGE")
-					? {
-							msg: "sent http message",
-							message: `${jsonStringifyCompat(message).substring(0, 100)}...`,
-						}
-					: { msg: "sent http message" },
-			);
-
-			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({
-				msg: "failed to send message, added to queue",
-				error,
-			});
-
-			// Assuming the socket is disconnected and will be reconnected soon
-			//
-			// Will attempt to resend soon
-			if (!opts?.ephemeral) {
-				this.#messageQueue.unshift(message);
-			}
-		}
-	}
-
-	async #parseMessage(data: ConnMessage): Promise<protocol.ToClient> {
-		invariant(this.#transport, "transport must be defined");
-
-		// Decode base64 since SSE sends raw strings
-		if (encodingIsBinary(this.#encoding) && "sse" in this.#transport) {
-			if (typeof data === "string") {
-				const binaryString = atob(data);
-				data = new Uint8Array(
-					[...binaryString].map((char) => char.charCodeAt(0)),
-				);
-			} else {
-				throw new errors.InternalError(
-					`Expected data to be a string for SSE, got ${data}.`,
-				);
-			}
-		}
+	async #parseMessage(data: ConnMessage): Promise<{
+		body:
+			| { tag: "Init"; val: { actorId: string; connectionId: string } }
+			| {
+					tag: "Error";
+					val: {
+						group: string;
+						code: string;
+						message: string;
+						metadata: unknown;
+						actionId: bigint | null;
+					};
+			  }
+			| { tag: "ActionResponse"; val: { id: bigint; output: unknown } }
+			| { tag: "Event"; val: { name: string; args: unknown } };
+	}> {
+		invariant(this.#websocket, "websocket must be defined");
 
 		const buffer = await inputDataToBuffer(data);
 
@@ -920,6 +836,59 @@ enc
 			this.#encoding,
 			buffer,
 			TO_CLIENT_VERSIONED,
+			ToClientSchema,
+			// JSON: values are already the correct type
+			(msg): ToClientJson => msg as ToClientJson,
+			// BARE: need to decode ArrayBuffer fields back to unknown
+			(msg): any => {
+				if (msg.body.tag === "Error") {
+					return {
+						body: {
+							tag: "Error",
+							val: {
+								group: msg.body.val.group,
+								code: msg.body.val.code,
+								message: msg.body.val.message,
+								metadata: msg.body.val.metadata
+									? cbor.decode(
+											new Uint8Array(
+												msg.body.val.metadata,
+											),
+										)
+									: null,
+								actionId: msg.body.val.actionId,
+							},
+						},
+					};
+				} else if (msg.body.tag === "ActionResponse") {
+					return {
+						body: {
+							tag: "ActionResponse",
+							val: {
+								id: msg.body.val.id,
+								output: cbor.decode(
+									new Uint8Array(msg.body.val.output),
+								),
+							},
+						},
+					};
+				} else if (msg.body.tag === "Event") {
+					return {
+						body: {
+							tag: "Event",
+							val: {
+								name: msg.body.val.name,
+								args: cbor.decode(
+									new Uint8Array(msg.body.val.args),
+								),
+							},
+						},
+					};
+				} else {
+					// Init has no ArrayBuffer fields
+					return msg;
+				}
+			},
 		);
 	}
 
@@ -964,13 +933,11 @@ enc
 		// Remove from registry
 		this.#client[ACTOR_CONNS_SYMBOL].delete(this);
 
-		// Disconnect transport cleanly
-		if (!this.#transport) {
-			// Nothing to do
-		} else if ("websocket" in this.#transport) {
+		// Disconnect websocket cleanly
+		if (this.#websocket) {
 			logger().debug("closing ws");
 
-			const ws = this.#transport.websocket;
+			const ws = this.#websocket;
 			// Check if WebSocket is already closed or closing
 			if (
 				ws.readyState === 2 /* CLOSING */ ||
@@ -986,41 +953,8 @@ enc
 				ws.close(1000, "Normal closure");
 				await promise;
 			}
-		} else if ("sse" in this.#transport) {
-			logger().debug("closing sse");
-
-			// Send close request to server for SSE connections
-			if (this.#connectionId && this.#connectionToken) {
-				try {
-					await sendHttpRequest({
-						url: "http://actor/connections/close",
-						method: "POST",
-						headers: {
-							[HEADER_CONN_ID]: this.#connectionId,
-							[HEADER_CONN_TOKEN]: this.#connectionToken,
-						},
-						encoding: this.#encoding,
-						skipParseResponse: true,
-						customFetch: this.#driver.sendRequest.bind(
-							this.#driver,
-							this.#actorId!,
-						),
-						requestVersionedDataHandler: TO_SERVER_VERSIONED,
-						responseVersionedDataHandler: TO_CLIENT_VERSIONED,
-					});
-				} catch (error) {
-					// Ignore errors when closing - connection may already be closed
-					logger().warn({
-						msg: "failed to send close request",
-						error,
-					});
-				}
-			}
-			this.#transport.sse.close();
-		} else {
-			assertUnreachable(this.#transport);
 		}
-		this.#transport = undefined;
+		this.#websocket = undefined;
 	}
 
 	#sendSubscription(eventName: string, subscribe: boolean) {

package/src/client/actor-handle.ts

@@ -15,6 +15,12 @@ import {
 	HTTP_ACTION_REQUEST_VERSIONED,
 	HTTP_ACTION_RESPONSE_VERSIONED,
 } from "@/schemas/client-protocol/versioned";
+import {
+	type HttpActionRequest as HttpActionRequestJson,
+	HttpActionRequestSchema,
+	type HttpActionResponse as HttpActionResponseJson,
+	HttpActionResponseSchema,
+} from "@/schemas/client-protocol-zod/mod";
 import { bufferToArrayBuffer } from "@/utils";
 import type { ActorDefinitionActions } from "./actor-common";
 import { type ActorConn, ActorConnRaw } from "./actor-conn";
@@ -100,8 +106,12 @@ export class ActorHandleRaw {
 				encoding: this.#encoding,
 			});
 			const responseData = await sendHttpRequest<
-				protocol.HttpActionRequest,
-				protocol.HttpActionResponse
+				protocol.HttpActionRequest, // Bare type
+				protocol.HttpActionResponse, // Bare type
+				HttpActionRequestJson, // Json type
+				HttpActionResponseJson, // Json type
+				unknown[], // Request type (the args array)
+				Response // Response type (the output value)
 			>({
 				url: `http://actor/action/${encodeURIComponent(opts.name)}`,
 				method: "POST",
@@ -111,9 +121,7 @@ export class ActorHandleRaw {
 						? { [HEADER_CONN_PARAMS]: JSON.stringify(this.#params) }
 						: {}),
 				},
-				body: {
-					args: bufferToArrayBuffer(cbor.encode(opts.args)),
-				} satisfies protocol.HttpActionRequest,
+				body: opts.args,
 				encoding: this.#encoding,
 				customFetch: this.#driver.sendRequest.bind(
 					this.#driver,
@@ -122,9 +130,24 @@ export class ActorHandleRaw {
 				signal: opts?.signal,
 				requestVersionedDataHandler: HTTP_ACTION_REQUEST_VERSIONED,
 				responseVersionedDataHandler: HTTP_ACTION_RESPONSE_VERSIONED,
+				requestZodSchema: HttpActionRequestSchema,
+				responseZodSchema: HttpActionResponseSchema,
+				// JSON Request: args is the raw value
+				requestToJson: (args): HttpActionRequestJson => ({
+					args,
+				}),
+				// BARE Request: args needs to be CBOR-encoded
+				requestToBare: (args): protocol.HttpActionRequest => ({
+					args: bufferToArrayBuffer(cbor.encode(args)),
+				}),
+				// JSON Response: output is the raw value
+				responseFromJson: (json): Response => json.output as Response,
+				// BARE Response: output is ArrayBuffer that needs CBOR-decoding
+				responseFromBare: (bare): Response =>
+					cbor.decode(new Uint8Array(bare.output)) as Response,
 			});
 
-			return cbor.decode(new Uint8Array(responseData.output));
+			return responseData;
 		} catch (err) {
 			// Standardize to ClientActorError instead of the native backend error
 			const { group, code, message, metadata } = deconstructError(