rivetkit 2.0.27 → 2.0.28

package/src/schemas/client-protocol/versioned.ts

@@ -1,16 +1,11 @@
-import {
-	createVersionedDataHandler,
-	type MigrationFn,
-} from "@/common/versioned-data";
-import type * as v1 from "../../../dist/schemas/client-protocol/v1";
+import { createVersionedDataHandler } from "vbare";
+import * as v1 from "../../../dist/schemas/client-protocol/v1";
 import * as v2 from "../../../dist/schemas/client-protocol/v2";
 
 export const CURRENT_VERSION = 2;
 
-const migrations = new Map<number, MigrationFn<any, any>>();
-
-// Migration from v1 to v2: Remove connectionToken from Init message
-migrations.set(1, (v1Data: v1.ToClient): v2.ToClient => {
+// Converter from v1 to v2: Remove connectionToken from Init message
+const v1ToV2 = (v1Data: v1.ToClient): v2.ToClient => {
 	// Handle Init message specifically to remove connectionToken
 	if (v1Data.body.tag === "Init") {
 		const { actorId, connectionId } = v1Data.body.val as v1.Init;
@@ -26,58 +21,194 @@ migrations.set(1, (v1Data: v1.ToClient): v2.ToClient => {
 	}
 	// All other messages are unchanged
 	return v1Data as unknown as v2.ToClient;
-});
+};
+
+// Converter from v2 to v1: Add empty connectionToken to Init message
+const v2ToV1 = (v2Data: v2.ToClient): v1.ToClient => {
+	// Handle Init message specifically to add connectionToken
+	if (v2Data.body.tag === "Init") {
+		const { actorId, connectionId } = v2Data.body.val;
+		return {
+			body: {
+				tag: "Init",
+				val: {
+					actorId,
+					connectionId,
+					connectionToken: "", // Add empty connectionToken for v1 compatibility
+				},
+			},
+		};
+	}
+	// All other messages are unchanged
+	return v2Data as unknown as v1.ToClient;
+};
 
 export const TO_SERVER_VERSIONED = createVersionedDataHandler<v2.ToServer>({
-	currentVersion: CURRENT_VERSION,
-	migrations,
-	serializeVersion: (data) => v2.encodeToServer(data),
-	deserializeVersion: (bytes) => v2.decodeToServer(bytes),
+	deserializeVersion: (bytes, version) => {
+		switch (version) {
+			case 1:
+				return v1.decodeToServer(bytes);
+			case 2:
+				return v2.decodeToServer(bytes);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	serializeVersion: (data, version) => {
+		switch (version) {
+			case 1:
+				return v1.encodeToServer(data as v1.ToServer);
+			case 2:
+				return v2.encodeToServer(data as v2.ToServer);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	deserializeConverters: () => [v1ToV2],
+	serializeConverters: () => [v2ToV1],
 });
 
 export const TO_CLIENT_VERSIONED = createVersionedDataHandler<v2.ToClient>({
-	currentVersion: CURRENT_VERSION,
-	migrations,
-	serializeVersion: (data) => v2.encodeToClient(data),
-	deserializeVersion: (bytes) => v2.decodeToClient(bytes),
+	deserializeVersion: (bytes, version) => {
+		switch (version) {
+			case 1:
+				return v1.decodeToClient(bytes);
+			case 2:
+				return v2.decodeToClient(bytes);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	serializeVersion: (data, version) => {
+		switch (version) {
+			case 1:
+				return v1.encodeToClient(data as v1.ToClient);
+			case 2:
+				return v2.encodeToClient(data as v2.ToClient);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	deserializeConverters: () => [v1ToV2],
+	serializeConverters: () => [v2ToV1],
 });
 
 export const HTTP_ACTION_REQUEST_VERSIONED =
 	createVersionedDataHandler<v2.HttpActionRequest>({
-		currentVersion: CURRENT_VERSION,
-		migrations,
-		serializeVersion: (data) => v2.encodeHttpActionRequest(data),
-		deserializeVersion: (bytes) => v2.decodeHttpActionRequest(bytes),
+		deserializeVersion: (bytes, version) => {
+			switch (version) {
+				case 1:
+					return v1.decodeHttpActionRequest(bytes);
+				case 2:
+					return v2.decodeHttpActionRequest(bytes);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		serializeVersion: (data, version) => {
+			switch (version) {
+				case 1:
+					return v1.encodeHttpActionRequest(
+						data as v1.HttpActionRequest,
+					);
+				case 2:
+					return v2.encodeHttpActionRequest(
+						data as v2.HttpActionRequest,
+					);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		deserializeConverters: () => [],
+		serializeConverters: () => [],
 	});
 
 export const HTTP_ACTION_RESPONSE_VERSIONED =
 	createVersionedDataHandler<v2.HttpActionResponse>({
-		currentVersion: CURRENT_VERSION,
-		migrations,
-		serializeVersion: (data) => v2.encodeHttpActionResponse(data),
-		deserializeVersion: (bytes) => v2.decodeHttpActionResponse(bytes),
+		deserializeVersion: (bytes, version) => {
+			switch (version) {
+				case 1:
+					return v1.decodeHttpActionResponse(bytes);
+				case 2:
+					return v2.decodeHttpActionResponse(bytes);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		serializeVersion: (data, version) => {
+			switch (version) {
+				case 1:
+					return v1.encodeHttpActionResponse(
+						data as v1.HttpActionResponse,
+					);
+				case 2:
+					return v2.encodeHttpActionResponse(
+						data as v2.HttpActionResponse,
+					);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		deserializeConverters: () => [],
+		serializeConverters: () => [],
 	});
 
 export const HTTP_RESPONSE_ERROR_VERSIONED =
 	createVersionedDataHandler<v2.HttpResponseError>({
-		currentVersion: CURRENT_VERSION,
-		migrations,
-		serializeVersion: (data) => v2.encodeHttpResponseError(data),
-		deserializeVersion: (bytes) => v2.decodeHttpResponseError(bytes),
-	});
-
-export const HTTP_RESOLVE_REQUEST_VERSIONED =
-	createVersionedDataHandler<v2.HttpResolveRequest>({
-		currentVersion: CURRENT_VERSION,
-		migrations,
-		serializeVersion: (_) => new Uint8Array(),
-		deserializeVersion: (bytes) => null,
+		deserializeVersion: (bytes, version) => {
+			switch (version) {
+				case 1:
+					return v1.decodeHttpResponseError(bytes);
+				case 2:
+					return v2.decodeHttpResponseError(bytes);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		serializeVersion: (data, version) => {
+			switch (version) {
+				case 1:
+					return v1.encodeHttpResponseError(
+						data as v1.HttpResponseError,
+					);
+				case 2:
+					return v2.encodeHttpResponseError(
+						data as v2.HttpResponseError,
+					);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		deserializeConverters: () => [],
+		serializeConverters: () => [],
 	});
 
 export const HTTP_RESOLVE_RESPONSE_VERSIONED =
 	createVersionedDataHandler<v2.HttpResolveResponse>({
-		currentVersion: CURRENT_VERSION,
-		migrations,
-		serializeVersion: (data) => v2.encodeHttpResolveResponse(data),
-		deserializeVersion: (bytes) => v2.decodeHttpResolveResponse(bytes),
+		deserializeVersion: (bytes, version) => {
+			switch (version) {
+				case 1:
+					return v1.decodeHttpResolveResponse(bytes);
+				case 2:
+					return v2.decodeHttpResolveResponse(bytes);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		serializeVersion: (data, version) => {
+			switch (version) {
+				case 1:
+					return v1.encodeHttpResolveResponse(
+						data as v1.HttpResolveResponse,
+					);
+				case 2:
+					return v2.encodeHttpResolveResponse(
+						data as v2.HttpResolveResponse,
+					);
+				default:
+					throw new Error(`Unknown version ${version}`);
+			}
+		},
+		deserializeConverters: () => [],
+		serializeConverters: () => [],
 	});

package/src/schemas/file-system-driver/mod.ts

@@ -1 +1 @@
-export * from "../../../dist/schemas/file-system-driver/v2";
+export * from "../../../dist/schemas/file-system-driver/v3";

package/src/schemas/file-system-driver/versioned.ts

@@ -1,51 +1,135 @@
-import {
-	createVersionedDataHandler,
-	type MigrationFn,
-} from "@/common/versioned-data";
+import { createVersionedDataHandler } from "vbare";
 import { bufferToArrayBuffer } from "@/utils";
-import type * as v1 from "../../../dist/schemas/file-system-driver/v1";
+import * as v1 from "../../../dist/schemas/file-system-driver/v1";
 import * as v2 from "../../../dist/schemas/file-system-driver/v2";
+import * as v3 from "../../../dist/schemas/file-system-driver/v3";
 
-export const CURRENT_VERSION = 2;
-
-const migrations = new Map<number, MigrationFn<any, any>>([
-	[
-		2,
-		(v1State: v1.ActorState): v2.ActorState => {
-			// Create a new kvStorage list with the legacy persist data
-			const kvStorage: v2.ActorKvEntry[] = [];
-
-			// Store the legacy persist data under key [1]
-			if (v1State.persistedData) {
-				// Key [1] as Uint8Array
-				const key = new Uint8Array([1]);
-				kvStorage.push({
-					key: bufferToArrayBuffer(key),
-					value: v1State.persistedData,
-				});
-			}
-
-			return {
-				actorId: v1State.actorId,
-				name: v1State.name,
-				key: v1State.key,
-				kvStorage,
-				createdAt: v1State.createdAt,
-			};
-		},
-	],
-]);
-
-export const ACTOR_STATE_VERSIONED = createVersionedDataHandler<v2.ActorState>({
-	currentVersion: CURRENT_VERSION,
-	migrations,
-	serializeVersion: (data) => v2.encodeActorState(data),
-	deserializeVersion: (bytes) => v2.decodeActorState(bytes),
+export const CURRENT_VERSION = 3;
+
+// Converter from v1 to v2
+const v1ToV2 = (v1State: v1.ActorState): v2.ActorState => {
+	// Create a new kvStorage list with the legacy persist data
+	const kvStorage: v2.ActorKvEntry[] = [];
+
+	// Store the legacy persist data under key [1]
+	if (v1State.persistedData) {
+		// Key [1] as Uint8Array
+		const key = new Uint8Array([1]);
+		kvStorage.push({
+			key: bufferToArrayBuffer(key),
+			value: v1State.persistedData,
+		});
+	}
+
+	return {
+		actorId: v1State.actorId,
+		name: v1State.name,
+		key: v1State.key,
+		kvStorage,
+		createdAt: v1State.createdAt,
+	};
+};
+
+// Converter from v2 to v3
+const v2ToV3 = (v2State: v2.ActorState): v3.ActorState => {
+	// Migrate from v2 to v3 by adding the new optional timestamp fields
+	return {
+		actorId: v2State.actorId,
+		name: v2State.name,
+		key: v2State.key,
+		kvStorage: v2State.kvStorage,
+		createdAt: v2State.createdAt,
+		startTs: null,
+		connectableTs: null,
+		sleepTs: null,
+		destroyTs: null,
+	};
+};
+
+// Converter from v3 to v2
+const v3ToV2 = (v3State: v3.ActorState): v2.ActorState => {
+	// Downgrade from v3 to v2 by removing the timestamp fields
+	return {
+		actorId: v3State.actorId,
+		name: v3State.name,
+		key: v3State.key,
+		kvStorage: v3State.kvStorage,
+		createdAt: v3State.createdAt,
+	};
+};
+
+// Converter from v2 to v1
+const v2ToV1 = (v2State: v2.ActorState): v1.ActorState => {
+	// Downgrade from v2 to v1 by converting kvStorage back to persistedData
+	// Find the persist data entry (key [1])
+	const persistDataEntry = v2State.kvStorage.find((entry) => {
+		const key = new Uint8Array(entry.key);
+		return key.length === 1 && key[0] === 1;
+	});
+
+	return {
+		actorId: v2State.actorId,
+		name: v2State.name,
+		key: v2State.key,
+		persistedData: persistDataEntry?.value || new ArrayBuffer(0),
+		createdAt: v2State.createdAt,
+	};
+};
+
+export const ACTOR_STATE_VERSIONED = createVersionedDataHandler<v3.ActorState>({
+	deserializeVersion: (bytes, version) => {
+		switch (version) {
+			case 1:
+				return v1.decodeActorState(bytes);
+			case 2:
+				return v2.decodeActorState(bytes);
+			case 3:
+				return v3.decodeActorState(bytes);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	serializeVersion: (data, version) => {
+		switch (version) {
+			case 1:
+				return v1.encodeActorState(data as v1.ActorState);
+			case 2:
+				return v2.encodeActorState(data as v2.ActorState);
+			case 3:
+				return v3.encodeActorState(data as v3.ActorState);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	deserializeConverters: () => [v1ToV2, v2ToV3],
+	serializeConverters: () => [v3ToV2, v2ToV1],
 });
 
-export const ACTOR_ALARM_VERSIONED = createVersionedDataHandler<v2.ActorAlarm>({
-	currentVersion: CURRENT_VERSION,
-	migrations,
-	serializeVersion: (data) => v2.encodeActorAlarm(data),
-	deserializeVersion: (bytes) => v2.decodeActorAlarm(bytes),
+export const ACTOR_ALARM_VERSIONED = createVersionedDataHandler<v3.ActorAlarm>({
+	deserializeVersion: (bytes, version) => {
+		switch (version) {
+			case 1:
+				return v1.decodeActorAlarm(bytes);
+			case 2:
+				return v2.decodeActorAlarm(bytes);
+			case 3:
+				return v3.decodeActorAlarm(bytes);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	serializeVersion: (data, version) => {
+		switch (version) {
+			case 1:
+				return v1.encodeActorAlarm(data as v1.ActorAlarm);
+			case 2:
+				return v2.encodeActorAlarm(data as v2.ActorAlarm);
+			case 3:
+				return v3.encodeActorAlarm(data as v3.ActorAlarm);
+			default:
+				throw new Error(`Unknown version ${version}`);
+		}
+	},
+	deserializeConverters: () => [],
+	serializeConverters: () => [],
 });

package/src/serde.ts

@@ -2,7 +2,7 @@ import * as cbor from "cbor-x";
 import invariant from "invariant";
 import type { z } from "zod";
 import { assertUnreachable } from "@/common/utils";
-import type { VersionedDataHandler } from "@/common/versioned-data";
+import type { VersionedDataHandler } from "vbare";
 import type { Encoding } from "@/mod";
 import { jsonParseCompat, jsonStringifyCompat } from "./actor/protocol/serde";
 
@@ -57,6 +57,7 @@ export function serializeWithEncoding<TBare, TJson, T = TBare>(
 	encoding: Encoding,
 	value: T,
 	versionedDataHandler: VersionedDataHandler<TBare> | undefined,
+	version: number | undefined,
 	zodSchema: z.ZodType<TJson>,
 	toJson: (value: T) => TJson,
 	toBare: (value: T) => TBare,
@@ -75,8 +76,14 @@ export function serializeWithEncoding<TBare, TJson, T = TBare>(
 				"VersionedDataHandler is required for 'bare' encoding",
 			);
 		}
+		if (version === undefined) {
+			throw new Error("version is required for 'bare' encoding");
+		}
 		const bareValue = toBare(value);
-		return versionedDataHandler.serializeWithEmbeddedVersion(bareValue);
+		return versionedDataHandler.serializeWithEmbeddedVersion(
+			bareValue,
+			version,
+		);
 	} else {
 		assertUnreachable(encoding);
 	}

package/src/test/config.ts

@@ -1,16 +1,17 @@
 import { z } from "zod";
 import { RunnerConfigSchema } from "@/registry/run-config";
 
-export const ConfigSchema = RunnerConfigSchema.removeDefault()
-	.extend({
-		hostname: z
-			.string()
-			.optional()
-			.default(process.env.HOSTNAME ?? "127.0.0.1"),
-		port: z
-			.number()
-			.optional()
-			.default(Number.parseInt(process.env.PORT ?? "8080")),
-	})
-	.default({});
+const ConfigSchemaBase = RunnerConfigSchema.removeDefault().extend({
+	hostname: z
+		.string()
+		.optional()
+		.default(process.env.HOSTNAME ?? "127.0.0.1"),
+	port: z
+		.number()
+		.optional()
+		.default(Number.parseInt(process.env.PORT ?? "8080")),
+});
+export const ConfigSchema = ConfigSchemaBase.default(() =>
+	ConfigSchemaBase.parse({}),
+);
 export type InputConfig = z.input<typeof ConfigSchema>;

package/src/test/mod.ts

@@ -1,122 +1,96 @@
 import { createServer } from "node:net";
-import { serve as honoServe, type ServerType } from "@hono/node-server";
+import { serve as honoServe } from "@hono/node-server";
 import { createNodeWebSocket } from "@hono/node-ws";
+import invariant from "invariant";
 import { type TestContext, vi } from "vitest";
-import { ClientConfigSchema } from "@/client/config";
 import { type Client, createClient } from "@/client/mod";
-import { chooseDefaultDriver } from "@/drivers/default";
 import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
-import {
-	configureInspectorAccessToken,
-	getInspectorUrl,
-} from "@/inspector/utils";
+import { configureInspectorAccessToken } from "@/inspector/utils";
 import { createManagerRouter } from "@/manager/router";
-import { createClientWithDriver } from "@/mod";
-import type { Registry } from "@/registry/mod";
+import { createClientWithDriver, type Registry, type RunConfig } from "@/mod";
+import { ClientConfigSchema } from "@/client/config";
 import { RunnerConfigSchema } from "@/registry/run-config";
-import { ConfigSchema, type InputConfig } from "./config";
 import { logger } from "./log";
 
-async function serve(
-	registry: Registry<any>,
-	inputConfig?: InputConfig,
-): Promise<ServerType> {
-	// Configure default configuration
-	inputConfig ??= {};
+export interface SetupTestResult<A extends Registry<any>> {
+	client: Client<A>;
+}
+
+// Must use `TestContext` since global hooks do not work when running concurrently
+export async function setupTest<A extends Registry<any>>(
+	c: TestContext,
+	registry: A,
+): Promise<SetupTestResult<A>> {
+	// Force enable test mode
+	registry.config.test.enabled = true;
 
-	const config = ConfigSchema.parse(inputConfig);
+	// Create driver
+	const driver = await createFileSystemOrMemoryDriver(
+		true,
+		`/tmp/rivetkit-test-${crypto.randomUUID()}`,
+	);
 
+	// Build driver config
+	// biome-ignore lint/style/useConst: Assigned later
 	let upgradeWebSocket: any;
-	if (!config.getUpgradeWebSocket) {
-		config.getUpgradeWebSocket = () => upgradeWebSocket!;
-	}
+	const config: RunConfig = RunnerConfigSchema.parse({
+		driver,
+		getUpgradeWebSocket: () => upgradeWebSocket!,
+		inspector: {
+			enabled: true,
+			token: () => "token",
+		},
+	});
 
 	// Create router
-	const runConfig = RunnerConfigSchema.parse(inputConfig);
-	const driver =
-		inputConfig.driver ?? (await createFileSystemOrMemoryDriver(false));
 	const managerDriver = driver.manager(registry.config, config);
-	const client = createClientWithDriver(
+	const internalClient = createClientWithDriver(
 		managerDriver,
 		ClientConfigSchema.parse({}),
 	);
 	configureInspectorAccessToken(config, managerDriver);
 	const { router } = createManagerRouter(
 		registry.config,
-		runConfig,
+		config,
 		managerDriver,
 		driver,
-		client,
+		internalClient,
 	);
 
 	// Inject WebSocket
 	const nodeWebSocket = createNodeWebSocket({ app: router });
 	upgradeWebSocket = nodeWebSocket.upgradeWebSocket;
 
+	// Start server
+	const port = await getPort();
 	const server = honoServe({
 		fetch: router.fetch,
-		hostname: config.hostname,
-		port: config.port,
+		hostname: "127.0.0.1",
+		port,
 	});
+	invariant(
+		nodeWebSocket.injectWebSocket !== undefined,
+		"should have injectWebSocket",
+	);
 	nodeWebSocket.injectWebSocket(server);
+	const endpoint = `http://127.0.0.1:${port}`;
 
-	logger().info({
-		msg: "rivetkit started",
-		hostname: config.hostname,
-		port: config.port,
-		definitions: Object.keys(registry.config.use).length,
-	});
+	logger().info({ msg: "test server listening", port });
 
-	return server;
-}
-
-export interface SetupTestResult<A extends Registry<any>> {
-	client: Client<A>;
-	mockDriver: {
-		actorDriver: {
-			setCreateVarsContext: (ctx: any) => void;
-		};
-	};
-}
-
-// Must use `TestContext` since global hooks do not work when running concurrently
-export async function setupTest<A extends Registry<any>>(
-	c: TestContext,
-	registry: A,
-): Promise<SetupTestResult<A>> {
-	vi.useFakeTimers();
-
-	// Set up mock driver for testing createVars context
-	const mockDriverContext: any = {};
-	const setDriverContextFn = (ctx: any) => {
-		mockDriverContext.current = ctx;
-	};
-
-	// We don't need to modify the driver context anymore since we're testing with the actual context
+	// Cleanup on test finish
+	c.onTestFinished(async () => {
+		await new Promise((resolve) => server.close(() => resolve(undefined)));
+	});
 
-	// Start server with a random port
-	const port = await getPort();
-	const server = await serve(registry, { port });
-	c.onTestFinished(
-		async () =>
-			await new Promise((resolve) => server.close(() => resolve())),
-	);
+	// Create client
+	const client = createClient<A>({
+		endpoint,
+		namespace: "default",
+		runnerName: "default",
+	});
+	c.onTestFinished(async () => await client.dispose());
 
-	throw "TODO: Fix engine port";
-
-	// // TODO: Figure out how to make this the correct endpoint
-	// // Create client
-	// const client = createClient<A>(`http://127.0.0.1:${port}`);
-	// c.onTestFinished(async () => await client.dispose());
-	//
-	// return {
-	// 	client,
-	// 	mockDriver: {
-	// 		actorDriver: {
-	// 			setCreateVarsContext: setDriverContextFn,
-	// 		},
-	// 	},
-	// };
+	return { client };
 }
 
 export async function getPort(): Promise<number> {