✨ Collect WebSocket resource events

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: bdibon

packages/browser-core/src/browser/webSocketObservable.spec.ts

@@ -1,5 +1,6 @@
 import { registerCleanupTask } from '../../test'
 import type { Subscription } from '../tools/observable'
+import { setAllowUntrustedEvents } from './addEventListener'
 import type { WebSocketContext } from './webSocketObservable'
 import { initWebSocketObservable, resetWebSocketObservable } from './webSocketObservable'
 
@@ -80,7 +81,8 @@ describe('webSocketObservable', () => {
   })
 
   function startTracking() {
-    subscription = initWebSocketObservable({ allowUntrustedEvents: true }).subscribe((context) => {
+    setAllowUntrustedEvents(true)
+    subscription = initWebSocketObservable().subscribe((context) => {
       contexts.push(context)
     })
   }
@@ -322,35 +324,4 @@ describe('webSocketObservable', () => {
       })
     })
   })
-
-  describe('with conflicting allowUntrustedEvents policies across callers', () => {
-    it('does not emit open or message-in for untrusted events when the customer disallows them', () => {
-      initWebSocketObservable({ allowUntrustedEvents: true })
-      subscription = initWebSocketObservable({ allowUntrustedEvents: false }).subscribe((context) => {
-        contexts.push(context)
-      })
-
-      const ws = new windowAsWebSocketHost.WebSocket('wss://example.com/socket')
-      ws.simulateOpen()
-      ws.simulateMessage('hello')
-
-      expect(getContexts('connecting').length).toBe(1)
-      expect(getContexts('open').length).toBe(0)
-      expect(getContexts('message-in').length).toBe(0)
-    })
-
-    it('emits open and message-in for untrusted events when every caller allows them', () => {
-      initWebSocketObservable({ allowUntrustedEvents: true })
-      subscription = initWebSocketObservable({ allowUntrustedEvents: true }).subscribe((context) => {
-        contexts.push(context)
-      })
-
-      const ws = new windowAsWebSocketHost.WebSocket('wss://example.com/socket')
-      ws.simulateOpen()
-      ws.simulateMessage('hello')
-
-      expect(getContexts('open').length).toBe(1)
-      expect(getContexts('message-in').length).toBe(1)
-    })
-  })
 })

packages/browser-core/src/browser/webSocketObservable.ts

@@ -6,10 +6,6 @@ import { instrumentConstructor, instrumentMethod } from '../tools/instrumentMeth
 import { Observable } from '../tools/observable'
 import { addEventListener } from './addEventListener'
 
-interface WebSocketObservableConfiguration {
-  allowUntrustedEvents?: boolean | undefined
-}
-
 type GlobalWithWebSocket = GlobalObject & { WebSocket: typeof WebSocket }
 
 function isGlobalWithWebSocket(global: GlobalObject): global is GlobalWithWebSocket {
@@ -64,18 +60,7 @@ export type WebSocketContext =
 
 let webSocketObservable: Observable<WebSocketContext> | undefined
 
-// The singleton WebSocket observable applies the latest caller's allowUntrustedEvents policy so
-// that the customer's configuration overrides an early call (e.g. from bufferedData) that opts
-// in before the customer config is parsed.
-let allowUntrustedEvents: boolean | undefined
-
-export function initWebSocketObservable(
-  configuration: WebSocketObservableConfiguration = {}
-): Observable<WebSocketContext> {
-  if (configuration.allowUntrustedEvents !== undefined) {
-    allowUntrustedEvents = configuration.allowUntrustedEvents
-  }
-
+export function initWebSocketObservable(): Observable<WebSocketContext> {
   if (!webSocketObservable) {
     webSocketObservable = createWebSocketObservable()
   }
@@ -143,23 +128,23 @@ function attachInstanceListeners(
   observable: Observable<WebSocketContext>,
   stopListeners: Array<() => void>
 ) {
-  const { stop: stopOpen } = addEventListener({ allowUntrustedEvents }, instance, 'open', () => {
+  const { stop: stopOpen } = addEventListener(instance, 'open', () => {
     observable.notify({
       state: 'open',
       instance,
       openClocks: clocksNow(),
       protocol: instance.protocol || '',
     })
   })
-  const { stop: stopMessage } = addEventListener({ allowUntrustedEvents }, instance, 'message', (event) => {
+  const { stop: stopMessage } = addEventListener(instance, 'message', (event) => {
     observable.notify({
       state: 'message-in',
       instance,
       size: computePayloadSize(event.data),
       at: clocksNow(),
     })
   })
-  const { stop: stopClose } = addEventListener({ allowUntrustedEvents }, instance, 'close', (event) => {
+  const { stop: stopClose } = addEventListener(instance, 'close', (event) => {
     observable.notify({
       state: 'closed',
       instance,
@@ -196,5 +181,4 @@ function computePayloadSize(data: unknown): number {
  */
 export function resetWebSocketObservable() {
   webSocketObservable = undefined
-  allowUntrustedEvents = undefined
 }

packages/browser-rum-core/src/boot/startRum.ts

@@ -17,13 +17,16 @@ import {
   startGlobalContext,
   startUserContext,
   startTabContext,
+  isExperimentalFeatureEnabled,
+  ExperimentalFeature,
 } from '@datadog/browser-core'
 import { createDOMMutationObservable } from '../browser/domMutationObservable'
 import { createWindowOpenObservable } from '../browser/windowOpenObservable'
 import { startInternalContext } from '../domain/contexts/internalContext'
 import { LifeCycle, LifeCycleEventType } from '../domain/lifeCycle'
 import { startViewHistory } from '../domain/contexts/viewHistory'
 import { startRequestCollection } from '../domain/requestCollection'
+import { startWebSocketCollection } from '../domain/webSocketCollection'
 import { startActionCollection } from '../domain/action/actionCollection'
 import { startErrorCollection } from '../domain/error/errorCollection'
 import { startResourceCollection } from '../domain/resource/resourceCollection'
@@ -230,6 +233,15 @@ export function startRumEventCollection(
 
   const vitalCollection = startVitalCollection(lifeCycle, pageStateHistory)
 
+  if (isExperimentalFeatureEnabled(ExperimentalFeature.TRACK_WEB_SOCKETS)) {
+    const webSocketCollection = startWebSocketCollection(
+      lifeCycle,
+      viewHistory,
+      vitalCollection.addDurationVital
+    )
+    cleanupTasks.push(webSocketCollection.stop)
+  }
+
   const internalContext = startInternalContext(
     configuration.applicationId,
     sessionManager,

packages/browser-rum-core/src/domain/lifeCycle.ts

@@ -4,6 +4,7 @@ import { AbstractLifeCycle } from '@datadog/browser-core'
 import type { RumEventDomainContext } from '../domainContext.types'
 import type { RawRumEvent, AssembledRumEvent } from '../rawRumEvent.types'
 import type { RequestCompleteEvent, RequestStartEvent } from './requestCollection'
+import type { WebSocketCompleteEvent } from './webSocketCollection'
 import type { AutoAction } from './action/actionCollection'
 import type { ViewEvent, ViewCreatedEvent, ViewEndedEvent, BeforeViewUpdateEvent } from './view/trackViews'
 import type { DurationVitalStart } from './vital/vitalCollection'
@@ -22,6 +23,7 @@ export const enum LifeCycleEventType {
   AFTER_VIEW_ENDED,
   REQUEST_STARTED,
   REQUEST_COMPLETED,
+  WEBSOCKET_COMPLETED,
 
   // The SESSION_EXPIRED lifecycle event has been introduced to represent when a session has expired
   // and trigger cleanup tasks related to this, prior to renewing the session. Its implementation is
@@ -67,6 +69,7 @@ declare const LifeCycleEventTypeAsConst: {
   AFTER_VIEW_ENDED: LifeCycleEventType.AFTER_VIEW_ENDED
   REQUEST_STARTED: LifeCycleEventType.REQUEST_STARTED
   REQUEST_COMPLETED: LifeCycleEventType.REQUEST_COMPLETED
+  WEBSOCKET_COMPLETED: LifeCycleEventType.WEBSOCKET_COMPLETED
   SESSION_EXPIRED: LifeCycleEventType.SESSION_EXPIRED
   SESSION_RENEWED: LifeCycleEventType.SESSION_RENEWED
   PAGE_MAY_EXIT: LifeCycleEventType.PAGE_MAY_EXIT
@@ -89,6 +92,7 @@ export interface LifeCycleEventMap {
   [LifeCycleEventTypeAsConst.AFTER_VIEW_ENDED]: ViewEndedEvent
   [LifeCycleEventTypeAsConst.REQUEST_STARTED]: RequestStartEvent
   [LifeCycleEventTypeAsConst.REQUEST_COMPLETED]: RequestCompleteEvent
+  [LifeCycleEventTypeAsConst.WEBSOCKET_COMPLETED]: WebSocketCompleteEvent
   [LifeCycleEventTypeAsConst.SESSION_EXPIRED]: void
   [LifeCycleEventTypeAsConst.SESSION_RENEWED]: void
   [LifeCycleEventTypeAsConst.PAGE_MAY_EXIT]: PageMayExitEvent

packages/browser-rum-core/src/domain/resource/resourceCollection.ts

@@ -1,45 +1,47 @@
-import type { Duration } from '@datadog/js-core/time'
-import { toServerDuration, relativeToClocks } from '@datadog/js-core/time'
 import {
+  addTelemetryDebug,
   combine,
-  generateUUID,
   createTaskQueue,
+  display,
+  generateUUID,
+  matchList,
   mockable,
+  RequestType,
+  ResourceType,
   runOnReadyState,
-  matchList,
   safeTruncate,
-  display,
-  addTelemetryDebug,
-  RequestType,
   setTimeout,
 } from '@datadog/browser-core'
-import type { MatchHeader, RumConfiguration } from '../configuration'
-import { RumPerformanceEntryType, createPerformanceObservable } from '../../browser/performanceObservable'
+import type { Duration } from '@datadog/js-core/time'
+import { elapsed, relativeToClocks, toServerDuration } from '@datadog/js-core/time'
+import { createPerformanceObservable, RumPerformanceEntryType } from '../../browser/performanceObservable'
+import { getNavigationEntry } from '../../browser/performanceUtils'
 import type { RumResourceEventDomainContext } from '../../domainContext.types'
 import type { NetworkHeaders, RawRumResourceEvent, ResourceRequest, ResourceResponse } from '../../rawRumEvent.types'
 import { RumEventType } from '../../rawRumEvent.types'
-import type { RawRumEventCollectedData, LifeCycle } from '../lifeCycle'
+import type { MatchHeader, RumConfiguration } from '../configuration'
+import { startEventTracker } from '../eventTracker'
+import { extractRegexMatch } from '../extractRegexMatch'
+import type { LifeCycle, RawRumEventCollectedData } from '../lifeCycle'
 import { LifeCycleEventType } from '../lifeCycle'
 import type { RequestCompleteEvent } from '../requestCollection'
-import { createSpanIdentifier } from '../tracing/identifier'
 import { getDocumentTraceId } from '../tracing/getDocumentTraceId'
-import { getNavigationEntry } from '../../browser/performanceUtils'
-import { startEventTracker } from '../eventTracker'
-import { extractRegexMatch } from '../extractRegexMatch'
+import { createSpanIdentifier } from '../tracing/identifier'
+import type { WebSocketCompleteEvent } from '../webSocketCollection'
+import type { GraphQlMetadata } from './graphql'
+import { extractGraphQlMetadata, findGraphQlConfiguration } from './graphql'
+import { createRequestRegistry } from './requestRegistry'
+import type { ResourceLikeEntry } from './resourceUtils'
 import {
+  computeResourceEntryDeliveryType,
   computeResourceEntryDetails,
   computeResourceEntryDuration,
-  computeResourceEntryType,
-  computeResourceEntrySize,
   computeResourceEntryProtocol,
-  computeResourceEntryDeliveryType,
+  computeResourceEntrySize,
+  computeResourceEntryType,
   isResourceEntryRequestType,
   sanitizeIfLongDataUrl,
 } from './resourceUtils'
-import type { ResourceLikeEntry } from './resourceUtils'
-import { createRequestRegistry } from './requestRegistry'
-import type { GraphQlMetadata } from './graphql'
-import { extractGraphQlMetadata, findGraphQlConfiguration } from './graphql'
 import type { ManualResourceData } from './trackManualResources'
 import { trackManualResources } from './trackManualResources'
 
@@ -51,6 +53,9 @@ export function startResourceCollection(lifeCycle: LifeCycle, configuration: Rum
   const taskQueue = mockable(createTaskQueue)()
   const requestRegistry = createRequestRegistry(lifeCycle)
 
+  lifeCycle.subscribe(LifeCycleEventType.WEBSOCKET_COMPLETED, (event: WebSocketCompleteEvent) => {
+    handleResource(() => assembleWebSocketResource(event))
+  })
   const performanceResourceSubscription = createPerformanceObservable({
     type: RumPerformanceEntryType.RESOURCE,
     buffered: true,
@@ -108,6 +113,55 @@ export function startResourceCollection(lifeCycle: LifeCycle, configuration: Rum
   }
 }
 
+function assembleWebSocketResource(
+  event: WebSocketCompleteEvent
+): RawRumEventCollectedData<RawRumResourceEvent> | undefined {
+  const duration = elapsed(event.startClocks.timeStamp, event.endClocks.timeStamp)
+
+  const rawRumEvent: RawRumResourceEvent = {
+    date: event.startClocks.timeStamp,
+    type: RumEventType.RESOURCE,
+    resource: {
+      id: generateUUID(),
+      type: ResourceType.WEBSOCKET,
+      url: event.url,
+      duration: toServerDuration(duration),
+      websocket: {
+        connection_id: event.connectionId,
+        handshake_succeeded: event.handshakeSucceeded,
+        start_time: event.startClocks.timeStamp,
+        end_time: event.endClocks.timeStamp,
+        start_view_id: event.startViewId,
+        end_view_id: event.endViewId,
+        tracking_end_reason: event.trackingEndReason,
+        close_code: event.closeCode,
+        close_reason: event.closeReason,
+        was_clean: event.wasClean,
+        messages_in: event.messagesIn,
+        messages_out: event.messagesOut,
+        time_to_first_message_in: event.firstMessageInOffset,
+        time_to_first_message_out: event.firstMessageOutOffset,
+        last_message_in_at: event.lastMessageInAt,
+        longest_inbound_silence: event.longestInboundSilence,
+        inbound_idle_duration_before_close: event.inboundIdleDurationBeforeClose,
+        buffered_amount_max: event.bufferedAmountMax,
+        protocol: event.protocol,
+        setup_duration: event.setupDuration,
+      },
+    },
+    _dd: {
+      discarded: false,
+    },
+  }
+
+  return {
+    startClocks: event.startClocks,
+    duration,
+    rawRumEvent,
+    domainContext: {},
+  }
+}
+
 function assembleResource(
   entry: ResourceLikeEntry,
   request: RequestCompleteEvent | undefined,