dd-trace 5.103.0 → 5.105.0

package/packages/datadog-instrumentations/src/mocha/utils.js

@@ -503,13 +503,38 @@ function getTestFinishInfo (test, status, config, error) {
   }
 }
 
-function getOnTestEndHandler (config) {
+function getOnTestEndHandler (config, finalAttemptHandlers) {
   return async function (test) {
     if (test._ddShouldSkipEfdRetry) {
       return
     }
     const ctx = getTestContext(test)
     const status = getTestStatus(test)
+    const shouldFinishTest = ctx && (!getAfterEachHooks(test).length || (test._ddIsDisabled && !test._ddIsAttemptToFix))
+    let testFinishInfo
+    let isFinalAttempt = false
+
+    // If there are afterEach to be run, we don't finish the test yet.
+    // Disabled tests (marked pending by us) are finished immediately without waiting for afterEach hooks.
+    // In older mocha versions, pending tests don't run afterEach hooks, so we can't rely on
+    // getOnHookEndHandler to finish the test. This mirrors Jest's approach where the skip handler
+    // directly sets finalStatus without waiting for hooks
+    if (!ctx && test.isPending()) {
+      test._ddIsFinalAttempt = true
+      isFinalAttempt = true
+    }
+
+    if (shouldFinishTest) {
+      testFinishInfo = getTestFinishInfo(test, status, config, ctx.err || test.err)
+      if (testFinishInfo.finalStatus !== undefined) {
+        test._ddIsFinalAttempt = true
+        isFinalAttempt = true
+      }
+    }
+
+    if (isFinalAttempt) {
+      finalAttemptHandlers?.onStart?.(test)
+    }
 
     // After finishing it might take a bit for the snapshot to be handled.
     // This means that tests retried with DI are BREAKPOINT_HIT_GRACE_PERIOD_MS slower at least.
@@ -521,13 +546,7 @@ function getOnTestEndHandler (config) {
       })
     }
 
-    // If there are afterEach to be run, we don't finish the test yet.
-    // Disabled tests (marked pending by us) are finished immediately without waiting for afterEach hooks.
-    // In older mocha versions, pending tests don't run afterEach hooks, so we can't rely on
-    // getOnHookEndHandler to finish the test. This mirrors Jest's approach where the skip handler
-    // directly sets finalStatus without waiting for hooks
-    if (ctx && (!getAfterEachHooks(test).length || (test._ddIsDisabled && !test._ddIsAttemptToFix))) {
-      const testFinishInfo = getTestFinishInfo(test, status, config, ctx.err || test.err)
+    if (shouldFinishTest) {
       testFinishCh.publish({
         status,
         hasBeenRetried: isMochaRetry(test),
@@ -536,6 +555,10 @@ function getOnTestEndHandler (config) {
         ...ctx.currentStore,
       })
     }
+
+    if (isFinalAttempt) {
+      finalAttemptHandlers?.onFinish?.(test)
+    }
   }
 }
 
@@ -552,6 +575,9 @@ function getOnHookEndHandler (config) {
         // skip to avoid double-publishing
         if (ctx && (!test._ddIsDisabled || test._ddIsAttemptToFix)) {
           const testFinishInfo = getTestFinishInfo(test, status, config, ctx.err || test.err)
+          if (testFinishInfo.finalStatus !== undefined) {
+            test._ddIsFinalAttempt = true
+          }
           testFinishCh.publish({
             status,
             hasBeenRetried: isMochaRetry(test),
@@ -583,6 +609,20 @@ function getOnFailHandler (isMain, config) {
         testContext.err = err
         errorCh.runStores(testContext, () => {})
         const testFinishInfo = getTestFinishInfo(test, 'fail', config, err)
+        // ATR never retries hook failures: this.retries(N) is set in runnableWrapper
+        // which only runs when the test function executes — hooks bypass that path,
+        // so _retries stays at -1 and getIsLastRetry returns false, leaving finalStatus
+        // undefined. We must also mark the attempt final when no clone-based retry
+        // mechanism (EFD original, EFD clone, ATF) has queued further attempts.
+        const noCloneRetries = !test._ddIsEfdRetry &&
+          !((test._ddIsNew || test._ddIsModified) && config.isEarlyFlakeDetectionEnabled) &&
+          !test._ddIsAttemptToFix
+        if (testFinishInfo.finalStatus !== undefined || noCloneRetries) {
+          test._ddIsFinalAttempt = true
+        }
+        // test.state is never set to 'failed' for hook failures (Mocha marks the hook,
+        // not the test). Flag it so finishRootSuiteForFile can compute the correct status.
+        test._ddHookFailed = true
         testFinishCh.publish({
           status: 'fail',
           hasBeenRetried: isMochaRetry(test),

package/packages/datadog-instrumentations/src/mongodb-core.js

@@ -174,7 +174,7 @@ function instrument (operation, command, instance, args, server, ns, ops, option
     name,
   }
   return startCh.runStores(ctx, () => {
-    args[index] = shimmer.wrapFunction(callback, callback => function (err, res) {
+    args[index] = shimmer.wrapCallback(callback, callback => function (err, res) {
       if (err) {
         ctx.error = err
         errorCh.publish(ctx)

package/packages/datadog-instrumentations/src/mongoose.js

@@ -125,21 +125,19 @@ addHook({
                     const resolve = args[0]
                     const reject = args[1]
 
-                    args[0] = shimmer.wrapFunction(resolve, resolve => function wrappedResolve (...args) {
-                      finishCh.publish(ctx)
-
-                      if (resolve) {
+                    if (typeof resolve === 'function') {
+                      args[0] = shimmer.wrapCallback(resolve, resolve => function wrappedResolve (...args) {
+                        finishCh.publish(ctx)
                         return resolve.apply(this, args)
-                      }
-                    })
-
-                    args[1] = shimmer.wrapFunction(reject, reject => function wrappedReject (...args) {
-                      finishCh.publish(ctx)
+                      })
+                    }
 
-                      if (reject) {
+                    if (typeof reject === 'function') {
+                      args[1] = shimmer.wrapCallback(reject, reject => function wrappedReject (...args) {
+                        finishCh.publish(ctx)
                         return reject.apply(this, args)
-                      }
-                    })
+                      })
+                    }
 
                     return originalThen.apply(this, args)
                   }

package/packages/datadog-instrumentations/src/mysql.js

@@ -29,7 +29,7 @@ addHook({ name: 'mysql', file: 'lib/Connection.js', versions: ['>=2'] }, Connect
 
         if (res._callback) {
           const cb = res._callback
-          res._callback = shimmer.wrapFunction(cb, cb => function (error, result) {
+          res._callback = shimmer.wrapCallback(cb, cb => function (error, result) {
             if (error) {
               ctx.error = error
               errorCh.publish(ctx)
@@ -86,7 +86,7 @@ addHook({ name: 'mysql', file: 'lib/Pool.js', versions: ['>=2'] }, Pool => {
     return startPoolQueryCh.runStores(ctx, () => {
       const cb = args[args.length - 1]
       if (typeof cb === 'function') {
-        args[args.length - 1] = shimmer.wrapFunction(cb, cb => function (...args) {
+        args[args.length - 1] = shimmer.wrapCallback(cb, cb => function (...args) {
           return finishPoolQueryCh.runStores(ctx, cb, this, ...args)
         })
       }

package/packages/datadog-instrumentations/src/mysql2.js

@@ -176,7 +176,7 @@ function wrapConnection (Connection, version) {
         if (typeof this.onResult === 'function') {
           const onResult = this.onResult
 
-          this.onResult = shimmer.wrapFunction(onResult, onResult => function (error) {
+          this.onResult = shimmer.wrapCallback(onResult, onResult => function (error) {
             if (error) {
               ctx.error = error
               errorCh.publish(ctx)

package/packages/datadog-instrumentations/src/nats.js

@@ -0,0 +1,182 @@
+'use strict'
+
+// Shimmer required: NATS consumer paths need argument modification — the user's
+// `opts.callback` is wrapped before being handed to SubscriptionImpl, and the
+// returned subscription's async iterator is wrapped so iterator-style consumers
+// get receive events. Orchestrion can only wrap method calls, not arguments
+// or returned iterables.
+
+const shimmer = require('../../datadog-shimmer')
+const { addHook, channel } = require('./helpers/instrument')
+
+const publishStartCh = channel('apm:nats:publish:start')
+const publishFinishCh = channel('apm:nats:publish:finish')
+const publishErrorCh = channel('apm:nats:publish:error')
+
+const consumeStartCh = channel('apm:nats:consume:start')
+const consumeFinishCh = channel('apm:nats:consume:finish')
+const consumeErrorCh = channel('apm:nats:consume:error')
+
+// Tracks connections that are currently inside a `request`/`requestMany` call
+// so the nested `this.publish(...)` they issue short-circuits without creating
+// a second producer span (the outer request wrap already created one and
+// injected headers — the inner publish would double-count it). A WeakSet avoids
+// changing the shape of the user's connection object.
+const requestsInFlight = new WeakSet()
+
+// Captured from the `lib/headers.js` hook below. The nats-core package always
+// imports `./headers` from `lib/nats.js`, so by the time we wrap `publish` the
+// reference is set. No defensive checks needed at call sites.
+let createHeaders
+
+addHook({ name: '@nats-io/nats-core', versions: ['>=3.0.0'], file: 'lib/headers.js' }, exports => {
+  createHeaders = exports.headers
+  return exports
+})
+
+// transport-node re-exports nats-core internals — the passthrough hook ensures
+// the package name is registered so `withVersions('nats', '@nats-io/transport-node', ...)`
+// can resolve it in plugin tests.
+addHook({ name: '@nats-io/transport-node', versions: ['>=3.0.0'] }, exports => exports)
+
+function wrapSyncProducer (original, type) {
+  return function (subject, data, options) {
+    if (!publishStartCh.hasSubscribers) {
+      return original.apply(this, arguments)
+    }
+    const opts = { ...options }
+    const ctx = { type, subject, data, options: opts, connection: this, createHeaders }
+    return publishStartCh.runStores(ctx, () => {
+      try {
+        return original.call(this, subject, data, opts)
+      } catch (err) {
+        ctx.error = err
+        publishErrorCh.publish(ctx)
+        throw err
+      } finally {
+        publishFinishCh.publish(ctx)
+      }
+    })
+  }
+}
+
+// publish is also wrapped by `wrapSyncProducer`, but request/requestMany call
+// `this.publish(...)` internally. Set a marker on the connection so the inner
+// publish wrap short-circuits — see `wrapPublish`.
+function wrapAsyncProducer (original, type) {
+  return function (subject, data, options) {
+    if (!publishStartCh.hasSubscribers) {
+      return original.apply(this, arguments)
+    }
+    const opts = { ...options }
+    const ctx = { type, subject, data, options: opts, connection: this, createHeaders }
+    return publishStartCh.runStores(ctx, () => {
+      requestsInFlight.add(this)
+      let promise
+      try {
+        // `request`/`requestMany` never throw synchronously — they wrap their own
+        // input validation in a try/catch that returns `Promise.reject`.
+        promise = original.call(this, subject, data, opts)
+      } finally {
+        // The nested `this.publish(...)` runs during the synchronous body of
+        // request/requestMany, so clearing the marker as soon as the call
+        // returns is sufficient — the promise resolution happens later.
+        requestsInFlight.delete(this)
+      }
+      return Promise.resolve(promise).then(
+        result => {
+          ctx.result = result
+          publishFinishCh.publish(ctx)
+          return result
+        },
+        err => {
+          ctx.error = err
+          publishErrorCh.publish(ctx)
+          publishFinishCh.publish(ctx)
+          throw err
+        }
+      )
+    })
+  }
+}
+
+function wrapPublish (original) {
+  const wrapped = wrapSyncProducer(original, 'publish')
+  return function (subject, data, options) {
+    // Called from inside request/requestMany — the outer wrap already produced
+    // a span and injected headers; running the inner wrap would double-count.
+    if (requestsInFlight.has(this)) {
+      return original.apply(this, arguments)
+    }
+    return wrapped.apply(this, arguments)
+  }
+}
+
+function wrapSubscribeCallback (userCallback, subject, connection) {
+  return function (err, message) {
+    if (!message || err) {
+      return userCallback.call(this, err, message)
+    }
+    const ctx = { subject, message, connection }
+    return consumeStartCh.runStores(ctx, () => {
+      try {
+        return userCallback.call(this, err, message)
+      } catch (e) {
+        ctx.error = e
+        consumeErrorCh.publish(ctx)
+        throw e
+      } finally {
+        consumeFinishCh.publish(ctx)
+      }
+    })
+  }
+}
+
+// Iterator-style consumers don't expose a delivery callback we can wrap, so
+// the consume span represents the moment of receipt only — it starts and
+// finishes before the value is yielded to user code, and the user's loop
+// body is not parented under the span.
+function wrapAsyncIteratorFactory (asyncIterator, subject, connection) {
+  return function () {
+    const iterator = asyncIterator.apply(this, arguments)
+    iterator.next = shimmer.wrapCallback(iterator.next, next => function () {
+      return next.apply(this, arguments).then(result => {
+        if (result && !result.done && result.value) {
+          const ctx = { subject, message: result.value, connection }
+          consumeStartCh.runStores(ctx, () => {
+            consumeFinishCh.publish(ctx)
+          })
+        }
+        return result
+      })
+    })
+    return iterator
+  }
+}
+
+addHook({ name: '@nats-io/nats-core', versions: ['>=3.0.0'], file: 'lib/nats.js' }, exports => {
+  const proto = exports.NatsConnectionImpl.prototype
+
+  shimmer.wrap(proto, 'publish', wrapPublish)
+  shimmer.wrap(proto, 'request', request => wrapAsyncProducer(request, 'request'))
+  shimmer.wrap(proto, 'requestMany', requestMany => wrapAsyncProducer(requestMany, 'requestMany'))
+
+  shimmer.wrap(proto, 'subscribe', subscribe => function (subject, opts) {
+    if (!consumeStartCh.hasSubscribers) {
+      return subscribe.apply(this, arguments)
+    }
+
+    const userOpts = opts ?? {}
+    if (typeof userOpts.callback === 'function') {
+      arguments[1] = { ...userOpts, callback: wrapSubscribeCallback(userOpts.callback, subject, this) }
+      return subscribe.apply(this, arguments)
+    }
+
+    const sub = subscribe.apply(this, arguments)
+    shimmer.wrap(sub, Symbol.asyncIterator, asyncIterator =>
+      wrapAsyncIteratorFactory(asyncIterator, subject, this))
+    return sub
+  })
+
+  return exports
+})

package/packages/datadog-instrumentations/src/nyc.js

@@ -2,7 +2,12 @@
 
 const shimmer = require('../../datadog-shimmer')
 const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
-const { setupSettingsCachePath } = require('../../dd-trace/src/ci-visibility/test-optimization-cache')
+const {
+  readCoverageBackfillFromCache,
+  readCoverageBackfillRootDirFromCache,
+  setupSettingsCachePath,
+} = require('../../dd-trace/src/ci-visibility/test-optimization-cache')
+const { applySkippedCoverageToCoverage } = require('../../dd-trace/src/plugins/util/test')
 const { addHook, channel } = require('./helpers/instrument')
 
 const codeCoverageWrapCh = channel('ci:nyc:wrap')
@@ -16,6 +21,38 @@ addHook({
   // when dd-trace fetches library configuration
   setupSettingsCachePath()
 
+  if (nycPackage.prototype.getCoverageMapFromAllCoverageFiles) {
+    // Some test frameworks receive skipped-suite coverage in the test process, but nyc merges reports later in the nyc
+    // process. Reuse the settings cache path as the process handoff so nyc can backfill skipped files before reporting.
+    shimmer.wrap(
+      nycPackage.prototype,
+      'getCoverageMapFromAllCoverageFiles',
+      getCoverageMapFromAllCoverageFiles => function (...args) {
+        const coverageMap = getCoverageMapFromAllCoverageFiles.apply(this, args)
+        const applyCoverageBackfill = (resolvedCoverageMap) => {
+          try {
+            if (!resolvedCoverageMap) {
+              return resolvedCoverageMap
+            }
+            applySkippedCoverageToCoverage(
+              resolvedCoverageMap,
+              readCoverageBackfillFromCache(),
+              readCoverageBackfillRootDirFromCache() || this.cwd
+            )
+          } catch {
+            // Do not break nyc's report generation if the cached backfill is stale or malformed.
+          }
+          return resolvedCoverageMap
+        }
+
+        if (coverageMap && typeof coverageMap.then === 'function') {
+          return coverageMap.then(applyCoverageBackfill)
+        }
+        return applyCoverageBackfill(coverageMap)
+      }
+    )
+  }
+
   // `wrap` is an async function
   shimmer.wrap(nycPackage.prototype, 'wrap', wrap => function (...args) {
     // Only relevant if the config `all` is set to true (for untested code coverage)

package/packages/datadog-instrumentations/src/openai.js

@@ -3,6 +3,7 @@
 const dc = require('dc-polyfill')
 const shimmer = require('../../datadog-shimmer')
 const { addHook } = require('./helpers/instrument')
+const aiGuard = require('./helpers/openai-ai-guard')
 
 const ch = dc.tracingChannel('apm:openai:request')
 const onStreamedChunkCh = dc.channel('apm:openai:request:chunk')
@@ -216,15 +217,20 @@ for (const extension of extensions) {
 
       for (const methodName of methods) {
         shimmer.wrap(targetPrototype, methodName, methodFn => function (...args) {
-          if (!ch.start.hasSubscribers) {
+          if (!ch.start.hasSubscribers && !aiGuard.hasSubscribers()) {
             return methodFn.apply(this, args)
           }
-
           // The OpenAI library lets you set `stream: true` on the options arg to any method
           // However, we only want to handle streamed responses in specific cases
           // chat.completions and completions
           const stream = streamedResponse && getOption(args, 'stream', false)
 
+          const guard = aiGuard.createGuard(baseResource, args[0], stream)
+
+          if (!ch.start.hasSubscribers && !guard) {
+            return methodFn.apply(this, args)
+          }
+
           const client = this._client || this.client
 
           const ctx = {
@@ -249,7 +255,7 @@ for (const extension of extensions) {
                   const parsedPromise = origApiPromParse.apply(this, args)
                     .then(body => Promise.all([this.responsePromise, body]))
 
-                  return handleUnwrappedAPIPromise(parsedPromise, ctx, stream)
+                  return handleUnwrappedAPIPromise(parsedPromise, ctx, stream, guard)
                 })
 
                 return unwrappedPromise
@@ -262,9 +268,11 @@ for (const extension of extensions) {
               const parsedPromise = origApiPromParse.apply(this, args)
                 .then(body => Promise.all([this.responsePromise, body]))
 
-              return handleUnwrappedAPIPromise(parsedPromise, ctx, stream)
+              return handleUnwrappedAPIPromise(parsedPromise, ctx, stream, guard)
             })
 
+            if (guard) aiGuard.wrapAsResponse(apiProm, guard)
+
             ch.end.publish(ctx)
 
             return apiProm
@@ -276,8 +284,10 @@ for (const extension of extensions) {
   }
 }
 
-function handleUnwrappedAPIPromise (apiProm, ctx, stream) {
-  return apiProm
+function handleUnwrappedAPIPromise (apiProm, ctx, stream, guard) {
+  const guardedApiProm = guard ? aiGuard.gateParse(apiProm, guard) : apiProm
+
+  return guardedApiProm
     .then(([{ response, options }, body]) => {
       if (stream) {
         if (body.iterator) {
@@ -287,22 +297,27 @@ function handleUnwrappedAPIPromise (apiProm, ctx, stream) {
             body.response.body, Symbol.asyncIterator, wrapStreamIterator(response, options, ctx)
           )
         }
-      } else {
-        finish(ctx, {
-          headers: response.headers,
-          data: body,
-          request: {
-            path: response.url,
-            method: options.method,
-          },
-        })
+        return body
       }
 
-      return body
+      finish(ctx, {
+        headers: response.headers,
+        data: body,
+        request: {
+          path: response.url,
+          method: options.method,
+        },
+      })
+
+      if (!guard) return body
+
+      return aiGuard.evaluateOutput(guard, body).then(() => body)
     })
     .catch(error => {
-      finish(ctx, undefined, error)
-
+      // ctx.result is set inside finish(); if absent, finish never ran (sync throw in
+      // success branch, before-model block, or openai error) — record the error now.
+      // If finish already ran successfully (after-model block), don't double-publish.
+      if (!ctx.result) finish(ctx, undefined, error)
       throw error
     })
 }

package/packages/datadog-instrumentations/src/oracledb.js

@@ -22,7 +22,7 @@ function finish (ctx) {
 
 addHook({ name: 'oracledb', versions: ['>=5'], file: 'lib/oracledb.js' }, oracledb => {
   shimmer.wrap(oracledb.Connection.prototype, 'execute', execute => {
-    return function wrappedExecute (dbQuery, ...args) {
+    return function wrappedExecute (dbQuery) {
       if (!startChannel.hasSubscribers) {
         return execute.apply(this, arguments)
       }
@@ -72,6 +72,11 @@ addHook({ name: 'oracledb', versions: ['>=5'], file: 'lib/oracledb.js' }, oracle
       }
 
       return startChannel.runStores(ctx, () => {
+        // bindStart is skipped when tracing is suppressed (legacy store is `noop`),
+        // leaving ctx.injected unset — do not overwrite the caller's SQL argument.
+        if (ctx.injected !== undefined) {
+          arguments[0] = ctx.injected
+        }
         try {
           let result = execute.apply(this, arguments)
 

package/packages/datadog-instrumentations/src/pg.js

@@ -193,7 +193,7 @@ function wrapPoolQuery (query) {
       }
 
       if (typeof cb === 'function') {
-        args[args.length - 1] = shimmer.wrapFunction(cb, cb => function (...args) {
+        args[args.length - 1] = shimmer.wrapCallback(cb, cb => function (...args) {
           finish(ctx)
           return cb.apply(this, args)
         })

package/packages/datadog-instrumentations/src/pino.js

@@ -6,7 +6,16 @@ const {
   addHook,
 } = require('./helpers/instrument')
 
+/**
+ * @param {string} symbol
+ * @param {(original: Function) => Function} wrapper
+ * @param {Function} pino
+ */
 function wrapPino (symbol, wrapper, pino) {
+  /**
+   * @param {unknown[]} args
+   * @returns {unknown}
+   */
   return function pinoWithTrace (...args) {
     const instance = pino.apply(this, args)
 
@@ -22,15 +31,18 @@ function wrapPino (symbol, wrapper, pino) {
 }
 
 function wrapAsJson (asJson) {
-  const ch = channel('apm:pino:log')
+  const jsonCh = channel('apm:pino:log:json')
   return function asJsonWithTrace (obj, msg, num, time) {
     obj = arguments[0] = obj || {}
 
-    const payload = { message: obj }
-    ch.publish(payload)
-    arguments[0] = payload.message
+    // Caller-provided `dd` wins -- skip the splice so a bespoke `dd` survives.
+    if (!jsonCh.hasSubscribers || Object.hasOwn(obj, 'dd')) {
+      return asJson.apply(this, arguments)
+    }
 
-    return asJson.apply(this, arguments)
+    const payload = { line: asJson.apply(this, arguments) }
+    jsonCh.publish(payload)
+    return payload.line
   }
 }