From 69c7ad8855604cd0a15e11cb08c12b820cf3c52c Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Mon, 13 Apr 2026 13:06:54 +0200 Subject: [PATCH] fix(WebSocketFrame): allow for broader connection types --- .../experimental/frames/websocket-frame.ts | 16 ++-- .../sources/interceptor-source.ts | 2 +- src/core/handlers/WebSocketHandler.ts | 11 ++- src/core/ws/utils/attachWebSocketLogger.ts | 88 +++++++++++-------- 4 files changed, 70 insertions(+), 47 deletions(-) diff --git a/src/core/experimental/frames/websocket-frame.ts b/src/core/experimental/frames/websocket-frame.ts index 73fb10b08..34dea24a3 100644 --- a/src/core/experimental/frames/websocket-frame.ts +++ b/src/core/experimental/frames/websocket-frame.ts @@ -1,9 +1,9 @@ import { TypedEvent } from 'rettime' -import { type WebSocketConnectionData } from '@mswjs/interceptors/WebSocket' import { kConnect, kAutoConnect, type WebSocketHandler, + type WebSocketConnectionFoo, } from '../../handlers/WebSocketHandler' import { NetworkFrame, @@ -16,8 +16,10 @@ import { import { devUtils } from '../../utils/internal/devUtils' import { HandlersController, AnyHandler } from '../handlers-controller' -export interface WebSocketNetworkFrameOptions { - connection: WebSocketConnectionData +export interface WebSocketNetworkFrameOptions< + Connection extends WebSocketConnectionFoo, +> { + connection: Connection } export type WebSocketNetworkFrameEventMap = { @@ -68,14 +70,16 @@ class UnhandledWebSocketExceptionEvent< } } -export abstract class WebSocketNetworkFrame extends NetworkFrame< +export abstract class WebSocketNetworkFrame< + Connection extends WebSocketConnectionFoo, +> extends NetworkFrame< 'ws', { - connection: WebSocketConnectionData + connection: Connection }, WebSocketNetworkFrameEventMap > { - constructor(options: WebSocketNetworkFrameOptions) { + constructor(options: WebSocketNetworkFrameOptions) { super('ws', { connection: options.connection, }) diff --git a/src/core/experimental/sources/interceptor-source.ts b/src/core/experimental/sources/interceptor-source.ts index edbb2d45d..784c25d3d 100644 --- a/src/core/experimental/sources/interceptor-source.ts +++ b/src/core/experimental/sources/interceptor-source.ts @@ -154,7 +154,7 @@ class InterceptorHttpNetworkFrame extends HttpNetworkFrame { } } -class InterceptorWebSocketNetworkFrame extends WebSocketNetworkFrame { +class InterceptorWebSocketNetworkFrame extends WebSocketNetworkFrame { constructor(args: { connection: WebSocketConnectionData }) { super({ connection: args.connection }) } diff --git a/src/core/handlers/WebSocketHandler.ts b/src/core/handlers/WebSocketHandler.ts index d0247b9e3..6ce9317d1 100644 --- a/src/core/handlers/WebSocketHandler.ts +++ b/src/core/handlers/WebSocketHandler.ts @@ -22,10 +22,13 @@ export type WebSocketHandlerEventMap = { connection: [args: WebSocketHandlerConnection] } -export interface WebSocketHandlerConnection { +export interface WebSocketConnectionFoo { client: WebSocketClientConnectionProtocol server: WebSocketServerConnectionProtocol info: WebSocketConnectionData['info'] +} + +export interface WebSocketHandlerConnection extends WebSocketConnectionFoo { params: PathParams } @@ -97,7 +100,7 @@ export class WebSocketHandler { } public async run( - connection: WebSocketConnectionData, + connection: WebSocketConnectionFoo, resolutionContext?: WebSocketResolutionContext, ): Promise { const parsedResult = this.parse({ @@ -160,7 +163,9 @@ export class WebSocketHandler { return this[kEmitter].emit('connection', connection) } - public log(connection: WebSocketConnectionData): () => void { + public log( + connection: WebSocketConnectionData | WebSocketConnectionFoo, + ): () => void { return attachWebSocketLogger(connection) } diff --git a/src/core/ws/utils/attachWebSocketLogger.ts b/src/core/ws/utils/attachWebSocketLogger.ts index 0afba3298..ba6afc5fc 100644 --- a/src/core/ws/utils/attachWebSocketLogger.ts +++ b/src/core/ws/utils/attachWebSocketLogger.ts @@ -1,13 +1,15 @@ import type { - WebSocketClientConnection, - WebSocketConnectionData, WebSocketData, + WebSocketConnectionData, + WebSocketClientConnection, + WebSocketClientConnectionProtocol, } from '@mswjs/interceptors/WebSocket' import { devUtils } from '../../utils/internal/devUtils' import { getTimestamp } from '../../utils/logging/getTimestamp' import { toPublicUrl } from '../../utils/request/toPublicUrl' import { getMessageLength } from './getMessageLength' import { getPublicData } from './getPublicData' +import type { WebSocketConnectionFoo } from '../../handlers/WebSocketHandler' export const colors = { system: '#3b82f6', @@ -17,7 +19,7 @@ export const colors = { } export function attachWebSocketLogger( - connection: WebSocketConnectionData, + connection: WebSocketConnectionData | WebSocketConnectionFoo, ): () => void { const { client, server } = connection const controller = new AbortController() @@ -48,31 +50,36 @@ export function attachWebSocketLogger( ) // Log client errors (connection closures due to errors). - client.socket.addEventListener( - 'error', - (event) => { - logClientError(event) - }, - { signal: controller.signal }, - ) + if ('socket' in client) { + client.socket.addEventListener( + 'error', + (event) => { + logClientError(event) + }, + { signal: controller.signal }, + ) + } const { send: originalClientSend } = client client.send = new Proxy(client.send, { apply(target, thisArg, args) { const [data] = args const messageEvent = new MessageEvent('message', { data }) - Object.defineProperties(messageEvent, { - currentTarget: { - enumerable: true, - writable: false, - value: client.socket, - }, - target: { - enumerable: true, - writable: false, - value: client.socket, - }, - }) + + if ('socket' in client) { + Object.defineProperties(messageEvent, { + currentTarget: { + enumerable: true, + writable: false, + value: client.socket, + }, + target: { + enumerable: true, + writable: false, + value: client.socket, + }, + }) + } queueMicrotask(() => { logIncomingMockedClientMessage(messageEvent) @@ -102,18 +109,21 @@ export function attachWebSocketLogger( apply(target, thisArg, args) { const [data] = args const messageEvent = new MessageEvent('message', { data }) - Object.defineProperties(messageEvent, { - currentTarget: { - enumerable: true, - writable: false, - value: server.socket, - }, - target: { - enumerable: true, - writable: false, - value: server.socket, - }, - }) + + if ('socket' in server) { + Object.defineProperties(messageEvent, { + currentTarget: { + enumerable: true, + writable: false, + value: server.socket, + }, + target: { + enumerable: true, + writable: false, + value: server.socket, + }, + }) + } logOutgoingMockedClientMessage(messageEvent) @@ -142,7 +152,9 @@ export function attachWebSocketLogger( * that intercepted this connection. This helps you see * what handlers observe this connection. */ -export function logConnectionOpen(client: WebSocketClientConnection) { +export function logConnectionOpen( + client: WebSocketClientConnectionProtocol | WebSocketClientConnection, +) { const publicUrl = toPublicUrl(client.url) console.groupCollapsed( @@ -150,8 +162,10 @@ export function logConnectionOpen(client: WebSocketClientConnection) { `color:${colors.system}`, 'color:inherit', ) - // eslint-disable-next-line no-console - console.log('Client:', client.socket) + if ('socket' in client) { + // eslint-disable-next-line no-console + console.log('Client:', client.socket) + } console.groupEnd() }