Ensure notify is patched when a client is cloned

clone-client doesn't copy over monkey patched methods on to new clones¹,
so we have to hook into the cloning process to monkey patch them

¹ This is actually a good thing as monkey patched methods will usually
work by holding a reference to an old client to call the un-patched
method. If we copied them over then they would still reference the
original client, not the clone

packages/in-flight/src/in-flight.js

@@ -1,47 +1,19 @@
 const cuid = require('@bugsnag/cuid')
+const clone = require('@bugsnag/core/lib/clone-client')
 
 const FLUSH_POLL_INTERVAL_MS = 50
 const inFlightRequests = new Map()
 
 const noop = () => {}
 
+// when a client is cloned, make sure to patch the clone's notify method too
+// we don't need to patch delivery when a client is cloned because the
+// original client's delivery method will be copied over to the clone
+clone.registerCallback(patchNotify)
+
 module.exports = {
   trackInFlight (client) {
-    const originalNotify = client._notify
-
-    client._notify = function (event, onError, callback = noop) {
-      const id = cuid()
-      inFlightRequests.set(id, true)
-
-      const _callback = function () {
-        inFlightRequests.delete(id)
-        callback.apply(null, arguments)
-      }
-
-      client._depth += 1
-
-      try {
-        originalNotify.call(client, event, onError, _callback)
-      } finally {
-        client._depth -= 1
-      }
-    }
-
-    const patchDelivery = (delivery) => {
-      const originalSendSession = delivery.sendSession
-
-      delivery.sendSession = function (session, callback = noop) {
-        const id = cuid()
-        inFlightRequests.set(id, true)
-
-        const _callback = function () {
-          inFlightRequests.delete(id)
-          callback.apply(null, arguments)
-        }
-
-        originalSendSession.call(delivery, session, _callback)
-      }
-    }
+    patchNotify(client)
 
     let delivery = client._delivery
     patchDelivery(delivery)
@@ -85,3 +57,52 @@ module.exports = {
     })
   }
 }
+
+// patch a client's _notify method to track in-flight requests
+// we patch _notify directly to track requests as early as possible and use the
+// "post report" delivery callback to know when a request finishes
+function patchNotify (client) {
+  const originalNotify = client._notify
+
+  client._notify = function (event, onError, callback = noop) {
+    const id = cuid()
+    inFlightRequests.set(id, true)
+
+    const _callback = function () {
+      inFlightRequests.delete(id)
+      callback.apply(null, arguments)
+    }
+
+    client._depth += 1
+
+    try {
+      originalNotify.call(client, event, onError, _callback)
+    } finally {
+      client._depth -= 1
+    }
+  }
+}
+
+// patch a delivery delegate's sendSession method to track in-flight requests
+// we do this on the delivery delegate because the client object doesn't
+// actually deliver sessions, a session delegate does
+// we can't patch the session delegate either because it will deliver sessions
+// in a way that makes sense on the platform, e.g. on node sessions are batched
+// into 1 request made every x seconds
+// therefore the only thing that knows when a session request is started and
+// when it finishes is the delivery delegate itself
+function patchDelivery (delivery) {
+  const originalSendSession = delivery.sendSession
+
+  delivery.sendSession = function (session, callback = noop) {
+    const id = cuid()
+    inFlightRequests.set(id, true)
+
+    const _callback = function () {
+      inFlightRequests.delete(id)
+      callback.apply(null, arguments)
+    }
+
+    originalSendSession.call(delivery, session, _callback)
+  }
+}

packages/in-flight/test/in-flight.test.ts

@@ -1,3 +1,4 @@
+import clone from '@bugsnag/core/lib/clone-client'
 import Client, { EventDeliveryPayload, SessionDeliveryPayload } from '@bugsnag/core/client'
 
 // The in-flight package has module level state which can leak between tests
@@ -43,6 +44,44 @@ describe('@bugsnag/in-flight', () => {
     expect(sendSession).not.toHaveBeenCalled()
   })
 
+  it('can track in-flight events after a client is cloned', () => {
+    const client = new Client({ apiKey: 'AN_API_KEY' })
+
+    // eslint thinks this is never reassigned, but it clearly is
+    let cloned: Client // eslint-disable-line prefer-const
+
+    const payloads: EventDeliveryPayload[] = []
+    const sendSession = jest.fn()
+
+    client._setDelivery(() => ({
+      sendEvent: (payload, cb) => {
+        expect(cloned._depth).toBe(2)
+        payloads.push(payload)
+        cb()
+      },
+      sendSession
+    }))
+
+    bugsnagInFlight.trackInFlight(client)
+
+    expect(payloads.length).toBe(0)
+
+    const onError = jest.fn()
+    const callback = jest.fn()
+
+    cloned = clone(client)
+
+    expect(cloned._depth).toBe(1)
+
+    cloned.notify(new Error('xyz'), onError, callback)
+
+    expect(cloned._depth).toBe(1)
+    expect(onError).toHaveBeenCalledTimes(1)
+    expect(callback).toHaveBeenCalledTimes(1)
+    expect(payloads.length).toBe(1)
+    expect(sendSession).not.toHaveBeenCalled()
+  })
+
   it('tracks in-flight sessions', () => {
     const client = new Client({ apiKey: 'AN_API_KEY' })
     const payloads: SessionDeliveryPayload[] = []
@@ -84,6 +123,49 @@ describe('@bugsnag/in-flight', () => {
     expect(client._sessionDelegate.resumeSession).not.toHaveBeenCalled()
   })
 
+  it('tracks in-flight sessions after a client has been cloned', () => {
+    const client = new Client({ apiKey: 'AN_API_KEY' })
+    const payloads: SessionDeliveryPayload[] = []
+    const sendEvent = jest.fn()
+    const callback = jest.fn()
+
+    client._sessionDelegate = {
+      startSession: jest.fn(function (client, session) {
+        client._delivery.sendSession(session, callback)
+
+        return client
+      }),
+      pauseSession: jest.fn(),
+      resumeSession: jest.fn()
+    }
+
+    client._setDelivery(() => ({
+      sendEvent,
+      sendSession: (payload, cb) => {
+        payloads.push(payload)
+        cb()
+      }
+    }))
+
+    bugsnagInFlight.trackInFlight(client)
+
+    expect(payloads.length).toBe(0)
+    expect(callback).not.toHaveBeenCalled()
+    expect(client._sessionDelegate.startSession).not.toHaveBeenCalled()
+    expect(client._sessionDelegate.pauseSession).not.toHaveBeenCalled()
+    expect(client._sessionDelegate.resumeSession).not.toHaveBeenCalled()
+
+    const cloned = clone(client)
+
+    cloned.startSession()
+
+    expect(payloads.length).toBe(1)
+    expect(callback).toHaveBeenCalledTimes(1)
+    expect(cloned._sessionDelegate.startSession).toHaveBeenCalledTimes(1)
+    expect(cloned._sessionDelegate.pauseSession).not.toHaveBeenCalled()
+    expect(cloned._sessionDelegate.resumeSession).not.toHaveBeenCalled()
+  })
+
   it('tracks all in-flight requests', () => {
     const client = new Client({ apiKey: 'AN_API_KEY' })
     const eventPayloads: EventDeliveryPayload[] = []