dd-trace 5.30.0 → 5.32.0

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

@@ -28,6 +28,42 @@ const newTasks = new WeakSet()
 const switchedStatuses = new WeakSet()
 const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')
 
+const BREAKPOINT_HIT_GRACE_PERIOD_MS = 400
+
+function waitForHitProbe () {
+  return new Promise(resolve => {
+    setTimeout(() => {
+      resolve()
+    }, BREAKPOINT_HIT_GRACE_PERIOD_MS)
+  })
+}
+
+function getProvidedContext () {
+  try {
+    const {
+      _ddIsEarlyFlakeDetectionEnabled,
+      _ddIsDiEnabled,
+      _ddKnownTests: knownTests,
+      _ddEarlyFlakeDetectionNumRetries: numRepeats
+    } = globalThis.__vitest_worker__.providedContext
+
+    return {
+      isDiEnabled: _ddIsDiEnabled,
+      isEarlyFlakeDetectionEnabled: _ddIsEarlyFlakeDetectionEnabled,
+      knownTests,
+      numRepeats
+    }
+  } catch (e) {
+    log.error('Vitest workers could not parse provided context, so some features will not work.')
+    return {
+      isDiEnabled: false,
+      isEarlyFlakeDetectionEnabled: false,
+      knownTests: {},
+      numRepeats: 0
+    }
+  }
+}
+
 function isReporterPackage (vitestPackage) {
   return vitestPackage.B?.name === 'BaseSequencer'
 }
@@ -143,7 +179,9 @@ function getSortWrapper (sort) {
       const knownTestsResponse = await getChannelPromise(knownTestsCh)
       if (!knownTestsResponse.err) {
         knownTests = knownTestsResponse.knownTests
-        const testFilepaths = await this.ctx.getTestFilepaths()
+        const getFilePaths = this.ctx.getTestFilepaths || this.ctx._globTestFilepaths
+
+        const testFilepaths = await getFilePaths.call(this.ctx)
 
         isEarlyFlakeDetectionFaultyCh.publish({
           knownTests: knownTests.vitest || {},
@@ -253,29 +291,26 @@ addHook({
   // `onBeforeRunTask` is run before any repetition or attempt is run
   shimmer.wrap(VitestTestRunner.prototype, 'onBeforeRunTask', onBeforeRunTask => async function (task) {
     const testName = getTestName(task)
-    try {
-      const {
-        _ddKnownTests: knownTests,
-        _ddIsEarlyFlakeDetectionEnabled: isEarlyFlakeDetectionEnabled,
-        _ddEarlyFlakeDetectionNumRetries: numRepeats
-      } = globalThis.__vitest_worker__.providedContext
-
-      if (isEarlyFlakeDetectionEnabled) {
-        isNewTestCh.publish({
-          knownTests,
-          testSuiteAbsolutePath: task.file.filepath,
-          testName,
-          onDone: (isNew) => {
-            if (isNew) {
-              task.repeats = numRepeats
-              newTasks.add(task)
-              taskToStatuses.set(task, [])
-            }
+
+    const {
+      knownTests,
+      isEarlyFlakeDetectionEnabled,
+      numRepeats
+    } = getProvidedContext()
+
+    if (isEarlyFlakeDetectionEnabled) {
+      isNewTestCh.publish({
+        knownTests,
+        testSuiteAbsolutePath: task.file.filepath,
+        testName,
+        onDone: (isNew) => {
+          if (isNew) {
+            task.repeats = numRepeats
+            newTasks.add(task)
+            taskToStatuses.set(task, [])
           }
-        })
-      }
-    } catch (e) {
-      log.error('Vitest workers could not parse known tests, so Early Flake Detection will not work.')
+        }
+      })
     }
 
     return onBeforeRunTask.apply(this, arguments)
@@ -283,9 +318,7 @@ addHook({
 
   // `onAfterRunTask` is run after all repetitions or attempts are run
   shimmer.wrap(VitestTestRunner.prototype, 'onAfterRunTask', onAfterRunTask => async function (task) {
-    const {
-      _ddIsEarlyFlakeDetectionEnabled: isEarlyFlakeDetectionEnabled
-    } = globalThis.__vitest_worker__.providedContext
+    const { isEarlyFlakeDetectionEnabled } = getProvidedContext()
 
     if (isEarlyFlakeDetectionEnabled && taskToStatuses.has(task)) {
       const statuses = taskToStatuses.get(task)
@@ -309,43 +342,40 @@ addHook({
     }
     const testName = getTestName(task)
     let isNew = false
-    let isEarlyFlakeDetectionEnabled = false
-    let isDiEnabled = false
-
-    try {
-      const {
-        _ddIsEarlyFlakeDetectionEnabled,
-        _ddIsDiEnabled
-      } = globalThis.__vitest_worker__.providedContext
 
-      isEarlyFlakeDetectionEnabled = _ddIsEarlyFlakeDetectionEnabled
-      isDiEnabled = _ddIsDiEnabled
+    const {
+      isEarlyFlakeDetectionEnabled,
+      isDiEnabled
+    } = getProvidedContext()
 
-      if (isEarlyFlakeDetectionEnabled) {
-        isNew = newTasks.has(task)
-      }
-    } catch (e) {
-      log.error('Vitest workers could not parse known tests, so Early Flake Detection will not work.')
+    if (isEarlyFlakeDetectionEnabled) {
+      isNew = newTasks.has(task)
     }
+
     const { retry: numAttempt, repeats: numRepetition } = retryInfo
 
     // We finish the previous test here because we know it has failed already
     if (numAttempt > 0) {
-      const probe = {}
+      const shouldWaitForHitProbe = isDiEnabled && numAttempt > 1
+      if (shouldWaitForHitProbe) {
+        await waitForHitProbe()
+      }
+
+      const promises = {}
+      const shouldSetProbe = isDiEnabled && numAttempt === 1
       const asyncResource = taskToAsync.get(task)
       const testError = task.result?.errors?.[0]
       if (asyncResource) {
         asyncResource.runInAsyncScope(() => {
           testErrorCh.publish({
             error: testError,
-            willBeRetried: true,
-            probe,
-            isDiEnabled
+            shouldSetProbe,
+            promises
           })
         })
         // We wait for the probe to be set
-        if (probe.setProbePromise) {
-          await probe.setProbePromise
+        if (promises.setProbePromise) {
+          await promises.setProbePromise
         }
       }
     }
@@ -401,7 +431,8 @@ addHook({
         testName,
         testSuiteAbsolutePath: task.file.filepath,
         isRetry: numAttempt > 0 || numRepetition > 0,
-        isNew
+        isNew,
+        mightHitProbe: isDiEnabled && numAttempt > 0
       })
     })
     return onBeforeTryTask.apply(this, arguments)
@@ -418,6 +449,12 @@ addHook({
       const status = getVitestTestStatus(task, retryCount)
       const asyncResource = taskToAsync.get(task)
 
+      const { isDiEnabled } = getProvidedContext()
+
+      if (isDiEnabled && retryCount > 1) {
+        await waitForHitProbe()
+      }
+
       if (asyncResource) {
         // We don't finish here because the test might fail in a later hook (afterEach)
         asyncResource.runInAsyncScope(() => {
@@ -457,15 +494,6 @@ addHook({
   return vitestPackage
 })
 
-addHook({
-  name: 'vitest',
-  versions: ['>=2.1.0'],
-  filePattern: 'dist/chunks/RandomSequencer.*'
-}, (randomSequencerPackage) => {
-  shimmer.wrap(randomSequencerPackage.B.prototype, 'sort', getSortWrapper)
-  return randomSequencerPackage
-})
-
 addHook({
   name: 'vitest',
   versions: ['>=2.0.5 <2.1.0'],
@@ -478,6 +506,24 @@ addHook({
   return vitestPackage
 })
 
+addHook({
+  name: 'vitest',
+  versions: ['>=2.1.0 <3.0.0'],
+  filePattern: 'dist/chunks/RandomSequencer.*'
+}, (randomSequencerPackage) => {
+  shimmer.wrap(randomSequencerPackage.B.prototype, 'sort', getSortWrapper)
+  return randomSequencerPackage
+})
+
+addHook({
+  name: 'vitest',
+  versions: ['>=3.0.0'],
+  filePattern: 'dist/chunks/resolveConfig.*'
+}, (randomSequencerPackage) => {
+  shimmer.wrap(randomSequencerPackage.B.prototype, 'sort', getSortWrapper)
+  return randomSequencerPackage
+})
+
 // Can't specify file because compiled vitest includes hashes in their files
 addHook({
   name: 'vitest',
@@ -498,15 +544,17 @@ addHook({
   versions: ['>=1.6.0'],
   file: 'dist/index.js'
 }, (vitestPackage, frameworkVersion) => {
-  shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
+  shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPaths) {
     let testSuiteError = null
     if (!testSuiteStartCh.hasSubscribers) {
       return startTests.apply(this, arguments)
     }
+    // From >=3.0.1, the first arguments changes from a string to an object containing the filepath
+    const testSuiteAbsolutePath = testPaths[0]?.filepath || testPaths[0]
 
     const testSuiteAsyncResource = new AsyncResource('bound-anonymous-fn')
     testSuiteAsyncResource.runInAsyncScope(() => {
-      testSuiteStartCh.publish({ testSuiteAbsolutePath: testPath[0], frameworkVersion })
+      testSuiteStartCh.publish({ testSuiteAbsolutePath, frameworkVersion })
     })
     const startTestsResponse = await startTests.apply(this, arguments)
 

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

@@ -0,0 +1,49 @@
+'use strict'
+
+const { channel, addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+const names = ['vm', 'node:vm']
+
+const runScriptStartChannel = channel('datadog:vm:run-script:start')
+const sourceTextModuleStartChannel = channel('datadog:vm:source-text-module:start')
+
+addHook({ name: names }, function (vm) {
+  vm.Script = class extends vm.Script {
+    constructor (code) {
+      super(...arguments)
+
+      if (runScriptStartChannel.hasSubscribers && code) {
+        runScriptStartChannel.publish({ code })
+      }
+    }
+  }
+
+  if (vm.SourceTextModule && typeof vm.SourceTextModule === 'function') {
+    vm.SourceTextModule = class extends vm.SourceTextModule {
+      constructor (code) {
+        super(...arguments)
+
+        if (sourceTextModuleStartChannel.hasSubscribers && code) {
+          sourceTextModuleStartChannel.publish({ code })
+        }
+      }
+    }
+  }
+
+  shimmer.wrap(vm, 'runInContext', wrapVMMethod)
+  shimmer.wrap(vm, 'runInNewContext', wrapVMMethod)
+  shimmer.wrap(vm, 'runInThisContext', wrapVMMethod)
+  shimmer.wrap(vm, 'compileFunction', wrapVMMethod)
+
+  return vm
+})
+
+function wrapVMMethod (original) {
+  return function wrappedVMMethod (code) {
+    if (runScriptStartChannel.hasSubscribers && code) {
+      runScriptStartChannel.publish({ code })
+    }
+
+    return original.apply(this, arguments)
+  }
+}

package/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime.js

@@ -0,0 +1,295 @@
+'use strict'
+
+const BaseAwsSdkPlugin = require('../base')
+const log = require('../../../dd-trace/src/log')
+
+const PROVIDER = {
+  AI21: 'AI21',
+  AMAZON: 'AMAZON',
+  ANTHROPIC: 'ANTHROPIC',
+  COHERE: 'COHERE',
+  META: 'META',
+  STABILITY: 'STABILITY',
+  MISTRAL: 'MISTRAL'
+}
+
+const enabledOperations = ['invokeModel']
+
+class BedrockRuntime extends BaseAwsSdkPlugin {
+  static get id () { return 'bedrock runtime' }
+
+  isEnabled (request) {
+    const operation = request.operation
+    if (!enabledOperations.includes(operation)) {
+      return false
+    }
+
+    return super.isEnabled(request)
+  }
+
+  generateTags (params, operation, response) {
+    let tags = {}
+    let modelName = ''
+    let modelProvider = ''
+    const modelMeta = params.modelId.split('.')
+    if (modelMeta.length === 2) {
+      [modelProvider, modelName] = modelMeta
+      modelProvider = modelProvider.toUpperCase()
+    } else {
+      [, modelProvider, modelName] = modelMeta
+      modelProvider = modelProvider.toUpperCase()
+    }
+
+    const shouldSetChoiceIds = modelProvider === PROVIDER.COHERE && !modelName.includes('embed')
+
+    const requestParams = extractRequestParams(params, modelProvider)
+    const textAndResponseReason = extractTextAndResponseReason(response, modelProvider, modelName, shouldSetChoiceIds)
+
+    tags = buildTagsFromParams(requestParams, textAndResponseReason, modelProvider, modelName, operation)
+
+    return tags
+  }
+}
+
+class Generation {
+  constructor ({ message = '', finishReason = '', choiceId = '' } = {}) {
+    // stringify message as it could be a single generated message as well as a list of embeddings
+    this.message = typeof message === 'string' ? message : JSON.stringify(message) || ''
+    this.finishReason = finishReason || ''
+    this.choiceId = choiceId || undefined
+  }
+}
+
+class RequestParams {
+  constructor ({
+    prompt = '',
+    temperature = undefined,
+    topP = undefined,
+    maxTokens = undefined,
+    stopSequences = [],
+    inputType = '',
+    truncate = '',
+    stream = '',
+    n = undefined
+  } = {}) {
+    // TODO: set a truncation limit to prompt
+    // stringify prompt as it could be a single prompt as well as a list of message objects
+    this.prompt = typeof prompt === 'string' ? prompt : JSON.stringify(prompt) || ''
+    this.temperature = temperature !== undefined ? temperature : undefined
+    this.topP = topP !== undefined ? topP : undefined
+    this.maxTokens = maxTokens !== undefined ? maxTokens : undefined
+    this.stopSequences = stopSequences || []
+    this.inputType = inputType || ''
+    this.truncate = truncate || ''
+    this.stream = stream || ''
+    this.n = n !== undefined ? n : undefined
+  }
+}
+
+function extractRequestParams (params, provider) {
+  const requestBody = JSON.parse(params.body)
+  const modelId = params.modelId
+
+  switch (provider) {
+    case PROVIDER.AI21: {
+      let userPrompt = requestBody.prompt
+      if (modelId.includes('jamba')) {
+        for (const message of requestBody.messages) {
+          if (message.role === 'user') {
+            userPrompt = message.content // Return the content of the most recent user message
+          }
+        }
+      }
+      return new RequestParams({
+        prompt: userPrompt,
+        temperature: requestBody.temperature,
+        topP: requestBody.top_p,
+        maxTokens: requestBody.max_tokens,
+        stopSequences: requestBody.stop_sequences
+      })
+    }
+    case PROVIDER.AMAZON: {
+      if (modelId.includes('embed')) {
+        return new RequestParams({ prompt: requestBody.inputText })
+      }
+      const textGenerationConfig = requestBody.textGenerationConfig || {}
+      return new RequestParams({
+        prompt: requestBody.inputText,
+        temperature: textGenerationConfig.temperature,
+        topP: textGenerationConfig.topP,
+        maxTokens: textGenerationConfig.maxTokenCount,
+        stopSequences: textGenerationConfig.stopSequences
+      })
+    }
+    case PROVIDER.ANTHROPIC: {
+      const prompt = requestBody.prompt || requestBody.messages
+      return new RequestParams({
+        prompt,
+        temperature: requestBody.temperature,
+        topP: requestBody.top_p,
+        maxTokens: requestBody.max_tokens_to_sample,
+        stopSequences: requestBody.stop_sequences
+      })
+    }
+    case PROVIDER.COHERE: {
+      if (modelId.includes('embed')) {
+        return new RequestParams({
+          prompt: requestBody.texts,
+          inputType: requestBody.input_type,
+          truncate: requestBody.truncate
+        })
+      }
+      return new RequestParams({
+        prompt: requestBody.prompt,
+        temperature: requestBody.temperature,
+        topP: requestBody.p,
+        maxTokens: requestBody.max_tokens,
+        stopSequences: requestBody.stop_sequences,
+        stream: requestBody.stream,
+        n: requestBody.num_generations
+      })
+    }
+    case PROVIDER.META: {
+      return new RequestParams({
+        prompt: requestBody.prompt,
+        temperature: requestBody.temperature,
+        topP: requestBody.top_p,
+        maxTokens: requestBody.max_gen_len
+      })
+    }
+    case PROVIDER.MISTRAL: {
+      return new RequestParams({
+        prompt: requestBody.prompt,
+        temperature: requestBody.temperature,
+        topP: requestBody.top_p,
+        maxTokens: requestBody.max_tokens,
+        stopSequences: requestBody.stop,
+        topK: requestBody.top_k
+      })
+    }
+    case PROVIDER.STABILITY: {
+      return new RequestParams()
+    }
+    default: {
+      return new RequestParams()
+    }
+  }
+}
+
+function extractTextAndResponseReason (response, provider, modelName, shouldSetChoiceIds) {
+  const body = JSON.parse(Buffer.from(response.body).toString('utf8'))
+
+  try {
+    switch (provider) {
+      case PROVIDER.AI21: {
+        if (modelName.includes('jamba')) {
+          const generations = body.choices || []
+          if (generations.length > 0) {
+            const generation = generations[0]
+            return new Generation({
+              message: generation.message,
+              finishReason: generation.finish_reason,
+              choiceId: shouldSetChoiceIds ? generation.id : undefined
+            })
+          }
+        }
+        const completions = body.completions || []
+        if (completions.length > 0) {
+          const completion = completions[0]
+          return new Generation({
+            message: completion.data?.text,
+            finishReason: completion?.finishReason,
+            choiceId: shouldSetChoiceIds ? completion?.id : undefined
+          })
+        }
+        return new Generation()
+      }
+      case PROVIDER.AMAZON: {
+        if (modelName.includes('embed')) {
+          return new Generation({ message: body.embedding })
+        }
+        const results = body.results || []
+        if (results.length > 0) {
+          const result = results[0]
+          return new Generation({ message: result.outputText, finishReason: result.completionReason })
+        }
+        break
+      }
+      case PROVIDER.ANTHROPIC: {
+        return new Generation({ message: body.completion || body.content, finishReason: body.stop_reason })
+      }
+      case PROVIDER.COHERE: {
+        if (modelName.includes('embed')) {
+          const embeddings = body.embeddings || [[]]
+          if (embeddings.length > 0) {
+            return new Generation({ message: embeddings[0] })
+          }
+        }
+        const generations = body.generations || []
+        if (generations.length > 0) {
+          const generation = generations[0]
+          return new Generation({
+            message: generation.text,
+            finishReason: generation.finish_reason,
+            choiceId: shouldSetChoiceIds ? generation.id : undefined
+          })
+        }
+        break
+      }
+      case PROVIDER.META: {
+        return new Generation({ message: body.generation, finishReason: body.stop_reason })
+      }
+      case PROVIDER.MISTRAL: {
+        const mistralGenerations = body.outputs || []
+        if (mistralGenerations.length > 0) {
+          const generation = mistralGenerations[0]
+          return new Generation({ message: generation.text, finishReason: generation.stop_reason })
+        }
+        break
+      }
+      case PROVIDER.STABILITY: {
+        return new Generation()
+      }
+      default: {
+        return new Generation()
+      }
+    }
+  } catch (error) {
+    log.warn('Unable to extract text/finishReason from response body. Defaulting to empty text/finishReason.')
+    return new Generation()
+  }
+
+  return new Generation()
+}
+
+function buildTagsFromParams (requestParams, textAndResponseReason, modelProvider, modelName, operation) {
+  const tags = {}
+
+  // add request tags
+  tags['resource.name'] = operation
+  tags['aws.bedrock.request.model'] = modelName
+  tags['aws.bedrock.request.model_provider'] = modelProvider
+  tags['aws.bedrock.request.prompt'] = requestParams.prompt
+  tags['aws.bedrock.request.temperature'] = requestParams.temperature
+  tags['aws.bedrock.request.top_p'] = requestParams.topP
+  tags['aws.bedrock.request.max_tokens'] = requestParams.maxTokens
+  tags['aws.bedrock.request.stop_sequences'] = requestParams.stopSequences
+  tags['aws.bedrock.request.input_type'] = requestParams.inputType
+  tags['aws.bedrock.request.truncate'] = requestParams.truncate
+  tags['aws.bedrock.request.stream'] = requestParams.stream
+  tags['aws.bedrock.request.n'] = requestParams.n
+
+  // add response tags
+  if (modelName.includes('embed')) {
+    tags['aws.bedrock.response.embedding_length'] = textAndResponseReason.message.length
+  }
+  if (textAndResponseReason.choiceId) {
+    tags['aws.bedrock.response.choices.id'] = textAndResponseReason.choiceId
+  }
+  tags['aws.bedrock.response.choices.text'] = textAndResponseReason.message
+  tags['aws.bedrock.response.choices.finish_reason'] = textAndResponseReason.finishReason
+
+  return tags
+}
+
+module.exports = BedrockRuntime

package/packages/datadog-plugin-aws-sdk/src/services/index.js

@@ -12,4 +12,5 @@ exports.sns = require('./sns')
 exports.sqs = require('./sqs')
 exports.states = require('./states')
 exports.stepfunctions = require('./stepfunctions')
+exports.bedrockruntime = require('./bedrockruntime')
 exports.default = require('./default')