diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.test.ts b/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.test.ts new file mode 100644 index 0000000000..bffe444461 --- /dev/null +++ b/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); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts b/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts index fec29c5dd4..bcecbca5fb 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/router-websocket-endpoints.ts +++ b/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(); + // 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"; } diff --git a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/raw-websocket.ts b/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/raw-websocket.ts index 04a624e2ee..e92c5044a0 100644 --- a/rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/raw-websocket.ts +++ b/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((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((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(); + }); }); }