dd-trace 5.132.0 → 5.134.0

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentation-registry.js

@@ -1,16 +1,38 @@
 'use strict'
 
 // Activated rewrites publish their module name to the plugin manager after evaluation.
-// Only pure Orchestrion integrations that need this activation set activate: true.
+// Integrations can also provide synchronous runtime setup before plugin activation.
+// Setup must not load rewrite targets: re-entrant activations of the same group are dropped.
+// activate: true covers every rewrite target in the entry; activate.modules narrows it to the packages
+// that should trigger activation, leaving the rest of the entry rewrite-only.
 // Other entries use hooks or another activation path, or do not need activation from a rewrite.
+/**
+ * @typedef {object} Activation
+ * @property {string} moduleName
+ * @property {string} [version]
+ *
+ * @typedef {object} ActivationConfig
+ * @property {string[]} [modules]
+ * @property {((activation: Activation) => void)} [setup]
+ *
+ * @typedef {object} InstrumentationRegistryEntry
+ * @property {true|ActivationConfig} [activate]
+ * @property {Array<{ module: { name: string } }>} instrumentations
+ */
+
+/** @satisfies {InstrumentationRegistryEntry[]} */
 const registry = [
   { instrumentations: require('./instrumentations/ai') },
   { activate: true, instrumentations: require('./instrumentations/azure-cosmos') },
   { instrumentations: require('./instrumentations/azure-durable-functions') },
   { activate: true, instrumentations: require('./instrumentations/bullmq') },
-  { instrumentations: require('./instrumentations/claude-agent-sdk') },
+  {
+    activate: { setup: () => require('../../claude-agent-sdk') },
+    instrumentations: require('./instrumentations/claude-agent-sdk'),
+  },
   { instrumentations: require('./instrumentations/graphql') },
   { instrumentations: require('./instrumentations/graphql-jit') },
+  { instrumentations: require('./instrumentations/jest') },
   { activate: true, instrumentations: require('./instrumentations/langchain') },
   { activate: true, instrumentations: require('./instrumentations/langgraph') },
   { activate: true, instrumentations: require('./instrumentations/mercurius') },
@@ -18,15 +40,68 @@ const registry = [
   { instrumentations: require('./instrumentations/openai-agents') },
   { instrumentations: require('./instrumentations/playwright') },
   { activate: true, instrumentations: require('./instrumentations/postgres') },
+  { instrumentations: require('./instrumentations/vitest') },
   { instrumentations: require('./instrumentations/webdriverio') },
-  { instrumentations: require('./instrumentations/aws-durable-execution-sdk-js') },
+  {
+    activate: { setup: () => require('../../aws-durable-execution-sdk-js') },
+    instrumentations: require('./instrumentations/aws-durable-execution-sdk-js'),
+  },
   { activate: true, instrumentations: require('./instrumentations/supabase') },
 ]
 
 const activatedModules = new Set()
+/** @type {Map<string, (activation: Activation) => void>} */
+const activationSetups = new Map()
 for (const { activate, instrumentations } of registry) {
-  if (!activate) continue
-  for (const { module } of instrumentations) activatedModules.add(module.name)
+  if (activate === undefined) continue
+  if (activate !== true && (typeof activate !== 'object' || activate === null || Array.isArray(activate))) {
+    throw new TypeError('Instrumentation registry activate must be true or an activation config object')
+  }
+  const config = activate === true ? undefined : activate
+  const setup = config?.setup
+  if (setup !== undefined && typeof setup !== 'function') {
+    throw new TypeError('Instrumentation registry activate.setup must be a function')
+  }
+
+  const rewriteModules = new Set(instrumentations.map(({ module }) => module.name))
+  for (const moduleName of resolveActivationModules(config?.modules, rewriteModules)) {
+    if (activatedModules.has(moduleName) && activationSetups.get(moduleName) !== setup) {
+      throw new TypeError(`Instrumentation registry activation module "${moduleName}" has conflicting setup functions`)
+    }
+    activatedModules.add(moduleName)
+    if (setup) activationSetups.set(moduleName, setup)
+  }
+}
+
+/**
+ * Activation covers every rewrite target in the entry unless it names the packages to activate. An explicit
+ * list keeps helper packages that share the entry rewrite-only.
+ * @param {string[]|undefined} modules
+ * @param {Set<string>} rewriteModules
+ * @returns {Set<string>}
+ */
+function resolveActivationModules (modules, rewriteModules) {
+  if (modules === undefined) return rewriteModules
+  if (!Array.isArray(modules) || modules.length === 0) {
+    throw new TypeError('Instrumentation registry activate.modules must be a non-empty array')
+  }
+
+  const entryModules = new Set()
+  for (const moduleName of modules) {
+    if (typeof moduleName !== 'string' || moduleName.length === 0) {
+      throw new TypeError('Instrumentation registry activate.modules must contain non-empty strings')
+    }
+    if (!rewriteModules.has(moduleName)) {
+      throw new TypeError(
+        `Instrumentation registry activation module "${moduleName}" is not a rewrite target in its entry`
+      )
+    }
+    if (entryModules.has(moduleName)) {
+      throw new TypeError(`Instrumentation registry activation module "${moduleName}" is duplicated in its entry`)
+    }
+    entryModules.add(moduleName)
+  }
+  return entryModules
 }
 
 const instrumentations = registry.flatMap(entry => entry.instrumentations)
@@ -38,4 +113,12 @@ function isRewriteActivationEnabled (moduleName) {
   return activatedModules.has(moduleName)
 }
 
-module.exports = { isRewriteActivationEnabled, instrumentations, registry }
+/**
+ * @param {string} moduleName
+ * @returns {((activation: Activation) => void)|undefined}
+ */
+function getActivationSetup (moduleName) {
+  return activationSetups.get(moduleName)
+}
+
+module.exports = { getActivationSetup, isRewriteActivationEnabled, instrumentations, registry }

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/graphql-jit.js

@@ -26,11 +26,10 @@ const COMPILE_DEFERRED_FIELD = 'FunctionDeclaration[id.name="compileDeferredFiel
  */
 
 /**
- * @param {GraphqlJitInstrumentation[]} instrumentations
  * @param {string} versionRange
  * @param {string} filePath
  */
-function addInstrumentations (instrumentations, versionRange, filePath) {
+function addInstrumentations (versionRange, filePath) {
   const moduleDefinition = { name: 'graphql-jit', versionRange, filePath }
   instrumentations.push(
     {
@@ -75,9 +74,9 @@ function addInstrumentations (instrumentations, versionRange, filePath) {
 const instrumentations = []
 // graphql-jit 0.7 and 0.8.0 use this layout and remain compatible with GraphQL 15 and 16.
 // GraphQL 17 compatibility starts at 0.8.1.
-addInstrumentations(instrumentations, '>=0.7.0 <0.8.5 || >=0.8.7 <0.9.0', 'dist/execution.js')
-addInstrumentations(instrumentations, '>=0.8.5 <0.8.7', 'dist/cjs/execution.js')
-addInstrumentations(instrumentations, '>=0.8.5 <0.8.7', 'dist/esm/execution.js')
-addInstrumentations(instrumentations, '>=0.8.7 <0.9.0', 'dist/execution.mjs')
+addInstrumentations('>=0.7.0 <0.8.5 || >=0.8.7 <0.9.0', 'dist/execution.js')
+addInstrumentations('>=0.8.5 <0.8.7', 'dist/cjs/execution.js')
+addInstrumentations('>=0.8.5 <0.8.7', 'dist/esm/execution.js')
+addInstrumentations('>=0.8.7 <0.9.0', 'dist/execution.mjs')
 
 module.exports = instrumentations

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

@@ -0,0 +1,12 @@
+'use strict'
+
+module.exports = [
+  ['>=24.8.0 <30.0.0', 'build/runTest.js'],
+  ['>=30.0.0', 'build/index.js'],
+  ['>=30.0.0', 'build/testWorker.js'],
+].map(([versionRange, filePath]) => ({
+  module: { name: 'jest-runner', versionRange, filePath },
+  // Wrap the returned promise so completion covers both successful and failed teardown.
+  functionQuery: { functionName: 'runTestInternal', kind: 'Sync' },
+  channelName: 'runTestInternal',
+}))

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

@@ -1,36 +1,53 @@
 'use strict'
 
-const handlers = ['src/index.js', 'cjs/src/index.js'].map(filePath => ({
-  module: {
-    name: 'postgres',
-    versionRange: '>=3.0.0',
-    filePath,
-  },
-  astQuery: 'Program',
-  transform: 'postgresQueryHandlers',
-  channelName: 'query',
-}))
+const instrumentations = []
 
-const preparations = ['src/connection.js', 'cjs/src/connection.js'].map(filePath => ({
-  module: {
-    name: 'postgres',
-    versionRange: '>=3.0.0',
-    filePath,
-  },
-  astQuery: 'Program',
-  transform: 'postgresQueryPreparation',
-  channelName: 'query:prepare',
-}))
+for (const filePath of ['src/index.js', 'cjs/src/index.js']) {
+  instrumentations.push({
+    module: {
+      name: 'postgres',
+      versionRange: '>=3.0.0',
+      filePath,
+    },
+    astQuery: 'Program',
+    transform: 'postgresQueryHandlers',
+    channelName: 'query',
+  })
+}
 
-const lifecycles = ['src/query.js', 'cjs/src/query.js'].map(filePath => ({
-  module: {
-    name: 'postgres',
-    versionRange: '>=3.0.0',
-    filePath,
-  },
-  astQuery: 'Program',
-  transform: 'postgresQueryLifecycle',
-  channelName: 'query',
-}))
+for (const filePath of ['src/connection.js', 'cjs/src/connection.js']) {
+  instrumentations.push({
+    module: {
+      name: 'postgres',
+      versionRange: '>=3.0.0',
+      filePath,
+    },
+    astQuery: 'Program',
+    transform: 'postgresQueryAcquire',
+    channelName: 'query:acquire',
+  }, {
+    module: {
+      name: 'postgres',
+      versionRange: '>=3.0.0',
+      filePath,
+    },
+    astQuery: 'Program',
+    transform: 'postgresQueryPreparation',
+    channelName: 'query:prepare',
+  })
+}
 
-module.exports = [...handlers, ...preparations, ...lifecycles]
+for (const filePath of ['src/query.js', 'cjs/src/query.js']) {
+  instrumentations.push({
+    module: {
+      name: 'postgres',
+      versionRange: '>=3.0.0',
+      filePath,
+    },
+    astQuery: 'Program',
+    transform: 'postgresQueryLifecycle',
+    channelName: 'query',
+  })
+}
+
+module.exports = instrumentations

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

@@ -0,0 +1,8 @@
+'use strict'
+
+module.exports = ['dist/index.js', 'dist/chunk-hooks.js'].map(filePath => ({
+  module: { name: '@vitest/runner', versionRange: '>=1.6.0 <4.1.0', filePath },
+  astQuery: 'FunctionDeclaration[id.name="createSuiteCollector"] ' +
+    'VariableDeclarator[id.name="task"] > FunctionExpression',
+  channelName: 'createTestTask',
+}))

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

@@ -506,9 +506,8 @@ module.exports = [
       versionRange: '>=9.0.0',
       filePath: 'build/index.js',
     },
-    astQuery: 'TryStatement:has(CatchClause IfStatement[test.operator=">"]' +
-      '[test.left.object.name="retries"][test.left.property.name="limit"]' +
-      '[test.right.object.name="retries"][test.right.property.name="attempts"]) > BlockStatement',
+    astQuery: 'FunctionDeclaration[id.name="executeAsync"] TryStatement:has(CatchClause ' +
+      'CallExpression[callee.object.name="executeAsync"][callee.property.name="call"]) > BlockStatement',
     channelName: 'executeAsync',
     transform: 'awaitContextCallback',
     transformOptions: {
@@ -521,10 +520,8 @@ module.exports = [
       versionRange: '>=9.0.0',
       filePath: 'build/index.js',
     },
-    astQuery: 'FunctionDeclaration[id.name="executeAsync"] CatchClause ' +
-      'IfStatement[test.operator=">"][test.left.object.name="retries"]' +
-      '[test.left.property.name="limit"][test.right.object.name="retries"]' +
-      '[test.right.property.name="attempts"]',
+    astQuery: 'FunctionDeclaration[id.name="executeAsync"] CatchClause > BlockStatement > ' +
+      'IfStatement:has(CallExpression[callee.object.name="executeAsync"][callee.property.name="call"])',
     channelName: 'executeAsync',
     transform: 'awaitContextCallback',
     transformOptions: {
@@ -540,7 +537,10 @@ module.exports = [
     },
     astQuery: 'AssignmentExpression[left.object.object.name="Spec"]' +
       '[left.object.property.name="prototype"][left.property.name="execute"] > FunctionExpression',
+    // A method name keeps the original function expression, so the start context
+    // carries `self` (the legacy Spec) that the WebdriverIO Jasmine lifecycle reads.
     functionQuery: {
+      methodName: 'execute',
       kind: 'Sync',
     },
     channelName: 'Spec_execute',

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

@@ -41,6 +41,8 @@
   "@supabase/realtime-js/dist/module/RealtimeChannel.js": "@supabase/realtime-js",
   "@supabase/storage-js/dist/index.cjs": "@supabase/storage-js",
   "@supabase/storage-js/dist/index.mjs": "@supabase/storage-js",
+  "@vitest/runner/dist/chunk-hooks.js": "@vitest/runner",
+  "@vitest/runner/dist/index.js": "@vitest/runner",
   "@wdio/cli/build/index.js": "@wdio/cli",
   "@wdio/config/build/node/index.js": "@wdio/config",
   "@wdio/jasmine-framework/build/index.js": "@wdio/jasmine-framework",
@@ -76,6 +78,9 @@
   "graphql/validation/validate.js": "graphql",
   "graphql/validation/validate.mjs": "graphql",
   "jasmine-core/lib/jasmine-core/jasmine.js": "jasmine-core",
+  "jest-runner/build/index.js": "jest-runner",
+  "jest-runner/build/runTest.js": "jest-runner",
+  "jest-runner/build/testWorker.js": "jest-runner",
   "mercurius/index.js": "mercurius",
   "playwright-core/lib/coreBundle.js": "playwright-core",
   "playwright/lib/common/config.js": "playwright",

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

@@ -12,11 +12,37 @@ const POSTGRES_SIMPLE_DBM_QUERIES = '__ddTracePostgresSimpleDbmQueries'
 const POSTGRES_STATEMENT = 'this.string ?? (this.tagged && this.strings?.length !== 1 ? undefined : this.strings?.[0])'
 
 module.exports = {
+  postgresQueryAcquire,
   postgresQueryHandlers,
   postgresQueryLifecycle,
   postgresQueryPreparation,
 }
 
+/**
+ * @param {object} state
+ * @param {import('estree').Program} program
+ */
+function postgresQueryAcquire (state, program) {
+  const channelVariable = injectPostgresTracingChannel(state, program)
+  const executeFunctions = query(program, 'FunctionDeclaration[id.name="execute"]')
+  assert(executeFunctions.length === 1, 'postgresQueryAcquire: unexpected execute function count')
+
+  const execute = executeFunctions[0]
+  const queryParameter = execute.params[0]
+  assert(queryParameter?.type === 'Identifier', 'postgresQueryAcquire: execute query parameter changed')
+  assert(execute.body.type === 'BlockStatement', 'postgresQueryAcquire: execute body changed')
+
+  const body = execute.body.body
+  const executeIndex = body.findIndex(statement => statement.type === 'TryStatement')
+  assert(executeIndex !== -1, 'postgresQueryAcquire: execute try block changed')
+
+  body.splice(executeIndex, 0, ...parse(`
+    if (${channelVariable}.start.hasSubscribers) {
+      ${channelVariable}.start.publish({ query: ${queryParameter.name} });
+    }
+  `).body)
+}
+
 /**
  * Starts a Postgres.js query at the pool handler that owns its connection options.
  *
@@ -154,6 +180,7 @@ function wrapPostgresHandler (handler, queryParameter, channelVariable) {
   assert(handler.body.type === 'BlockStatement', 'postgresQueryHandlers: handler body changed')
   assert(!handler.async && !handler.generator, 'postgresQueryHandlers: unsupported handler kind')
 
+  // Direct handoffs publish before the handler returns, so they record zero pool wait.
   const originalBody = handler.body.body
   const wrapperBody = parse(`
     function wrapper () {
@@ -167,12 +194,18 @@ function wrapPostgresHandler (handler, queryParameter, channelVariable) {
           __ddTraceContext.host = ${POSTGRES_OPTIONS}.host[0];
           __ddTraceContext.port = ${POSTGRES_OPTIONS}.port[0];
         }
-        return ${channelVariable}.start.runStores(__ddTraceContext, () => {});
+        const __ddTraceResult = ${channelVariable}.start.runStores(__ddTraceContext, () => {});
+        if (!__ddTraceContext.poolWaitFinished) {
+          __ddTraceContext.poolWaitStartMs = performance.now();
+        }
+        return __ddTraceResult;
       }
     }
   `).body[0].body.body
 
-  wrapperBody[0].consequent.body.at(-1).argument.arguments[1].body.body = originalBody
+  const tracingCalls = query(wrapperBody[0], 'CallExpression[callee.property.name="runStores"]')
+  assert(tracingCalls.length === 1, 'postgresQueryHandlers: runStores call changed')
+  tracingCalls[0].arguments[1].body.body = originalBody
   handler.body.body = [...wrapperBody, ...originalBody]
 }
 
@@ -272,8 +305,13 @@ function findPostgresQueryIdentifier (program) {
 
   const declarations = query(program, 'VariableDeclarator[id.type="ObjectPattern"]')
   for (const declaration of declarations) {
-    if (declaration.init?.type !== 'CallExpression' || declaration.init.callee.name !== 'require') continue
-    if (declaration.init.arguments[0]?.value !== './query.js') continue
+    if (
+      declaration.init?.type !== 'CallExpression' ||
+      declaration.init.callee.name !== 'require' ||
+      declaration.init.arguments[0]?.value !== './query.js'
+    ) {
+      continue
+    }
 
     for (const property of declaration.id.properties) {
       if (property.type === 'Property' && property.key.name === 'Query' && property.value.type === 'Identifier') {

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

@@ -140,8 +140,7 @@ function awaitContextCallback (state, node, _parent, ancestry) {
     const generatedCallback = createAwaitedContextCallback(
       state,
       insertionTarget,
-      callbackAncestry,
-      'awaitContextCallback'
+      callbackAncestry
     )
     if (!generatedCallback) return
 
@@ -158,8 +157,7 @@ function awaitContextCallback (state, node, _parent, ancestry) {
   const generatedCallback = createAwaitedContextCallback(
     state,
     node.consequent,
-    ancestry,
-    'awaitContextCallback'
+    ancestry
   )
   if (!generatedCallback) return
 
@@ -181,18 +179,17 @@ function awaitContextCallback (state, node, _parent, ancestry) {
  * @param {Parameters<typeof awaitContextCallback>[0]} state
  * @param {import('estree').BlockStatement} insertionTarget
  * @param {import('estree').Node[]} ancestry
- * @param {string} transformName
  * @returns {{
  *   callbackStatements: import('estree').Statement[],
  *   callbackVariable: string
  * }|undefined}
  */
-function createAwaitedContextCallback (state, insertionTarget, ancestry, transformName) {
+function createAwaitedContextCallback (state, insertionTarget, ancestry) {
   const { callbackArgumentNames = [], callbackName, callbackThis = false } = state.transformOptions ?? {}
 
-  assert(identifierPattern.test(callbackName), `${transformName}: callbackName must be an identifier`)
+  assert(identifierPattern.test(callbackName), 'awaitContextCallback: callbackName must be an identifier')
   assert(Array.isArray(callbackArgumentNames) && callbackArgumentNames.every(name => identifierPattern.test(name)),
-    `${transformName}: callbackArgumentNames must be identifiers`)
+    'awaitContextCallback: callbackArgumentNames must be identifiers')
 
   let enclosingFunction
   let hasTraceWrapper = false
@@ -215,8 +212,8 @@ function createAwaitedContextCallback (state, insertionTarget, ancestry, transfo
     hasTraceWrapper ||= hasContextBinding && hasTracedBinding
   }
 
-  assert(enclosingFunction?.async, `${transformName}: expected an enclosing async function`)
-  assert(hasTraceWrapper, `${transformName}: expected an enclosing trace wrapper`)
+  assert(enclosingFunction?.async, 'awaitContextCallback: expected an enclosing async function')
+  assert(hasTraceWrapper, 'awaitContextCallback: expected an enclosing trace wrapper')
 
   const callbackVariable = `__apm$${callbackName}`
   if (query(insertionTarget, `[id.name="${callbackVariable}"]`).length > 0) return

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

@@ -14,9 +14,7 @@ const connectionOptionsCache = new WeakMap()
 
 function wrapRedis (Redis) {
   shimmer.wrap(Redis.prototype, 'sendCommand', sendCommand => function (command, stream) {
-    if (!startCh.hasSubscribers) return sendCommand.call(this, command, stream)
-
-    if (!command?.promise) return sendCommand.call(this, command, stream)
+    if (!startCh.hasSubscribers || !command?.promise) return sendCommand.call(this, command, stream)
 
     const options = this.options || {}
     let connectionOptions = connectionOptionsCache.get(this)
@@ -32,7 +30,7 @@ function wrapRedis (Redis) {
       connectionName: options.connectionName,
     }
     return startCh.runStores(ctx, () => {
-      command.promise.then(() => finish(finishCh, errorCh, ctx), err => finish(finishCh, errorCh, ctx, err))
+      command.promise.then(() => finish(ctx), err => finish(ctx, err))
 
       return sendCommand.call(this, command, stream)
     })
@@ -51,7 +49,11 @@ addHook({ name: 'ioredis', versions: ['>=4.11.0 <5'], file: 'built/redis/index.j
 
 addHook({ name: 'ioredis', versions: ['>=5'] }, wrapRedis)
 
-function finish (finishCh, errorCh, ctx, error) {
+/**
+ * @param {Record<string, unknown> & { error?: Error }} ctx
+ * @param {Error} [error]
+ */
+function finish (ctx, error) {
   if (error) {
     ctx.error = error
     errorCh.publish(ctx)

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

@@ -12,9 +12,7 @@ const errorCh = channel('apm:iovalkey:command:error')
 
 addHook({ name: 'iovalkey', versions: ['>=0.0.1'] }, Valkey => {
   shimmer.wrap(Valkey.prototype, 'sendCommand', sendCommand => function (command, stream) {
-    if (!startCh.hasSubscribers) return sendCommand.apply(this, arguments)
-
-    if (!command?.promise) return sendCommand.apply(this, arguments)
+    if (!startCh.hasSubscribers || !command?.promise) return sendCommand.apply(this, arguments)
 
     const options = this.options || {}
     const connectionName = options.connectionName
@@ -22,7 +20,7 @@ addHook({ name: 'iovalkey', versions: ['>=0.0.1'] }, Valkey => {
 
     const ctx = { command: command.name, args: command.args, connectionOptions, connectionName }
     return startCh.runStores(ctx, () => {
-      command.promise.then(() => finish(finishCh, errorCh, ctx), err => finish(finishCh, errorCh, ctx, err))
+      command.promise.then(() => finish(ctx), err => finish(ctx, err))
 
       try {
         return sendCommand.apply(this, arguments)
@@ -37,7 +35,11 @@ addHook({ name: 'iovalkey', versions: ['>=0.0.1'] }, Valkey => {
   return Valkey
 })
 
-function finish (finishCh, errorCh, ctx, error) {
+/**
+ * @param {Record<string, unknown> & { error?: Error }} ctx
+ * @param {Error} [error]
+ */
+function finish (ctx, error) {
   if (error) {
     ctx.error = error
     errorCh.publish(ctx)