@rivetkit/cloudflare-workers 2.0.2 → 2.0.4-rc.1

package/src/actor-handler-do.ts

@@ -1,62 +1,59 @@
 import { DurableObject, env } from "cloudflare:workers";
-import type {
-	ActorKey,
-	ActorRouter,
-	Registry,
-	RunConfig,
-} from "@rivetkit/core";
-import {
-	createActorRouter,
-	createClientWithDriver,
-	createInlineClientDriver,
-} from "@rivetkit/core";
-import { serializeEmptyPersistData } from "@rivetkit/core/driver-helpers";
 import type { ExecutionContext } from "hono";
 import invariant from "invariant";
+import type { ActorKey, ActorRouter, Registry, RegistryConfig } from "rivetkit";
+import { createActorRouter, createClientWithDriver } from "rivetkit";
+import type { ActorDriver, ManagerDriver } from "rivetkit/driver-helpers";
+import { getInitialActorKvState } from "rivetkit/driver-helpers";
+import type { GetUpgradeWebSocket } from "rivetkit/utils";
+import { stringifyError } from "rivetkit/utils";
 import {
+	ActorGlobalState,
 	CloudflareDurableObjectGlobalState,
 	createCloudflareActorsActorDriverBuilder,
 } from "./actor-driver";
+import { buildActorId, parseActorId } from "./actor-id";
+import { kvGet, kvPut } from "./actor-kv";
+import { GLOBAL_KV_KEYS } from "./global-kv";
 import type { Bindings } from "./handler";
+import { getCloudflareAmbientEnv } from "./handler";
 import { logger } from "./log";
-
-export const KEYS = {
-	NAME: "rivetkit:name",
-	KEY: "rivetkit:key",
-	PERSIST_DATA: "rivetkit:data",
-};
+import { CloudflareActorsManagerDriver } from "./manager-driver";
 
 export interface ActorHandlerInterface extends DurableObject {
-	initialize(req: ActorInitRequest): Promise<void>;
+	create(req: ActorInitRequest): Promise<ActorInitResponse>;
+	getMetadata(): Promise<
+		| {
+				actorId: string;
+				name: string;
+				key: ActorKey;
+				destroying: boolean;
+		  }
+		| undefined
+	>;
+	managerKvGet(key: Uint8Array): Promise<Uint8Array | null>;
 }
 
 export interface ActorInitRequest {
 	name: string;
 	key: ActorKey;
 	input?: unknown;
+	allowExisting: boolean;
 }
-
-interface InitializedData {
-	name: string;
-	key: ActorKey;
-}
+export type ActorInitResponse =
+	| { success: { actorId: string; created: boolean } }
+	| { error: { actorAlreadyExists: true } };
 
 export type DurableObjectConstructor = new (
 	...args: ConstructorParameters<typeof DurableObject<Bindings>>
 ) => DurableObject<Bindings>;
 
-interface LoadedActor {
-	actorRouter: ActorRouter;
-}
-
 export function createActorDurableObject(
 	registry: Registry<any>,
-	rootRunConfig: RunConfig,
+	getUpgradeWebSocket: GetUpgradeWebSocket,
 ): DurableObjectConstructor {
 	const globalState = new CloudflareDurableObjectGlobalState();
-
-	// Configure to use the runner role instead of server role
-	const runConfig = Object.assign({}, rootRunConfig, { role: "runner" });
+	const parsedConfig = registry.parseConfig();
 
 	/**
 	 * Startup steps:
@@ -68,46 +65,104 @@ export function createActorDurableObject(
 		extends DurableObject<Bindings>
 		implements ActorHandlerInterface
 	{
-		#initialized?: InitializedData;
-		#initializedPromise?: PromiseWithResolvers<void>;
-
-		#actor?: LoadedActor;
+		/**
+		 * This holds a strong reference to ActorGlobalState.
+		 * CloudflareDurableObjectGlobalState holds a weak reference so we can
+		 * access it elsewhere.
+		 **/
+		#state: ActorGlobalState;
+
+		constructor(
+			...args: ConstructorParameters<typeof DurableObject<Bindings>>
+		) {
+			super(...args);
+
+			// Initialize SQL table for key-value storage
+			//
+			// We do this instead of using the native KV storage so we can store blob keys. The native CF KV API only supports string keys.
+			this.ctx.storage.sql.exec(`
+				CREATE TABLE IF NOT EXISTS _rivetkit_kv_storage(
+					key BLOB PRIMARY KEY,
+					value BLOB
+				);
+			`);
+
+			// Initialize SQL table for actor metadata
+			//
+			// id always equals 1 in order to ensure that there's always exactly 1 row in this table
+			this.ctx.storage.sql.exec(`
+				CREATE TABLE IF NOT EXISTS _rivetkit_metadata(
+					id INTEGER PRIMARY KEY CHECK (id = 1),
+					name TEXT NOT NULL,
+					key TEXT NOT NULL,
+					destroyed INTEGER DEFAULT 0,
+					generation INTEGER DEFAULT 0
+				);
+			`);
+
+			// Get or create the actor state from the global WeakMap
+			const state = globalState.getActorState(this.ctx);
+			if (state) {
+				this.#state = state;
+			} else {
+				this.#state = new ActorGlobalState();
+				globalState.setActorState(this.ctx, this.#state);
+			}
+		}
 
-		async #loadActor(): Promise<LoadedActor> {
-			// Wait for init
-			if (!this.#initialized) {
-				// Wait for init
-				if (this.#initializedPromise) {
-					await this.#initializedPromise.promise;
-				} else {
-					this.#initializedPromise = Promise.withResolvers();
-					const res = await this.ctx.storage.get([
-						KEYS.NAME,
-						KEYS.KEY,
-						KEYS.PERSIST_DATA,
-					]);
-					if (res.get(KEYS.PERSIST_DATA)) {
-						const name = res.get(KEYS.NAME) as string;
-						if (!name) throw new Error("missing actor name");
-						const key = res.get(KEYS.KEY) as ActorKey;
-						if (!key) throw new Error("missing actor key");
-
-						logger().debug("already initialized", { name, key });
-
-						this.#initialized = { name, key };
-						this.#initializedPromise.resolve();
+		async #loadActor() {
+			invariant(this.#state, "State should be initialized");
+
+			// Check if initialized
+			if (!this.#state.initialized) {
+				// Query SQL for initialization data
+				const cursor = this.ctx.storage.sql.exec(
+					"SELECT name, key, destroyed, generation FROM _rivetkit_metadata WHERE id = 1",
+				);
+				const result = cursor.raw().next();
+
+				if (!result.done && result.value) {
+					const name = result.value[0] as string;
+					const key = JSON.parse(
+						result.value[1] as string,
+					) as ActorKey;
+					const destroyed = result.value[2] as number;
+					const generation = result.value[3] as number;
+
+					// Only initialize if not destroyed
+					if (!destroyed) {
+						logger().debug({
+							msg: "already initialized",
+							name,
+							key,
+							generation,
+						});
+
+						this.#state.initialized = { name, key, generation };
 					} else {
-						logger().debug("waiting to initialize");
+						logger().debug("actor is destroyed, cannot load");
+						throw new Error("Actor is destroyed");
 					}
+				} else {
+					logger().debug("not initialized");
+					throw new Error("Actor is not initialized");
 				}
 			}
 
 			// Check if already loaded
-			if (this.#actor) {
-				return this.#actor;
+			if (this.#state.actor) {
+				// Assert that the cached actor has the correct generation
+				// This will catch any cases where #state.actor has a stale generation
+				invariant(
+					!this.#state.initialized ||
+						this.#state.actor.generation ===
+							this.#state.initialized.generation,
+					`Stale actor cached: actor generation ${this.#state.actor.generation} != initialized generation ${this.#state.initialized?.generation}. This should not happen.`,
+				);
+				return this.#state.actor;
 			}
 
-			if (!this.#initialized) throw new Error("Not initialized");
+			if (!this.#state.initialized) throw new Error("Not initialized");
 
 			// Register DO with global state first
 			// HACK: This leaks the DO context, but DO does not provide a native way
@@ -116,93 +171,253 @@ export function createActorDurableObject(
 			const actorId = this.ctx.id.toString();
 			globalState.setDOState(actorId, { ctx: this.ctx, env: env });
 
-			// Configure actor driver
-			invariant(runConfig.driver, "runConfig.driver");
-			runConfig.driver.actor =
-				createCloudflareActorsActorDriverBuilder(globalState);
-
-			// Create manager driver (we need this for the actor router)
-			const managerDriver = runConfig.driver.manager(
-				registry.config,
-				runConfig,
-			);
+			// Create manager driver
+			const managerDriver = new CloudflareActorsManagerDriver();
 
 			// Create inline client
-			const inlineClient = createClientWithDriver(
-				createInlineClientDriver(managerDriver),
-			);
+			const inlineClient = createClientWithDriver(managerDriver);
+
+			// Create actor driver builder
+			const actorDriverBuilder =
+				createCloudflareActorsActorDriverBuilder(globalState);
 
 			// Create actor driver
-			const actorDriver = runConfig.driver.actor(
-				registry.config,
-				runConfig,
+			const actorDriver = actorDriverBuilder(
+				parsedConfig,
 				managerDriver,
 				inlineClient,
 			);
 
 			// Create actor router
-			const actorRouter = createActorRouter(runConfig, actorDriver);
+			const actorRouter = createActorRouter(
+				parsedConfig,
+				actorDriver,
+				getUpgradeWebSocket,
+				registry.config.test?.enabled ?? false,
+			);
 
-			// Save actor
-			this.#actor = {
+			// Save actor with generation
+			this.#state.actor = {
 				actorRouter,
+				actorDriver,
+				generation: this.#state.initialized.generation,
 			};
 
+			// Build actor ID with generation for loading
+			const actorIdWithGen = buildActorId(
+				actorId,
+				this.#state.initialized.generation,
+			);
+
 			// Initialize the actor instance with proper metadata
 			// This ensures the actor driver knows about this actor
-			await actorDriver.loadActor(actorId);
+			await actorDriver.loadActor(actorIdWithGen);
 
-			return this.#actor;
+			return this.#state.actor;
 		}
 
-		/** RPC called by the service that creates the DO to initialize it. */
-		async initialize(req: ActorInitRequest) {
-			// TODO: Need to add this to a core promise that needs to be resolved before start
+		/** RPC called to get actor metadata without creating it */
+		async getMetadata(): Promise<
+			| {
+					actorId: string;
+					name: string;
+					key: ActorKey;
+					destroying: boolean;
+			  }
+			| undefined
+		> {
+			// Query the metadata
+			const cursor = this.ctx.storage.sql.exec(
+				"SELECT name, key, destroyed, generation FROM _rivetkit_metadata WHERE id = 1",
+			);
+			const result = cursor.raw().next();
+
+			if (!result.done && result.value) {
+				const name = result.value[0] as string;
+				const key = JSON.parse(result.value[1] as string) as ActorKey;
+				const destroyed = result.value[2] as number;
+				const generation = result.value[3] as number;
+
+				// Check if destroyed
+				if (destroyed) {
+					logger().debug({
+						msg: "getMetadata: actor is destroyed",
+						name,
+						key,
+						generation,
+					});
+					return undefined;
+				}
 
-			await this.ctx.storage.put({
-				[KEYS.NAME]: req.name,
-				[KEYS.KEY]: req.key,
-				[KEYS.PERSIST_DATA]: serializeEmptyPersistData(req.input),
+				// Build actor ID with generation
+				const doId = this.ctx.id.toString();
+				const actorId = buildActorId(doId, generation);
+				const destroying =
+					globalState.getActorState(this.ctx)?.destroying ?? false;
+
+				logger().debug({
+					msg: "getMetadata: found actor metadata",
+					actorId,
+					name,
+					key,
+					generation,
+					destroying,
+				});
+
+				return { actorId, name, key, destroying };
+			}
+
+			logger().debug({
+				msg: "getMetadata: no metadata found",
 			});
-			this.#initialized = {
+			return undefined;
+		}
+
+		/** RPC called by ManagerDriver.kvGet to read from KV. */
+		async managerKvGet(key: Uint8Array): Promise<Uint8Array | null> {
+			return kvGet(this.ctx.storage.sql, key);
+		}
+
+		/** RPC called by the manager to create a DO. Can optionally allow existing actors. */
+		async create(req: ActorInitRequest): Promise<ActorInitResponse> {
+			// Check if actor exists
+			const checkCursor = this.ctx.storage.sql.exec(
+				"SELECT destroyed, generation FROM _rivetkit_metadata WHERE id = 1",
+			);
+			const checkResult = checkCursor.raw().next();
+
+			let created = false;
+			let generation = 0;
+
+			if (!checkResult.done && checkResult.value) {
+				const destroyed = checkResult.value[0] as number;
+				generation = checkResult.value[1] as number;
+
+				if (!destroyed) {
+					// Actor exists and is not destroyed
+					if (!req.allowExisting) {
+						// Fail if not allowing existing actors
+						logger().debug({
+							msg: "create failed: actor already exists",
+							name: req.name,
+							key: req.key,
+							generation,
+						});
+						return { error: { actorAlreadyExists: true } };
+					}
+
+					// Return existing actor
+					logger().debug({
+						msg: "actor already exists",
+						key: req.key,
+						generation,
+					});
+					const doId = this.ctx.id.toString();
+					const actorId = buildActorId(doId, generation);
+					return { success: { actorId, created: false } };
+				}
+
+				// Actor exists but is destroyed - resurrect with incremented generation
+				generation = generation + 1;
+				created = true;
+
+				// Clear stale actor from previous generation
+				// This is necessary because the DO instance may still be in memory
+				// with the old #state.actor field from before the destroy
+				if (this.#state) {
+					this.#state.actor = undefined;
+				}
+
+				logger().debug({
+					msg: "resurrecting destroyed actor",
+					key: req.key,
+					oldGeneration: generation - 1,
+					newGeneration: generation,
+				});
+			} else {
+				// No actor exists - will create with generation 0
+				generation = 0;
+				created = true;
+				logger().debug({
+					msg: "creating new actor",
+					key: req.key,
+					generation,
+				});
+			}
+
+			// Perform upsert - either inserts new or updates destroyed actor
+			this.ctx.storage.sql.exec(
+				`INSERT INTO _rivetkit_metadata (id, name, key, destroyed, generation)
+				VALUES (1, ?, ?, 0, ?)
+				ON CONFLICT(id) DO UPDATE SET
+					name = excluded.name,
+					key = excluded.key,
+					destroyed = 0,
+					generation = excluded.generation`,
+				req.name,
+				JSON.stringify(req.key),
+				generation,
+			);
+
+			this.#state.initialized = {
 				name: req.name,
 				key: req.key,
+				generation,
 			};
 
-			logger().debug("initialized actor", { key: req.key });
+			// Build actor ID with generation
+			const doId = this.ctx.id.toString();
+			const actorId = buildActorId(doId, generation);
+
+			// Initialize storage and update KV when created or resurrected
+			if (created) {
+				// Initialize persist data in KV storage
+				initializeActorKvStorage(this.ctx.storage.sql, req.input);
+
+				// Update metadata in the background
+				const env = getCloudflareAmbientEnv();
+				const actorData = { name: req.name, key: req.key, generation };
+				this.ctx.waitUntil(
+					env.ACTOR_KV.put(
+						GLOBAL_KV_KEYS.actorMetadata(actorId),
+						JSON.stringify(actorData),
+					),
+				);
+			}
 
-			// Preemptively actor so the lifecycle hooks are called
+			// Preemptively load actor so the lifecycle hooks are called
 			await this.#loadActor();
+
+			logger().debug({
+				msg: created
+					? "actor created/resurrected"
+					: "returning existing actor",
+				actorId,
+				created,
+				generation,
+			});
+
+			return { success: { actorId, created } };
 		}
 
 		async fetch(request: Request): Promise<Response> {
-			const { actorRouter } = await this.#loadActor();
+			const { actorRouter, generation } = await this.#loadActor();
+
+			// Build actor ID with generation
+			const doId = this.ctx.id.toString();
+			const actorId = buildActorId(doId, generation);
 
-			const actorId = this.ctx.id.toString();
 			return await actorRouter.fetch(request, {
 				actorId,
 			});
 		}
 
 		async alarm(): Promise<void> {
-			await this.#loadActor();
-			const actorId = this.ctx.id.toString();
+			const { actorDriver, generation } = await this.#loadActor();
 
-			// Get the actor driver
-			invariant(runConfig.driver, "runConfig.driver");
-			const managerDriver = runConfig.driver.manager(
-				registry.config,
-				runConfig,
-			);
-			const inlineClient = createClientWithDriver(
-				createInlineClientDriver(managerDriver),
-			);
-			const actorDriver = runConfig.driver.actor(
-				registry.config,
-				runConfig,
-				managerDriver,
-				inlineClient,
-			);
+			// Build actor ID with generation
+			const doId = this.ctx.id.toString();
+			const actorId = buildActorId(doId, generation);
 
 			// Load the actor instance and trigger alarm
 			const actor = await actorDriver.loadActor(actorId);
@@ -210,3 +425,13 @@ export function createActorDurableObject(
 		}
 	};
 }
+
+function initializeActorKvStorage(
+	sql: SqlStorage,
+	input: unknown | undefined,
+): void {
+	const initialKvState = getInitialActorKvState(input);
+	for (const [key, value] of initialKvState) {
+		kvPut(sql, key, value);
+	}
+}

package/src/actor-id.ts

@@ -0,0 +1,38 @@
+/**
+ * Actor ID utilities for managing actor IDs with generation tracking.
+ *
+ * Actor IDs are formatted as: `{doId}:{generation}`
+ * This allows tracking actor resurrection and preventing stale references.
+ */
+
+/**
+ * Build an actor ID from a Durable Object ID and generation number.
+ * @param doId The Durable Object ID
+ * @param generation The generation number (increments on resurrection)
+ * @returns The formatted actor ID
+ */
+export function buildActorId(doId: string, generation: number): string {
+	return `${doId}:${generation}`;
+}
+
+/**
+ * Parse an actor ID into its components.
+ * @param actorId The actor ID to parse
+ * @returns A tuple of [doId, generation]
+ * @throws Error if the actor ID format is invalid
+ */
+export function parseActorId(actorId: string): [string, number] {
+	const parts = actorId.split(":");
+	if (parts.length !== 2) {
+		throw new Error(`Invalid actor ID format: ${actorId}`);
+	}
+
+	const [doId, generationStr] = parts;
+	const generation = parseInt(generationStr, 10);
+
+	if (Number.isNaN(generation)) {
+		throw new Error(`Invalid generation number in actor ID: ${actorId}`);
+	}
+
+	return [doId, generation];
+}

package/src/actor-kv.ts

@@ -0,0 +1,71 @@
+export function kvGet(sql: SqlStorage, key: Uint8Array): Uint8Array | null {
+	const cursor = sql.exec(
+		"SELECT value FROM _rivetkit_kv_storage WHERE key = ?",
+		key,
+	);
+	const result = cursor.raw().next();
+
+	if (!result.done && result.value) {
+		return toUint8Array(result.value[0]);
+	}
+	return null;
+}
+
+export function kvPut(
+	sql: SqlStorage,
+	key: Uint8Array,
+	value: Uint8Array,
+): void {
+	sql.exec(
+		"INSERT OR REPLACE INTO _rivetkit_kv_storage (key, value) VALUES (?, ?)",
+		key,
+		value,
+	);
+}
+
+export function kvDelete(sql: SqlStorage, key: Uint8Array): void {
+	sql.exec("DELETE FROM _rivetkit_kv_storage WHERE key = ?", key);
+}
+
+export function kvListPrefix(
+	sql: SqlStorage,
+	prefix: Uint8Array,
+): [Uint8Array, Uint8Array][] {
+	const cursor = sql.exec("SELECT key, value FROM _rivetkit_kv_storage");
+	const entries: [Uint8Array, Uint8Array][] = [];
+
+	for (const row of cursor.raw()) {
+		const key = toUint8Array(row[0]);
+		const value = toUint8Array(row[1]);
+
+		// Check if key starts with prefix
+		if (hasPrefix(key, prefix)) {
+			entries.push([key, value]);
+		}
+	}
+
+	return entries;
+}
+
+// Helper function to convert SqlStorageValue to Uint8Array
+function toUint8Array(
+	value: string | number | ArrayBuffer | Uint8Array | null,
+): Uint8Array {
+	if (value instanceof Uint8Array) {
+		return value;
+	}
+	if (value instanceof ArrayBuffer) {
+		return new Uint8Array(value);
+	}
+	throw new Error(
+		`Unexpected SQL value type: ${typeof value} (${value?.constructor?.name})`,
+	);
+}
+
+function hasPrefix(arr: Uint8Array, prefix: Uint8Array): boolean {
+	if (prefix.length > arr.length) return false;
+	for (let i = 0; i < prefix.length; i++) {
+		if (arr[i] !== prefix[i]) return false;
+	}
+	return true;
+}

package/src/config.ts

@@ -1,12 +1,24 @@
-import { RunConfigSchema } from "@rivetkit/core/driver-helpers";
-import type { Hono } from "hono";
-import { z } from "zod";
-
-export const ConfigSchema = RunConfigSchema.removeDefault()
-	.omit({ driver: true, getUpgradeWebSocket: true })
-	.extend({
-		app: z.custom<Hono>().optional(),
-	})
-	.default({});
+import type { Client } from "rivetkit";
+import { z } from "zod/v4";
+
+const ConfigSchemaBase = z.object({
+	/** Path that the Rivet manager API will be mounted. */
+	managerPath: z.string().optional().default("/api/rivet"),
+
+	/** Runner key for authentication. */
+	runnerKey: z.string().optional(),
+
+	/** Disable the welcome message. */
+	noWelcome: z.boolean().optional().default(false),
+
+	fetch: z
+		.custom<
+			ExportedHandlerFetchHandler<{ RIVET: Client<any> }, unknown>
+		>()
+		.optional(),
+});
+export const ConfigSchema = ConfigSchemaBase.default(() =>
+	ConfigSchemaBase.parse({}),
+);
 export type InputConfig = z.input<typeof ConfigSchema>;
 export type Config = z.infer<typeof ConfigSchema>;

package/src/global-kv.ts

@@ -0,0 +1,6 @@
+/** KV keys for using Workers KV to store actor metadata globally. */
+export const GLOBAL_KV_KEYS = {
+	actorMetadata: (actorId: string): string => {
+		return `actor:${actorId}:metadata`;
+	},
+};