fix(ws): handle removing untracked event listeners (#7419)

crysmags · web-flow · commit 76e93d7934c6 · 2026-02-03T13:37:43.000-05:00
diff --git a/packages/datadog-instrumentations/src/ws.js b/packages/datadog-instrumentations/src/ws.js
@@ -95,7 +95,7 @@ function wrapListener (originalOn) {
 function removeListener (originalOff) {
   return function (eventName, handler) {
     if (eventName === 'message') {
-      const wrappedHandler = eventHandlerMap.get(handler)
+      const wrappedHandler = eventHandlerMap.get(handler) || handler
       return originalOff.call(this, eventName, wrappedHandler)
     }
     return originalOff.apply(this, arguments)
diff --git a/packages/datadog-plugin-ws/test/index.spec.js b/packages/datadog-plugin-ws/test/index.spec.js
@@ -95,6 +95,26 @@ describe('Plugin', () => {
           })
         })
 
+        it('should handle removing a listener that was never added', (done) => {
+          wsServer.on('connection', (ws) => {
+            connectionReceived = true
+            ws.send('test message')
+          })
+
+          const neverAddedHandler = () => {
+            throw new Error('this should never be called')
+          }
+
+          client.on('message', (msg) => {
+            assert.strictEqual(msg.toString(), 'test message')
+            done()
+          })
+
+          assert.doesNotThrow(() => {
+            client.off('message', neverAddedHandler)
+          })
+        })
+
         it('should do automatic instrumentation for server connections', done => {
           connectionReceived = false
 

Original file line number	Diff line number	Diff line change
`@@ -95,7 +95,7 @@ function wrapListener (originalOn) {`
`95`	`95`	`function removeListener (originalOff) {`
`96`	`96`	`return function (eventName, handler) {`
`97`	`97`	`if (eventName === 'message') {`
`98`		`- const wrappedHandler = eventHandlerMap.get(handler)`
	`98`	`+ const wrappedHandler = eventHandlerMap.get(handler) \|\| handler`
`99`	`99`	`return originalOff.call(this, eventName, wrappedHandler)`
`100`	`100`	`}`
`101`	`101`	`return originalOff.apply(this, arguments)`