dd-trace 5.132.0 → 5.134.0

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

@@ -1,6 +1,9 @@
 'use strict'
 
 const { AsyncResource } = require('node:async_hooks')
+const { realpathSync } = require('node:fs')
+const { resolve, sep } = require('node:path')
+const { fileURLToPath } = require('node:url')
 
 const { createCoverageMap } = require('../../../../vendor/dist/istanbul-lib-coverage')
 const satisfies = require('../../../../vendor/dist/semifies')
@@ -16,6 +19,7 @@ const { EMPTY_EFD_RETRY_POLICY } = require('../../../dd-trace/src/ci-visibility/
 const { writeCoverageBackfillToCache } = require('../../../dd-trace/src/ci-visibility/test-optimization-cache')
 const log = require('../../../dd-trace/src/log')
 const { getEnvironmentVariable } = require('../../../dd-trace/src/config/helper')
+const { getCallSites } = require('../../../dd-trace/src/plugins/util/stacktrace')
 const {
   getTestSuitePath,
   MOCHA_WORKER_TELEMETRY_PAYLOAD_CODE,
@@ -88,6 +92,9 @@ const runnerRecoveryStates = new WeakMap()
 const runnersWithPendingCoverageReset = new WeakSet()
 const parallelRunners = new WeakSet()
 const wrappedRunnerEmitPrototypes = new WeakSet()
+const rootSuitesWithHookFiles = new WeakSet()
+const ROOT_HOOK_NAMES = ['_beforeAll', '_afterAll', '_beforeEach', '_afterEach']
+const ROOT_SUITE_ARRAY_NAMES = ['tests', 'suites', '_onlyTests', '_onlySuites', ...ROOT_HOOK_NAMES]
 let hasWarnedDeprecatedMochaVersion = false
 
 const unskippableSuites = []
@@ -169,6 +176,40 @@ function haveRootTestsFinished (rootTests) {
   return true
 }
 
+/**
+ * @param {{ file: string }} test
+ */
+function startRootSuite (test) {
+  const ctx = testFileToSuiteCtx.get(test.file)
+  if (ctx?._pendingRootStart) {
+    ctx._pendingRootStart = false
+    testSuiteStartCh.runStores(ctx, () => {})
+  }
+}
+
+/**
+ * @param {string} file
+ * @param {'pass' | 'fail' | 'skip'} status
+ * @param {object} [runner] Defers coverage reset until this runner's suite-end reporters finish.
+ */
+function finishRootSuite (file, status, runner) {
+  const ctx = testFileToSuiteCtx.get(file)
+  if (!ctx) {
+    log.warn('No ctx found for suite', file)
+    return
+  }
+
+  if (global.__coverage__) {
+    const coverageFiles = getCoveredFilesFromCoverage(global.__coverage__)
+    testSuiteCodeCoverageCh.publish({ coverageFiles, suiteFile: file })
+    mergeCoverage(global.__coverage__, originalCoverageMap)
+    if (runner) runnersWithPendingCoverageReset.add(runner)
+    else resetCoverage(global.__coverage__)
+  }
+
+  testSuiteFinishCh.publish({ status, ...ctx.currentStore }, () => {})
+}
+
 function getSuitePath (suite) {
   return getTestSuitePath(suite.file, process.cwd())
 }
@@ -931,6 +972,23 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
       return
     }
 
+    if (runner._opts?.cleanReferencesAfterRun === false) {
+      const rootSuite = runner.suite
+      const originalArrays = ROOT_SUITE_ARRAY_NAMES.map(name => rootSuite[name])
+      // Restore definitions before completion callbacks can start another run, including with the plugin disabled.
+      runner.prependOnceListener('end', () => {
+        for (let index = 0; index < ROOT_SUITE_ARRAY_NAMES.length; index++) {
+          rootSuite[ROOT_SUITE_ARRAY_NAMES[index]] = originalArrays[index]
+        }
+      })
+    }
+
+    // Root hooks must be selected before .only() removes other files from the test tree.
+    for (const name of ROOT_HOOK_NAMES) {
+      runner.suite[name] = runner.suite[name].filter(hook => !hook.file ||
+        !suitesToSkip.includes(getSuitePath(hook)) || unskippableSuites.includes(hook.file))
+    }
+
     // We remove the suites that we skip through ITR
     // Mocha normally applies exclusivity after this asynchronous configuration step.
     if (typeof runner.suite.hasOnly === 'function') {
@@ -940,17 +998,28 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     }
     const numTestsToRun = runner.grepTotal(runner.suite)
     const filteredSuites = getFilteredSuites(runner.suite.suites)
+    // Root tests share a Mocha suite across files, but TIA still selects them by file.
+    const filteredRootTests = getFilteredSuites(runner.suite.tests)
     const { suitesToRun, suitesToSkipForRun } = filteredSuites
 
-    isSuitesSkipped = suitesToRun.length !== runner.suite.suites.length
+    isSuitesSkipped = suitesToRun.length !== runner.suite.suites.length ||
+      filteredRootTests.suitesToRun.length !== runner.suite.tests.length
 
     log.debug('%d out of %d suites are going to run.', suitesToRun.length, runner.suite.suites.length)
 
     runner.suite.suites = suitesToRun
+    runner.suite.tests = filteredRootTests.suitesToRun
+    // Mocha reapplies exclusivity when the delayed runner resumes.
+    if (runner.suite._onlyTests.length) {
+      runner.suite._onlyTests = runner.suite.tests
+    }
+    if (runner.suite._onlySuites.length) {
+      runner.suite._onlySuites = runner.suite._onlySuites.filter(suite => suitesToRun.includes(suite))
+    }
     areAllSuitesSkipped = isSuitesSkipped && numTestsToRun > 0 && runner.grepTotal(runner.suite) === 0
 
-    skippedSuites = [...filteredSuites.skippedSuites]
-    suitesToSkip = suitesToSkipForRun
+    skippedSuites = [...new Set([...filteredSuites.skippedSuites, ...filteredRootTests.skippedSuites])]
+    suitesToSkip = [...new Set([...suitesToSkipForRun, ...filteredRootTests.suitesToSkipForRun])]
     skippedSuitesCoverage = getSkippedSuitesCoverageForRun()
     writeCoverageBackfillToCache(skippedSuitesCoverage, getCoverageRootDir())
 
@@ -1095,11 +1164,51 @@ function runMochaRunner (run, runner, args) {
 /**
  * @param {Function} Mocha
  * @param {string} frameworkVersion
+ * @param {import('./common').MochaHookMeta} hookMeta
  * @returns {Function}
  */
-function wrapMochaRun (Mocha, frameworkVersion) {
+function wrapMochaRun (Mocha, frameworkVersion, hookMeta) {
   warnDeprecatedMochaVersion(frameworkVersion)
 
+  // Shimmer is required to attach listeners to the root suite created by each Mocha instance.
+  shimmer.wrap(Mocha.prototype, 'ui', ui => function (...args) {
+    const result = ui.apply(this, args)
+    const suite = this.suite
+    if (!rootSuitesWithHookFiles.has(suite)) {
+      rootSuitesWithHookFiles.add(suite)
+      let file
+      suite.prependListener('pre-require', (context, testFile) => {
+        file = testFile
+      })
+      suite.on('post-require', () => {
+        file = undefined
+      })
+      // Hook objects are created at runtime; inspect the caller while registration is still on the stack.
+      shimmer.wrap(suite, '_createHook', createHook => function createHookWithFile (...args) {
+        const hook = createHook.apply(this, args)
+        if (file && !hook.file) {
+          try {
+            const { moduleBaseDir } = hookMeta
+            const mochaPathPrefix = moduleBaseDir && moduleBaseDir + sep
+            for (const site of getCallSites(createHookWithFile)) {
+              const source = site.getFileName()
+              if (!source || source.startsWith('node:')) continue
+              const sourceFile = source.startsWith('file:') ? fileURLToPath(source) : source
+              if (mochaPathPrefix && sourceFile.startsWith(mochaPathPrefix)) continue
+              // Helpers register shared setup; only direct declarations belong to the loading test file.
+              if (sourceFile === file || realpathSync(sourceFile) === realpathSync(file)) hook.file = file
+              break
+            }
+          } catch (error) {
+            log.debug('Error capturing Mocha root hook source: %s', error)
+          }
+        }
+        return hook
+      })
+    }
+    return result
+  })
+
   // Shimmer is required because run must return its Runner while execution is paused and resumed after configuration.
   shimmer.wrap(Mocha.prototype, 'run', run => function (...args) {
     // Workers do not need to request any data, just run the tests
@@ -1123,7 +1232,7 @@ function wrapMochaRun (Mocha, frameworkVersion) {
     this.files.forEach((path) => {
       const isUnskippable = isMarkedAsUnskippable({ path })
       if (isUnskippable) {
-        unskippableSuites.push(path)
+        unskippableSuites.push(resolve(path))
       }
     })
 
@@ -1200,8 +1309,7 @@ addHook({
   filePattern: String.raw`lib/runner\.(?:c?js)$`,
 }, function (runnerPackage, frameworkVersion) {
   const Runner = runnerPackage.Runner ?? runnerPackage.default ?? runnerPackage
-  if (typeof Runner !== 'function') return
-  if (patched.has(Runner)) return
+  if (typeof Runner !== 'function' || patched.has(Runner)) return
 
   patched.add(Runner)
   wrapRunnerEmit(Runner)
@@ -1226,12 +1334,12 @@ addHook({
       runnerFrameworkErrors.delete(this)
     }
 
-    const { suitesByTestFile, numSuitesByTestFile } = getSuitesByTestFile(this.suite)
+    let { suitesByTestFile, numSuitesByTestFile } = getSuitesByTestFile(this.suite)
     // Root-level tests (direct children of root, no describe wrapper) keyed by file.
     // Populated during the root 'suite' event so the normal finish path can include them
     // in mixed-file status calculation.
     const rootTestsByFile = new Map()
-    // Counts how many original tests per pure-root file still need their final attempt.
+    // Counts how many selected tests per pure-root file still need their final attempt.
     // Hits zero when the last test's lifecycle completes, triggering the suite finish.
     const rootPendingCountByFile = new Map()
     const rootFinalizationPendingCountByFile = new Map()
@@ -1304,33 +1412,18 @@ addHook({
       }
       rootPendingCountByFile.delete(file)
 
-      const ctx = testFileToSuiteCtx.get(file)
-      if (!ctx) {
-        log.warn('No ctx found for suite', file)
-        return
-      }
-
       const rootTests = rootTestsByFile.get(file) || []
-      const status = getRootSuiteStatus(rootTests)
-
-      if (global.__coverage__) {
-        const coverageFiles = getCoveredFilesFromCoverage(global.__coverage__)
-        testSuiteCodeCoverageCh.publish({ coverageFiles, suiteFile: file })
-        mergeCoverage(global.__coverage__, originalCoverageMap)
-        resetCoverage(global.__coverage__)
-      }
-
-      testSuiteFinishCh.publish({ status, ...ctx.currentStore }, () => {})
+      finishRootSuite(file, getRootSuiteStatus(rootTests))
     }
 
-    function finishRootSuiteFallbackForFile (file) {
+    function finishRootSuiteFallbackForFile (file, runner) {
       const ctx = testFileToSuiteCtx.get(file)
       if (!ctx || !rootPendingCountByFile.has(file)) return
 
       const rootTests = rootTestsByFile.get(file) || []
       const status = haveRootTestsFinished(rootTests) ? getRootSuiteStatus(rootTests) : 'fail'
       rootPendingCountByFile.delete(file)
-      testSuiteFinishCh.publish({ status, ...ctx.currentStore }, () => {})
+      finishRootSuite(file, status, runner)
     }
 
     function finishRootSuiteAfterFinalAttempt (test) {
@@ -1361,13 +1454,7 @@ addHook({
     this.prependListener('test', getOnTestHandler(true))
     // Reporters are registered before Runner#run, so prepend this after the test handler
     // to start the suite first and preserve the test event if a reporter throws.
-    this.prependListener('test', function (test) {
-      const ctx = testFileToSuiteCtx.get(test.file)
-      if (ctx?._pendingRootStart) {
-        ctx._pendingRootStart = false
-        testSuiteStartCh.runStores(ctx, () => {})
-      }
-    })
+    this.prependListener('test', startRootSuite)
 
     // Reporters are registered before this run wrapper. Prepend terminal handlers
     // so a reporter error cannot strand a test span that Mocha already completed.
@@ -1407,8 +1494,16 @@ addHook({
     this.prependListener('fail', getOnFailHandler(true, config))
 
     this.prependListener('pending', getOnPendingHandler())
+    // Static skips emit 'pending' without 'test', so their suite must start here too.
+    this.prependListener('pending', startRootSuite)
 
     this.prependListener('suite', function (suite) {
+      if (suite.root) {
+        // TIA and .only() have now selected the suites that can finish during this run.
+        const selectedSuites = getSuitesByTestFile(suite)
+        suitesByTestFile = selectedSuites.suitesByTestFile
+        numSuitesByTestFile = selectedSuites.numSuitesByTestFile
+      }
       if (suite.root || !suite.tests.length) {
         // This branch can be triggered when we have top level it(...) inside test files.
         // In that case, they all (even if they are from different files) are going to be
@@ -1418,9 +1513,14 @@ addHook({
         // below. Since 'suite' fires for root first, in the mixed case the ctx is created
         // here and the describe-based handler finds it already set.
         if (suite.root && suite.tests.length > 0) {
+          const { _grep: grep, _invert: invert } = this
+          const lastIndex = grep.lastIndex
+          const selectedTests = suite.tests.filter(test => grep.test(test.fullTitle()) !== !!invert)
+          // Inspecting a stateful grep must not change which tests Mocha subsequently selects.
+          if (grep.lastIndex !== lastIndex) grep.lastIndex = lastIndex
           const files = new Set(suite.tests.map(test => test.file).filter(Boolean))
           for (const file of files) {
-            const testsForFile = suite.tests.filter(t => t.file === file)
+            const testsForFile = selectedTests.filter(t => t.file === file)
             rootTestsByFile.set(file, testsForFile)
             // Only track the countdown for pure root-level files.
             // Mixed files are finished by the normal 'suite end' path.
@@ -1461,10 +1561,10 @@ addHook({
     this.prependListener('suite end', function (suite) {
       if (suite.root) {
         // Normal case: pure root-level files are finished by the 'test end' / 'hook end'
-        // listeners via finishRootSuiteForFile. Two edge cases remain here:
+        // listeners via finishRootSuiteForFile. Fallback cases remain here:
         //
-        // 1. All-pending: no 'test' event fired, _pendingRootStart is still true.
-        //    Start and immediately finish with 'skip'.
+        // 1. Unstarted file: no 'test' or 'pending' event fired (e.g. all tests
+        //    were filtered out). Start and immediately finish with 'skip'.
         //
         // 2. Aborted mid-run (e.g. a beforeEach hook failure): Mocha skips remaining
         //    tests and jumps straight to 'suite end'. rootPendingCountByFile still has
@@ -1491,7 +1591,7 @@ addHook({
               continue
             }
 
-            finishRootSuiteFallbackForFile(test.file)
+            finishRootSuiteFallbackForFile(test.file, this)
           }
         }
         return
@@ -1584,10 +1684,7 @@ addHook({
   file: 'src/WorkerHandler.js',
 }, (workerHandlerPackage) => {
   shimmer.wrap(workerHandlerPackage.prototype, 'exec', exec => function (_, path) {
-    if (!testFinishCh.hasSubscribers) {
-      return exec.apply(this, arguments)
-    }
-    if (!path?.length) {
+    if (!testFinishCh.hasSubscribers || !path?.length) {
       return exec.apply(this, arguments)
     }
     const [testSuiteAbsolutePath] = path

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

@@ -43,6 +43,9 @@ const originalEfdFns = new WeakMap()
 const originalFns = new WeakMap()
 const originalPendingByTest = new WeakMap()
 const originalRetriesByTest = new WeakMap()
+let automaticRetryTests = new WeakSet()
+let automaticRetryAttempts = new WeakMap()
+let nativeRetryTests = new WeakSet()
 const testToStartLine = new WeakMap()
 const testFileToSuiteCtx = new Map()
 const datadogRetryOriginals = new WeakMap()
@@ -346,6 +349,9 @@ function resetRunState (rootSuite) {
   efdRetryCountByTestFullName.clear()
   efdSlowAbortedTests.clear()
   suiteHookFailureSuites.clear()
+  automaticRetryTests = new WeakSet()
+  automaticRetryAttempts = new WeakMap()
+  nativeRetryTests = new WeakSet()
   dynamicAtrRetryCountByTest = new WeakMap()
   attemptToFixExecutions.clear()
   loggedAttemptToFixTests.clear()
@@ -423,7 +429,7 @@ function isMochaRetry (test) {
 }
 
 function getIsLastRetry (test) {
-  return test._currentRetry === test._retries
+  return test._currentRetry >= test._retries
 }
 
 function getTestFullName (test) {
@@ -534,6 +540,18 @@ function inheritDatadogPropertiesFromRetriedTest (test) {
 
 function runnableWrapper (RunnablePackage, libraryConfig) {
   const Runnable = RunnablePackage.Runnable ?? RunnablePackage
+  // Mocha 12 can load its ESM class through require() without the ESM rewriter.
+  // Intercept the exported prototype to observe setters in both injection modes.
+  shimmer.wrap(Runnable.prototype, 'retries', retries => function () {
+    if (arguments.length && automaticRetryTests.delete(this)) {
+      const originalTest = this._retriedTest || this
+      // afterEach can change this runnable after Mocha has already copied its budget into a queued clone.
+      if (!automaticRetryAttempts.get(originalTest)?.has(this._currentRetry + 1)) {
+        nativeRetryTests.add(originalTest)
+      }
+    }
+    return retries.apply(this, arguments)
+  })
   shimmer.wrap(Runnable.prototype, 'run', run => function (...args) {
     if (!testFinishCh.hasSubscribers) {
       return run.apply(this, args)
@@ -559,11 +577,12 @@ function runnableWrapper (RunnablePackage, libraryConfig) {
           return onRunnableFinished.apply(this, arguments)
         }
       }
-    } else if (libraryConfig?.isFlakyTestRetriesEnabled) {
+    } else if (libraryConfig?.isFlakyTestRetriesEnabled && !nativeRetryTests.has(test._retriedTest || test)) {
       if (!originalRetriesByTest.has(test)) {
         const originalRetries = originalRetriesByTest.get(test._retriedTest) ?? test.retries()
         originalRetriesByTest.set(test, originalRetries)
       }
+      automaticRetryTests.delete(test)
       if (libraryConfig.isDynamicAtrEnabled) {
         // Dynamic ATR: set the max possible retries initially.
         // The actual duration-based count is computed after the first attempt.
@@ -581,6 +600,7 @@ function runnableWrapper (RunnablePackage, libraryConfig) {
       } else {
         test.retries(libraryConfig.flakyTestRetriesCount)
       }
+      automaticRetryTests.add(test)
     }
 
     if (isTestHook || this.type === 'test') {
@@ -833,6 +853,7 @@ function getTestFinishInfo (test, status, config, error) {
 
   const isAttemptToFixRetry = test._ddIsAttemptToFix && testStatuses.length > 1
   const isAtrRetry = config.isFlakyTestRetriesEnabled &&
+    automaticRetryAttempts.get(test._retriedTest || test)?.has(test._currentRetry) &&
     !test._ddIsAttemptToFix &&
     !test._ddIsEfdRetry
   const isFinalAttempt = isTerminalHookFailure || status !== 'fail' || test._currentRetry >= test._retries
@@ -899,6 +920,7 @@ function getOnTestEndHandler (config, finalAttemptHandlers) {
     const shouldFinishTest = ctx && (
       test._ddReporterStartFailed ||
       test._ddReporterTerminalFailed ||
+      !test.fn ||
       !getAfterEachHooks(test).length ||
       (test._ddIsDisabled && !test._ddIsAttemptToFix)
     )
@@ -906,8 +928,8 @@ function getOnTestEndHandler (config, finalAttemptHandlers) {
     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.
+    // If there are afterEach hooks to run, we don't finish the test yet.
+    // Static skips have no function and never execute hooks. Disabled tests also finish immediately.
     // 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
@@ -1177,9 +1199,20 @@ function getOnFailHandler (isMain, config) {
 
 function getOnTestRetryHandler (config) {
   return function (test, err) {
+    if (automaticRetryTests.has(test)) {
+      // Mocha has already queued the clone; later budget changes cannot change its scheduling owner.
+      const originalTest = test._retriedTest || test
+      let attempts = automaticRetryAttempts.get(originalTest)
+      if (!attempts) {
+        attempts = new Set()
+        automaticRetryAttempts.set(originalTest, attempts)
+      }
+      attempts.add(test._currentRetry + 1)
+    }
     const isFirstAttempt = test._currentRetry === 0
     const isDynamicAtrTest = config.isDynamicAtrEnabled &&
       config.isFlakyTestRetriesEnabled &&
+      automaticRetryTests.has(test) &&
       !test._ddIsAttemptToFix &&
       !isEarlyFlakeDetectionTest(test, config)
     if (isDynamicAtrTest && isFirstAttempt) {
@@ -1216,6 +1249,7 @@ function publishTestRetry (test, err, config) {
     const willBeRetried = test._currentRetry < test._retries
     const isAtrRetry = !isFirstAttempt &&
       config.isFlakyTestRetriesEnabled &&
+      automaticRetryAttempts.get(test._retriedTest || test)?.has(test._currentRetry) &&
       !test._ddIsAttemptToFix &&
       !test._ddIsEfdRetry
     const promises = {}

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

@@ -443,8 +443,7 @@ addHook({
   filePattern: String.raw`lib/runner\.(?:c?js)$`,
 }, function (runnerPackage) {
   const Runner = runnerPackage.Runner ?? runnerPackage.default ?? runnerPackage
-  if (typeof Runner !== 'function') return
-  if (patchedRunners.has(Runner)) return
+  if (typeof Runner !== 'function' || patchedRunners.has(Runner)) return
 
   patchedRunners.add(Runner)
   shimmer.wrap(Runner.prototype, 'runTests', runTests => getRunTestsWrapper(runTests, config))

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

@@ -45,7 +45,7 @@ addHook({ name: 'mongodb', versions: ['>=4.6.0 <6.4.0'], file: 'lib/cmap/connect
 
 addHook({ name: 'mongodb', versions: ['>=6.4.0'], file: 'lib/cmap/connection.js' }, Connection => {
   const proto = Connection.Connection.prototype
-  shimmer.wrap(proto, 'command', command => wrapConnectionCommand(command, 'command', undefined, instrumentPromise))
+  shimmer.wrap(proto, 'command', command => wrapConnectionCommand(command, 'command', instrumentPromise))
 })
 
 addHook({ name: 'mongodb', versions: ['>=3.3 <4'], file: 'lib/core/wireprotocol/index.js' }, wp => wrapWp(wp))
@@ -95,8 +95,8 @@ function wrapUnifiedCommand (command, operation, name) {
   }
 }
 
-function wrapConnectionCommand (command, operation, name, instrumentFn = instrument) {
-  const opts = { name }
+function wrapConnectionCommand (command, operation, instrumentFn = instrument) {
+  const opts = { name: undefined }
   return function (ns, ops) {
     if (!startCh.hasSubscribers) {
       return command.apply(this, arguments)

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

@@ -83,7 +83,7 @@ addHook({
 
         let callbackWrapped = false
 
-        const wrapCallbackIfExist = (args, ctx) => {
+        const wrapCallbackIfExist = (args) => {
           const lastArgumentIndex = args.length - 1
 
           if (typeof args[lastArgumentIndex] === 'function') {
@@ -106,7 +106,7 @@ addHook({
         }
 
         return startCh.runStores(ctx, () => {
-          wrapCallbackIfExist(args, ctx)
+          wrapCallbackIfExist(args)
 
           // A callback query executes synchronously inside `method.apply`, so the
           // exec scope (and the deferred driver call it spawns) must be entered
@@ -121,7 +121,7 @@ addHook({
             shimmer.wrap(res, 'exec', originalExec => {
               return function wrappedExec (...args) {
                 if (!callbackWrapped) {
-                  wrapCallbackIfExist(args, ctx)
+                  wrapCallbackIfExist(args)
                 }
 
                 const execResult = execCh.runStores(ctx, () => originalExec.apply(this, args))

package/packages/datadog-instrumentations/src/nova-sonic/audio.js

@@ -0,0 +1,173 @@
+'use strict'
+
+const { LLMOBS_AUDIO_ACCUMULATE_MAX_BYTES } = require('../../../dd-trace/src/llmobs/constants/audio')
+
+// Leave room for the WAV header when the retained PCM is encoded.
+const MAX_AUDIO = LLMOBS_AUDIO_ACCUMULATE_MAX_BYTES - 128
+const RATES = new Set([8000, 16_000, 24_000, 48_000])
+const INVALID_BASE64 = /[^A-Za-z0-9+/]/
+
+/** @param {object} configuration */
+function pcmRate (configuration = {}) {
+  return configuration.mediaType === 'audio/lpcm' && configuration.sampleSizeBits === 16 &&
+    configuration.channelCount === 1 && configuration.encoding === 'base64' &&
+    RATES.has(configuration.sampleRateHertz)
+    ? configuration.sampleRateHertz
+    : 0
+}
+
+/** @param {string} encoded */
+function decodedLength (encoded) {
+  if (typeof encoded !== 'string' || encoded.length % 4) return
+  const padding = encoded.endsWith('==') ? 2 : encoded.endsWith('=') ? 1 : 0
+  if (INVALID_BASE64.test(encoded.slice(0, encoded.length - padding))) return
+  return encoded.length / 4 * 3 - padding
+}
+
+/** A bounded, connection-wide sample timeline. Content-name changes do not reset its origin. */
+class InputAudio {
+  rate = 0
+  valid = true
+  total = 0
+  anchor
+  #chunks = []
+  #bytes = 0
+
+  /** @param {object} configuration */
+  configure (configuration) {
+    const rate = pcmRate(configuration)
+    if (!rate || (this.total && rate !== this.rate)) this.valid = false
+    this.rate = rate
+  }
+
+  /**
+   * @param {string} encoded
+   * @param {number} now Milliseconds at SDK consumption, not an acknowledgement from the provider.
+   */
+  append (encoded, now) {
+    const length = decodedLength(encoded)
+    if (length === undefined) {
+      this.valid = false
+      this.clear()
+      return
+    }
+    if (length % 2) this.valid = false
+    this.total += length
+    if (this.rate) this.anchor = now - this.total * 1000 / (2 * this.rate)
+    if (!this.valid || !length) return
+    // Decode from a base64 quartet boundary, retaining at most two alignment bytes
+    // beyond the cap instead of allocating the entire incoming frame.
+    const start = Math.max(0, length - MAX_AUDIO)
+    const retained = Buffer.from(encoded.slice(Math.floor(start / 3) * 4), 'base64').subarray(start % 3)
+    this.#chunks.push(retained)
+    this.#bytes += retained.length
+    while (this.#bytes > MAX_AUDIO) {
+      const first = this.#chunks[0]
+      const excess = this.#bytes - MAX_AUDIO
+      if (first.length <= excess) {
+        this.#chunks.shift()
+        this.#bytes -= first.length
+      } else {
+        this.#chunks[0] = Buffer.from(first.subarray(excess))
+        this.#bytes -= excess
+      }
+    }
+  }
+
+  /**
+   * @param {number} startMs
+   * @param {number} endMs
+   */
+  clip (startMs, endMs) {
+    if (!this.valid || !this.rate || !Number.isFinite(startMs) || !Number.isFinite(endMs) || startMs >= endMs) return
+    const start = Math.floor(startMs * this.rate / 1000) * 2
+    const end = Math.floor(endMs * this.rate / 1000) * 2
+    const base = this.total - this.#bytes
+    if (start < base || end > this.total) return
+    const clip = Buffer.allocUnsafe(end - start)
+    let offset = base
+    let written = 0
+    for (const chunk of this.#chunks) {
+      const chunkEnd = offset + chunk.length
+      if (chunkEnd > start) {
+        written += chunk.copy(clip, written, Math.max(start - offset, 0), Math.min(end - offset, chunk.length))
+      }
+      offset = chunkEnd
+      if (offset >= end) break
+    }
+    return clip
+  }
+
+  clear () {
+    this.#chunks = []
+    this.#bytes = 0
+  }
+}
+
+/** Ordered PCM playback projection for one response, including genuine underrun silence. */
+class OutputAudio {
+  rate = 0
+  start
+  end
+  interrupted
+  generatedBytes = 0
+  timingValid = true
+  omittedReason
+  #chunks = []
+  #bytes = 0
+
+  /**
+   * @param {string} encoded
+   * @param {object} configuration
+   * @param {number} now
+   */
+  append (encoded, configuration, now) {
+    const length = decodedLength(encoded)
+    if (length !== undefined) this.generatedBytes += length
+    // Late chunks cannot resurrect playback or invalidate the clip already cut at interruption.
+    if (this.interrupted !== undefined) return
+    const rate = pcmRate(configuration)
+    if (length === undefined || !rate || length % 2 || (this.rate && this.rate !== rate)) {
+      this.timingValid = false
+      this.omittedReason = 'invalid_audio'
+      this.#chunks = []
+      this.#bytes = 0
+      return
+    }
+    if (!length) return
+    this.rate ||= rate
+    this.start ??= now
+    const previousEnd = this.end ?? now
+    const gap = Math.floor(Math.max(0, now - previousEnd) * rate / 1000) * 2
+    this.end = previousEnd + (gap + length) * 1000 / (2 * rate)
+    if (!this.omittedReason && this.#bytes + gap + length <= MAX_AUDIO) {
+      if (gap) this.#chunks.push(Buffer.alloc(gap))
+      this.#chunks.push(Buffer.from(encoded, 'base64'))
+      this.#bytes += gap + length
+    } else {
+      this.omittedReason ||= 'retention_limit'
+      this.#chunks = []
+      this.#bytes = 0
+    }
+  }
+
+  /** @param {number} now */
+  interrupt (now) {
+    if (this.interrupted !== undefined) return
+    this.interrupted = now
+    if (this.start === undefined || this.end === undefined || !this.rate) return
+    const frames = Math.floor(Math.max(0, now - this.start) * this.rate / 1000)
+    this.end = Math.min(this.end, this.start + frames * 1000 / this.rate)
+    if (this.#bytes > frames * 2) {
+      const clip = Buffer.concat(this.#chunks, this.#bytes).subarray(0, frames * 2)
+      this.#chunks = [Buffer.from(clip)]
+      this.#bytes = clip.length
+    }
+  }
+
+  pcm () {
+    if (!this.omittedReason && this.#bytes) return Buffer.concat(this.#chunks, this.#bytes)
+  }
+}
+
+module.exports = { InputAudio, OutputAudio }