dd-trace 5.43.0 → 5.44.0

package/packages/dd-trace/src/llmobs/plugins/vertexai.js

@@ -0,0 +1,196 @@
+'use strict'
+
+const LLMObsPlugin = require('./base')
+const {
+  extractModel,
+  extractSystemInstructions
+} = require('../../../../datadog-plugin-google-cloud-vertexai/src/utils')
+
+class VertexAILLMObsPlugin extends LLMObsPlugin {
+  static get id () { return 'vertexai' } // used for llmobs telemetry
+  static get prefix () {
+    return 'tracing:apm:vertexai:request'
+  }
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    const history = ctx.instance?.historyInternal || []
+    ctx.history = history
+
+    return {
+      kind: 'llm',
+      modelName: extractModel(ctx.instance),
+      modelProvider: 'google',
+      name: ctx.resource
+    }
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    const { instance, result, request } = ctx
+    const history = ctx.history || []
+    const systemInstructions = extractSystemInstructions(instance)
+
+    const metadata = getMetadata(instance)
+    const inputMessages = extractInputMessages(request, history, systemInstructions)
+    const outputMessages = extractOutputMessages(result)
+    const metrics = extractMetrics(result)
+
+    this._tagger.tagLLMIO(span, inputMessages, outputMessages)
+    this._tagger.tagMetadata(span, metadata)
+    this._tagger.tagMetrics(span, metrics)
+  }
+}
+
+function getMetadata (instance) {
+  const metadata = {}
+
+  const modelConfig = instance.generationConfig
+  if (!modelConfig) return metadata
+
+  for (const [parameter, parameterKey] of [
+    ['temperature', 'temperature'],
+    ['maxOutputTokens', 'max_output_tokens'],
+    ['candidateCount', 'candidate_count'],
+    ['topP', 'top_p'],
+    ['topK', 'top_k']
+  ]) {
+    if (modelConfig[parameter]) {
+      metadata[parameterKey] = modelConfig[parameter]
+    }
+  }
+
+  return metadata
+}
+
+function extractInputMessages (request, history, systemInstructions) {
+  const contents = typeof request === 'string' || Array.isArray(request) ? request : request.contents
+  const messages = []
+
+  if (systemInstructions) {
+    for (const instruction of systemInstructions) {
+      messages.push({ content: instruction || '', role: 'system' })
+    }
+  }
+
+  for (const content of history) {
+    messages.push(...extractMessagesFromContent(content))
+  }
+
+  if (typeof contents === 'string') {
+    messages.push({ content: contents })
+    return messages
+  }
+
+  if (isPart(contents)) {
+    messages.push(extractMessageFromPart(contents))
+    return messages
+  }
+
+  if (!Array.isArray(contents)) {
+    messages.push({
+      content: '[Non-array content object: ' +
+      `${(typeof contents.toString === 'function' ? contents.toString() : String(contents))}]`
+    })
+    return messages
+  }
+
+  for (const content of contents) {
+    if (typeof content === 'string') {
+      messages.push({ content })
+      continue
+    }
+
+    if (isPart(content)) {
+      messages.push(extractMessageFromPart(content))
+      continue
+    }
+
+    messages.push(...extractMessagesFromContent(content))
+  }
+
+  return messages
+}
+
+function extractOutputMessages (result) {
+  if (!result) return [{ content: '' }]
+  const { response } = result
+
+  if (!response) return [{ content: '' }]
+
+  const outputMessages = []
+  const candidates = response.candidates || []
+  for (const candidate of candidates) {
+    const content = candidate.content || ''
+    outputMessages.push(...extractMessagesFromContent(content))
+  }
+
+  return outputMessages
+}
+
+function extractMessagesFromContent (content) {
+  const messages = []
+
+  const role = content.role || ''
+  const parts = content.parts || []
+  if (parts == null || parts.length === 0 || !Array.isArray(parts)) {
+    const message = {
+      content:
+      `[Non-text content object: ${(typeof content.toString === 'function' ? content.toString() : String(content))}]`
+    }
+    if (role) message.role = role
+    messages.push(message)
+    return messages
+  }
+
+  for (const part of parts) {
+    const message = extractMessageFromPart(part, role)
+    messages.push(message)
+  }
+
+  return messages
+}
+
+function extractMessageFromPart (part, role) {
+  const text = part.text || ''
+  const functionCall = part.functionCall
+  const functionResponse = part.functionResponse
+
+  const message = { content: text }
+  if (role) message.role = role
+  if (functionCall) {
+    message.toolCalls = [{
+      name: functionCall.name,
+      arguments: functionCall.args
+    }]
+  }
+  if (functionResponse) {
+    message.content = `[tool result: ${functionResponse.response}]`
+  }
+
+  return message
+}
+
+function extractMetrics (result) {
+  if (!result) return {}
+  const { response } = result
+
+  if (!response) return {}
+
+  const tokenCounts = response.usageMetadata
+  const metrics = {}
+  if (tokenCounts) {
+    metrics.inputTokens = tokenCounts.promptTokenCount
+    metrics.outputTokens = tokenCounts.candidatesTokenCount
+    metrics.totalTokens = tokenCounts.totalTokenCount
+  }
+
+  return metrics
+}
+
+function isPart (part) {
+  return part.text || part.functionCall || part.functionResponse
+}
+
+module.exports = VertexAILLMObsPlugin

package/packages/dd-trace/src/llmobs/sdk.js

@@ -430,6 +430,7 @@ class LLMObs extends NoopLLMObs {
       modelProvider,
       sessionId,
       mlApp,
+      _decorator,
       ...spanOptions
     } = options
 
@@ -438,6 +439,7 @@ class LLMObs extends NoopLLMObs {
       modelName,
       modelProvider,
       sessionId,
+      _decorator,
       spanOptions
     }
   }

package/packages/dd-trace/src/llmobs/span_processor.js

@@ -7,6 +7,7 @@ const {
   METADATA,
   INPUT_MESSAGES,
   INPUT_VALUE,
+  INTEGRATION,
   OUTPUT_MESSAGES,
   INPUT_DOCUMENTS,
   OUTPUT_DOCUMENTS,
@@ -26,6 +27,8 @@ const {
   ERROR_STACK
 } = require('../constants')
 
+const telemetry = require('./telemetry')
+
 const LLMObsTagger = require('./tagger')
 
 const tracerVersion = require('../../../../package.json').version
@@ -48,6 +51,7 @@ class LLMObsSpanProcessor {
 
     try {
       const formattedEvent = this.format(span)
+      telemetry.incrementLLMObsSpanFinishedCount(span)
       this._writer.append(formattedEvent)
     } catch (e) {
       // this should be a rare case
@@ -186,6 +190,8 @@ class LLMObsSpanProcessor {
     const errType = span.context()._tags[ERROR_TYPE] || error?.name
     if (errType) tags.error_type = errType
     if (sessionId) tags.session_id = sessionId
+    const integration = LLMObsTagger.tagMap.get(span)?.[INTEGRATION]
+    if (integration) tags.integration = integration
     const existingTags = LLMObsTagger.tagMap.get(span)?.[TAGS] || {}
     if (existingTags) tags = { ...tags, ...existingTags }
     return Object.entries(tags).map(([key, value]) => `${key}:${value ?? ''}`)

package/packages/dd-trace/src/llmobs/tagger.js

@@ -22,7 +22,9 @@ const {
   ROOT_PARENT_ID,
   INPUT_TOKENS_METRIC_KEY,
   OUTPUT_TOKENS_METRIC_KEY,
-  TOTAL_TOKENS_METRIC_KEY
+  TOTAL_TOKENS_METRIC_KEY,
+  INTEGRATION,
+  DECORATOR
 } = require('./constants/tags')
 
 // global registry of LLMObs spans
@@ -51,7 +53,9 @@ class LLMObsTagger {
     mlApp,
     parent,
     kind,
-    name
+    name,
+    integration,
+    _decorator
   } = {}) {
     if (!this._config.llmobs.enabled) return
     if (!kind) return // do not register it in the map if it doesn't have an llmobs span kind
@@ -66,6 +70,8 @@ class LLMObsTagger {
 
     sessionId = sessionId || registry.get(parent)?.[SESSION_ID]
     if (sessionId) this._setTag(span, SESSION_ID, sessionId)
+    if (integration) this._setTag(span, INTEGRATION, integration)
+    if (_decorator) this._setTag(span, DECORATOR, _decorator)
 
     if (!mlApp) mlApp = registry.get(parent)?.[ML_APP] || this._config.llmobs.mlApp
     this._setTag(span, ML_APP, mlApp)

package/packages/dd-trace/src/llmobs/telemetry.js

@@ -1,12 +1,93 @@
 'use strict'
 
+const {
+  SPAN_KIND,
+  MODEL_PROVIDER,
+  PARENT_ID_KEY,
+  SESSION_ID,
+  ROOT_PARENT_ID,
+  INTEGRATION,
+  DECORATOR
+} = require('./constants/tags')
+
+const ERROR_TYPE = require('../constants')
+
 const telemetryMetrics = require('../telemetry/metrics')
+
+const LLMObsTagger = require('./tagger')
+
 const llmobsMetrics = telemetryMetrics.manager.namespace('mlobs')
 
+function extractIntegrationFromTags (tags) {
+  if (!Array.isArray(tags)) return null
+  const integrationTag = tags.find(tag => tag.startsWith('integration:'))
+  if (!integrationTag) return null
+  return integrationTag.split(':')[1] || null
+}
+
+function extractTagsFromSpanEvent (event) {
+  const spanKind = event.meta?.['span.kind'] || ''
+  const integration = extractIntegrationFromTags(event.tags)
+  const error = event.status === 'error'
+  const autoinstrumented = integration != null
+
+  return {
+    span_kind: spanKind,
+    autoinstrumented: Number(autoinstrumented),
+    error: error ? 1 : 0,
+    integration: integration || 'N/A'
+  }
+}
+
 function incrementLLMObsSpanStartCount (tags, value = 1) {
   llmobsMetrics.count('span.start', tags).inc(value)
 }
 
+function incrementLLMObsSpanFinishedCount (span, value = 1) {
+  const mlObsTags = LLMObsTagger.tagMap.get(span)
+  const spanTags = span.context()._tags
+
+  const isRootSpan = mlObsTags[PARENT_ID_KEY] === ROOT_PARENT_ID
+  const hasSessionId = mlObsTags[SESSION_ID] != null
+  const integration = mlObsTags[INTEGRATION]
+  const autoInstrumented = integration != null
+  const decorator = !!mlObsTags[DECORATOR]
+  const spanKind = mlObsTags[SPAN_KIND]
+  const modelProvider = mlObsTags[MODEL_PROVIDER]
+  const error = spanTags.error || spanTags[ERROR_TYPE]
+
+  const tags = {
+    autoinstrumented: Number(autoInstrumented),
+    has_session_id: Number(hasSessionId),
+    is_root_span: Number(isRootSpan),
+    span_kind: spanKind,
+    integration: integration || 'N/A',
+    error: error ? 1 : 0
+  }
+  if (!autoInstrumented) {
+    tags.decorator = Number(decorator)
+  }
+  if (modelProvider) {
+    tags.model_provider = modelProvider
+  }
+
+  llmobsMetrics.count('span.finished', tags).inc(value)
+}
+
+function recordLLMObsRawSpanSize (event, rawEventSize) {
+  const tags = extractTagsFromSpanEvent(event)
+  llmobsMetrics.distribution('span.raw_size', tags).track(rawEventSize)
+}
+
+function recordLLMObsSpanSize (event, eventSize, shouldTruncate) {
+  const tags = extractTagsFromSpanEvent(event)
+  tags.truncated = Number(shouldTruncate)
+  llmobsMetrics.distribution('span.size', tags).track(eventSize)
+}
+
 module.exports = {
-  incrementLLMObsSpanStartCount
+  incrementLLMObsSpanStartCount,
+  incrementLLMObsSpanFinishedCount,
+  recordLLMObsRawSpanSize,
+  recordLLMObsSpanSize
 }

package/packages/dd-trace/src/llmobs/writers/spans/base.js

@@ -4,6 +4,7 @@ const { EVP_EVENT_SIZE_LIMIT, EVP_PAYLOAD_SIZE_LIMIT } = require('../../constant
 const { DROPPED_VALUE_TEXT } = require('../../constants/text')
 const { DROPPED_IO_COLLECTION_ERROR } = require('../../constants/tags')
 const BaseWriter = require('../base')
+const telemetry = require('../../telemetry')
 const logger = require('../../../log')
 
 const tracerVersion = require('../../../../../../package.json').version
@@ -18,11 +19,19 @@ class LLMObsSpanWriter extends BaseWriter {
 
   append (event) {
     const eventSizeBytes = Buffer.from(JSON.stringify(event)).byteLength
-    if (eventSizeBytes > EVP_EVENT_SIZE_LIMIT) {
+    telemetry.recordLLMObsRawSpanSize(event, eventSizeBytes)
+
+    const shouldTruncate = eventSizeBytes > EVP_EVENT_SIZE_LIMIT
+    let processedEventSizeBytes = eventSizeBytes
+
+    if (shouldTruncate) {
       logger.warn(`Dropping event input/output because its size (${eventSizeBytes}) exceeds the 1MB event size limit`)
       event = this._truncateSpanEvent(event)
+      processedEventSizeBytes = Buffer.from(JSON.stringify(event)).byteLength
     }
 
+    telemetry.recordLLMObsSpanSize(event, processedEventSizeBytes, shouldTruncate)
+
     if (this._bufferSize + eventSizeBytes > EVP_PAYLOAD_SIZE_LIMIT) {
       logger.debug('Flusing queue because queing next event will exceed EvP payload limit')
       this.flush()

package/packages/dd-trace/src/plugin_manager.js

@@ -28,9 +28,6 @@ loadChannel.subscribe(({ name }) => {
   maybeEnable(plugins[name])
 })
 
-// Always enabled
-maybeEnable(require('../../datadog-plugin-dd-trace-api/src'))
-
 function maybeEnable (Plugin) {
   if (!Plugin || typeof Plugin !== 'function') return
   if (!pluginClasses[Plugin.id]) {

package/packages/dd-trace/src/plugins/ci_plugin.js

@@ -28,9 +28,7 @@ const {
   DI_DEBUG_ERROR_SNAPSHOT_ID_SUFFIX,
   DI_DEBUG_ERROR_FILE_SUFFIX,
   DI_DEBUG_ERROR_LINE_SUFFIX,
-  DD_CAPABILITIES_EARLY_FLAKE_DETECTION,
-  DD_CAPABILITIES_AUTO_TEST_RETRIES,
-  DD_CAPABILITIES_TEST_IMPACT_ANALYSIS
+  getLibraryCapabilitiesTags
 } = require('./util/test')
 const Plugin = require('./plugin')
 const { COMPONENT } = require('../constants')
@@ -41,23 +39,17 @@ const {
   TELEMETRY_EVENT_CREATED,
   TELEMETRY_ITR_SKIPPED
 } = require('../ci-visibility/telemetry')
-const { CI_PROVIDER_NAME, GIT_REPOSITORY_URL, GIT_COMMIT_SHA, GIT_BRANCH, CI_WORKSPACE_PATH } = require('./util/tags')
+const {
+  CI_PROVIDER_NAME,
+  GIT_REPOSITORY_URL,
+  GIT_COMMIT_SHA,
+  GIT_BRANCH,
+  CI_WORKSPACE_PATH,
+  GIT_COMMIT_MESSAGE
+} = require('./util/tags')
 const { OS_VERSION, OS_PLATFORM, OS_ARCHITECTURE, RUNTIME_NAME, RUNTIME_VERSION } = require('./util/env')
 const getDiClient = require('../ci-visibility/dynamic-instrumentation')
 
-const UNSUPPORTED_TIA_FRAMEWORKS = ['playwright', 'vitest']
-const UNSUPPORTED_TIA_FRAMEWORKS_PARALLEL_MODE = ['cucumber', 'mocha']
-
-function isTiaSupported (testFramework, isParallel) {
-  if (UNSUPPORTED_TIA_FRAMEWORKS.includes(testFramework)) {
-    return false
-  }
-  if (isParallel && UNSUPPORTED_TIA_FRAMEWORKS_PARALLEL_MODE.includes(testFramework)) {
-    return false
-  }
-  return true
-}
-
 module.exports = class CiPlugin extends Plugin {
   constructor (...args) {
     super(...args)
@@ -75,17 +67,13 @@ module.exports = class CiPlugin extends Plugin {
         } else {
           this.libraryConfig = libraryConfig
         }
-        const { isItrEnabled, isEarlyFlakeDetectionEnabled, isFlakyTestRetriesEnabled } = this.libraryConfig || {}
+
+        const libraryCapabilitiesTags = getLibraryCapabilitiesTags(this.constructor.id, isParallel)
         const metadataTags = {
           test: {
-            [DD_CAPABILITIES_TEST_IMPACT_ANALYSIS]: isItrEnabled ? 'true' : 'false',
-            [DD_CAPABILITIES_EARLY_FLAKE_DETECTION]: isEarlyFlakeDetectionEnabled ? 'true' : 'false',
-            [DD_CAPABILITIES_AUTO_TEST_RETRIES]: isFlakyTestRetriesEnabled ? 'true' : 'false'
+            ...libraryCapabilitiesTags
           }
         }
-        if (!isTiaSupported(this.constructor.id, isParallel)) {
-          metadataTags.test[DD_CAPABILITIES_TEST_IMPACT_ANALYSIS] = undefined
-        }
         this.tracer._exporter.addMetadataTags(metadataTags)
         onDone({ err, libraryConfig })
       })
@@ -251,7 +239,8 @@ module.exports = class CiPlugin extends Plugin {
       [RUNTIME_VERSION]: runtimeVersion,
       [GIT_BRANCH]: branch,
       [CI_PROVIDER_NAME]: ciProviderName,
-      [CI_WORKSPACE_PATH]: repositoryRoot
+      [CI_WORKSPACE_PATH]: repositoryRoot,
+      [GIT_COMMIT_MESSAGE]: commitMessage
     } = this.testEnvironmentMetadata
 
     this.repositoryRoot = repositoryRoot || process.cwd()
@@ -269,7 +258,8 @@ module.exports = class CiPlugin extends Plugin {
       runtimeName,
       runtimeVersion,
       branch,
-      testLevel: 'suite'
+      testLevel: 'suite',
+      commitMessage
     }
   }
 

package/packages/dd-trace/src/plugins/database.js

@@ -63,7 +63,7 @@ class DatabasePlugin extends StoragePlugin {
     return tracerService
   }
 
-  createDbmComment (span, serviceName, isPreparedStatement = false) {
+  createDbmComment (span, serviceName, disableFullMode = false) {
     const mode = this.config.dbmPropagationMode
     const dbmService = this.getDbmServiceName(span, serviceName)
 
@@ -73,7 +73,7 @@ class DatabasePlugin extends StoragePlugin {
 
     const servicePropagation = this.createDBMPropagationCommentService(dbmService, span)
 
-    if (isPreparedStatement || mode === 'service') {
+    if (disableFullMode || mode === 'service') {
       return servicePropagation
     } else if (mode === 'full') {
       span.setTag('_dd.dbm_trace_injected', 'true')
@@ -82,8 +82,8 @@ class DatabasePlugin extends StoragePlugin {
     }
   }
 
-  injectDbmQuery (span, query, serviceName, isPreparedStatement = false) {
-    const dbmTraceComment = this.createDbmComment(span, serviceName, isPreparedStatement)
+  injectDbmQuery (span, query, serviceName, disableFullMode = false) {
+    const dbmTraceComment = this.createDbmComment(span, serviceName, disableFullMode)
 
     if (!dbmTraceComment) {
       return query

package/packages/dd-trace/src/plugins/plugin.js

@@ -18,10 +18,12 @@ class Subscription {
   }
 
   enable () {
+    // TODO: Once Node.js v18.6.0 is no longer supported, we should use `dc.subscribe(event, handler)` instead
     this._channel.subscribe(this._handler)
   }
 
   disable () {
+    // TODO: Once Node.js v18.6.0 is no longer supported, we should use `dc.unsubscribe(event, handler)` instead
     this._channel.unsubscribe(this._handler)
   }
 }

package/packages/dd-trace/src/plugins/util/test.js

@@ -58,6 +58,7 @@ const TEST_IS_RETRY = 'test.is_retry'
 const TEST_EARLY_FLAKE_ENABLED = 'test.early_flake.enabled'
 const TEST_EARLY_FLAKE_ABORT_REASON = 'test.early_flake.abort_reason'
 const TEST_RETRY_REASON = 'test.retry_reason'
+const TEST_HAS_FAILED_ALL_RETRIES = 'test.has_failed_all_retries'
 
 const CI_APP_ORIGIN = 'ciapp-test'
 
@@ -103,6 +104,12 @@ const EFD_TEST_NAME_REGEX = new RegExp(EFD_STRING + ' \\(#\\d+\\): ', 'g')
 const DD_CAPABILITIES_TEST_IMPACT_ANALYSIS = '_dd.library_capabilities.test_impact_analysis'
 const DD_CAPABILITIES_EARLY_FLAKE_DETECTION = '_dd.library_capabilities.early_flake_detection'
 const DD_CAPABILITIES_AUTO_TEST_RETRIES = '_dd.library_capabilities.auto_test_retries'
+const DD_CAPABILITIES_TEST_MANAGEMENT_QUARANTINE = '_dd.library_capabilities.test_management.quarantine'
+const DD_CAPABILITIES_TEST_MANAGEMENT_DISABLE = '_dd.library_capabilities.test_management.disable'
+const DD_CAPABILITIES_TEST_MANAGEMENT_ATTEMPT_TO_FIX = '_dd.library_capabilities.test_management.attempt_to_fix'
+const UNSUPPORTED_TIA_FRAMEWORKS = ['playwright', 'vitest']
+const UNSUPPORTED_TIA_FRAMEWORKS_PARALLEL_MODE = ['cucumber', 'mocha']
+const UNSUPPORTED_ATTEMPT_TO_FIX_FRAMEWORKS_PARALLEL_MODE = ['mocha']
 
 const TEST_LEVEL_EVENT_TYPES = [
   'test',
@@ -120,9 +127,16 @@ const DI_DEBUG_ERROR_SNAPSHOT_ID_SUFFIX = 'snapshot_id'
 const DI_DEBUG_ERROR_FILE_SUFFIX = 'file'
 const DI_DEBUG_ERROR_LINE_SUFFIX = 'line'
 
+// Test Management tags
+const TEST_MANAGEMENT_IS_ATTEMPT_TO_FIX = 'test.test_management.is_attempt_to_fix'
 const TEST_MANAGEMENT_IS_DISABLED = 'test.test_management.is_test_disabled'
 const TEST_MANAGEMENT_IS_QUARANTINED = 'test.test_management.is_quarantined'
 const TEST_MANAGEMENT_ENABLED = 'test.test_management.enabled'
+const TEST_MANAGEMENT_ATTEMPT_TO_FIX_PASSED = 'test.test_management.attempt_to_fix_passed'
+
+// Test Management utils strings
+const ATTEMPT_TO_FIX_STRING = "Retried by Datadog's Test Management"
+const ATTEMPT_TEST_NAME_REGEX = new RegExp(ATTEMPT_TO_FIX_STRING + ' \\(#\\d+\\): ', 'g')
 
 module.exports = {
   TEST_CODE_OWNERS,
@@ -155,6 +169,7 @@ module.exports = {
   TEST_EARLY_FLAKE_ENABLED,
   TEST_EARLY_FLAKE_ABORT_REASON,
   TEST_RETRY_REASON,
+  TEST_HAS_FAILED_ALL_RETRIES,
   getTestEnvironmentMetadata,
   getTestParametersString,
   finishAllTraceSpans,
@@ -192,7 +207,9 @@ module.exports = {
   EFD_STRING,
   EFD_TEST_NAME_REGEX,
   removeEfdStringFromTestName,
+  removeAttemptToFixStringFromTestName,
   addEfdStringToTestName,
+  addAttemptToFixStringToTestName,
   getIsFaultyEarlyFlakeDetection,
   TEST_BROWSER_DRIVER,
   TEST_BROWSER_DRIVER_VERSION,
@@ -202,6 +219,9 @@ module.exports = {
   DD_CAPABILITIES_TEST_IMPACT_ANALYSIS,
   DD_CAPABILITIES_EARLY_FLAKE_DETECTION,
   DD_CAPABILITIES_AUTO_TEST_RETRIES,
+  DD_CAPABILITIES_TEST_MANAGEMENT_QUARANTINE,
+  DD_CAPABILITIES_TEST_MANAGEMENT_DISABLE,
+  DD_CAPABILITIES_TEST_MANAGEMENT_ATTEMPT_TO_FIX,
   TEST_LEVEL_EVENT_TYPES,
   getNumFromKnownTests,
   getFileAndLineNumberFromError,
@@ -212,9 +232,12 @@ module.exports = {
   DI_DEBUG_ERROR_LINE_SUFFIX,
   getFormattedError,
   DD_TEST_IS_USER_PROVIDED_SERVICE,
+  TEST_MANAGEMENT_IS_ATTEMPT_TO_FIX,
   TEST_MANAGEMENT_IS_DISABLED,
   TEST_MANAGEMENT_IS_QUARANTINED,
-  TEST_MANAGEMENT_ENABLED
+  TEST_MANAGEMENT_ENABLED,
+  TEST_MANAGEMENT_ATTEMPT_TO_FIX_PASSED,
+  getLibraryCapabilitiesTags
 }
 
 // Returns pkg manager and its version, separated by '-', e.g. npm-8.15.0 or yarn-1.22.19
@@ -622,10 +645,18 @@ function addEfdStringToTestName (testName, numAttempt) {
   return `${EFD_STRING} (#${numAttempt}): ${testName}`
 }
 
+function addAttemptToFixStringToTestName (testName, numAttempt) {
+  return `${ATTEMPT_TO_FIX_STRING} (#${numAttempt}): ${testName}`
+}
+
 function removeEfdStringFromTestName (testName) {
   return testName.replace(EFD_TEST_NAME_REGEX, '')
 }
 
+function removeAttemptToFixStringFromTestName (testName) {
+  return testName.replace(ATTEMPT_TEST_NAME_REGEX, '')
+}
+
 function getIsFaultyEarlyFlakeDetection (projectSuites, testsBySuiteName, faultyThresholdPercentage) {
   let newSuites = 0
   for (const suite of projectSuites) {
@@ -718,3 +749,33 @@ function getFormattedError (error, repositoryRoot) {
 
   return newError
 }
+
+function getLibraryCapabilitiesTags (testFramework, isParallel) {
+  function isTiaSupported (testFramework, isParallel) {
+    if (UNSUPPORTED_TIA_FRAMEWORKS.includes(testFramework)) {
+      return false
+    }
+    if (isParallel && UNSUPPORTED_TIA_FRAMEWORKS_PARALLEL_MODE.includes(testFramework)) {
+      return false
+    }
+    return true
+  }
+
+  function isAttemptToFixSupported (testFramework, isParallel) {
+    if (isParallel && UNSUPPORTED_ATTEMPT_TO_FIX_FRAMEWORKS_PARALLEL_MODE.includes(testFramework)) {
+      return false
+    }
+    return true
+  }
+
+  return {
+    [DD_CAPABILITIES_TEST_IMPACT_ANALYSIS]: isTiaSupported(testFramework, isParallel) ? '1' : undefined,
+    [DD_CAPABILITIES_EARLY_FLAKE_DETECTION]: '1',
+    [DD_CAPABILITIES_AUTO_TEST_RETRIES]: '1',
+    [DD_CAPABILITIES_TEST_MANAGEMENT_QUARANTINE]: '1',
+    [DD_CAPABILITIES_TEST_MANAGEMENT_DISABLE]: '1',
+    [DD_CAPABILITIES_TEST_MANAGEMENT_ATTEMPT_TO_FIX]: isAttemptToFixSupported(testFramework, isParallel)
+      ? '2'
+      : undefined
+  }
+}