diff --git a/src/core/handlers/WebSocketHandler.ts b/src/core/handlers/WebSocketHandler.ts index 7e18d9c56..580a1e03f 100644 --- a/src/core/handlers/WebSocketHandler.ts +++ b/src/core/handlers/WebSocketHandler.ts @@ -29,6 +29,10 @@ export interface WebSocketHandlerConnection { params: PathParams } +export type WebSocketHandlerConnectionTransformer = ( + connection: WebSocketHandlerConnection, +) => WebSocketHandlerConnection + export interface WebSocketResolutionContext { baseUrl?: string [kAutoConnect]?: boolean @@ -49,7 +53,10 @@ export class WebSocketHandler { protected [kEmitter]: Emitter - constructor(protected readonly url: Path) { + constructor( + protected readonly url: Path, + private readonly transform?: WebSocketHandlerConnectionTransformer, + ) { this.id = createRequestId() this[kEmitter] = new Emitter() @@ -113,10 +120,15 @@ export class WebSocketHandler { return null } - const resolvedConnection: WebSocketHandlerConnection = { - ...connection, - params: parsedResult.match.params || {}, - } + const resolvedConnection = this.transform + ? this.transform({ + ...connection, + params: parsedResult.match.params || {}, + }) + : ({ + ...connection, + params: parsedResult.match.params || {}, + } as WebSocketHandlerConnection) if (resolutionContext?.[kAutoConnect] ?? true) { if (this[kConnect](resolvedConnection)) { diff --git a/src/core/ws.test.ts b/src/core/ws.test.ts index 0f4eff759..fdfe23444 100644 --- a/src/core/ws.test.ts +++ b/src/core/ws.test.ts @@ -1,5 +1,6 @@ // @vitest-environment node-websocket import { ws } from './ws' +import { createTestWebSocketConnection } from '../../test/support/ws-test-utils' it('exports the "link()" method', () => { expect(ws).toHaveProperty('link') @@ -19,3 +20,39 @@ it('throws an error when given a non-path argument to "ws.link()"', () => { ws.link(2), ).toThrow('Expected a WebSocket server URL to be a valid path but got number') }) + +it('supports transforming the connection in the link()', async () => { + const service = ws.link('ws://localhost', { + transform(connection) { + return { + ...connection, + params: { + roomId: 'room-1', + }, + } + }, + }) + const listener = vi.fn() + + const handler = service.addEventListener('connection', listener) + + const result = await handler.run( + createTestWebSocketConnection('ws://localhost'), + ) + + expect(result).toEqual( + expect.objectContaining({ + params: { + roomId: 'room-1', + }, + }), + ) + expect(listener).toHaveBeenCalledOnce() + expect(listener).toHaveBeenCalledWith( + expect.objectContaining({ + params: { + roomId: 'room-1', + }, + }), + ) +}) diff --git a/src/core/ws.ts b/src/core/ws.ts index 0c99ee823..e09e2c66a 100644 --- a/src/core/ws.ts +++ b/src/core/ws.ts @@ -8,6 +8,7 @@ import { WebSocketHandler, kEmitter, type WebSocketHandlerEventMap, + type WebSocketHandlerConnectionTransformer, } from './handlers/WebSocketHandler' import { hasRefCounted } from './utils/internal/hasRefCounted' import { @@ -32,6 +33,10 @@ export type WebSocketEventListener< EventType extends keyof WebSocketHandlerEventMap, > = (...args: WebSocketHandlerEventMap[EventType]) => void +export interface WebSocketLinkOptions { + transform?: WebSocketHandlerConnectionTransformer +} + export type WebSocketLink = { /** * A set of all WebSocket clients connected @@ -97,7 +102,10 @@ export type WebSocketLink = { * client.send('hello from server!') * }) */ -function createWebSocketLinkHandler(url: Path): WebSocketLink { +function createWebSocketLinkHandler( + url: Path, + options: WebSocketLinkOptions = {}, +): WebSocketLink { invariant(url, 'Expected a WebSocket server URL but got undefined') invariant( @@ -124,7 +132,7 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink { return clientManager.clients }, addEventListener(event, listener) { - const webSocketHandler = new WebSocketHandler(url) + const webSocketHandler = new WebSocketHandler(url, options.transform) // Add the connection event listener for when the // handler matches and emits a connection event. diff --git a/test/typings/ws.test-d.ts b/test/typings/ws.test-d.ts index a72fd14f5..a5f76269e 100644 --- a/test/typings/ws.test-d.ts +++ b/test/typings/ws.test-d.ts @@ -34,6 +34,20 @@ it('supports "connection" event listener', () => { }) }) +it('supports "connection" transform option', () => { + const link = ws.link('ws://localhost', { + transform(connection) { + return { + ...connection, + } + }, + }) + + link.addEventListener('connection', (connection) => { + expectTypeOf(connection).toEqualTypeOf() + }) +}) + it('errors on arbitrary event names passed to the link', () => { const link = ws.link('ws://localhost')