fix(rivetkit): fix routing of /websocket

rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.test.ts

@@ -0,0 +1,54 @@
+import { describe, expect, test } from "vitest";
+import {
+	PATH_WEBSOCKET_BASE,
+	PATH_WEBSOCKET_PREFIX,
+} from "@/common/actor-router-consts";
+
+/**
+ * Unit tests for WebSocket path routing logic.
+ *
+ * These tests verify the path matching behavior in routeWebSocket
+ * without needing a full actor setup.
+ *
+ * NOTE: The driver-file-system end-to-end tests pass because the driver
+ * correctly strips query parameters before calling routeWebSocket
+ * (see FileSystemManagerDriver.openWebSocket). However, the bug still
+ * exists in routeWebSocket itself and could be triggered by other callers
+ * (e.g., engine driver's runnerWebSocket which passes requestPath directly).
+ */
+describe("websocket path routing", () => {
+	// Helper that replicates the routing logic from routeWebSocket
+	// After fix: strips query params before comparing
+	function matchesWebSocketPath(requestPath: string): boolean {
+		const requestPathWithoutQuery = requestPath.split("?")[0];
+		return (
+			requestPathWithoutQuery === PATH_WEBSOCKET_BASE ||
+			requestPathWithoutQuery.startsWith(PATH_WEBSOCKET_PREFIX)
+		);
+	}
+
+	test("should match base websocket path without query", () => {
+		expect(matchesWebSocketPath("/websocket")).toBe(true);
+	});
+
+	test("should match websocket path with trailing slash", () => {
+		expect(matchesWebSocketPath("/websocket/")).toBe(true);
+	});
+
+	test("should match websocket path with subpath", () => {
+		expect(matchesWebSocketPath("/websocket/foo")).toBe(true);
+		expect(matchesWebSocketPath("/websocket/foo/bar")).toBe(true);
+	});
+
+	test("should match websocket path with subpath and query", () => {
+		// This works because "/websocket/foo?query" starts with "/websocket/"
+		expect(matchesWebSocketPath("/websocket/foo?query=value")).toBe(true);
+	});
+
+	// FIX: Query parameters are now stripped before routing comparison.
+	// This ensures /websocket?query correctly routes to the websocket handler.
+	test("should match base websocket path with query parameters", () => {
+		expect(matchesWebSocketPath("/websocket?token=abc")).toBe(true);
+		expect(matchesWebSocketPath("/websocket?foo=bar&baz=123")).toBe(true);
+	});
+});

rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts

@@ -88,10 +88,15 @@ export async function routeWebSocket(
 		// Promise used to wait for the websocket close in `disconnect`
 		const closePromiseResolvers = promiseWithResolvers<void>();
 
+		// Strip query parameters from requestPath for routing purposes.
+		// This handles paths like "/websocket?query=value" which should route
+		// to the raw websocket handler.
+		const requestPathWithoutQuery = requestPath.split("?")[0];
+
 		// Route WebSocket & create driver
 		let handler: WebSocketHandler;
 		let connDriver: ConnDriver;
-		if (requestPath === PATH_CONNECT) {
+		if (requestPathWithoutQuery === PATH_CONNECT) {
 			const { driver, setWebSocket } = createWebSocketDriver(
 				isHibernatable
 					? { gatewayId: gatewayId!, requestId: requestId! }
@@ -103,8 +108,8 @@ export async function routeWebSocket(
 			handler = handleWebSocketConnect.bind(undefined, setWebSocket);
 			connDriver = driver;
 		} else if (
-			requestPath === PATH_WEBSOCKET_BASE ||
-			requestPath.startsWith(PATH_WEBSOCKET_PREFIX)
+			requestPathWithoutQuery === PATH_WEBSOCKET_BASE ||
+			requestPathWithoutQuery.startsWith(PATH_WEBSOCKET_PREFIX)
 		) {
 			const { driver, setWebSocket } = createRawWebSocketDriver(
 				isHibernatable
@@ -114,7 +119,7 @@ export async function routeWebSocket(
 			);
 			handler = handleRawWebSocket.bind(undefined, setWebSocket);
 			connDriver = driver;
-		} else if (requestPath === PATH_INSPECTOR_CONNECT) {
+		} else if (requestPathWithoutQuery === PATH_INSPECTOR_CONNECT) {
 			if (!actor.inspectorToken) {
 				throw "WebSocket Inspector Unauthorized: actor does not provide inspector access";
 			}

rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/raw-websocket.ts

@@ -469,5 +469,45 @@ export function runRawWebSocketTests(driverTestConfig: DriverTestConfig) {
 
 			ws.close();
 		});
+
+		test("should handle query parameters on base websocket path (no subpath)", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const actor = client.rawWebSocketActor.getOrCreate([
+				"base-path-query-params",
+			]);
+
+			// Test WebSocket with ONLY query parameters on the base path
+			// This tests the case where path is "/websocket?foo=bar" without trailing slash
+			const ws = await actor.websocket("?token=secret&session=123");
+
+			await new Promise<void>((resolve, reject) => {
+				ws.addEventListener("open", () => resolve(), { once: true });
+				ws.addEventListener("error", reject);
+				ws.addEventListener("close", (evt: any) => {
+					reject(new Error(`WebSocket closed: code=${evt.code} reason=${evt.reason}`));
+				});
+			});
+
+			const requestInfoPromise = new Promise<any>((resolve, reject) => {
+				ws.addEventListener("message", (event: any) => {
+					const data = JSON.parse(event.data as string);
+					if (data.type === "requestInfo") {
+						resolve(data);
+					}
+				});
+				ws.addEventListener("close", reject);
+			});
+
+			// Send request to get the request info
+			ws.send(JSON.stringify({ type: "getRequestInfo" }));
+
+			const requestInfo = await requestInfoPromise;
+
+			// Verify query parameters were preserved even on base websocket path
+			expect(requestInfo.url).toContain("token=secret");
+			expect(requestInfo.url).toContain("session=123");
+
+			ws.close();
+		});
 	});
 }