fix(core): fix websocket protocol for rivetkit

Author: NathanFlurry

clients/openapi/openapi.json

@@ -24,7 +24,15 @@
             "schema": {
               "type": "string"
             },
-            "required": true,
+            "required": false,
+            "name": "actor_ids",
+            "in": "query"
+          },
+          {
+            "schema": {
+              "type": "string"
+            },
+            "required": false,
             "name": "key",
             "in": "query"
           }

packages/cloudflare-workers/src/websocket.ts

@@ -65,7 +65,10 @@ export const upgradeWebSocket: UpgradeWebSocket<
 	const protocols = c.req.header("Sec-WebSocket-Protocol");
 	if (
 		typeof protocols === "string" &&
-		protocols.includes(WS_PROTOCOL_STANDARD)
+		protocols
+			.split(",")
+			.map((x) => x.trim())
+			.includes(WS_PROTOCOL_STANDARD)
 	) {
 		headers["Sec-WebSocket-Protocol"] = WS_PROTOCOL_STANDARD;
 	}

packages/rivetkit/package.json

@@ -163,7 +163,7 @@
     "@bare-ts/lib": "~0.3.0",
     "@hono/standard-validator": "^0.1.3",
     "@hono/zod-openapi": "^0.19.10",
-    "@rivetkit/engine-runner": "https://pkg.pr.new/rivet-dev/engine/@rivetkit/engine-runner@b72b232",
+    "@rivetkit/engine-runner": "https://pkg.pr.new/rivet-dev/engine/@rivetkit/engine-runner@a0639d0",
     "@rivetkit/fast-json-patch": "^3.1.2",
     "cbor-x": "^1.6.0",
     "hono": "^4.7.0",

packages/rivetkit/src/actor/router.ts

@@ -23,6 +23,9 @@ import {
 	HEADER_ENCODING,
 	PATH_CONNECT_WEBSOCKET,
 	PATH_RAW_WEBSOCKET_PREFIX,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_TOKEN,
 } from "@/common/actor-router-consts";
 import {
 	handleRouteError,
@@ -81,8 +84,23 @@ export function createActorRouter(
 		const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
 		if (upgradeWebSocket) {
 			return upgradeWebSocket(async (c) => {
-				const encodingRaw = c.req.header(HEADER_ENCODING);
-				const connParamsRaw = c.req.header(HEADER_CONN_PARAMS);
+				// Parse configuration from Sec-WebSocket-Protocol header
+				const protocols = c.req.header("sec-websocket-protocol");
+				let encodingRaw: string | undefined;
+				let connParamsRaw: string | undefined;
+
+				if (protocols) {
+					const protocolList = protocols.split(",").map((p) => p.trim());
+					for (const protocol of protocolList) {
+						if (protocol.startsWith(WS_PROTOCOL_ENCODING)) {
+							encodingRaw = protocol.substring(WS_PROTOCOL_ENCODING.length);
+						} else if (protocol.startsWith(WS_PROTOCOL_CONN_PARAMS)) {
+							connParamsRaw = decodeURIComponent(
+								protocol.substring(WS_PROTOCOL_CONN_PARAMS.length),
+							);
+						}
+					}
+				}
 
 				const encoding = EncodingSchema.parse(encodingRaw);
 				const connParams = connParamsRaw
@@ -172,8 +190,26 @@ export function createActorRouter(
 		const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
 		if (upgradeWebSocket) {
 			return upgradeWebSocket(async (c) => {
-				const encodingRaw = c.req.header(HEADER_ENCODING);
-				const connParamsRaw = c.req.header(HEADER_CONN_PARAMS);
+				// Parse configuration from Sec-WebSocket-Protocol header
+				const protocols = c.req.header("sec-websocket-protocol");
+				let encodingRaw: string | undefined;
+				let connParamsRaw: string | undefined;
+				let token: string | undefined;
+
+				if (protocols) {
+					const protocolList = protocols.split(",").map((p) => p.trim());
+					for (const protocol of protocolList) {
+						if (protocol.startsWith(WS_PROTOCOL_ENCODING)) {
+							encodingRaw = protocol.substring(WS_PROTOCOL_ENCODING.length);
+						} else if (protocol.startsWith(WS_PROTOCOL_CONN_PARAMS)) {
+							connParamsRaw = decodeURIComponent(
+								protocol.substring(WS_PROTOCOL_CONN_PARAMS.length),
+							);
+						} else if (protocol.startsWith(WS_PROTOCOL_TOKEN)) {
+							token = protocol.substring(WS_PROTOCOL_TOKEN.length);
+						}
+					}
+				}
 
 				const encoding = EncodingSchema.parse(encodingRaw);
 				const connParams = connParamsRaw

packages/rivetkit/src/drivers/engine/actor-driver.ts

@@ -22,15 +22,16 @@ import type { Client } from "@/client/client";
 import {
 	PATH_CONNECT_WEBSOCKET,
 	PATH_RAW_WEBSOCKET_PREFIX,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_TOKEN,
 } from "@/common/actor-router-consts";
 import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
 import { getLogger } from "@/common/log";
 import type { UniversalWebSocket } from "@/common/websocket-interface";
 import {
 	type ActorDriver,
 	type AnyActorInstance,
-	HEADER_CONN_PARAMS,
-	HEADER_ENCODING,
 	type ManagerDriver,
 	serializeEmptyPersistData,
 } from "@/driver-helpers/mod";
@@ -293,9 +294,26 @@ export class EngineActorDriver implements ActorDriver {
 
 		const url = new URL(request.url);
 
-		// Parse headers
-		const encodingRaw = request.headers.get(HEADER_ENCODING);
-		const connParamsRaw = request.headers.get(HEADER_CONN_PARAMS);
+		// Parse configuration from Sec-WebSocket-Protocol header
+		const protocols = request.headers.get("sec-websocket-protocol");
+		let encodingRaw: string | undefined;
+		let connParamsRaw: string | undefined;
+		let token: string | undefined;
+
+		if (protocols) {
+			const protocolList = protocols.split(",").map((p) => p.trim());
+			for (const protocol of protocolList) {
+				if (protocol.startsWith(WS_PROTOCOL_ENCODING)) {
+					encodingRaw = protocol.substring(WS_PROTOCOL_ENCODING.length);
+				} else if (protocol.startsWith(WS_PROTOCOL_CONN_PARAMS)) {
+					connParamsRaw = decodeURIComponent(
+						protocol.substring(WS_PROTOCOL_CONN_PARAMS.length),
+					);
+				} else if (protocol.startsWith(WS_PROTOCOL_TOKEN)) {
+					token = protocol.substring(WS_PROTOCOL_TOKEN.length);
+				}
+			}
+		}
 
 		const encoding = EncodingSchema.parse(encodingRaw);
 		const connParams = connParamsRaw ? JSON.parse(connParamsRaw) : undefined;