fix(core): fix websocket protocol for rivetkit

Author: NathanFlurry

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,10 +163,10 @@
     "@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@9676fe4",
+    "@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.9.8",
+    "hono": "^4.7.0",
     "invariant": "^2.2.4",
     "nanoevents": "^9.1.0",
     "on-change": "^5.0.1",

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

@@ -24,15 +24,14 @@ import {
 	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";
@@ -309,18 +308,24 @@ export class EngineActorDriver implements ActorDriver {
 		if (protocols === null)
 			throw new Error(`Missing sec-websocket-protocol header`);
 
+		// Parse configuration from Sec-WebSocket-Protocol header
+		const protocols = request.headers.get("sec-websocket-protocol");
 		let encodingRaw: string | undefined;
 		let connParamsRaw: string | undefined;
-
-		// Parse 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),
-				);
+		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);
+				}
 			}
 		}