dd-trace 6.7.0 → 6.9.0

package/packages/datadog-instrumentations/src/grpc/client.js

@@ -13,6 +13,8 @@ const errorChannel = channel('apm:grpc:client:request:error')
 const finishChannel = channel('apm:grpc:client:request:finish')
 const emitChannel = channel('apm:grpc:client:request:emit')
 
+/** @typedef {{ length: number, [index: number]: unknown } & Iterable<unknown>} ArgumentsLike */
+
 function createWrapMakeRequest (type, hasPeer = false) {
   const metadataIndex = type === types.client_stream || type === types.bidi ? 3 : 4
 
@@ -186,9 +188,9 @@ function callMethod (client, method, args, path, metadata, type, hasPeer = false
  * one at `index` when the user did not pass their own.
  *
  * @param {object} client
- * @param {ArrayLike<unknown>} args
+ * @param {ArgumentsLike} args
  * @param {number} index
- * @returns {{ metadata: object | undefined, args: ArrayLike<unknown> }}
+ * @returns {{ metadata: object | undefined, args: ArgumentsLike }}
  */
 function resolveMetadata (client, args, index) {
   const grpc = client && getGrpc(client)

package/packages/datadog-instrumentations/src/helpers/callback-instrumentor.js

@@ -3,6 +3,8 @@
 const shimmer = require('../../../datadog-shimmer')
 const { channel } = require('./instrument')
 
+/** @typedef {{ length: number, [index: number]: unknown } & Iterable<unknown>} ArgumentsLike */
+
 /**
  * Create a shimmer-compatible instrumentor for callback-style APIs whose work is offloaded to the
  * libuv worker thread pool (e.g. zlib.gzip, crypto.pbkdf2, dns.lookup). Builds a set of three
@@ -22,7 +24,7 @@ const { channel } = require('./instrument')
  * @param {boolean} [options.captureResult] set `ctx.result` to the callback's first
  *   non-error argument before publishing `:finish`. Plugins that tag spans from the call's
  *   return value (e.g. the DNS lookup plugin) rely on this.
- * @returns {(buildContext: (thisArg: unknown, args: IArguments) => object | undefined) =>
+ * @returns {(buildContext: (thisArg: unknown, args: ArgumentsLike) => object | undefined) =>
  *   (fn: Function) => Function}
  */
 function createCallbackInstrumentor (prefix, { captureResult = false } = {}) {

package/packages/datadog-instrumentations/src/helpers/extract-package-and-module-path.js

@@ -2,11 +2,21 @@
 
 const NM = 'node_modules/'
 
+/**
+ * @typedef {object} PackageAndModulePath
+ * @property {string|null} pkg
+ * @property {string|null} path
+ * @property {string} [pkgJson]
+ */
+
 /**
  * For a given full path to a module,
  *   return the package name it belongs to and the local path to the module
  *   input: '/foo/node_modules/@co/stuff/foo/bar/baz.js'
  *   output: { pkg: '@co/stuff', path: 'foo/bar/baz.js',  pkgJson: '/foo/node_modules/@co/stuff/package.json' }
+ *
+ * @param {string} fullPath
+ * @returns {PackageAndModulePath}
  */
 module.exports = function extractPackageAndModulePath (fullPath) {
   const nm = fullPath.lastIndexOf(NM)

package/packages/datadog-instrumentations/src/helpers/hook.js

@@ -6,6 +6,7 @@ const iitm = require('../../../dd-trace/src/iitm')
 const ritm = require('../../../dd-trace/src/ritm')
 const log = require('../../../dd-trace/src/log')
 const requirePackageJson = require('../../../dd-trace/src/require-package-json')
+const { isNodeBuiltinModuleName } = require('./shared-utils')
 
 /**
  * @param {string} moduleBaseDir
@@ -82,7 +83,7 @@ function Hook (modules, hookOptions, onrequire) {
       moduleVersion ||= getVersion(moduleBaseDir)
     } catch (error) {
       log.error('Error getting version for "%s": %s', moduleName, error.message, error)
-      return
+      return moduleExports
     }
 
     if (
@@ -101,7 +102,23 @@ function Hook (modules, hookOptions, onrequire) {
       }
     }
 
-    const newExports = wrappedOnrequire(moduleExports, moduleName, moduleBaseDir, moduleVersion, isIitm)
+    // A builtin's named ESM exports are a second view of what `default` already points at, and
+    // `patched` cannot relate them because it keys on the object the hook ran against.
+    let newExports
+    if (defaultWrapResult && isNodeBuiltinModuleName(moduleName)) {
+      // Accessor exports are read, not skipped: `replaceGetter` leaves Node 20's `fs.opendir`
+      // an accessor, and Node's ESM facade resolved every export before handing us this view.
+      for (const key of Object.keys(defaultWrapResult)) {
+        const wrapped = defaultWrapResult[key]
+        const existing = moduleExports[key]
+        if (existing !== wrapped && existing !== undefined) {
+          moduleExports[key] = wrapped
+        }
+      }
+      newExports = moduleExports
+    } else {
+      newExports = wrappedOnrequire(moduleExports, moduleName, moduleBaseDir, moduleVersion, isIitm)
+    }
 
     if (defaultWrapResult && defaultExportAliases) {
       newExports.default = defaultWrapResult

package/packages/datadog-instrumentations/src/helpers/hooks.js

@@ -58,6 +58,10 @@ module.exports = {
   '@smithy/core': () => require('../aws-sdk'),
   '@smithy/smithy-client': () => require('../aws-sdk'),
   '@vitest/runner': { esmFirst: true, fn: () => require('../vitest') },
+  '@wdio/cli': { esmFirst: true, fn: () => require('../webdriverio') },
+  '@wdio/jasmine-framework': { esmFirst: true, fn: () => require('../webdriverio') },
+  '@wdio/local-runner': { esmFirst: true, fn: () => require('../webdriverio') },
+  '@wdio/utils': { esmFirst: true, fn: () => require('../webdriverio') },
   aerospike: () => require('../aerospike'),
   ai: () => require('../ai'),
   amqp10: () => require('../amqp10'),

package/packages/datadog-instrumentations/src/helpers/register.js

@@ -57,19 +57,23 @@ const alreadyLoggedIncompatibleIntegrations = new Set()
 // Always disable prefixed and unprefixed node modules if one is disabled.
 if (disabledInstrumentations.size) {
   const builtinsSet = new Set(builtinModules)
+  const disabledBuiltinCounterparts = []
   for (const name of disabledInstrumentations) {
     const hasPrefix = name.startsWith('node:')
     if (hasPrefix || builtinsSet.has(name)) {
       if (hasPrefix) {
         const unprefixedName = name.slice(5)
         if (!disabledInstrumentations.has(unprefixedName)) {
-          disabledInstrumentations.add(unprefixedName)
+          disabledBuiltinCounterparts.push(unprefixedName)
         }
       } else if (!disabledInstrumentations.has(`node:${name}`)) {
-        disabledInstrumentations.add(`node:${name}`)
+        disabledBuiltinCounterparts.push(`node:${name}`)
       }
     }
   }
+  for (const name of disabledBuiltinCounterparts) {
+    disabledInstrumentations.add(name)
+  }
   builtinsSet.clear()
 }
 
@@ -163,7 +167,9 @@ function logAbortedIntegrations () {
   for (const [nameVersion, success] of instrumentedIntegrationsSuccess) {
     // Only ever log a single version of an integration, even if it is loaded later.
     if (!success && !alreadyLoggedIncompatibleIntegrations.has(nameVersion)) {
-      const [name, version] = nameVersion.split('@')
+      const lastAtPosition = nameVersion.lastIndexOf('@')
+      const name = nameVersion.slice(0, lastAtPosition)
+      const version = nameVersion.slice(lastAtPosition + 1)
       telemetry('abort.integration', [
         `integration:${name}`,
         `integration_version:${version}`,

package/packages/datadog-instrumentations/src/helpers/rewriter/index.js

@@ -5,8 +5,9 @@ const { join } = require('path')
 const { pathToFileURL } = require('url')
 const log = require('../../../../dd-trace/src/log')
 const { create } = require('../../../../../vendor/dist/@apm-js-collab/code-transformer')
-const { waitForAsyncEnd } = require('./transforms')
 const instrumentations = require('./instrumentations')
+const { getRewriteTarget } = require('./targets')
+const { waitForAsyncEnd } = require('./transforms')
 
 // `dc-polyfill` is referenced from injected `require()` (CJS) and `import`
 // (ESM) statements that the transformer splices into the rewritten module.
@@ -39,20 +40,26 @@ for (const matcher of [matcherCjs, matcherEsm]) {
 
 // Keep the marker split: source-map scanners can read a contiguous token in
 // string literals as this file's own inline map.
+// eslint-disable-next-line unicorn/no-useless-concat -- Keep the source-map marker non-contiguous.
 const SOURCE_MAP_PREFIX = '//# sourceMapping' + 'URL=data:application/json;base64,'
 
-function rewrite (content, filename, format) {
+/**
+ * @param {string|Buffer|ArrayBuffer|Uint8Array} content
+ * @param {string} filename
+ * @param {string} [format]
+ * @param {{ moduleName: string, filePath: string }} [target]
+ * @returns {string|Buffer|ArrayBuffer|Uint8Array}
+ */
+function rewrite (content, filename, format, target) {
   if (!content) return content
-  if (!filename.includes('node_modules')) return content
+
+  target ||= getRewriteTarget(filename)
+  if (!target) return content
 
   filename = filename.replace('file://', '')
 
   const moduleType = format === 'module' ? 'esm' : 'cjs'
-  const [modulePath] = filename.split('/node_modules/').reverse()
-  const moduleParts = modulePath.split('/')
-  const splitIndex = moduleParts[0].startsWith('@') ? 2 : 1
-  const moduleName = moduleParts.slice(0, splitIndex).join('/')
-  const filePath = moduleParts.slice(splitIndex).join('/')
+  const { moduleName, filePath } = target
   const version = getVersion(filename, filePath)
 
   if (disabled.has(moduleName)) return content
@@ -63,8 +70,10 @@ function rewrite (content, filename, format) {
   if (!transformer) return content
 
   try {
+    const source = getSourceText(content)
+
     // TODO: pass existing sourcemap as input for remapping
-    const { code, map } = transformer.transform(content, moduleType)
+    const { code, map } = transformer.transform(source, moduleType)
 
     if (!map) return code
 
@@ -78,6 +87,22 @@ function rewrite (content, filename, format) {
   return content
 }
 
+/** @typedef {{ buffer: ArrayBuffer | SharedArrayBuffer, byteLength: number, byteOffset: number }} BufferView */
+
+/**
+ * Convert the source representations accepted by Node.js loader hooks to text.
+ *
+ * @param {string | ArrayBuffer | BufferView} source
+ * @returns {string}
+ */
+function getSourceText (source) {
+  if (typeof source === 'string') return source
+  if (ArrayBuffer.isView(source)) {
+    return Buffer.from(source.buffer, source.byteOffset, source.byteLength).toString('utf8')
+  }
+  return Buffer.from(source).toString('utf8')
+}
+
 function disable (instrumentation) {
   disabled.add(instrumentation)
 }

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/index.js

@@ -12,5 +12,6 @@ module.exports = [
   ...require('./modelcontextprotocol-sdk'),
   ...require('./openai-agents'),
   ...require('./playwright'),
+  ...require('./webdriverio'),
   ...require('./aws-durable-execution-sdk-js'),
 ]

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/webdriverio.js

@@ -0,0 +1,175 @@
+'use strict'
+
+module.exports = [
+  {
+    module: {
+      name: '@wdio/cli',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'Launcher',
+      methodName: '_startInstance',
+      kind: 'Async',
+    },
+    channelName: 'Launcher_startInstance',
+  },
+  {
+    module: {
+      name: '@wdio/local-runner',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'LocalRunner',
+      methodName: 'run',
+      kind: 'Async',
+    },
+    channelName: 'LocalRunner_run',
+  },
+  {
+    module: {
+      name: '@wdio/local-runner',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'LocalRunner',
+      methodName: 'shutdown',
+      kind: 'Async',
+    },
+    channelName: 'LocalRunner_shutdown',
+  },
+  {
+    module: {
+      name: '@wdio/local-runner',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    astQuery: 'VariableDeclarator[id.name="LocalRunner"] > ClassExpression > ClassBody > ' +
+      'MethodDefinition[key.name="shutdown"] ReturnStatement > ' +
+      'CallExpression[callee.object.name="promise"][callee.property.name="then"], ' +
+      'ClassDeclaration[id.name="LocalRunner"] > ClassBody > MethodDefinition[key.name="shutdown"] ReturnStatement > ' +
+      'CallExpression[callee.object.name="promise"][callee.property.name="then"]',
+    channelName: 'LocalRunner_shutdown',
+    transform: 'waitForAsyncEnd',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'JasmineAdapter',
+      methodName: 'init',
+      kind: 'Async',
+    },
+    channelName: 'JasmineAdapter_init',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    astQuery: 'VariableDeclarator[id.name="JasmineAdapter"] > ClassExpression > ClassBody > ' +
+      'MethodDefinition[key.name="init"] ReturnStatement > ' +
+      'CallExpression[callee.object.name="promise"][callee.property.name="then"], ' +
+      'ClassDeclaration[id.name="JasmineAdapter"] > ClassBody > MethodDefinition[key.name="init"] ReturnStatement > ' +
+      'CallExpression[callee.object.name="promise"][callee.property.name="then"]',
+    channelName: 'JasmineAdapter_init',
+    transform: 'waitForAsyncEnd',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'JasmineAdapter',
+      methodName: 'run',
+      kind: 'Async',
+    },
+    channelName: 'JasmineAdapter_run',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    astQuery: 'VariableDeclarator[id.name="JasmineAdapter"] > ClassExpression > ClassBody > ' +
+      'MethodDefinition[key.name="run"] ReturnStatement > ' +
+      'CallExpression[callee.object.name="promise"][callee.property.name="then"], ' +
+      'ClassDeclaration[id.name="JasmineAdapter"] > ClassBody > MethodDefinition[key.name="run"] ReturnStatement > ' +
+      'CallExpression[callee.object.name="promise"][callee.property.name="then"]',
+    channelName: 'JasmineAdapter_run',
+    transform: 'waitForAsyncEnd',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'JasmineReporter',
+      methodName: 'specDone',
+      kind: 'Sync',
+    },
+    channelName: 'JasmineReporter_specDone',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'JasmineReporter',
+      methodName: 'specStarted',
+      kind: 'Sync',
+    },
+    channelName: 'JasmineReporter_specStarted',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'JasmineReporter',
+      methodName: 'suiteStarted',
+      kind: 'Sync',
+    },
+    channelName: 'JasmineReporter_suiteStarted',
+  },
+  {
+    module: {
+      name: '@wdio/jasmine-framework',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      className: 'JasmineReporter',
+      methodName: 'suiteDone',
+      kind: 'Sync',
+    },
+    channelName: 'JasmineReporter_suiteDone',
+  },
+  {
+    module: {
+      name: '@wdio/utils',
+      versionRange: '>=9.0.0',
+      filePath: 'build/index.js',
+    },
+    functionQuery: {
+      expressionName: 'testFrameworkFnWrapper',
+      kind: 'Async',
+    },
+    channelName: 'testFrameworkFnWrapper',
+  },
+]

package/packages/datadog-instrumentations/src/helpers/rewriter/loader.mjs

@@ -1,4 +1,8 @@
-import { rewrite } from './index.js'
+import * as Module from 'module'
+
+const require = Module.createRequire(import.meta.url)
+const { getRewriteTarget } = require('./targets.js')
+let rewriter
 
 async function load (url, context, nextLoad) {
   const result = await nextLoad(url, context)
@@ -12,6 +16,11 @@ function loadSync (url, context, nextLoad) {
   return rewriteResult(result, url, context)
 }
 
+/**
+ * @param {{ format?: string, source?: unknown }} result
+ * @param {string} url
+ * @param {{ format?: string }} context
+ */
 function rewriteResult (result, url, context) {
   const format = result.format || context.format
 
@@ -19,7 +28,16 @@ function rewriteResult (result, url, context) {
   // double-instruments CommonJS entrypoints loaded through sync hooks.
   if (format === 'commonjs') return result
 
-  result.source = rewrite(result.source, url, format)
+  if (result.source) {
+    const target = getRewriteTarget(url)
+    if (target) {
+      if (!rewriter) {
+        rewriter = require('./index.js')
+      }
+
+      result.source = rewriter.rewrite(result.source, url, format, target)
+    }
+  }
 
   return result
 }

package/packages/datadog-instrumentations/src/helpers/rewriter/targets.js

@@ -0,0 +1,32 @@
+'use strict'
+
+// Generated by scripts/generate-rewriter-targets.js from ./instrumentations. Do not edit by hand.
+const targets = require('./targets.json')
+
+const NODE_MODULES = '/node_modules/'
+
+/**
+ * @typedef {object} RewriteTarget
+ * @property {string} moduleName
+ * @property {string} filePath
+ */
+
+/**
+ * @param {string} filename
+ * @returns {RewriteTarget|undefined}
+ */
+function getRewriteTarget (filename) {
+  const nodeModulesIndex = filename.lastIndexOf(NODE_MODULES)
+  if (nodeModulesIndex === -1) return
+
+  const modulePath = filename.slice(nodeModulesIndex + NODE_MODULES.length)
+  const moduleName = targets[modulePath]
+  if (typeof moduleName !== 'string') return
+
+  return {
+    moduleName,
+    filePath: modulePath.slice(moduleName.length + 1),
+  }
+}
+
+module.exports = { getRewriteTarget }

package/packages/datadog-instrumentations/src/helpers/rewriter/targets.json

@@ -0,0 +1,61 @@
+{
+  "@anthropic-ai/claude-agent-sdk/sdk.mjs": "@anthropic-ai/claude-agent-sdk",
+  "@aws/durable-execution-sdk-js/dist-cjs/index.js": "@aws/durable-execution-sdk-js",
+  "@aws/durable-execution-sdk-js/dist/index.mjs": "@aws/durable-execution-sdk-js",
+  "@azure/cosmos/dist/browser/plugins/Plugin.js": "@azure/cosmos",
+  "@azure/cosmos/dist/commonjs/plugins/Plugin.js": "@azure/cosmos",
+  "@azure/cosmos/dist/esm/plugins/Plugin.js": "@azure/cosmos",
+  "@azure/cosmos/dist/react-native/plugins/Plugin.js": "@azure/cosmos",
+  "@graphql-tools/executor/cjs/execution/execute.js": "@graphql-tools/executor",
+  "@graphql-tools/executor/cjs/execution/normalizedExecutor.js": "@graphql-tools/executor",
+  "@graphql-tools/executor/esm/execution/execute.js": "@graphql-tools/executor",
+  "@graphql-tools/executor/esm/execution/normalizedExecutor.js": "@graphql-tools/executor",
+  "@langchain/core/dist/embeddings.cjs": "@langchain/core",
+  "@langchain/core/dist/embeddings.js": "@langchain/core",
+  "@langchain/core/dist/language_models/chat_models.cjs": "@langchain/core",
+  "@langchain/core/dist/language_models/chat_models.js": "@langchain/core",
+  "@langchain/core/dist/language_models/llms.cjs": "@langchain/core",
+  "@langchain/core/dist/language_models/llms.js": "@langchain/core",
+  "@langchain/core/dist/runnables/base.cjs": "@langchain/core",
+  "@langchain/core/dist/runnables/base.js": "@langchain/core",
+  "@langchain/core/dist/tools/index.cjs": "@langchain/core",
+  "@langchain/core/dist/tools/index.js": "@langchain/core",
+  "@langchain/core/dist/vectorstores.cjs": "@langchain/core",
+  "@langchain/core/dist/vectorstores.js": "@langchain/core",
+  "@langchain/langgraph/dist/pregel/index.cjs": "@langchain/langgraph",
+  "@langchain/langgraph/dist/pregel/index.js": "@langchain/langgraph",
+  "@modelcontextprotocol/sdk/dist/cjs/client/index.js": "@modelcontextprotocol/sdk",
+  "@modelcontextprotocol/sdk/dist/cjs/server/mcp.js": "@modelcontextprotocol/sdk",
+  "@modelcontextprotocol/sdk/dist/esm/client/index.js": "@modelcontextprotocol/sdk",
+  "@modelcontextprotocol/sdk/dist/esm/server/mcp.js": "@modelcontextprotocol/sdk",
+  "@openai/agents-openai/dist/openaiChatCompletionsModel.js": "@openai/agents-openai",
+  "@openai/agents-openai/dist/openaiChatCompletionsModel.mjs": "@openai/agents-openai",
+  "@wdio/cli/build/index.js": "@wdio/cli",
+  "@wdio/jasmine-framework/build/index.js": "@wdio/jasmine-framework",
+  "@wdio/local-runner/build/index.js": "@wdio/local-runner",
+  "@wdio/utils/build/index.js": "@wdio/utils",
+  "ai/dist/index.js": "ai",
+  "ai/dist/index.mjs": "ai",
+  "bullmq/dist/cjs/classes/flow-producer.js": "bullmq",
+  "bullmq/dist/cjs/classes/queue.js": "bullmq",
+  "bullmq/dist/cjs/classes/worker.js": "bullmq",
+  "bullmq/dist/esm/classes/flow-producer.js": "bullmq",
+  "bullmq/dist/esm/classes/queue.js": "bullmq",
+  "bullmq/dist/esm/classes/worker.js": "bullmq",
+  "graphql/execution/execute.js": "graphql",
+  "graphql/execution/execute.mjs": "graphql",
+  "graphql/language/parser.js": "graphql",
+  "graphql/language/parser.mjs": "graphql",
+  "graphql/language/printer.js": "graphql",
+  "graphql/language/printer.mjs": "graphql",
+  "graphql/language/visitor.js": "graphql",
+  "graphql/language/visitor.mjs": "graphql",
+  "graphql/utilities/index.js": "graphql",
+  "graphql/utilities/index.mjs": "graphql",
+  "graphql/validation/validate.js": "graphql",
+  "graphql/validation/validate.mjs": "graphql",
+  "mercurius/index.js": "mercurius",
+  "playwright-core/lib/coreBundle.js": "playwright-core",
+  "playwright/lib/index.js": "playwright",
+  "playwright/lib/runner/index.js": "playwright"
+}

package/packages/datadog-instrumentations/src/helpers/rewriter/transforms.js

@@ -17,8 +17,8 @@ const { parse, query } = require('./compiler')
 module.exports = { waitForAsyncEnd }
 
 /**
- * Injects a wait for `ctx.asyncEndPromise` into a generated `tracePromise`
- * wrapper's native-Promise fulfillment handler.
+ * Injects settlement-specific asyncEnd waits into a generated `tracePromise`
+ * wrapper's native-Promise handlers.
  *
  * @param {object} _state
  * @param {import('estree').CallExpression} node
@@ -26,36 +26,68 @@ module.exports = { waitForAsyncEnd }
  */
 function waitForAsyncEnd (_state, node) {
   const onFulfilled = node.arguments[0]
-  const statements = onFulfilled?.body?.body
+  const onRejected = node.arguments[1]
 
-  if (!statements || query(onFulfilled.body, '[id.name=__apm$asyncEndPromise]').length > 0) {
+  if (!onFulfilled?.body || !onRejected?.body) {
     return
   }
 
-  const returnIndex = statements.findIndex(statement =>
-    statement.type === 'ReturnStatement' && statement.argument
+  injectAsyncEndCallbackWait(onFulfilled.body, 'ReturnStatement', 'resolveCallback')
+  injectAsyncEndCallbackWait(onRejected.body, 'ThrowStatement', 'rejectCallback')
+}
+
+/**
+ * Injects a settlement-specific callback wait before an exit.
+ *
+ * @param {import('estree').BlockStatement} body
+ * @param {'ReturnStatement'|'ThrowStatement'} exitType
+ * @param {'resolveCallback'|'rejectCallback'} callbackProperty
+ * @returns {void}
+ */
+function injectAsyncEndCallbackWait (body, exitType, callbackProperty) {
+  const callbackVariable = `__apm$${callbackProperty}`
+  if (query(body, `[id.name=${callbackVariable}]`).length > 0) {
+    return
+  }
+
+  const exitIndex = body.body.findIndex(statement =>
+    statement.type === exitType && statement.argument
   )
 
-  // The generated fulfillment handler always ends in a return; a miss means the
+  // The generated settlement handlers always end in a return or throw; a miss means the
   // upstream template changed and the caller's try/catch falls back to the
   // unwrapped source.
-  assert(returnIndex !== -1, 'waitForAsyncEnd: no return statement to wait on')
+  assert(exitIndex !== -1, `waitForAsyncEnd: no ${exitType} to wait on`)
 
+  // This runs inside tracePromise's native-Promise settlement handler. The
+  // Promise adapts subscriber callback completion to that existing chain.
   const waitStatements = parse(`
     function wrapper () {
-      const __apm$asyncEndPromise = __apm$ctx.asyncEndPromise;
-      if (__apm$asyncEndPromise && typeof __apm$asyncEndPromise.then === 'function') {
-        return __apm$asyncEndPromise.then(() => __apm$result, () => __apm$result);
+      const ${callbackVariable} = __apm$ctx.${callbackProperty};
+      if (typeof ${callbackVariable} === 'function') {
+        return new Promise(${callbackVariable}).then(() => __apm$result, () => __apm$result);
       }
     }
   `).body[0].body.body
 
-  // Resolve to whatever the fulfillment handler returns (its return argument),
-  // so a subscriber that reassigned `__apm$ctx.result` in `asyncEnd` still wins.
-  const returnArgument = statements[returnIndex].argument
-  const { arguments: onSettled } = waitStatements[1].consequent.body[0].argument
-  onSettled[0].body = clone(returnArgument)
-  onSettled[1].body = clone(returnArgument)
+  const exitArgument = body.body[exitIndex].argument
+  const callbackIf = waitStatements[1]
+  const { arguments: onCallbackSettled } = callbackIf.consequent.body[0].argument
+
+  if (exitType === 'ThrowStatement') {
+    for (const handler of onCallbackSettled) {
+      handler.body = {
+        type: 'BlockStatement',
+        body: [{
+          type: 'ThrowStatement',
+          argument: clone(exitArgument),
+        }],
+      }
+    }
+  } else {
+    onCallbackSettled[0].body = clone(exitArgument)
+    onCallbackSettled[1].body = clone(exitArgument)
+  }
 
-  statements.splice(returnIndex, 0, ...waitStatements)
+  body.body.splice(exitIndex, 0, ...waitStatements)
 }