feat: support ws.link connection transform

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<WebSocketHandlerEventMap>
 
-  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)) {

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',
+      },
+    }),
+  )
+})

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.

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<WebSocketHandlerConnection>()
+  })
+})
+
 it('errors on arbitrary event names passed to the link', () => {
   const link = ws.link('ws://localhost')