dd-trace 6.0.0 → 6.1.0

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

@@ -1,8 +1,5 @@
 'use strict'
 
-// Capture real timers at module load time, before any test can install fake timers.
-const realSetTimeout = setTimeout
-
 const { performance } = require('node:perf_hooks')
 
 const shimmer = require('../../datadog-shimmer')
@@ -18,6 +15,7 @@ const {
   testFinishTimeCh,
   testPassCh,
   testErrorCh,
+  testDiWaitCh,
   testSkipCh,
   testFnCh,
   testSuiteStartCh,
@@ -32,39 +30,36 @@ const {
   getVitestTestProperties,
 } = require('./vitest-util')
 
-const BREAKPOINT_HIT_GRACE_PERIOD_MS = 400
-
 const taskToCtx = new WeakMap()
 const taskToTestProperties = new WeakMap()
 const taskToStatuses = new WeakMap()
 const taskToReportedErrorCount = new WeakMap()
 const attemptToFixTaskToStatuses = new WeakMap()
+const fileToHasConcurrentTests = new WeakMap()
 const originalHookFns = new WeakMap()
 const newTasks = new WeakSet()
 const dynamicNameTasks = new WeakSet()
 const disabledTasks = new WeakSet()
 const quarantinedTasks = new WeakSet()
 const attemptToFixTasks = new WeakSet()
+const attemptToFixRetryTasks = new WeakSet()
 const modifiedTasks = new WeakSet()
+const efdRetryTasks = new WeakSet()
 const efdDeterminedRetries = new WeakMap()
 const efdSlowAbortedTasks = new WeakSet()
 const efdExecutionStartByTask = new WeakMap()
 const efdSkippedRetryResults = new WeakMap()
 const attemptToFixExecutions = new Map()
 const loggedAttemptToFixTests = new Set()
-let isRetryReasonEfd = false
-let isRetryReasonAttemptToFix = false
 const switchedStatuses = new WeakSet()
 let vitestGetFn = null
 let vitestSetFn = null
 let vitestGetHooks = null
 
 function waitForHitProbe () {
-  return new Promise(resolve => {
-    realSetTimeout(() => {
-      resolve()
-    }, BREAKPOINT_HIT_GRACE_PERIOD_MS)
-  })
+  const promises = {}
+  testDiWaitCh.publish({ promises })
+  return promises.hitBreakpointPromise
 }
 
 function getVitestTestStatus (test, retryCount) {
@@ -138,6 +133,70 @@ function wrapBeforeEachCleanupResult (task, result) {
   return result
 }
 
+/**
+ * Returns whether a Vitest task tree includes any concurrent task.
+ *
+ * @param {Array<{ type?: string, concurrent?: boolean, tasks?: object[] }>|undefined} tasks
+ * @returns {boolean}
+ */
+function hasConcurrentTask (tasks) {
+  if (!tasks) return false
+
+  for (const task of tasks) {
+    if (task.concurrent === true) return true
+    if (hasConcurrentTask(task.tasks)) return true
+  }
+
+  return false
+}
+
+/**
+ * Returns whether a Vitest file includes any concurrent test.
+ *
+ * @param {{ tasks?: object[] }} file
+ * @returns {boolean}
+ */
+function hasConcurrentTests (file) {
+  const cached = fileToHasConcurrentTests.get(file)
+  if (cached !== undefined) return cached
+
+  const hasConcurrent = hasConcurrentTask(file.tasks)
+  fileToHasConcurrentTests.set(file, hasConcurrent)
+  return hasConcurrent
+}
+
+/**
+ * Gets the task associated with a Vitest hook invocation.
+ *
+ * @param {unknown[]} args
+ * @param {object} fallbackTask
+ * @returns {object}
+ */
+function getTaskFromHookArgs (args, fallbackTask) {
+  return args[0]?.task || fallbackTask
+}
+
+/**
+ * Wraps a Vitest hook so it runs inside the span context for the current test.
+ *
+ * @param {'beforeEach'|'afterEach'} hookType
+ * @param {Function} fn
+ * @param {object} fallbackTask
+ * @returns {Function}
+ */
+function wrapSuiteHookFn (hookType, fn, fallbackTask) {
+  return shimmer.wrapFunction(fn, fn => function (...args) {
+    const task = getTaskFromHookArgs(args, fallbackTask)
+    const result = testFnCh.runStores(taskToCtx.get(task), () => fn.apply(this, args))
+
+    if (hookType === 'beforeEach') {
+      return wrapBeforeEachCleanupResult(task, result)
+    }
+
+    return result
+  })
+}
+
 function wrapVitestTestRunner (VitestTestRunner) {
   // `onBeforeRunTask` is run before any repetition or attempt is run
   // `onBeforeRunTask` is an async function
@@ -163,7 +222,9 @@ function wrapVitestTestRunner (VitestTestRunner) {
         isQuarantined,
       } = testProperties
       if (isAttemptToFix) {
-        isRetryReasonAttemptToFix = task.repeats !== testManagementAttemptToFixRetries
+        if (task.repeats !== testManagementAttemptToFixRetries) {
+          attemptToFixRetryTasks.add(task)
+        }
         disableFrameworkRetries(task)
         task.repeats = testManagementAttemptToFixRetries
         attemptToFixTasks.add(task)
@@ -183,7 +244,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
 
     if (isImpactedTestsEnabled && testProperties.isModified) {
       if (isEarlyFlakeDetectionEnabled) {
-        isRetryReasonEfd = true
+        efdRetryTasks.add(task)
         disableFrameworkRetries(task)
         task.repeats = numRepeats
       }
@@ -193,7 +254,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
 
     if (isKnownTestsEnabled && testProperties.isNew && !attemptToFixTasks.has(task)) {
       if (isEarlyFlakeDetectionEnabled && !modifiedTasks.has(task)) {
-        isRetryReasonEfd = true
+        efdRetryTasks.add(task)
         disableFrameworkRetries(task)
         task.repeats = numRepeats
       }
@@ -267,6 +328,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
     }
 
     const { retry: numAttempt, repeats: numRepetition } = retryInfo
+    const isFailedTestReplayAllowed = !hasConcurrentTests(task.file)
     const isEfdManagedTask = isEarlyFlakeDetectionEnabled && taskToStatuses.has(task) && !attemptToFixTasks.has(task)
 
     if (isEfdManagedTask && numRepetition > 0 && !efdDeterminedRetries.has(task)) {
@@ -303,19 +365,20 @@ function wrapVitestTestRunner (VitestTestRunner) {
 
     // We finish the previous test here because we know it has failed already
     if (numAttempt > 0) {
-      const shouldWaitForHitProbe = isDiEnabled && numAttempt > 1
+      const shouldWaitForHitProbe = isDiEnabled && isFailedTestReplayAllowed && numAttempt > 1
       if (shouldWaitForHitProbe) {
         await waitForHitProbe()
       }
 
       const promises = {}
-      const shouldSetProbe = isDiEnabled && numAttempt === 1
+      const shouldSetProbe = isDiEnabled && isFailedTestReplayAllowed && numAttempt === 1
       const ctx = taskToCtx.get(task)
       const testError = getCurrentAttemptTestError(task, task.result?.errors)
       if (ctx) {
         testErrorCh.publish({
           error: testError,
           shouldSetProbe,
+          shouldWaitForHitProbe,
           promises,
           ...ctx.currentStore,
         })
@@ -377,18 +440,18 @@ function wrapVitestTestRunner (VitestTestRunner) {
 
     const isRetryReasonAtr = numAttempt > 0 &&
       isFlakyTestRetriesEnabledForTask(providedContext, task) &&
-      !isRetryReasonAttemptToFix &&
-      !isRetryReasonEfd
+      !attemptToFixRetryTasks.has(task) &&
+      !efdRetryTasks.has(task)
 
     const ctx = {
       testName,
       testSuiteAbsolutePath: task.file.filepath,
       isRetry: numAttempt > 0 || numRepetition > 0,
-      isRetryReasonEfd,
-      isRetryReasonAttemptToFix: isRetryReasonAttemptToFix && numRepetition > 0,
+      isRetryReasonEfd: efdRetryTasks.has(task),
+      isRetryReasonAttemptToFix: attemptToFixRetryTasks.has(task) && numRepetition > 0,
       isNew,
       hasDynamicName: dynamicNameTasks.has(task),
-      mightHitProbe: isDiEnabled && numAttempt > 0,
+      mightHitProbe: isDiEnabled && isFailedTestReplayAllowed && numAttempt > 0,
       isAttemptToFix: attemptToFixTasks.has(task),
       isDisabled: disabledTasks.has(task),
       isQuarantined: quarantinedTasks.has(task),
@@ -429,17 +492,9 @@ function wrapVitestTestRunner (VitestTestRunner) {
           if (!hookArray) continue
           for (let i = 0; i < hookArray.length; i++) {
             const currentFn = hookArray[i]
-            const originalFn = originalHookFns.get(currentFn) || currentFn
-            const wrappedFn = shimmer.wrapFunction(originalFn, fn => function (...args) {
-              const result = testFnCh.runStores(taskToCtx.get(task), () => fn.apply(this, args))
-
-              if (hookType === 'beforeEach') {
-                return wrapBeforeEachCleanupResult(task, result)
-              }
-
-              return result
-            })
-            originalHookFns.set(wrappedFn, originalFn)
+            if (originalHookFns.has(currentFn)) continue
+            const wrappedFn = wrapSuiteHookFn(hookType, currentFn, task)
+            originalHookFns.set(wrappedFn, currentFn)
             hookArray[i] = wrappedFn
           }
         }
@@ -468,6 +523,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
       const ctx = taskToCtx.get(task)
 
       const { isDiEnabled } = getProvidedContext()
+      const isFailedTestReplayAllowed = !hasConcurrentTests(task.file)
 
       if (efdSkippedRetryResults.has(task)) {
         task.result = efdSkippedRetryResults.get(task)
@@ -475,7 +531,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
         return result
       }
 
-      if (isDiEnabled && retryInfo.retry > 1) {
+      if (isDiEnabled && isFailedTestReplayAllowed && retryInfo.retry > 1) {
         await waitForHitProbe()
       }
 
@@ -601,6 +657,7 @@ addHook({
       testCommand: providedContext.testCommand,
       repositoryRoot: providedContext.repositoryRoot,
       codeOwnersEntries: providedContext.codeOwnersEntries,
+      testEnvironmentMetadata: providedContext.testEnvironmentMetadata,
     }
     testSuiteStartCh.runStores(testSuiteCtx, () => {})
     const startTestsResponse = await startTests.apply(this, arguments)
@@ -611,6 +668,7 @@ addHook({
     })
 
     const testTasks = getTypeTasks(startTestsResponse[0].tasks)
+    const testEventPromises = []
 
     // Only one test task per test, even if there are retries
     for (const task of testTasks) {
@@ -639,12 +697,17 @@ addHook({
           if (testCtx) {
             const isSkippedByTestManagement =
               !attemptToFixTasks.has(task) && (disabledTasks.has(task) || quarantinedTasks.has(task))
+            const promises = {}
             testPassCh.publish({
               task,
               finalStatus: isSkippedByTestManagement ? 'skip' : 'pass',
               earlyFlakeAbortReason: efdSlowAbortedTasks.has(task) ? 'slow' : undefined,
+              promises,
               ...testCtx.currentStore,
             })
+            if (promises.hitBreakpointPromise) {
+              testEventPromises.push(promises.hitBreakpointPromise)
+            }
           }
         } else if (state === 'fail' || isSwitchedStatus) {
           let hasFailedAllRetries = false
@@ -684,6 +747,7 @@ addHook({
 
           if (testCtx) {
             const isRetry = task.result?.retryCount > 0
+            const promises = {}
             // `duration` is the duration of all the retries, so it can't be used if there are retries
 
             let finalStatus
@@ -708,8 +772,12 @@ addHook({
               attemptToFixFailed,
               finalStatus,
               earlyFlakeAbortReason: efdSlowAbortedTasks.has(task) ? 'slow' : undefined,
+              promises,
               ...testCtx.currentStore,
             })
+            if (promises.hitBreakpointPromise) {
+              testEventPromises.push(promises.hitBreakpointPromise)
+            }
           }
           if (errors?.length) {
             testSuiteError = testError // we store the error to bubble it up to the suite
@@ -725,6 +793,8 @@ addHook({
       }
     }
 
+    await Promise.all(testEventPromises)
+
     const testSuiteResult = startTestsResponse[0].result
 
     if (testSuiteResult.errors?.length) { // Errors from root level hooks

package/packages/datadog-plugin-ai/src/index.js

@@ -5,7 +5,7 @@ const VercelAILLMObsPlugin = require('../../dd-trace/src/llmobs/plugins/ai')
 const VercelAITracingPlugin = require('./tracing')
 
 class VercelAIPlugin extends CompositePlugin {
-  static get id () { return 'ai' }
+  static id = 'ai'
   static get plugins () {
     return {
       llmobs: VercelAILLMObsPlugin,

package/packages/datadog-plugin-ai/src/tracing.js

@@ -1,10 +1,11 @@
 'use strict'
 
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
-const { getModelProvider } = require('./utils')
+const { getModelProvider, parseModelProvider } = require('./utils')
 
-class VercelAITracingPlugin extends TracingPlugin {
-  static id = 'ai'
+class DdTelemetryPlugin extends TracingPlugin {
+  static id = 'ai_tracing_dd_telemetry'
   static prefix = 'tracing:dd-trace:vercel-ai'
 
   bindStart (ctx) {
@@ -30,4 +31,66 @@ class VercelAITracingPlugin extends TracingPlugin {
   }
 }
 
+class VercelAiTelemetryPlugin extends TracingPlugin {
+  static id = 'ai_tracing_vercel_telemetry'
+  static prefix = 'tracing:ai:telemetry'
+
+  #streamedCalls = new Set()
+
+  constructor () {
+    super(...arguments)
+
+    this.addSub('dd-trace:vercel-ai:chunk', ({ ctx, chunk, done }) => {
+      ctx.streamConsumed = done
+    })
+  }
+
+  bindStart (ctx) {
+    const { type: name, event } = ctx
+    const model = event.modelId
+    const modelProvider = parseModelProvider(event.provider, model)
+
+    let isStream = this.#streamedCalls.has(event.callId)
+    if (name.includes('stream')) {
+      this.#streamedCalls.add(event.callId)
+      isStream = true
+
+      ctx.streamConsumed = false
+    }
+
+    ctx.isStream = isStream
+
+    this.startSpan(name, {
+      meta: {
+        'resource.name': event.functionId ?? name,
+        'ai.request.model': model,
+        'ai.request.model_provider': modelProvider,
+      },
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    // check if isStreamed and stream resolved
+    // this event will fire multiple times for the same channel
+    if (ctx.isStream && ctx.result?.stream && !ctx.streamConsumed) return
+
+    if (ctx.type?.includes('stream')) {
+      this.#streamedCalls.delete(ctx.event?.callId)
+    }
+
+    const span = ctx.currentStore?.span
+    span?.finish()
+  }
+}
+
+class VercelAITracingPlugin extends CompositePlugin {
+  static id = 'ai_tracing'
+  static plugins = {
+    dd: DdTelemetryPlugin,
+    ai: VercelAiTelemetryPlugin,
+  }
+}
+
 module.exports = VercelAITracingPlugin

package/packages/datadog-plugin-ai/src/utils.js

@@ -1,6 +1,6 @@
 'use strict'
 
-const { parseModelId } = require('../../datadog-plugin-aws-sdk/src/services/bedrockruntime/utils')
+const { parseModelId: parseBedrockModelId } = require('../../datadog-plugin-aws-sdk/src/services/bedrockruntime/utils')
 
 /**
  * Get the model provider from the span tags or attributes.
@@ -11,12 +11,24 @@ const { parseModelId } = require('../../datadog-plugin-aws-sdk/src/services/bedr
  */
 function getModelProvider (tags) {
   const modelProviderTag = tags['ai.model.provider']
-  const providerParts = modelProviderTag?.split('.')
+  const modelId = tags['ai.model.id']
+
+  return parseModelProvider(modelProviderTag, modelId)
+}
+
+/**
+ * Parse the model provider from the raw provider string.
+ *
+ * @param {string} rawProvider
+ * @param {string} modelId
+ * @returns {string}
+ */
+function parseModelProvider (rawProvider, modelId) {
+  const providerParts = rawProvider?.split('.')
   const provider = providerParts?.[0]
 
   if (provider === 'amazon-bedrock') {
-    const modelId = tags['ai.model.id']
-    const model = modelId && parseModelId(modelId)
+    const model = modelId && parseBedrockModelId(modelId)
     return model?.modelProvider ?? provider
   }
 
@@ -25,4 +37,5 @@ function getModelProvider (tags) {
 
 module.exports = {
   getModelProvider,
+  parseModelProvider,
 }

package/packages/datadog-plugin-aws-durable-execution-sdk-js/src/context.js

@@ -50,6 +50,7 @@ class BaseContextPlugin extends TracingPlugin {
 
     this.startSpan(spanName, {
       resource,
+      type: this.constructor.type,
       kind: this.constructor.kind,
       meta,
       metrics,

package/packages/datadog-plugin-aws-durable-execution-sdk-js/src/handler.js

@@ -52,6 +52,7 @@ class AwsDurableExecutionSdkJsHandlerPlugin extends TracingPlugin {
 
     this.startSpan(this.operationName(), {
       resource: handler?.name,
+      type: this.constructor.type,
       kind: this.constructor.kind,
       meta,
     }, ctx)

package/packages/datadog-plugin-child_process/src/index.js

@@ -97,11 +97,22 @@ class ChildProcessPlugin extends TracingPlugin {
   }
 
   asyncEnd (ctx) {
-    const { result } = ctx
+    const { result, error } = ctx
+    let exitCode
+
+    if (result !== null && typeof result === 'object') {
+      // util.promisify(execFile) resolves with a { stdout, stderr } object on
+      // success, where the exit code is 0.
+      exitCode = result.status ?? 0
+    } else if (result === undefined && error !== undefined) {
+      exitCode = error.status ?? error.code ?? 0
+    } else {
+      exitCode = result
+    }
 
     const span = ctx.currentStore?.span || this.activeSpan
 
-    span?.setTag('cmd.exit_code', `${result}`)
+    span?.setTag('cmd.exit_code', `${exitCode}`)
     span?.finish()
 
     return ctx.parentStore

package/packages/datadog-plugin-claude-agent-sdk/src/index.js

@@ -0,0 +1,31 @@
+'use strict'
+
+require('../../datadog-instrumentations/src/claude-agent-sdk')
+
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
+const ClaudeAgentSdkLLMObsPlugins = require('../../dd-trace/src/llmobs/plugins/claude-agent-sdk')
+const ClaudeAgentSdkTracingPlugins = require('./tracing')
+
+class ClaudeAgentSdkPlugin extends CompositePlugin {
+  static id = 'claude-agent-sdk'
+
+  static get plugins () {
+    const plugins = {}
+
+    // LLM Obs plugins must be registered before tracing plugins so that
+    // annotations are added to the span before it finishes.
+    // The tracing plugin uses `bindStart` vs the LLM Obs plugin's `start`,
+    // so the span is created in the tracing plugin before the LLM Obs one runs.
+    for (const Plugin of ClaudeAgentSdkLLMObsPlugins) {
+      plugins[Plugin.id] = Plugin
+    }
+
+    for (const Plugin of ClaudeAgentSdkTracingPlugins) {
+      plugins[Plugin.id] = Plugin
+    }
+
+    return plugins
+  }
+}
+
+module.exports = ClaudeAgentSdkPlugin

package/packages/datadog-plugin-claude-agent-sdk/src/tracing.js

@@ -0,0 +1,107 @@
+'use strict'
+
+const { storage } = require('../../datadog-core')
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+const { splitModel } = require('./util')
+
+class QueryTracingPlugin extends TracingPlugin {
+  static id = 'claude_agent_sdk_query'
+  static operation = 'turn'
+  static system = 'claude-agent-sdk'
+  static prefix = 'tracing:orchestrion:@anthropic-ai/claude-agent-sdk:query'
+
+  bindStart (ctx) {
+    this.startSpan('claude_agent_sdk.query', {
+      meta: { 'resource.name': 'claude_agent_sdk.query' },
+      startTime: ctx.startTime,
+    }, ctx)
+    ctx.runInContext = fn => storage('legacy').run(ctx.currentStore, fn)
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    if (!ctx.streamResolved) return
+
+    ctx.currentStore?.span?.finish(ctx.finishTime)
+  }
+}
+
+class StepTracingPlugin extends TracingPlugin {
+  static id = 'claude_agent_sdk_step'
+  static operation = 'step'
+  static system = 'claude-agent-sdk'
+  static prefix = 'tracing:apm:claude-agent-sdk:step'
+
+  bindStart (ctx) {
+    this.startSpan(`step-${ctx.stepIndex}`, {
+      meta: { 'resource.name': 'claude_agent_sdk.step' },
+      startTime: ctx.startTime,
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  end (ctx) {
+    ctx.currentStore?.span?.finish(ctx.finishTime)
+  }
+}
+
+class ToolTracingPlugin extends TracingPlugin {
+  static id = 'claude_agent_sdk_tool'
+  static operation = 'tool'
+  static system = 'claude-agent-sdk'
+  static prefix = 'tracing:apm:claude-agent-sdk:tool'
+
+  bindStart (ctx) {
+    const toolName = ctx.name || 'claude_agent_sdk.tool'
+
+    this.startSpan(toolName, {
+      meta: { 'resource.name': 'claude_agent_sdk.tool' },
+      startTime: ctx.startTime,
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  end (ctx) {
+    const span = ctx.currentStore?.span
+    if (ctx.error) this.addError(ctx.error, span)
+    span?.finish(ctx.finishTime)
+  }
+}
+
+class LlmTracingPlugin extends TracingPlugin {
+  static id = 'claude_agent_sdk_llm'
+  static operation = 'llm'
+  static system = 'claude-agent-sdk'
+  static prefix = 'tracing:apm:claude-agent-sdk:llm'
+
+  bindStart (ctx) {
+    const { model } = ctx
+
+    const { modelName, modelProvider } = splitModel(model)
+    const name = modelName || 'claude_agent_sdk.llm'
+
+    this.startSpan(name, {
+      meta: {
+        'resource.name': 'claude_agent_sdk.llm',
+        'claude-agent-sdk.request.model_name': modelName,
+        'claude-agent-sdk.request.model_provider': modelProvider,
+      },
+      startTime: ctx.startTime,
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  end (ctx) {
+    ctx.currentStore?.span?.finish(ctx.finishTime)
+  }
+}
+
+module.exports = [
+  QueryTracingPlugin,
+  StepTracingPlugin,
+  ToolTracingPlugin,
+  LlmTracingPlugin,
+]

package/packages/datadog-plugin-claude-agent-sdk/src/util.js

@@ -0,0 +1,15 @@
+'use strict'
+
+/**
+ * Parse a model ID string into the model name and provider for claude agent sdk provided model IDs
+ * @param {string} model model id
+ * @returns {{ modelName: string, modelProvider: string }}
+ */
+function splitModel (model) {
+  if (!model) return { modelName: undefined, modelProvider: 'anthropic' }
+  const idx = model.indexOf('/')
+  if (idx === -1) return { modelName: model, modelProvider: 'anthropic' }
+  return { modelName: model.slice(idx + 1), modelProvider: model.slice(0, idx) }
+}
+
+module.exports = { splitModel }