rivetkit 2.0.24-rc.1 → 2.0.25-rc.1

package/src/actor/protocol/old.ts

@@ -13,25 +13,24 @@ 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 } from "@/serde";
 import { assertUnreachable, bufferToArrayBuffer } from "../../utils";
-import { ActionContext } from "../action";
-import type { Conn } from "../conn";
-import type { ActorInstance } from "../instance";
-
-export const TransportSchema = z.enum(["websocket", "sse"]);
-
-/**
- * Transport mechanism used to communicate between client & actor.
- */
-export type Transport = z.infer<typeof TransportSchema>;
+import { CONN_SEND_MESSAGE_SYMBOL, type Conn } from "../conn/mod";
+import { ActionContext } from "../contexts/action";
+import type { ActorInstance } from "../instance/mod";
 
 interface MessageEventOpts {
 	encoding: Encoding;
 	maxIncomingMessageSize: number;
 }
 
-function getValueLength(value: InputData): number {
+export function getValueLength(value: InputData): number {
 	if (typeof value === "string") {
 		return value.length;
 	} else if (value instanceof Blob) {
@@ -70,7 +69,17 @@ export async function inputDataToBuffer(
 export async function parseMessage(
 	value: InputData,
 	opts: MessageEventOpts,
-): Promise<protocol.ToServer> {
+): Promise<{
+	body:
+		| {
+				tag: "ActionRequest";
+				val: { id: bigint; name: string; args: unknown };
+		  }
+		| {
+				tag: "SubscriptionRequest";
+				val: { eventName: string; subscribe: boolean };
+		  };
+}> {
 	// Validate value length
 	const length = getValueLength(value);
 	if (length > opts.maxIncomingMessageSize) {
@@ -88,7 +97,34 @@ export async function parseMessage(
 	}
 
 	// Deserialize message
-	return deserializeWithEncoding(opts.encoding, buffer, TO_SERVER_VERSIONED);
+	return deserializeWithEncoding(
+		opts.encoding,
+		buffer,
+		TO_SERVER_VERSIONED,
+		ToServerSchema,
+		// JSON: values are already the correct type
+		(json: ToServerJson): any => json,
+		// BARE: need to decode ArrayBuffer fields back to unknown
+		(bare: protocol.ToServer): any => {
+			if (bare.body.tag === "ActionRequest") {
+				return {
+					body: {
+						tag: "ActionRequest",
+						val: {
+							id: bare.body.val.id,
+							name: bare.body.val.name,
+							args: cbor.decode(
+								new Uint8Array(bare.body.val.args),
+							),
+						},
+					},
+				};
+			} else {
+				// SubscriptionRequest has no ArrayBuffer fields
+				return bare;
+			}
+		},
+	);
 }
 
 export interface ProcessMessageHandler<
@@ -122,7 +158,17 @@ export async function processMessage<
 	I,
 	DB extends AnyDatabaseProvider,
 >(
-	message: protocol.ToServer,
+	message: {
+		body:
+			| {
+					tag: "ActionRequest";
+					val: { id: bigint; name: string; args: unknown };
+			  }
+			| {
+					tag: "SubscriptionRequest";
+					val: { eventName: string; subscribe: boolean };
+			  };
+	},
 	actor: ActorInstance<S, CP, CS, V, I, DB>,
 	conn: Conn<S, CP, CS, V, I, DB>,
 	handler: ProcessMessageHandler<S, CP, CS, V, I, DB>,
@@ -138,10 +184,9 @@ export async function processMessage<
 				throw new errors.Unsupported("Action");
 			}
 
-			const { id, name, args: argsRaw } = message.body.val;
+			const { id, name, args } = message.body.val;
 			actionId = id;
 			actionName = name;
-			const args = cbor.decode(new Uint8Array(argsRaw));
 
 			actor.rLog.debug({
 				msg: "processing action request",
@@ -149,14 +194,15 @@ export async function processMessage<
 				actionName: name,
 			});
 
-			const ctx = new ActionContext<S, CP, CS, V, I, DB>(
-				actor.actorContext,
-				conn,
-			);
+			const ctx = new ActionContext<S, CP, CS, V, I, DB>(actor, conn);
 
 			// Process the action request and wait for the result
 			// This will wait for async actions to complete
-			const output = await handler.onExecuteAction(ctx, name, args);
+			const output = await handler.onExecuteAction(
+				ctx,
+				name,
+				args as unknown[],
+			);
 
 			actor.rLog.debug({
 				msg: "sending action response",
@@ -167,20 +213,31 @@ export async function processMessage<
 			});
 
 			// Send the response back to the client
-			conn._sendMessage(
-				new CachedSerializer<protocol.ToClient>(
-					{
+			conn[CONN_SEND_MESSAGE_SYMBOL](
+				new CachedSerializer(
+					output,
+					TO_CLIENT_VERSIONED,
+					ToClientSchema,
+					// JSON: output is the raw value
+					(value): ToClientJson => ({
 						body: {
-							tag: "ActionResponse",
+							tag: "ActionResponse" as const,
 							val: {
 								id: id,
-								output: bufferToArrayBuffer(
-									cbor.encode(output),
-								),
+								output: value,
 							},
 						},
-					},
-					TO_CLIENT_VERSIONED,
+					}),
+					// BARE/CBOR: output needs to be CBOR-encoded to ArrayBuffer
+					(value): protocol.ToClient => ({
+						body: {
+							tag: "ActionResponse" as const,
+							val: {
+								id: id,
+								output: bufferToArrayBuffer(cbor.encode(value)),
+							},
+						},
+					}),
 				),
 			);
 
@@ -236,23 +293,54 @@ export async function processMessage<
 		});
 
 		// Build response
-		conn._sendMessage(
-			new CachedSerializer<protocol.ToClient>(
-				{
+		const errorData = { group, code, message, metadata, actionId };
+		conn[CONN_SEND_MESSAGE_SYMBOL](
+			new CachedSerializer(
+				errorData,
+				TO_CLIENT_VERSIONED,
+				ToClientSchema,
+				// JSON: metadata is the raw value (keep as undefined if not present)
+				(value): ToClientJson => {
+					const val: any = {
+						group: value.group,
+						code: value.code,
+						message: value.message,
+						actionId:
+							value.actionId !== undefined
+								? value.actionId
+								: null,
+					};
+					if (value.metadata !== undefined) {
+						val.metadata = value.metadata;
+					}
+					return {
+						body: {
+							tag: "Error" as const,
+							val,
+						},
+					};
+				},
+				// BARE/CBOR: metadata needs to be CBOR-encoded to ArrayBuffer
+				// Note: protocol.Error expects `| null` for optional fields (BARE protocol)
+				(value): protocol.ToClient => ({
 					body: {
-						tag: "Error",
+						tag: "Error" as const,
 						val: {
-							group,
-							code,
-							message,
-							metadata: bufferToArrayBuffer(
-								cbor.encode(metadata),
-							),
-							actionId: actionId ?? null,
+							group: value.group,
+							code: value.code,
+							message: value.message,
+							metadata: value.metadata
+								? bufferToArrayBuffer(
+										cbor.encode(value.metadata),
+									)
+								: null,
+							actionId:
+								value.actionId !== undefined
+									? value.actionId
+									: null,
 						},
 					},
-				},
-				TO_CLIENT_VERSIONED,
+				}),
 			),
 		);
 

package/src/actor/protocol/serde.ts

@@ -22,14 +22,26 @@ export type Encoding = z.infer<typeof EncodingSchema>;
 /**
  * Helper class that helps serialize data without re-serializing for the same encoding.
  */
-export class CachedSerializer<T> {
+export class CachedSerializer<TBare, TJson, T = TBare> {
 	#data: T;
 	#cache = new Map<Encoding, OutputData>();
-	#versionedDataHandler: VersionedDataHandler<T>;
-
-	constructor(data: T, versionedDataHandler: VersionedDataHandler<T>) {
+	#versionedDataHandler: VersionedDataHandler<TBare>;
+	#zodSchema: z.ZodType<TJson>;
+	#toJson: (value: T) => TJson;
+	#toBare: (value: T) => TBare;
+
+	constructor(
+		data: T,
+		versionedDataHandler: VersionedDataHandler<TBare>,
+		zodSchema: z.ZodType<TJson>,
+		toJson: (value: T) => TJson,
+		toBare: (value: T) => TBare,
+	) {
 		this.#data = data;
 		this.#versionedDataHandler = versionedDataHandler;
+		this.#zodSchema = zodSchema;
+		this.#toJson = toJson;
+		this.#toBare = toBare;
 	}
 
 	public get rawData(): T {
@@ -45,6 +57,9 @@ export class CachedSerializer<T> {
 				encoding,
 				this.#data,
 				this.#versionedDataHandler,
+				this.#zodSchema,
+				this.#toJson,
+				this.#toBare,
 			);
 			this.#cache.set(encoding, serialized);
 			return serialized;