rivetkit 2.0.2 → 2.0.4

package/src/drivers/engine/config.ts

@@ -0,0 +1,31 @@
+import type { Hono } from "hono";
+import { z } from "zod";
+import { getEnvUniversal } from "@/utils";
+
+export const ConfigSchema = z
+	.object({
+		app: z.custom<Hono>().optional(),
+		endpoint: z
+			.string()
+			.default(
+				() => getEnvUniversal("RIVET_ENGINE") ?? "http://localhost:6420",
+			),
+		pegboardEndpoint: z.string().optional(),
+		namespace: z
+			.string()
+			.default(() => getEnvUniversal("RIVET_NAMESPACE") ?? "default"),
+		runnerName: z
+			.string()
+			.default(() => getEnvUniversal("RIVET_RUNNER") ?? "rivetkit"),
+		// TODO: Automatically attempt to determine key by common env vars (e.g. k8s pod name)
+		runnerKey: z
+			.string()
+			.default(
+				() => getEnvUniversal("RIVET_RUNNER_KEY") ?? crypto.randomUUID(),
+			),
+		totalSlots: z.number().default(100_000),
+	})
+	.default({});
+
+export type InputConfig = z.input<typeof ConfigSchema>;
+export type Config = z.infer<typeof ConfigSchema>;

package/src/drivers/engine/kv.ts

@@ -0,0 +1,3 @@
+export const KEYS = {
+	PERSIST_DATA: Uint8Array.from([1]),
+};

package/src/drivers/engine/log.ts

@@ -0,0 +1,5 @@
+import { getLogger } from "@/common/log";
+
+export function logger() {
+	return getLogger("driver-engine");
+}

package/src/drivers/engine/mod.ts

@@ -0,0 +1,35 @@
+import type { Client } from "@/client/client";
+import type { ManagerDriver } from "@/manager/driver";
+import type { RegistryConfig } from "@/registry/config";
+import type { DriverConfig, RunConfig } from "@/registry/run-config";
+import { RemoteManagerDriver } from "@/remote-manager-driver/mod";
+import { EngineActorDriver } from "./actor-driver";
+import { ConfigSchema, type InputConfig } from "./config";
+
+export { EngineActorDriver } from "./actor-driver";
+export { type Config, ConfigSchema, type InputConfig } from "./config";
+
+export function createEngineDriver(inputConfig?: InputConfig): DriverConfig {
+	const config = ConfigSchema.parse(inputConfig);
+
+	return {
+		name: "engine",
+		manager: (_registryConfig, runConfig) => {
+			return new RemoteManagerDriver(runConfig);
+		},
+		actor: (
+			registryConfig: RegistryConfig,
+			runConfig: RunConfig,
+			managerDriver: ManagerDriver,
+			inlineClient: Client<any>,
+		) => {
+			return new EngineActorDriver(
+				registryConfig,
+				runConfig,
+				managerDriver,
+				inlineClient,
+				config,
+			);
+		},
+	};
+}

package/src/drivers/file-system/actor.ts

@@ -0,0 +1,91 @@
+import type { GenericConnGlobalState } from "@/actor/generic-conn-driver";
+import { loggerWithoutContext } from "@/actor/log";
+import type { AnyClient } from "@/client/client";
+import type {
+	ActorDriver,
+	AnyActorInstance,
+	ManagerDriver,
+} from "@/driver-helpers/mod";
+import type { RegistryConfig, RunConfig } from "@/mod";
+import { bufferToArrayBuffer } from "@/utils";
+import type { FileSystemGlobalState } from "./global-state";
+
+export type ActorDriverContext = Record<never, never>;
+
+/**
+ * File System implementation of the Actor Driver
+ */
+export class FileSystemActorDriver implements ActorDriver {
+	#registryConfig: RegistryConfig;
+	#runConfig: RunConfig;
+	#managerDriver: ManagerDriver;
+	#inlineClient: AnyClient;
+	#state: FileSystemGlobalState;
+
+	constructor(
+		registryConfig: RegistryConfig,
+		runConfig: RunConfig,
+		managerDriver: ManagerDriver,
+		inlineClient: AnyClient,
+		state: FileSystemGlobalState,
+	) {
+		this.#registryConfig = registryConfig;
+		this.#runConfig = runConfig;
+		this.#managerDriver = managerDriver;
+		this.#inlineClient = inlineClient;
+		this.#state = state;
+	}
+
+	async loadActor(actorId: string): Promise<AnyActorInstance> {
+		return this.#state.startActor(
+			this.#registryConfig,
+			this.#runConfig,
+			this.#inlineClient,
+			this,
+			actorId,
+		);
+	}
+
+	getGenericConnGlobalState(actorId: string): GenericConnGlobalState {
+		return this.#state.getActorOrError(actorId).genericConnGlobalState;
+	}
+
+	/**
+	 * Get the current storage directory path
+	 */
+	get storagePath(): string {
+		return this.#state.storagePath;
+	}
+
+	getContext(_actorId: string): ActorDriverContext {
+		return {};
+	}
+
+	async readPersistedData(actorId: string): Promise<Uint8Array | undefined> {
+		return new Uint8Array(
+			(await this.#state.loadActorStateOrError(actorId)).persistedData,
+		);
+	}
+
+	async writePersistedData(actorId: string, data: Uint8Array): Promise<void> {
+		const state = await this.#state.loadActorStateOrError(actorId);
+
+		// Save state to disk
+		await this.#state.writeActor(actorId, {
+			...state,
+			persistedData: bufferToArrayBuffer(data),
+		});
+	}
+
+	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
+		await this.#state.setActorAlarm(actor.id, timestamp);
+	}
+
+	getDatabase(actorId: string): Promise<unknown | undefined> {
+		return this.#state.createDatabase(actorId);
+	}
+
+	sleep(actorId: string): Promise<void> {
+		return this.#state.sleepActor(actorId);
+	}
+}