rivetkit 2.0.1 → 2.0.3

package/src/common/versioned-data.ts

@@ -0,0 +1,95 @@
+export interface VersionedData<T> {
+	version: number;
+	data: T;
+}
+
+export type MigrationFn<TFrom, TTo> = (data: TFrom) => TTo;
+
+export interface VersionedDataConfig<T> {
+	currentVersion: number;
+	migrations: Map<number, MigrationFn<any, any>>;
+	serializeVersion: (data: T) => Uint8Array;
+	deserializeVersion: (bytes: Uint8Array) => T;
+}
+
+export class VersionedDataHandler<T> {
+	constructor(private config: VersionedDataConfig<T>) {}
+
+	serializeWithEmbeddedVersion(data: T): Uint8Array {
+		const versioned: VersionedData<Uint8Array> = {
+			version: this.config.currentVersion,
+			data: this.config.serializeVersion(data),
+		};
+
+		return this.embedVersion(versioned);
+	}
+
+	deserializeWithEmbeddedVersion(bytes: Uint8Array): T {
+		const versioned = this.extractVersion(bytes);
+		return this.deserialize(versioned.data, versioned.version);
+	}
+
+	serialize(data: T, version: number): Uint8Array {
+		return this.config.serializeVersion(data);
+	}
+
+	deserialize(bytes: Uint8Array, version: number): T {
+		if (version === this.config.currentVersion) {
+			return this.config.deserializeVersion(bytes);
+		}
+
+		if (version > this.config.currentVersion) {
+			throw new Error(
+				`Cannot decode data from version ${version}, current version is ${this.config.currentVersion}`,
+			);
+		}
+
+		let currentData: any = this.config.deserializeVersion(bytes);
+		let currentVersion = version;
+
+		while (currentVersion < this.config.currentVersion) {
+			const migration = this.config.migrations.get(currentVersion);
+			if (!migration) {
+				throw new Error(
+					`No migration found from version ${currentVersion} to ${currentVersion + 1}`,
+				);
+			}
+
+			currentData = migration(currentData);
+			currentVersion++;
+		}
+
+		return currentData;
+	}
+
+	private embedVersion(data: VersionedData<Uint8Array>): Uint8Array {
+		const versionBytes = new Uint8Array(4);
+		new DataView(versionBytes.buffer).setUint32(0, data.version, true);
+
+		const result = new Uint8Array(versionBytes.length + data.data.length);
+		result.set(versionBytes);
+		result.set(data.data, versionBytes.length);
+
+		return result;
+	}
+
+	private extractVersion(bytes: Uint8Array): VersionedData<Uint8Array> {
+		if (bytes.length < 4) {
+			throw new Error("Invalid versioned data: too short");
+		}
+
+		const version = new DataView(bytes.buffer, bytes.byteOffset).getUint32(
+			0,
+			true,
+		);
+		const data = bytes.slice(4);
+
+		return { version, data };
+	}
+}
+
+export function createVersionedDataHandler<T>(
+	config: VersionedDataConfig<T>,
+): VersionedDataHandler<T> {
+	return new VersionedDataHandler(config);
+}

package/src/common/websocket-interface.ts

@@ -0,0 +1,49 @@
+// Define minimal event interfaces to avoid conflicts between different WebSocket implementations
+export interface RivetEvent {
+	type: string;
+	target?: any;
+	currentTarget?: any;
+}
+
+export interface RivetMessageEvent extends RivetEvent {
+	data: any;
+}
+
+export interface RivetCloseEvent extends RivetEvent {
+	code: number;
+	reason: string;
+	wasClean: boolean;
+}
+
+/**
+ * Common WebSocket interface that can be implemented by different WebSocket-like classes
+ * This is compatible with the standard WebSocket API but allows for custom implementations
+ */
+export interface UniversalWebSocket {
+	// WebSocket readyState values
+	readonly CONNECTING: 0;
+	readonly OPEN: 1;
+	readonly CLOSING: 2;
+	readonly CLOSED: 3;
+
+	// Properties
+	readonly readyState: 0 | 1 | 2 | 3;
+	binaryType: "arraybuffer" | "blob";
+	readonly bufferedAmount: number;
+	readonly extensions: string;
+	readonly protocol: string;
+	readonly url: string;
+
+	// Methods
+	send(data: string | ArrayBufferLike | Blob | ArrayBufferView): void;
+	close(code?: number, reason?: string): void;
+	addEventListener(type: string, listener: (event: any) => void): void;
+	removeEventListener(type: string, listener: (event: any) => void): void;
+	dispatchEvent(event: RivetEvent): boolean;
+
+	// Event handlers (optional)
+	onopen?: ((event: RivetEvent) => void) | null;
+	onclose?: ((event: RivetCloseEvent) => void) | null;
+	onerror?: ((event: RivetEvent) => void) | null;
+	onmessage?: ((event: RivetMessageEvent) => void) | null;
+}

package/src/common/websocket.ts

@@ -0,0 +1,43 @@
+import { logger } from "@/client/log";
+
+// Global singleton promise that will be reused for subsequent calls
+let webSocketPromise: Promise<typeof WebSocket> | null = null;
+
+export async function importWebSocket(): Promise<typeof WebSocket> {
+	// Return existing promise if we already started loading
+	if (webSocketPromise !== null) {
+		return webSocketPromise;
+	}
+
+	// Create and store the promise
+	webSocketPromise = (async () => {
+		let _WebSocket: typeof WebSocket;
+
+		if (typeof WebSocket !== "undefined") {
+			// Browser environment
+			_WebSocket = WebSocket as unknown as typeof WebSocket;
+			logger().debug("using native websocket");
+		} else {
+			// Node.js environment
+			try {
+				const ws = await import("ws");
+				_WebSocket = ws.default as unknown as typeof WebSocket;
+				logger().debug("using websocket from npm");
+			} catch {
+				// WS not available
+				_WebSocket = class MockWebSocket {
+					constructor() {
+						throw new Error(
+							'WebSocket support requires installing the "ws" peer dependency.',
+						);
+					}
+				} as unknown as typeof WebSocket;
+				logger().debug("using mock websocket");
+			}
+		}
+
+		return _WebSocket;
+	})();
+
+	return webSocketPromise;
+}

package/src/driver-helpers/mod.ts

@@ -0,0 +1,22 @@
+export type { ActorDriver } from "@/actor/driver";
+export type { ActorInstance, AnyActorInstance } from "@/actor/instance";
+export {
+	HEADER_ACTOR_ID,
+	HEADER_ACTOR_QUERY,
+	HEADER_AUTH_DATA,
+	HEADER_CONN_ID,
+	HEADER_CONN_PARAMS,
+	HEADER_CONN_TOKEN,
+	HEADER_ENCODING,
+	HEADER_EXPOSE_INTERNAL_ERROR,
+} from "@/actor/router-endpoints";
+export type {
+	ActorOutput,
+	CreateInput,
+	GetForIdInput,
+	GetOrCreateWithKeyInput,
+	GetWithKeyInput,
+	ManagerDriver,
+} from "@/manager/driver";
+export { DriverConfigSchema, RunConfigSchema } from "@/registry/run-config";
+export { serializeEmptyPersistData } from "./utils";

package/src/driver-helpers/utils.ts

@@ -0,0 +1,17 @@
+import * as cbor from "cbor-x";
+import type * as schema from "@/schemas/actor-persist/mod";
+import { PERSISTED_ACTOR_VERSIONED } from "@/schemas/actor-persist/versioned";
+import { bufferToArrayBuffer } from "@/utils";
+
+export function serializeEmptyPersistData(
+	input: unknown | undefined,
+): Uint8Array {
+	const persistData: schema.PersistedActor = {
+		input: input !== undefined ? bufferToArrayBuffer(cbor.encode(input)) : null,
+		hasInitialized: false,
+		state: bufferToArrayBuffer(cbor.encode(undefined)),
+		connections: [],
+		scheduledEvents: [],
+	};
+	return PERSISTED_ACTOR_VERSIONED.serializeWithEmbeddedVersion(persistData);
+}

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

@@ -0,0 +1,7 @@
+import { getLogger } from "@/common/log";
+
+export const LOGGER_NAME = "test-suite";
+
+export function logger() {
+	return getLogger(LOGGER_NAME);
+}

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

@@ -0,0 +1,213 @@
+import { serve as honoServe } from "@hono/node-server";
+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 { createInlineClientDriver } from "@/inline-client-driver/mod";
+import { createManagerRouter } from "@/manager/router";
+import type { DriverConfig, Registry, RunConfig } from "@/mod";
+import { RunConfigSchema } from "@/registry/run-config";
+import { getPort } from "@/test/mod";
+import { runActionFeaturesTests } from "./tests/action-features";
+import { runActorAuthTests } from "./tests/actor-auth";
+import { runActorConnTests } from "./tests/actor-conn";
+import { runActorConnStateTests } from "./tests/actor-conn-state";
+import {
+	runActorDriverTests,
+	runActorDriverTestsWithTransport,
+} from "./tests/actor-driver";
+import { runActorErrorHandlingTests } from "./tests/actor-error-handling";
+import { runActorHandleTests } from "./tests/actor-handle";
+import { runActorInlineClientTests } from "./tests/actor-inline-client";
+import { runActorInspectorTests } from "./tests/actor-inspector";
+import { runActorMetadataTests } from "./tests/actor-metadata";
+import { runActorOnStateChangeTests } from "./tests/actor-onstatechange";
+import { runActorVarsTests } from "./tests/actor-vars";
+import { runManagerDriverTests } from "./tests/manager-driver";
+import { runRawHttpTests } from "./tests/raw-http";
+import { runRawHttpDirectRegistryTests } from "./tests/raw-http-direct-registry";
+import { runRawHttpRequestPropertiesTests } from "./tests/raw-http-request-properties";
+import { runRawWebSocketTests } from "./tests/raw-websocket";
+import { runRawWebSocketDirectRegistryTests } from "./tests/raw-websocket-direct-registry";
+import { runRequestAccessTests } from "./tests/request-access";
+
+export interface SkipTests {
+	schedule?: boolean;
+	sleep?: boolean;
+}
+
+export interface DriverTestConfig {
+	/** Deploys an registry and returns the connection endpoint. */
+	start(projectDir: string): Promise<DriverDeployOutput>;
+
+	/**
+	 * If we're testing with an external system, we should use real timers
+	 * instead of Vitest's mocked timers.
+	 **/
+	useRealTimers?: boolean;
+
+	/** Cloudflare Workers has some bugs with cleanup. */
+	HACK_skipCleanupNet?: boolean;
+
+	skip?: SkipTests;
+
+	transport?: Transport;
+
+	clientType: ClientType;
+
+	cleanup?: () => Promise<void>;
+}
+
+/**
+ * The type of client to run the test with.
+ *
+ * The logic for HTTP vs inline is very different, so this helps validate all behavior matches.
+ **/
+type ClientType = "http" | "inline";
+
+export interface DriverDeployOutput {
+	endpoint: string;
+
+	/** Cleans up the test. */
+	cleanup(): Promise<void>;
+}
+
+/** Runs all Vitest tests against the provided drivers. */
+export function runDriverTests(
+	driverTestConfigPartial: Omit<DriverTestConfig, "clientType" | "transport">,
+) {
+	for (const clientType of ["http", "inline"] as ClientType[]) {
+		const driverTestConfig: DriverTestConfig = {
+			...driverTestConfigPartial,
+			clientType,
+		};
+
+		describe(`client type (${clientType})`, () => {
+			runActorDriverTests(driverTestConfig);
+			runManagerDriverTests(driverTestConfig);
+
+			for (const transport of ["websocket", "sse"] as Transport[]) {
+				describe(`transport (${transport})`, () => {
+					runActorConnTests({
+						...driverTestConfig,
+						transport,
+					});
+
+					runActorConnStateTests({ ...driverTestConfig, transport });
+
+					runRequestAccessTests({ ...driverTestConfig, transport });
+
+					runActorDriverTestsWithTransport({ ...driverTestConfig, transport });
+				});
+			}
+
+			runActorHandleTests(driverTestConfig);
+
+			runActionFeaturesTests(driverTestConfig);
+
+			runActorVarsTests(driverTestConfig);
+
+			runActorMetadataTests(driverTestConfig);
+
+			runActorOnStateChangeTests(driverTestConfig);
+
+			runActorErrorHandlingTests(driverTestConfig);
+
+			runActorAuthTests(driverTestConfig);
+
+			runActorInlineClientTests(driverTestConfig);
+
+			runRawHttpTests(driverTestConfig);
+
+			runRawHttpRequestPropertiesTests(driverTestConfig);
+
+			runRawWebSocketTests(driverTestConfig);
+
+			runRawHttpDirectRegistryTests(driverTestConfig);
+
+			runRawWebSocketDirectRegistryTests(driverTestConfig);
+
+			runActorInspectorTests(driverTestConfig);
+		});
+	}
+}
+
+/**
+ * Helper function to adapt the drivers to the Node.js runtime for tests.
+ *
+ * This is helpful for drivers that run in-process as opposed to drivers that rely on external tools.
+ */
+export async function createTestRuntime(
+	registryPath: string,
+	driverFactory: (registry: Registry<any>) => Promise<{
+		driver: DriverConfig;
+		cleanup?: () => Promise<void>;
+	}>,
+): Promise<DriverDeployOutput> {
+	const {
+		mod: { registry },
+	} = await bundleRequire<{ registry: Registry<any> }>({
+		filepath: registryPath,
+	});
+
+	// TODO: Find a cleaner way of flagging an registry as test mode (ideally not in the config itself)
+	// Force enable test
+	registry.config.test.enabled = true;
+
+	// Build drivers
+	const { driver, cleanup: driverCleanup } = await driverFactory(registry);
+
+	// Build driver config
+	let injectWebSocket: NodeWebSocket["injectWebSocket"] | undefined;
+	let upgradeWebSocket: any;
+	const config: RunConfig = RunConfigSchema.parse({
+		driver,
+		getUpgradeWebSocket: () => upgradeWebSocket!,
+		inspector: {
+			enabled: true,
+			token: () => "token",
+		},
+	});
+
+	// Create router
+	const managerDriver = driver.manager(registry.config, config);
+	const inlineDriver = createInlineClientDriver(managerDriver);
+	const { router } = createManagerRouter(
+		registry.config,
+		config,
+		inlineDriver,
+		managerDriver,
+		false,
+	);
+
+	// Inject WebSocket
+	const nodeWebSocket = createNodeWebSocket({ app: router });
+	upgradeWebSocket = nodeWebSocket.upgradeWebSocket;
+	injectWebSocket = nodeWebSocket.injectWebSocket;
+
+	// Start server
+	const port = await getPort();
+	const server = honoServe({
+		fetch: router.fetch,
+		hostname: "127.0.0.1",
+		port,
+	});
+	invariant(injectWebSocket !== undefined, "should have injectWebSocket");
+	injectWebSocket(server);
+	const endpoint = `http://127.0.0.1:${port}`;
+
+	// Cleanup
+	const cleanup = async () => {
+		// Stop server
+		await new Promise((resolve) => server.close(() => resolve(undefined)));
+
+		// Extra cleanup
+		await driverCleanup?.();
+	};
+
+	return {
+		endpoint,
+		cleanup,
+	};
+}