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

package/src/manager-driver.ts

@@ -1,41 +1,31 @@
-import type { Context as HonoContext } from "hono";
-import type { Encoding } from "rivetkit";
+import type { Hono, Context as HonoContext } from "hono";
+import type { Encoding, RegistryConfig, UniversalWebSocket } from "rivetkit";
 import {
 	type ActorOutput,
 	type CreateInput,
 	type GetForIdInput,
 	type GetOrCreateWithKeyInput,
 	type GetWithKeyInput,
-	HEADER_AUTH_DATA,
-	HEADER_CONN_PARAMS,
-	HEADER_ENCODING,
-	HEADER_EXPOSE_INTERNAL_ERROR,
+	type ListActorsInput,
+	type ManagerDisplayInformation,
 	type ManagerDriver,
+	WS_PROTOCOL_ACTOR,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_STANDARD,
+	WS_PROTOCOL_TARGET,
 } from "rivetkit/driver-helpers";
-import { ActorAlreadyExists, InternalError } from "rivetkit/errors";
+import {
+	ActorDuplicateKey,
+	ActorNotFound,
+	InternalError,
+} from "rivetkit/errors";
+import { assertUnreachable } from "rivetkit/utils";
+import { parseActorId } from "./actor-id";
 import { getCloudflareAmbientEnv } from "./handler";
 import { logger } from "./log";
 import type { Bindings } from "./mod";
-import { serializeKey, serializeNameAndKey } from "./util";
-
-// Actor metadata structure
-interface ActorData {
-	name: string;
-	key: string[];
-}
-
-const KEYS = {
-	ACTOR: {
-		// Combined key for actor metadata (name and key)
-		metadata: (actorId: string) => `actor:${actorId}:metadata`,
-
-		// Key index function for actor lookup
-		keyIndex: (name: string, key: string[] = []) => {
-			// Use serializeKey for consistent handling of all keys
-			return `actor_key:${serializeKey(key)}`;
-		},
-	},
-};
+import { serializeNameAndKey } from "./util";
 
 const STANDARD_WEBSOCKET_HEADERS = [
 	"connection",
@@ -47,16 +37,24 @@ const STANDARD_WEBSOCKET_HEADERS = [
 ];
 
 export class CloudflareActorsManagerDriver implements ManagerDriver {
-	async sendRequest(actorId: string, actorRequest: Request): Promise<Response> {
+	async sendRequest(
+		actorId: string,
+		actorRequest: Request,
+	): Promise<Response> {
 		const env = getCloudflareAmbientEnv();
 
-		logger().debug("sending request to durable object", {
+		// Parse actor ID to get DO ID
+		const [doId] = parseActorId(actorId);
+
+		logger().debug({
+			msg: "sending request to durable object",
 			actorId,
+			doId,
 			method: actorRequest.method,
 			url: actorRequest.url,
 		});
 
-		const id = env.ACTOR_DO.idFromString(actorId);
+		const id = env.ACTOR_DO.idFromString(doId);
 		const stub = env.ACTOR_DO.get(id);
 
 		return await stub.fetch(actorRequest);
@@ -67,34 +65,45 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
-	): Promise<WebSocket> {
+	): Promise<UniversalWebSocket> {
 		const env = getCloudflareAmbientEnv();
 
-		logger().debug("opening websocket to durable object", { actorId, path });
+		// Parse actor ID to get DO ID
+		const [doId] = parseActorId(actorId);
+
+		logger().debug({
+			msg: "opening websocket to durable object",
+			actorId,
+			doId,
+			path,
+		});
 
 		// Make a fetch request to the Durable Object with WebSocket upgrade
-		const id = env.ACTOR_DO.idFromString(actorId);
+		const id = env.ACTOR_DO.idFromString(doId);
 		const stub = env.ACTOR_DO.get(id);
 
+		const protocols: string[] = [];
+		protocols.push(WS_PROTOCOL_STANDARD);
+		protocols.push(`${WS_PROTOCOL_TARGET}actor`);
+		protocols.push(`${WS_PROTOCOL_ACTOR}${encodeURIComponent(actorId)}`);
+		protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
+		if (params) {
+			protocols.push(
+				`${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`,
+			);
+		}
+
 		const headers: Record<string, string> = {
 			Upgrade: "websocket",
 			Connection: "Upgrade",
-			[HEADER_EXPOSE_INTERNAL_ERROR]: "true",
-			[HEADER_ENCODING]: encoding,
+			"sec-websocket-protocol": protocols.join(", "),
 		};
-		if (params) {
-			headers[HEADER_CONN_PARAMS] = JSON.stringify(params);
-		}
-		// HACK: See packages/drivers/cloudflare-workers/src/websocket.ts
-		headers["sec-websocket-protocol"] = "rivetkit";
 
 		// Use the path parameter to determine the URL
-		const url = `http://actor${path}`;
+		const normalizedPath = path.startsWith("/") ? path : `/${path}`;
+		const url = `http://actor${normalizedPath}`;
 
-		logger().debug("rewriting websocket url", {
-			from: path,
-			to: url,
-		});
+		logger().debug({ msg: "rewriting websocket url", from: path, to: url });
 
 		const response = await stub.fetch(url, {
 			headers,
@@ -103,11 +112,12 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 
 		if (!webSocket) {
 			throw new InternalError(
-				"missing websocket connection in response from DO",
+				`missing websocket connection in response from DO\n\nStatus: ${response.status}\nResponse: ${await response.text()}`,
 			);
 		}
 
-		logger().debug("durable object websocket connection open", {
+		logger().debug({
+			msg: "durable object websocket connection open",
 			actorId,
 		});
 
@@ -122,7 +132,11 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 			(webSocket as any).dispatchEvent(event);
 		}, 0);
 
-		return webSocket as unknown as WebSocket;
+		return webSocket as unknown as UniversalWebSocket;
+	}
+
+	async buildGatewayUrl(actorId: string): Promise<string> {
+		return `http://actor/gateway/${encodeURIComponent(actorId)}`;
 	}
 
 	async proxyRequest(
@@ -130,14 +144,21 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		actorRequest: Request,
 		actorId: string,
 	): Promise<Response> {
-		logger().debug("forwarding request to durable object", {
+		const env = getCloudflareAmbientEnv();
+
+		// Parse actor ID to get DO ID
+		const [doId] = parseActorId(actorId);
+
+		logger().debug({
+			msg: "forwarding request to durable object",
 			actorId,
+			doId,
 			method: actorRequest.method,
 			url: actorRequest.url,
 		});
 
-		const id = c.env.ACTOR_DO.idFromString(actorId);
-		const stub = c.env.ACTOR_DO.get(id);
+		const id = env.ACTOR_DO.idFromString(doId);
+		const stub = env.ACTOR_DO.get(id);
 
 		return await stub.fetch(actorRequest);
 	}
@@ -148,9 +169,9 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
-		authData: unknown,
 	): Promise<Response> {
-		logger().debug("forwarding websocket to durable object", {
+		logger().debug({
+			msg: "forwarding websocket to durable object",
 			actorId,
 			path,
 		});
@@ -166,7 +187,8 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		const newUrl = new URL(`http://actor${path}`);
 		const actorRequest = new Request(newUrl, c.req.raw);
 
-		logger().debug("rewriting websocket url", {
+		logger().debug({
+			msg: "rewriting websocket url",
 			from: c.req.url,
 			to: actorRequest.url,
 		});
@@ -175,49 +197,85 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		// HACK: Since we can't build a new request, we need to remove
 		// non-standard headers manually
 		const headerKeys: string[] = [];
-		actorRequest.headers.forEach((v, k) => headerKeys.push(k));
+		actorRequest.headers.forEach((v, k) => {
+			headerKeys.push(k);
+		});
 		for (const k of headerKeys) {
 			if (!STANDARD_WEBSOCKET_HEADERS.includes(k)) {
 				actorRequest.headers.delete(k);
 			}
 		}
 
-		// Add RivetKit headers
-		actorRequest.headers.set(HEADER_EXPOSE_INTERNAL_ERROR, "true");
-		actorRequest.headers.set(HEADER_ENCODING, encoding);
+		// Build protocols for WebSocket connection
+		const protocols: string[] = [];
+		protocols.push(WS_PROTOCOL_STANDARD);
+		protocols.push(`${WS_PROTOCOL_TARGET}actor`);
+		protocols.push(`${WS_PROTOCOL_ACTOR}${encodeURIComponent(actorId)}`);
+		protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
 		if (params) {
-			actorRequest.headers.set(HEADER_CONN_PARAMS, JSON.stringify(params));
-		}
-		if (authData) {
-			actorRequest.headers.set(HEADER_AUTH_DATA, JSON.stringify(authData));
+			protocols.push(
+				`${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`,
+			);
 		}
+		actorRequest.headers.set(
+			"sec-websocket-protocol",
+			protocols.join(", "),
+		);
 
-		const id = c.env.ACTOR_DO.idFromString(actorId);
-		const stub = c.env.ACTOR_DO.get(id);
+		// Parse actor ID to get DO ID
+		const env = getCloudflareAmbientEnv();
+		const [doId] = parseActorId(actorId);
+		const id = env.ACTOR_DO.idFromString(doId);
+		const stub = env.ACTOR_DO.get(id);
 
 		return await stub.fetch(actorRequest);
 	}
 
 	async getForId({
 		c,
+		name,
 		actorId,
-	}: GetForIdInput<{ Bindings: Bindings }>): Promise<ActorOutput | undefined> {
+	}: GetForIdInput<{ Bindings: Bindings }>): Promise<
+		ActorOutput | undefined
+	> {
 		const env = getCloudflareAmbientEnv();
 
-		// Get actor metadata from KV (combined name and key)
-		const actorData = (await env.ACTOR_KV.get(KEYS.ACTOR.metadata(actorId), {
-			type: "json",
-		})) as ActorData | null;
+		// Parse actor ID to get DO ID and expected generation
+		const [doId, expectedGeneration] = parseActorId(actorId);
+
+		// Get the Durable Object stub
+		const id = env.ACTOR_DO.idFromString(doId);
+		const stub = env.ACTOR_DO.get(id);
+
+		// Call the DO's getMetadata method
+		const result = await stub.getMetadata();
+
+		if (!result) {
+			logger().debug({
+				msg: "getForId: actor not found",
+				actorId,
+			});
+			return undefined;
+		}
 
-		// If the actor doesn't exist, return undefined
-		if (!actorData) {
+		// Check if the actor IDs match in order to check if the generation matches
+		if (result.actorId !== actorId) {
+			logger().debug({
+				msg: "getForId: generation mismatch",
+				requestedActorId: actorId,
+				actualActorId: result.actorId,
+			});
 			return undefined;
 		}
 
+		if (result.destroying) {
+			throw new ActorNotFound(actorId);
+		}
+
 		return {
-			actorId,
-			name: actorData.name,
-			key: actorData.key,
+			actorId: result.actorId,
+			name: result.name,
+			key: result.key,
 		};
 	}
 
@@ -230,44 +288,82 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 	> {
 		const env = getCloudflareAmbientEnv();
 
-		logger().debug("getWithKey: searching for actor", { name, key });
+		logger().debug({ msg: "getWithKey: searching for actor", name, key });
 
 		// Generate deterministic ID from the name and key
-		// This is aligned with how createActor generates IDs
 		const nameKeyString = serializeNameAndKey(name, key);
-		const actorId = env.ACTOR_DO.idFromName(nameKeyString).toString();
+		const doId = env.ACTOR_DO.idFromName(nameKeyString).toString();
 
-		// Check if the actor metadata exists
-		const actorData = await env.ACTOR_KV.get(KEYS.ACTOR.metadata(actorId), {
-			type: "json",
-		});
+		// Try to get the Durable Object to see if it exists
+		const id = env.ACTOR_DO.idFromString(doId);
+		const stub = env.ACTOR_DO.get(id);
 
-		if (!actorData) {
-			logger().debug("getWithKey: no actor found with matching name and key", {
+		// Check if actor exists without creating it
+		const result = await stub.getMetadata();
+
+		if (result) {
+			logger().debug({
+				msg: "getWithKey: found actor with matching name and key",
+				actorId: result.actorId,
+				name: result.name,
+				key: result.key,
+			});
+			return {
+				actorId: result.actorId,
+				name: result.name,
+				key: result.key,
+			};
+		} else {
+			logger().debug({
+				msg: "getWithKey: no actor found with matching name and key",
 				name,
 				key,
-				actorId,
+				doId,
 			});
 			return undefined;
 		}
+	}
 
-		logger().debug("getWithKey: found actor with matching name and key", {
-			actorId,
+	async getOrCreateWithKey({
+		c,
+		name,
+		key,
+		input,
+	}: GetOrCreateWithKeyInput<{ Bindings: Bindings }>): Promise<ActorOutput> {
+		const env = getCloudflareAmbientEnv();
+
+		// Create a deterministic ID from the actor name and key
+		// This ensures that actors with the same name and key will have the same ID
+		const nameKeyString = serializeNameAndKey(name, key);
+		const doId = env.ACTOR_DO.idFromName(nameKeyString);
+
+		// Get or create actor using the Durable Object's method
+		const actor = env.ACTOR_DO.get(doId);
+		const result = await actor.create({
 			name,
 			key,
+			input,
+			allowExisting: true,
 		});
-		return this.#buildActorOutput(c, actorId);
-	}
+		if ("success" in result) {
+			const { actorId, created } = result.success;
+			logger().debug({
+				msg: "getOrCreateWithKey result",
+				actorId,
+				name,
+				key,
+				created,
+			});
 
-	async getOrCreateWithKey(
-		input: GetOrCreateWithKeyInput,
-	): Promise<ActorOutput> {
-		// TODO: Prevent race condition here
-		const getOutput = await this.getWithKey(input);
-		if (getOutput) {
-			return getOutput;
+			return {
+				actorId,
+				name,
+				key,
+			};
+		} else if ("error" in result) {
+			throw new Error(`Error: ${JSON.stringify(result.error)}`);
 		} else {
-			return await this.createActor(input);
+			assertUnreachable(result);
 		}
 	}
 
@@ -279,62 +375,70 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 	}: CreateInput<{ Bindings: Bindings }>): Promise<ActorOutput> {
 		const env = getCloudflareAmbientEnv();
 
-		// Check if actor with the same name and key already exists
-		const existingActor = await this.getWithKey({ c, name, key });
-		if (existingActor) {
-			throw new ActorAlreadyExists(name, key);
-		}
-
 		// Create a deterministic ID from the actor name and key
 		// This ensures that actors with the same name and key will have the same ID
 		const nameKeyString = serializeNameAndKey(name, key);
 		const doId = env.ACTOR_DO.idFromName(nameKeyString);
-		const actorId = doId.toString();
 
-		// Init actor
+		// Create actor - this will fail if it already exists
 		const actor = env.ACTOR_DO.get(doId);
-		await actor.initialize({
+		const result = await actor.create({
 			name,
 			key,
 			input,
+			allowExisting: false,
 		});
 
-		// Store combined actor metadata (name and key)
-		const actorData: ActorData = { name, key };
-		await env.ACTOR_KV.put(
-			KEYS.ACTOR.metadata(actorId),
-			JSON.stringify(actorData),
-		);
+		if ("success" in result) {
+			const { actorId } = result.success;
+			return {
+				actorId,
+				name,
+				key,
+			};
+		} else if ("error" in result) {
+			if (result.error.actorAlreadyExists) {
+				throw new ActorDuplicateKey(name, key);
+			}
 
-		// Add to key index for lookups by name and key
-		await env.ACTOR_KV.put(KEYS.ACTOR.keyIndex(name, key), actorId);
+			throw new InternalError(
+				`Unknown error creating actor: ${JSON.stringify(result.error)}`,
+			);
+		} else {
+			assertUnreachable(result);
+		}
+	}
 
-		return {
-			actorId,
+	async listActors({ c, name }: ListActorsInput): Promise<ActorOutput[]> {
+		logger().warn({
+			msg: "listActors not fully implemented for Cloudflare Workers",
 			name,
-			key,
+		});
+		return [];
+	}
+
+	displayInformation(): ManagerDisplayInformation {
+		return {
+			properties: {
+				Driver: "Cloudflare Workers",
+			},
 		};
 	}
 
-	// Helper method to build actor output from an ID
-	async #buildActorOutput(
-		c: any,
-		actorId: string,
-	): Promise<ActorOutput | undefined> {
+	setGetUpgradeWebSocket(): void {
+		// No-op for Cloudflare Workers - WebSocket upgrades are handled by the DO
+	}
+
+	async kvGet(actorId: string, key: Uint8Array): Promise<string | null> {
 		const env = getCloudflareAmbientEnv();
 
-		const actorData = (await env.ACTOR_KV.get(KEYS.ACTOR.metadata(actorId), {
-			type: "json",
-		})) as ActorData | null;
+		// Parse actor ID to get DO ID
+		const [doId] = parseActorId(actorId);
 
-		if (!actorData) {
-			return undefined;
-		}
+		const id = env.ACTOR_DO.idFromString(doId);
+		const stub = env.ACTOR_DO.get(id);
 
-		return {
-			actorId,
-			name: actorData.name,
-			key: actorData.key,
-		};
+		const value = await stub.managerKvGet(key);
+		return value !== null ? new TextDecoder().decode(value) : null;
 	}
 }

package/src/mod.ts

@@ -1,3 +1,11 @@
+export type { Client } from "rivetkit";
 export type { DriverContext } from "./actor-driver";
+export { createActorDurableObject } from "./actor-handler-do";
 export type { InputConfig as Config } from "./config";
-export { type Bindings, createServer, createServerHandler } from "./handler";
+export {
+	type Bindings,
+	createHandler,
+	createInlineClient,
+	HandlerOutput,
+	InlineOutput,
+} from "./handler";

package/src/websocket.ts

@@ -4,6 +4,7 @@
 
 import type { UpgradeWebSocket, WSEvents, WSReadyState } from "hono/ws";
 import { defineWebSocketHelper, WSContext } from "hono/ws";
+import { WS_PROTOCOL_STANDARD } from "rivetkit/driver-helpers";
 
 // Based on https://github.com/honojs/hono/issues/1153#issuecomment-1767321332
 export const upgradeWebSocket: UpgradeWebSocket<
@@ -57,14 +58,24 @@ export const upgradeWebSocket: UpgradeWebSocket<
 	// we have to do this after `server.accept() is called`
 	events.onOpen?.(new Event("open"), wsContext);
 
+	// Build response headers
+	const headers: Record<string, string> = {};
+
+	// Set Sec-WebSocket-Protocol if does not exist
+	const protocols = c.req.header("Sec-WebSocket-Protocol");
+	if (
+		typeof protocols === "string" &&
+		protocols
+			.split(",")
+			.map((x) => x.trim())
+			.includes(WS_PROTOCOL_STANDARD)
+	) {
+		headers["Sec-WebSocket-Protocol"] = WS_PROTOCOL_STANDARD;
+	}
+
 	return new Response(null, {
 		status: 101,
-		headers: {
-			// HACK: Required in order for Cloudflare to not error with "Network connection lost"
-			//
-			// This bug undocumented. Cannot easily reproduce outside of RivetKit.
-			"Sec-WebSocket-Protocol": "rivetkit",
-		},
+		headers,
 		webSocket: client,
 	});
 });