rivetkit 2.0.8 → 2.0.10

package/src/actor/protocol/serde.ts

@@ -123,9 +123,81 @@ export function encodeDataToString(message: OutputData): string {
 	}
 }
 
+function base64DecodeToUint8Array(base64: string): Uint8Array {
+	// Check if Buffer is available (Node.js)
+	if (typeof Buffer !== "undefined") {
+		return new Uint8Array(Buffer.from(base64, "base64"));
+	}
+
+	// Browser environment - use atob
+	const binary = atob(base64);
+	const len = binary.length;
+	const bytes = new Uint8Array(len);
+	for (let i = 0; i < len; i++) {
+		bytes[i] = binary.charCodeAt(i);
+	}
+	return bytes;
+}
+
+function base64DecodeToArrayBuffer(base64: string): ArrayBuffer {
+	return base64DecodeToUint8Array(base64).buffer as ArrayBuffer;
+}
+
 /** Stringifies with compat for values that BARE & CBOR supports. */
 export function jsonStringifyCompat(input: any): string {
-	return JSON.stringify(input, (_key, value) =>
-		typeof value === "bigint" ? value.toString() : value,
-	);
+	return JSON.stringify(input, (_key, value) => {
+		if (typeof value === "bigint") {
+			return ["$BigInt", value.toString()];
+		} else if (value instanceof ArrayBuffer) {
+			return ["$ArrayBuffer", base64EncodeArrayBuffer(value)];
+		} else if (value instanceof Uint8Array) {
+			return ["$Uint8Array", base64EncodeUint8Array(value)];
+		}
+
+		// Escape user arrays that start with $ by prepending another $
+		if (
+			Array.isArray(value) &&
+			value.length === 2 &&
+			typeof value[0] === "string" &&
+			value[0].startsWith("$")
+		) {
+			return ["$" + value[0], value[1]];
+		}
+
+		return value;
+	});
+}
+
+/** Parses JSON with compat for values that BARE & CBOR supports. */
+export function jsonParseCompat(input: string): any {
+	return JSON.parse(input, (_key, value) => {
+		// Handle arrays with $ prefix
+		if (
+			Array.isArray(value) &&
+			value.length === 2 &&
+			typeof value[0] === "string" &&
+			value[0].startsWith("$")
+		) {
+			// Known special types
+			if (value[0] === "$BigInt") {
+				return BigInt(value[1]);
+			} else if (value[0] === "$ArrayBuffer") {
+				return base64DecodeToArrayBuffer(value[1]);
+			} else if (value[0] === "$Uint8Array") {
+				return base64DecodeToUint8Array(value[1]);
+			}
+
+			// Unescape user arrays that started with $ ($$foo -> $foo)
+			if (value[0].startsWith("$$")) {
+				return [value[0].substring(1), value[1]];
+			}
+
+			// Unknown type starting with $ - this is an error
+			throw new Error(
+				`Unknown JSON encoding type: ${value[0]}. This may indicate corrupted data or a version mismatch.`,
+			);
+		}
+
+		return value;
+	});
 }

package/src/actor/router-endpoints.ts

@@ -26,7 +26,7 @@ import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
 import { deconstructError, stringifyError } from "@/common/utils";
 import type { UniversalWebSocket } from "@/common/websocket-interface";
 import { HonoWebSocketAdapter } from "@/manager/hono-websocket-adapter";
-import type { RunConfig } from "@/registry/run-config";
+import type { RunnerConfig } from "@/registry/run-config";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import {
 	HTTP_ACTION_REQUEST_VERSIONED,
@@ -106,7 +106,7 @@ export interface WebSocketOpts {
  */
 export async function handleWebSocketConnect(
 	req: Request | undefined,
-	runConfig: RunConfig,
+	runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	actorId: string,
 	encoding: Encoding,
@@ -325,7 +325,7 @@ export async function handleWebSocketConnect(
  */
 export async function handleSseConnect(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	actorId: string,
 ) {
@@ -437,7 +437,7 @@ export async function handleSseConnect(
  */
 export async function handleAction(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	actionName: string,
 	actorId: string,
@@ -505,7 +505,7 @@ export async function handleAction(
  */
 export async function handleConnectionMessage(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	connId: string,
 	connToken: string,
@@ -542,7 +542,7 @@ export async function handleConnectionMessage(
 
 export async function handleConnectionClose(
 	c: HonoContext,
-	_runConfig: RunConfig,
+	_runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	connId: string,
 	connToken: string,

package/src/actor/router.ts

@@ -41,7 +41,7 @@ import {
 	createActorInspectorRouter,
 } from "@/inspector/actor";
 import { isInspectorEnabled, secureInspector } from "@/inspector/utils";
-import type { RunConfig } from "@/registry/run-config";
+import type { RunnerConfig } from "@/registry/run-config";
 import { ConnDriverKind } from "./conn-drivers";
 import type { ActorDriver } from "./driver";
 import { InternalError } from "./errors";
@@ -67,7 +67,7 @@ export type ActorRouter = Hono<{ Bindings: ActorRouterBindings }>;
  * Creates a router that runs on the partitioned instance.
  */
 export function createActorRouter(
-	runConfig: RunConfig,
+	runConfig: RunnerConfig,
 	actorDriver: ActorDriver,
 	isTest: boolean,
 ): ActorRouter {

package/src/client/actor-conn.ts

@@ -315,13 +315,34 @@ enc
 			});
 		});
 		ws.addEventListener("message", async (ev) => {
-			this.#handleOnMessage(ev.data);
+			try {
+				await this.#handleOnMessage(ev.data);
+			} catch (err) {
+				logger().error({
+					msg: "error in websocket message handler",
+					error: stringifyError(err),
+				});
+			}
 		});
 		ws.addEventListener("close", (ev) => {
-			this.#handleOnClose(ev);
+			try {
+				this.#handleOnClose(ev);
+			} catch (err) {
+				logger().error({
+					msg: "error in websocket close handler",
+					error: stringifyError(err),
+				});
+			}
 		});
 		ws.addEventListener("error", (_ev) => {
-			this.#handleOnError();
+			try {
+				this.#handleOnError();
+			} catch (err) {
+				logger().error({
+					msg: "error in websocket error handler",
+					error: stringifyError(err),
+				});
+			}
 		});
 	}
 

package/src/client/config.ts

@@ -1,38 +1,32 @@
 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;
+import { type GetUpgradeWebSocket, getEnvUniversal } from "@/utils";
 
 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 for Rivet Engine or RivetKit manager API. */
+	endpoint: z
+		.string()
+		.optional()
+		.transform(
+			(x) =>
+				x ??
+				getEnvUniversal("RIVET_ENGINE") ??
+				getEnvUniversal("RIVET_ENDPOINT"),
+		),
+
+	/** Token to use to authenticate with the API. */
 	token: z
 		.string()
 		.optional()
 		.transform((x) => x ?? getEnvUniversal("RIVET_TOKEN")),
 
-	totalSlots: z.number().optional(),
-
-	headers: z.record(z.string()).optional().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 to connect to. */
 	namespace: z
 		.string()
 		.default(() => getEnvUniversal("RIVET_NAMESPACE") ?? "default"),
 
+	/** Name of the runner. This is used to group together runners in to different pools. */
 	runnerName: z
 		.string()
 		.default(() => getEnvUniversal("RIVET_RUNNER") ?? "rivetkit"),
@@ -41,10 +35,9 @@ export const ClientConfigSchema = z.object({
 
 	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(...)`
+	headers: z.record(z.string()).optional().default({}),
+
+	// See RunConfig.getUpgradeWebSocket
 	getUpgradeWebSocket: z.custom<GetUpgradeWebSocket>().optional(),
 });
 

package/src/driver-helpers/mod.ts

@@ -1,5 +1,6 @@
 export type { ActorDriver } from "@/actor/driver";
 export type { ActorInstance, AnyActorInstance } from "@/actor/instance";
+export { generateRandomString } from "@/actor/utils";
 export {
 	ALLOWED_PUBLIC_HEADERS,
 	HEADER_ACTOR_ID,
@@ -31,5 +32,8 @@ export type {
 	ManagerDisplayInformation,
 	ManagerDriver,
 } from "@/manager/driver";
-export { DriverConfigSchema, RunConfigSchema } from "@/registry/run-config";
+export {
+	DriverConfigSchema,
+	RunnerConfigSchema as RunConfigSchema,
+} from "@/registry/run-config";
 export { serializeEmptyPersistData } from "./utils";

package/src/driver-test-suite/mod.ts

@@ -3,11 +3,16 @@ import { createNodeWebSocket, type NodeWebSocket } from "@hono/node-ws";
 import { bundleRequire } from "bundle-require";
 import invariant from "invariant";
 import { describe } from "vitest";
-import type { Transport } from "@/client/mod";
+import type { Encoding, Transport } from "@/client/mod";
 import { configureInspectorAccessToken } from "@/inspector/utils";
 import { createManagerRouter } from "@/manager/router";
-import type { DriverConfig, Registry, RunConfig } from "@/mod";
-import { RunConfigSchema } from "@/registry/run-config";
+import {
+	createClientWithDriver,
+	type DriverConfig,
+	type Registry,
+	type RunConfig,
+} from "@/mod";
+import { RunnerConfigSchema } from "@/registry/run-config";
 import { getPort } from "@/test/mod";
 import { logger } from "./log";
 import { runActionFeaturesTests } from "./tests/action-features";
@@ -55,6 +60,8 @@ export interface DriverTestConfig {
 
 	transport?: Transport;
 
+	encoding?: Encoding;
+
 	clientType: ClientType;
 
 	cleanup?: () => Promise<void>;
@@ -78,68 +85,81 @@ export interface DriverDeployOutput {
 
 /** Runs all Vitest tests against the provided drivers. */
 export function runDriverTests(
-	driverTestConfigPartial: Omit<DriverTestConfig, "clientType" | "transport">,
+	driverTestConfigPartial: Omit<
+		DriverTestConfig,
+		"clientType" | "transport" | "encoding"
+	>,
 ) {
 	const clientTypes: ClientType[] = driverTestConfigPartial.skip?.inline
 		? ["http"]
 		: ["http", "inline"];
 	for (const clientType of clientTypes) {
-		const driverTestConfig: DriverTestConfig = {
-			...driverTestConfigPartial,
-			clientType,
-		};
-
 		describe(`client type (${clientType})`, () => {
-			runActorDriverTests(driverTestConfig);
-			runManagerDriverTests(driverTestConfig);
+			const encodings: Encoding[] = ["bare", "cbor", "json"];
 
-			const transports: Transport[] = driverTestConfig.skip?.sse
-				? ["websocket"]
-				: ["websocket", "sse"];
-			for (const transport of transports) {
-				describe(`transport (${transport})`, () => {
-					runActorConnTests({
-						...driverTestConfig,
-						transport,
-					});
+			for (const encoding of encodings) {
+				describe(`encoding (${encoding})`, () => {
+					const driverTestConfig: DriverTestConfig = {
+						...driverTestConfigPartial,
+						clientType,
+						encoding,
+					};
 
-					runActorConnStateTests({ ...driverTestConfig, transport });
+					runActorDriverTests(driverTestConfig);
+					runManagerDriverTests(driverTestConfig);
 
-					runActorReconnectTests({ ...driverTestConfig, transport });
+					const transports: Transport[] = driverTestConfig.skip?.sse
+						? ["websocket"]
+						: ["websocket", "sse"];
+					for (const transport of transports) {
+						describe(`transport (${transport})`, () => {
+							runActorConnTests({
+								...driverTestConfig,
+								transport,
+							});
 
-					runRequestAccessTests({ ...driverTestConfig, transport });
+							runActorConnStateTests({ ...driverTestConfig, transport });
 
-					runActorDriverTestsWithTransport({ ...driverTestConfig, transport });
-				});
-			}
+							runActorReconnectTests({ ...driverTestConfig, transport });
 
-			runActorHandleTests(driverTestConfig);
+							runRequestAccessTests({ ...driverTestConfig, transport });
 
-			runActionFeaturesTests(driverTestConfig);
+							runActorDriverTestsWithTransport({
+								...driverTestConfig,
+								transport,
+							});
+						});
+					}
 
-			runActorVarsTests(driverTestConfig);
+					runActorHandleTests(driverTestConfig);
 
-			runActorMetadataTests(driverTestConfig);
+					runActionFeaturesTests(driverTestConfig);
 
-			runActorOnStateChangeTests(driverTestConfig);
+					runActorVarsTests(driverTestConfig);
 
-			runActorErrorHandlingTests(driverTestConfig);
+					runActorMetadataTests(driverTestConfig);
 
-			runActorInlineClientTests(driverTestConfig);
+					runActorOnStateChangeTests(driverTestConfig);
 
-			runRawHttpTests(driverTestConfig);
+					runActorErrorHandlingTests(driverTestConfig);
 
-			runRawHttpRequestPropertiesTests(driverTestConfig);
+					runActorInlineClientTests(driverTestConfig);
 
-			runRawWebSocketTests(driverTestConfig);
+					runRawHttpTests(driverTestConfig);
 
-			// TODO: re-expose this once we can have actor queries on the gateway
-			// runRawHttpDirectRegistryTests(driverTestConfig);
+					runRawHttpRequestPropertiesTests(driverTestConfig);
 
-			// TODO: re-expose this once we can have actor queries on the gateway
-			// runRawWebSocketDirectRegistryTests(driverTestConfig);
+					runRawWebSocketTests(driverTestConfig);
 
-			runActorInspectorTests(driverTestConfig);
+					// TODO: re-expose this once we can have actor queries on the gateway
+					// runRawHttpDirectRegistryTests(driverTestConfig);
+
+					// TODO: re-expose this once we can have actor queries on the gateway
+					// runRawWebSocketDirectRegistryTests(driverTestConfig);
+
+					runActorInspectorTests(driverTestConfig);
+				});
+			}
 		});
 	}
 }
@@ -199,7 +219,7 @@ export async function createTestRuntime(
 		// Build driver config
 		// biome-ignore lint/style/useConst: Assigned later
 		let upgradeWebSocket: any;
-		const config: RunConfig = RunConfigSchema.parse({
+		const config: RunConfig = RunnerConfigSchema.parse({
 			driver,
 			getUpgradeWebSocket: () => upgradeWebSocket!,
 			inspector: {
@@ -210,12 +230,14 @@ export async function createTestRuntime(
 
 		// Create router
 		const managerDriver = driver.manager(registry.config, config);
+		const client = createClientWithDriver(managerDriver);
 		configureInspectorAccessToken(config, managerDriver);
 		const { router } = createManagerRouter(
 			registry.config,
 			config,
 			managerDriver,
-			undefined,
+			driver,
+			client,
 		);
 
 		// Inject WebSocket

package/src/driver-test-suite/utils.ts

@@ -36,17 +36,20 @@ export async function setupDriverTest(
 			namespace,
 			runnerName,
 			transport: driverTestConfig.transport,
+			encoding: driverTestConfig.encoding,
 		});
 	} else if (driverTestConfig.clientType === "inline") {
 		// Use inline client from driver
 		const transport = driverTestConfig.transport ?? "websocket";
+		const encoding = driverTestConfig.encoding ?? "bare";
 		const managerDriver = createTestInlineClientDriver(
 			endpoint,
-			"bare",
+			encoding,
 			transport,
 		);
 		const runConfig = RunConfigSchema.parse({
 			transport: transport,
+			encoding: encoding,
 		});
 		client = createClientWithDriver(managerDriver, runConfig);
 	} else {

package/src/drivers/default.ts

@@ -2,15 +2,21 @@ import { UserError } from "@/actor/errors";
 import { loggerWithoutContext } from "@/actor/log";
 import { createEngineDriver } from "@/drivers/engine/mod";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
-import type { DriverConfig, RunConfig } from "@/registry/run-config";
+import type { DriverConfig, RunnerConfig } from "@/registry/run-config";
 
 /**
  * Chooses the appropriate driver based on the run configuration.
  */
-export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
+export function chooseDefaultDriver(runConfig: RunnerConfig): DriverConfig {
 	if (runConfig.endpoint && runConfig.driver) {
 		throw new UserError(
-			"Cannot specify both 'engine' and 'driver' in configuration",
+			"Cannot specify both 'endpoint' and 'driver' in configuration",
+		);
+	}
+
+	if (runConfig.runnerKind === "serverless" && !runConfig.endpoint) {
+		throw new UserError(
+			"Cannot use 'serverless' runnerKind without the 'endpoint' config set.",
 		);
 	}
 
@@ -18,16 +24,12 @@ export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig {
 		return runConfig.driver;
 	}
 
-	if (runConfig.endpoint) {
+	if (runConfig.endpoint || runConfig.token) {
 		loggerWithoutContext().debug({
 			msg: "using rivet engine driver",
 			endpoint: runConfig.endpoint,
 		});
-		// TODO: Add all properties from config
-		return createEngineDriver({
-			endpoint: runConfig.endpoint,
-			token: runConfig.token,
-		});
+		return createEngineDriver();
 	}
 
 	loggerWithoutContext().debug({ msg: "using default file system driver" });