Add websocket pong handler

Author: mxiao-cll

docs/components/transport-types/websocket-transport.md

@@ -158,6 +158,23 @@ The heartbeat interval is controlled by the `WS_HEARTBEAT_INTERVAL_MS` adapter s
 
 **Note:** The heartbeat only starts if the `heartbeat` handler is provided. If you don't need heartbeat functionality, simply omit the `heartbeat` handler.
 
+### Pong handler (optional)
+
+When using WebSocket protocol-level pings (via `connection.ping()`), the server will respond with a `pong` frame. The `pong` handler lets you react to those responses — for example, to track latency, log liveness confirmations, or reset a timeout.
+
+```typescript
+handlers: {
+  heartbeat: (connection) => {
+    connection.ping()
+  },
+  pong: (connection, data) => {
+    logger.debug('Received pong from server', { data: data.toString() })
+  },
+}
+```
+
+The `pong` handler receives the WebSocket connection and a `Buffer` containing any data the server included in the pong frame. It is optional — if omitted, pong frames are silently ignored.
+
 ### Retrieving and storing the response or errors
 
 As shown in the first example, the **handlers** object accepts a function called **message** which will be executed when Data Provider sends a message through the WS connection. It takes this message as its first argument and the adapter context as the second, and should build and return a list of response objects (_ProviderResult_) that will be stored in the response cache for the endpoint.

src/transports/websocket.ts

@@ -82,6 +82,15 @@ export interface WebSocketTransportConfig<T extends WebsocketTransportGenerics>
      */
     heartbeat?: (wsConnection: WebSocket, context: EndpointContext<T>) => Promise<void> | void
 
+    /**
+     * Handles when the WS receives a pong
+     *
+     * @param wsConnection - the WebSocket with an established connection
+     * @param data - the data received by the WS
+     * @returns void
+     */
+    pong?: (wsConnection: WebSocket, data: Buffer) => void
+
     /**
      * Handles when the websocket connection dispatches an error event
      * Optional to let the adapter handle the event in its own way if it decides to
@@ -371,6 +380,9 @@ export class WebSocketTransport<
     )
     connection.addEventListener('error', handlers.error)
     connection.addEventListener('close', handlers.close)
+    if (this.config.handlers.pong) {
+      connection.on('pong', (data) => this.config.handlers.pong?.(connection, data))
+    }
 
     return connection
   }

src/util/testing-utils.ts

@@ -491,9 +491,13 @@ export function setEnvVariables(envVariables: NodeJS.ProcessEnv): void {
 export const mockWebSocketProvider = (provider: typeof WebSocketClassProvider): void => {
   // Extend mock WebSocket class to bypass protocol headers error
   class MockWebSocket extends WebSocket {
+    // Separate listener map for ws-style .on()/.emit() (e.g. 'pong'), which mock-socket doesn't support
+    private wsListeners: Map<string, ((...args: unknown[]) => void)[]> = new Map()
+
     constructor(url: string, protocol: string | string[] | Record<string, string> | undefined) {
       super(url, protocol instanceof Object ? undefined : protocol)
     }
+
     // This is part of the 'ws' node library but not the common interface, but it's used in our WS transport
     removeAllListeners() {
       for (const eventType in this.listeners) {
@@ -502,6 +506,22 @@ export const mockWebSocketProvider = (provider: typeof WebSocketClassProvider):
           delete this.listeners[eventType]
         }
       }
+      this.wsListeners.clear()
+    }
+
+    // Ws-style EventEmitter API used for protocol-level events like 'pong'
+    on(event: string, listener: (...args: unknown[]) => void) {
+      const existing = this.wsListeners.get(event) ?? []
+      this.wsListeners.set(event, [...existing, listener])
+      return this
+    }
+
+    emit(event: string, ...args: unknown[]): boolean {
+      const listeners = this.wsListeners.get(event) ?? []
+      for (const l of listeners) {
+        l(...args)
+      }
+      return listeners.length > 0
     }
   }
 

test/transports/websocket.test.ts

@@ -64,6 +64,7 @@ const createAdapter = (
     connection: WebSocket,
     context: EndpointContext<WebSocketTypes>,
   ) => Promise<void> | void,
+  pongHandler?: (connection: WebSocket, data: Buffer) => void,
 ): Adapter => {
   const websocketTransport = new WebSocketTransport<WebSocketTypes>({
     url: () => ENDPOINT_URL,
@@ -96,6 +97,7 @@ const createAdapter = (
         ]
       },
       heartbeat: heartbeatHandler,
+      pong: pongHandler,
     },
     builders: {
       subscribeMessage: (params) => `S:${params.base}/${params.quote}`,
@@ -1313,6 +1315,93 @@ test.serial('does not heartbeat when handler throws an error', async (t) => {
   await t.context.clock.runToLastAsync()
 })
 
+test.serial('calls pong handler when a pong frame is received', async (t) => {
+  const base = 'ETH'
+  const quote = 'DOGE'
+
+  mockWebSocketProvider(WebSocketClassProvider)
+  const mockWsServer = new Server(ENDPOINT_URL, { mock: false })
+  let pongCallCount = 0
+
+  mockWsServer.on('connection', (socket) => {
+    socket.on('message', () => {
+      socket.send(
+        JSON.stringify({
+          pair: `${base}/${quote}`,
+          value: price,
+        }),
+      )
+    })
+  })
+
+  const websocketTransport = new WebSocketTransport<WebSocketTypes>({
+    url: () => ENDPOINT_URL,
+    options: () => ({ headers: { 'x-auth-token': 'token' } }),
+    handlers: {
+      message(message) {
+        if (!message.pair) {
+          return []
+        }
+        const [curBase, curQuote] = message.pair.split('/')
+        return [
+          {
+            params: { base: curBase, quote: curQuote },
+            response: {
+              data: { result: message.value },
+              result: message.value,
+              timestamps: { providerIndicatedTimeUnixMs: Date.now() },
+            },
+          },
+        ]
+      },
+      pong: () => {
+        pongCallCount++
+      },
+    },
+    builders: {
+      subscribeMessage: (params) => `S:${params.base}/${params.quote}`,
+      unsubscribeMessage: (params) => ({
+        request: 'unsubscribe',
+        pair: `${params.base}/${params.quote}`,
+      }),
+    },
+  })
+
+  const webSocketEndpoint = new AdapterEndpoint({
+    name: 'TEST',
+    transport: websocketTransport,
+    inputParameters,
+  })
+
+  const config = new AdapterConfig({}, { envDefaultOverrides: { BACKGROUND_EXECUTE_MS_WS } })
+
+  const adapter = new Adapter({
+    name: 'TEST',
+    defaultEndpoint: 'test',
+    endpoints: [webSocketEndpoint],
+    config,
+  })
+
+  const testAdapter = await TestAdapter.startWithMockedCache(adapter, t.context)
+
+  await testAdapter.startBackgroundExecuteThenGetResponse(t, {
+    requestData: { base, quote },
+    expectedResponse: {
+      data: { result: price },
+      result: price,
+      statusCode: 200,
+    },
+  })
+
+  websocketTransport.wsConnection?.emit('pong')
+
+  t.is(pongCallCount, 1)
+
+  testAdapter.api.close()
+  mockWsServer.close()
+  await t.context.clock.runToLastAsync()
+})
+
 test.serial(
   'increments failover counter on abnormal closure and passes it to url function',
   async (t) => {