@rivetkit/cloudflare-workers 2.0.6 → 2.0.7

package/src/actor-driver.ts

@@ -4,17 +4,14 @@ import type {
 	RegistryConfig,
 	RunConfig,
 } from "rivetkit";
-import {
-	createGenericConnDrivers,
-	GenericConnGlobalState,
-	lookupInRegistry,
-} from "rivetkit";
+import { lookupInRegistry } from "rivetkit";
 import type { Client } from "rivetkit/client";
 import type {
 	ActorDriver,
 	AnyActorInstance,
 	ManagerDriver,
 } from "rivetkit/driver-helpers";
+import { promiseWithResolvers } from "rivetkit/utils";
 import { KEYS } from "./actor-handler-do";
 
 interface DurableObjectGlobalState {
@@ -49,8 +46,8 @@ export interface DriverContext {
 // Actor handler to track running instances
 class ActorHandler {
 	actor?: AnyActorInstance;
-	actorPromise?: PromiseWithResolvers<void> = Promise.withResolvers();
-	genericConnGlobalState = new GenericConnGlobalState();
+	actorPromise?: ReturnType<typeof promiseWithResolvers<void>> =
+		promiseWithResolvers();
 }
 
 export class CloudflareActorsActorDriver implements ActorDriver {
@@ -114,11 +111,7 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		handler.actor = definition.instantiate();
 
 		// Start actor
-		const connDrivers = createGenericConnDrivers(
-			handler.genericConnGlobalState,
-		);
 		await handler.actor.start(
-			connDrivers,
 			this,
 			this.#inlineClient,
 			actorId,
@@ -134,14 +127,6 @@ export class CloudflareActorsActorDriver implements ActorDriver {
 		return handler.actor;
 	}
 
-	getGenericConnGlobalState(actorId: string): GenericConnGlobalState {
-		const handler = this.#actors.get(actorId);
-		if (!handler) {
-			throw new Error(`Actor ${actorId} not loaded`);
-		}
-		return handler.genericConnGlobalState;
-	}
-
 	getContext(actorId: string): DriverContext {
 		const state = this.#globalState.getDOState(actorId);
 		return { state: state.ctx };

package/src/actor-handler-do.ts

@@ -3,7 +3,12 @@ import type { ExecutionContext } from "hono";
 import invariant from "invariant";
 import type { ActorKey, ActorRouter, Registry, RunConfig } from "rivetkit";
 import { createActorRouter, createClientWithDriver } from "rivetkit";
-import { serializeEmptyPersistData } from "rivetkit/driver-helpers";
+import type { ActorDriver } from "rivetkit/driver-helpers";
+import {
+	type ManagerDriver,
+	serializeEmptyPersistData,
+} from "rivetkit/driver-helpers";
+import { promiseWithResolvers } from "rivetkit/utils";
 import {
 	CloudflareDurableObjectGlobalState,
 	createCloudflareActorsActorDriverBuilder,
@@ -38,6 +43,7 @@ export type DurableObjectConstructor = new (
 
 interface LoadedActor {
 	actorRouter: ActorRouter;
+	actorDriver: ActorDriver;
 }
 
 export function createActorDurableObject(
@@ -60,7 +66,7 @@ export function createActorDurableObject(
 		implements ActorHandlerInterface
 	{
 		#initialized?: InitializedData;
-		#initializedPromise?: PromiseWithResolvers<void>;
+		#initializedPromise?: ReturnType<typeof promiseWithResolvers<void>>;
 
 		#actor?: LoadedActor;
 
@@ -71,7 +77,7 @@ export function createActorDurableObject(
 				if (this.#initializedPromise) {
 					await this.#initializedPromise.promise;
 				} else {
-					this.#initializedPromise = Promise.withResolvers();
+					this.#initializedPromise = promiseWithResolvers();
 					const res = await this.ctx.storage.get([
 						KEYS.NAME,
 						KEYS.KEY,
@@ -118,6 +124,8 @@ export function createActorDurableObject(
 				runConfig,
 			);
 
+			configureInspectorAccessToken(registry.config, managerDriver);
+
 			// Create inline client
 			const inlineClient = createClientWithDriver(managerDriver);
 
@@ -135,6 +143,7 @@ export function createActorDurableObject(
 			// Save actor
 			this.#actor = {
 				actorRouter,
+				actorDriver,
 			};
 
 			// Initialize the actor instance with proper metadata
@@ -174,26 +183,18 @@ export function createActorDurableObject(
 		}
 
 		async alarm(): Promise<void> {
-			await this.#loadActor();
+			const { actorDriver } = await this.#loadActor();
 			const actorId = this.ctx.id.toString();
 
-			// Get the actor driver
-			invariant(runConfig.driver, "runConfig.driver");
-			const managerDriver = runConfig.driver.manager(
-				registry.config,
-				runConfig,
-			);
-			const inlineClient = createClientWithDriver(managerDriver);
-			const actorDriver = runConfig.driver.actor(
-				registry.config,
-				runConfig,
-				managerDriver,
-				inlineClient,
-			);
-
 			// Load the actor instance and trigger alarm
 			const actor = await actorDriver.loadActor(actorId);
 			await actor._onAlarm();
 		}
 	};
 }
+function configureInspectorAccessToken(
+	config: any,
+	managerDriver: ManagerDriver,
+) {
+	throw new Error("Function not implemented.");
+}

package/src/manager-driver.ts

@@ -6,11 +6,15 @@ import {
 	type GetForIdInput,
 	type GetOrCreateWithKeyInput,
 	type GetWithKeyInput,
-	HEADER_AUTH_DATA,
-	HEADER_CONN_PARAMS,
-	HEADER_ENCODING,
 	type ManagerDisplayInformation,
 	type ManagerDriver,
+	WS_PROTOCOL_ACTOR,
+	WS_PROTOCOL_CONN_ID,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_CONN_TOKEN,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_STANDARD,
+	WS_PROTOCOL_TARGET,
 } from "rivetkit/driver-helpers";
 import { ActorAlreadyExists, InternalError } from "rivetkit/errors";
 import { getCloudflareAmbientEnv } from "./handler";
@@ -68,6 +72,8 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
+		connId?: string,
+		connToken?: string,
 	): Promise<UniversalWebSocket> {
 		const env = getCloudflareAmbientEnv();
 
@@ -81,19 +87,32 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		const id = env.ACTOR_DO.idFromString(actorId);
 		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}${actorId}`);
+		protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
+		if (params) {
+			protocols.push(
+				`${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`,
+			);
+		}
+		if (connId) {
+			protocols.push(`${WS_PROTOCOL_CONN_ID}${connId}`);
+		}
+		if (connToken) {
+			protocols.push(`${WS_PROTOCOL_CONN_TOKEN}${connToken}`);
+		}
+
 		const headers: Record<string, string> = {
 			Upgrade: "websocket",
 			Connection: "Upgrade",
-			[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({ msg: "rewriting websocket url", from: path, to: url });
 
@@ -104,7 +123,7 @@ 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()}`,
 			);
 		}
 
@@ -151,7 +170,6 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 		actorId: string,
 		encoding: Encoding,
 		params: unknown,
-		authData: unknown,
 	): Promise<Response> {
 		logger().debug({
 			msg: "forwarding websocket to durable object",
@@ -187,14 +205,18 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
 			}
 		}
 
-		// Add RivetKit headers
-		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}${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);

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<
@@ -62,8 +63,14 @@ export const upgradeWebSocket: UpgradeWebSocket<
 
 	// Set Sec-WebSocket-Protocol if does not exist
 	const protocols = c.req.header("Sec-WebSocket-Protocol");
-	if (typeof protocols === "string" && protocols.includes("rivetkit")) {
-		headers["Sec-WebSocket-Protocol"] = "rivetkit";
+	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, {