dd-trace 6.12.0 → 6.14.0

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

@@ -5,7 +5,7 @@ const { storage } = require('../../datadog-core')
 const tags = require('../../../ext/tags')
 const formats = require('../../../ext/formats')
 const HTTP_HEADERS = formats.HTTP_HEADERS
-const log = require('../../dd-trace/src/log')
+const { getClientStatusValidator } = require('../../dd-trace/src/plugins/util/status-validator')
 const { buildClientHttpUrl } = require('../../dd-trace/src/plugins/util/url')
 const { stripQueryAndFragment } = require('../../dd-trace/src/util')
 const { CLIENT_PORT_KEY } = require('../../dd-trace/src/constants')
@@ -16,6 +16,8 @@ const {
   HTTP_RESPONSE_HEADERS,
 } = tags
 
+const UPGRADE_CHANNEL = 'apm:undici:request:upgrade'
+
 // WeakMap to store span context for native undici request objects
 const requestContexts = new WeakMap()
 
@@ -33,6 +35,7 @@ class UndiciPlugin extends HttpClientPlugin {
     this.addSub('undici:request:headers', this.#onNativeRequestHeaders.bind(this))
     this.addSub('undici:request:trailers', this.#onNativeRequestTrailers.bind(this))
     this.addSub('undici:request:error', this.#onNativeRequestError.bind(this))
+    this.addSub(UPGRADE_CHANNEL, this.#onNativeRequestUpgrade.bind(this))
   }
 
   // ===========================================
@@ -112,30 +115,15 @@ class UndiciPlugin extends HttpClientPlugin {
     // Store span context for request for later retrieval
     requestContexts.set(request, {
       span,
-      store,
-      uri,
       method,
     })
-
-    // Enter the span context
-    storage('legacy').enterWith({ ...store, span })
   }
 
   #onNativeRequestBodySent ({ request }) {
     const ctx = requestContexts.get(request)
     if (!ctx || ctx.method !== 'CONNECT') return
 
-    const { span, store } = ctx
-
-    this.config.hooks.request(span, null, null)
-
-    span.finish()
-
-    requestContexts.delete(request)
-
-    if (store) {
-      storage('legacy').enterWith(store)
-    }
+    this.#finishNativeRequest(request)
   }
 
   #onNativeRequestHeaders ({ request, response }) {
@@ -160,50 +148,48 @@ class UndiciPlugin extends HttpClientPlugin {
   }
 
   #onNativeRequestTrailers ({ request }) {
-    const ctx = requestContexts.get(request)
-    if (!ctx) return
-
-    const { span, store } = ctx
-
-    // Call the request hook if configured
-    this.config.hooks.request(span, null, null)
-
-    // Finish the span
-    span.finish()
+    this.#finishNativeRequest(request)
+  }
 
-    // Clean up
-    requestContexts.delete(request)
+  #onNativeRequestError ({ request, error }) {
+    this.#finishNativeRequest(request, error)
+  }
 
-    // Restore parent store
-    if (store) {
-      storage('legacy').enterWith(store)
-    }
+  /**
+   * @param {{ error?: Error, headers: unknown, request: object, statusCode: number }} context
+   */
+  #onNativeRequestUpgrade ({ error, headers, request, statusCode }) {
+    this.#onNativeRequestHeaders({ request, response: { headers, statusCode } })
+    this.#finishNativeRequest(request, error)
   }
 
-  #onNativeRequestError ({ request, error }) {
+  /**
+   * @param {object} request
+   * @param {Error} [error]
+   */
+  #finishNativeRequest (request, error) {
     const ctx = requestContexts.get(request)
     if (!ctx) return
 
-    const { span, store } = ctx
+    const { span } = ctx
 
     // Don't record AbortError as an error - it's user-initiated cancellation
     if (error && error.name !== 'AbortError') {
       span.setTag('error', error)
     }
 
-    // Call the request hook if configured
-    this.config.hooks.request(span, null, null)
+    const requestHook = this.config.hooks.request
+    if (requestHook !== noop) {
+      this.tracer.scope().activate(span, () => {
+        requestHook(span, null, null)
+      })
+    }
 
     // Finish the span
     span.finish()
 
     // Clean up
     requestContexts.delete(request)
-
-    // Restore parent store
-    if (store) {
-      storage('legacy').enterWith(store)
-    }
   }
 
   // ===========================================
@@ -302,7 +288,7 @@ function normalizeHeaders (headers) {
 }
 
 function normalizeConfig (config) {
-  const validateStatus = getStatusValidator(config)
+  const validateStatus = getClientStatusValidator(config)
   const hooks = getHooks(config)
 
   return {
@@ -312,19 +298,6 @@ function normalizeConfig (config) {
   }
 }
 
-function getStatusValidator (config) {
-  if (typeof config.validateStatus === 'function') {
-    return config.validateStatus
-  } else if (Object.hasOwn(config, 'validateStatus')) {
-    log.error('Expected `validateStatus` to be a function.')
-  }
-  return defaultValidateStatus
-}
-
-function defaultValidateStatus (code) {
-  return code < 400 || code >= 500
-}
-
 function getHooks (config) {
   const request = config.hooks?.request ?? noop
 

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

@@ -361,13 +361,13 @@ class VitestPlugin extends CiPlugin {
         [TEST_SOURCE_START]: testStartLine || 1,
         [TEST_STATUS]: 'skip',
         [TEST_FINAL_STATUS]: 'skip',
-        ...(isAttemptToFix ? { [TEST_MANAGEMENT_IS_ATTEMPT_TO_FIX]: 'true' } : {}),
-        ...(isDisabled ? { [TEST_MANAGEMENT_IS_DISABLED]: 'true' } : {}),
-        ...(isQuarantined ? { [TEST_MANAGEMENT_IS_QUARANTINED]: 'true' } : {}),
-        ...(isNew ? { [TEST_IS_NEW]: 'true' } : {}),
-        ...(isRumActive ? { [TEST_IS_RUM_ACTIVE]: 'true' } : {}),
-        ...(isTestFrameworkWorker ? { [TEST_IS_TEST_FRAMEWORK_WORKER]: 'true' } : {}),
       }
+      if (isAttemptToFix) extraTags[TEST_MANAGEMENT_IS_ATTEMPT_TO_FIX] = 'true'
+      if (isDisabled) extraTags[TEST_MANAGEMENT_IS_DISABLED] = 'true'
+      if (isQuarantined) extraTags[TEST_MANAGEMENT_IS_QUARANTINED] = 'true'
+      if (isNew) extraTags[TEST_IS_NEW] = 'true'
+      if (isRumActive) extraTags[TEST_IS_RUM_ACTIVE] = 'true'
+      if (isTestFrameworkWorker) extraTags[TEST_IS_TEST_FRAMEWORK_WORKER] = 'true'
       setBrowserTags(extraTags, {
         browserDriver,
         browserName,
@@ -523,7 +523,6 @@ class VitestPlugin extends CiPlugin {
           this.telemetry.ciVisEvent(TELEMETRY_CODE_COVERAGE_FINISHED, 'suite', { library: coverageLibrary })
           this.telemetry.distribution(TELEMETRY_CODE_COVERAGE_NUM_FILES, {}, relativeFiles.length)
         }
-        this.tracer._exporter.deferTestSuiteSpan?.(testSuiteSpan)
         testSuiteSpan.finish()
         finishAllTraceSpans(testSuiteSpan)
       }
@@ -567,7 +566,6 @@ class VitestPlugin extends CiPlugin {
     this.addSub('ci:vitest:session:finish', ({
       status,
       error,
-      isTestSessionFinalizationError,
       testCodeCoverageLinesTotal,
       isEarlyFlakeDetectionEnabled,
       isEarlyFlakeDetectionFaulty,
@@ -590,9 +588,6 @@ class VitestPlugin extends CiPlugin {
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
       if (error) {
-        if (isTestSessionFinalizationError) {
-          this.tracer._exporter.setDeferredTestSuiteError?.(error)
-        }
         this.testModuleSpan.setTag('error', error)
         this.testSessionSpan.setTag('error', error)
       }
@@ -622,7 +617,6 @@ class VitestPlugin extends CiPlugin {
       if (vitestPool) {
         this.testSessionSpan.setTag(VITEST_POOL, vitestPool)
       }
-      this.tracer._exporter.exportDeferredTestSuiteSpans?.()
       this.testModuleSpan.finish()
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'module')
       this.testSessionSpan.finish()

package/packages/dd-trace/src/aiguard/evaluation.js

@@ -1,18 +1,22 @@
 'use strict'
 
 const { HTTP_CLIENT_IP, HTTP_USERAGENT, NETWORK_CLIENT_IP } = require('../../../../ext/tags')
-const clone = require('../../../../vendor/dist/rfdc')({ proto: false, circles: false })
+const createRfdc = require('../../../../vendor/dist/rfdc')
+const clone = createRfdc({ proto: false, circles: false })
 const { USER_ID, USER_SESSION_ID } = require('../appsec/addresses')
 const { getActiveRequest } = require('../appsec/store')
 const { keepTrace } = require('../priority_sampler')
 const { extractIp } = require('../plugins/util/ip_extractor')
 const { AI_GUARD } = require('../standalone/product')
 const telemetryMetrics = require('../telemetry/metrics')
+const { normalizeRedactionReplacements, redactMessages } = require('./redaction')
 const TAGS = require('./tags')
 
 const ALLOW = 'ALLOW'
 
+/** @typedef {import('../../../../index').aiguard.ContentPart} ContentPart */
 /** @typedef {import('../../../../index').aiguard.Message} Message */
+/** @typedef {import('../../../../index').aiguard.RedactionReplacement} RedactionReplacement */
 /** @typedef {import('../opentracing/span')} Span */
 
 /**
@@ -24,19 +28,22 @@ const ALLOW = 'ALLOW'
  * @property {Record<string, unknown>} tagProbabilities
  * @property {boolean} hasTagProbabilities
  * @property {boolean} blockingEnabled
+ * @property {unknown} redactionReplacements
  */
 
 /**
  * @typedef {object} EvaluationOutcome
  * @property {{ action: string, reason: string|undefined, tags: Array<unknown>,
- *   tagProbabilities: Record<string, unknown>, sds: Array<unknown> }} result
+ *   tagProbabilities: Record<string, unknown>, sds: Array<unknown>, messages: Message[],
+ *   redactionReplacements: RedactionReplacement[] }} result
  * @property {boolean} shouldBlock
  * @property {boolean} hasTagProbabilities
+ * @property {{ enabled: boolean, applied: boolean, failures: number }} redaction
  */
 
 /**
  * @typedef {object} EvaluationMetaStruct
- * @property {Message[]} messages
+ * @property {Message[]} [messages]
  * @property {Array<unknown>} [attack_categories]
  * @property {Array<unknown>} [sds]
  * @property {Record<string, unknown>} [tag_probs]
@@ -45,6 +52,7 @@ const ALLOW = 'ALLOW'
 /**
  * @typedef {object} EvaluationReport
  * @property {Span} span
+ * @property {Message[]} messages
  * @property {EvaluationMetaStruct} metaStruct
  * @property {{ source: string, integration: string }} telemetryTags
  */
@@ -79,6 +87,7 @@ function parseEvaluationResponse (body) {
     tagProbabilities: isRecord(attributes.tag_probs) ? attributes.tag_probs : {},
     hasTagProbabilities: isRecord(attributes.tag_probs),
     blockingEnabled: attributes.is_blocking_enabled === true,
+    redactionReplacements: attributes.redaction_replacements,
   }
 }
 
@@ -96,11 +105,16 @@ function shouldBlockEvaluation (block, evaluation) {
 /**
  * Applies local evaluation policy and creates the result returned to the caller.
  *
+ * @param {Message[]} messages
  * @param {EvaluationResponse} evaluation
- * @param {boolean} block
+ * @param {{ block: boolean, redactionEnabled: boolean }} options
  * @returns {EvaluationOutcome}
  */
-function createEvaluationOutcome (evaluation, block) {
+function createEvaluationOutcome (messages, evaluation, options) {
+  const redaction = options.redactionEnabled
+    ? redactMessages(messages, evaluation.redactionReplacements)
+    : { messages, redacted: false, failures: 0 }
+
   return {
     result: {
       action: evaluation.action,
@@ -108,9 +122,16 @@ function createEvaluationOutcome (evaluation, block) {
       tags: evaluation.tags,
       tagProbabilities: evaluation.tagProbabilities,
       sds: evaluation.sdsFindings,
+      messages: redaction.messages,
+      redactionReplacements: normalizeRedactionReplacements(evaluation.redactionReplacements),
     },
-    shouldBlock: shouldBlockEvaluation(block, evaluation),
+    shouldBlock: shouldBlockEvaluation(options.block, evaluation),
     hasTagProbabilities: evaluation.hasTagProbabilities,
+    redaction: {
+      enabled: options.redactionEnabled,
+      applied: redaction.redacted,
+      failures: redaction.failures,
+    },
   }
 }
 
@@ -130,6 +151,7 @@ class EvaluationReporter {
   #config
   #maxMessagesLength
   #maxContentSize
+  #redactionEnabled
 
   /**
    * @param {import('../config/config-base')} config
@@ -138,6 +160,7 @@ class EvaluationReporter {
     this.#config = config
     this.#maxMessagesLength = config.experimental.aiguard.maxMessagesLength
     this.#maxContentSize = config.experimental.aiguard.maxContentSize
+    this.#redactionEnabled = config.experimental.aiguard.redactionEnabled
   }
 
   /**
@@ -150,19 +173,18 @@ class EvaluationReporter {
    */
   start (span, messages, options) {
     const telemetryTags = { source: options.source, integration: options.integration }
-    const last = messages.at(-1)
+    const evaluatedMessages = this.#redactionEnabled ? clone(messages) : messages
+    const last = evaluatedMessages.at(-1)
     const target = this.#isToolCall(last) ? 'tool' : 'prompt'
     span.setTag(TAGS.TARGET_TAG_KEY, target)
     if (target === 'tool') {
-      const name = this.#getToolName(last, messages)
+      const name = this.#getToolName(last, evaluatedMessages)
       if (name) {
         span.setTag(TAGS.TOOL_NAME_TAG_KEY, name)
       }
     }
 
-    const metaStruct = {
-      messages: this.#buildMessagesForMetaStruct(messages, telemetryTags),
-    }
+    const metaStruct = {}
     span.meta_struct = {
       [TAGS.META_STRUCT_KEY]: metaStruct,
     }
@@ -178,6 +200,7 @@ class EvaluationReporter {
 
     return {
       span,
+      messages: evaluatedMessages,
       metaStruct,
       telemetryTags,
     }
@@ -191,6 +214,7 @@ class EvaluationReporter {
    * @returns {void}
    */
   fail (report, errorType) {
+    report.metaStruct.messages = this.#buildMessagesForMetaStruct(report.messages, report.telemetryTags)
     aiguardMetrics.count(TAGS.TELEMETRY_REQUESTS, { error: true, ...report.telemetryTags }).inc(1)
     aiguardMetrics.count(TAGS.TELEMETRY_ERROR, { type: errorType, ...report.telemetryTags }).inc(1)
   }
@@ -203,18 +227,38 @@ class EvaluationReporter {
    * @returns {void}
    */
   finish (report, outcome) {
-    const { result, shouldBlock } = outcome
+    const { result, redaction, shouldBlock } = outcome
 
     if (result.tags.length > 0) report.metaStruct.attack_categories = result.tags
     if (result.sds.length > 0) report.metaStruct.sds = result.sds
     if (outcome.hasTagProbabilities) report.metaStruct.tag_probs = result.tagProbabilities
 
-    const requestTelemetryTags = {
-      action: result.action,
-      error: false,
-      block: shouldBlock,
-      ...report.telemetryTags,
+    if (redaction.enabled) {
+      report.span.setTag(TAGS.REDACTED_TAG_KEY, redaction.applied ? 'true' : 'false')
+      if (redaction.failures > 0) {
+        aiguardMetrics.count(TAGS.TELEMETRY_ERROR, {
+          type: TAGS.ERROR_TYPE_REDACTION,
+          ...report.telemetryTags,
+        }).inc(redaction.failures)
+      }
     }
+
+    report.metaStruct.messages = this.#buildMessagesForMetaStruct(result.messages, report.telemetryTags)
+
+    const requestTelemetryTags = redaction.enabled
+      ? {
+          action: result.action,
+          error: false,
+          block: shouldBlock,
+          redacted: redaction.applied,
+          ...report.telemetryTags,
+        }
+      : {
+          action: result.action,
+          error: false,
+          block: shouldBlock,
+          ...report.telemetryTags,
+        }
     aiguardMetrics.count(TAGS.TELEMETRY_REQUESTS, requestTelemetryTags).inc(1)
 
     report.span.setTag(TAGS.ACTION_TAG_KEY, result.action)
@@ -242,10 +286,7 @@ class EvaluationReporter {
     let contentTruncated = false
     for (let i = messages.length - size; i < messages.length; i++) {
       const message = clone(messages[i])
-      if (message.content?.length > this.#maxContentSize) {
-        contentTruncated = true
-        message.content = message.content.slice(0, this.#maxContentSize)
-      }
+      if (this.#truncateMessageContent(message)) contentTruncated = true
       result.push(message)
     }
     if (contentTruncated) {
@@ -254,6 +295,40 @@ class EvaluationReporter {
     return result
   }
 
+  /**
+   * Truncates text in a cloned message to one shared content-size limit.
+   *
+   * @param {{ content?: string|ContentPart[] }} message
+   * @returns {boolean}
+   */
+  #truncateMessageContent (message) {
+    const { content } = message
+    if (typeof content === 'string') {
+      if (content.length <= this.#maxContentSize) return false
+
+      message.content = content.slice(0, this.#maxContentSize)
+      return true
+    }
+
+    if (!Array.isArray(content)) return false
+
+    let remainingContentSize = this.#maxContentSize
+    let truncated = false
+    for (const part of content) {
+      const text = part?.text
+      if (typeof text !== 'string') continue
+
+      if (text.length > remainingContentSize) {
+        part.text = text.slice(0, remainingContentSize)
+        truncated = true
+        remainingContentSize = 0
+      } else {
+        remainingContentSize -= text.length
+      }
+    }
+    return truncated
+  }
+
   /**
    * Returns whether a message represents a tool call or tool output.
    *

package/packages/dd-trace/src/aiguard/integrations/vercel-ai.js

@@ -74,18 +74,17 @@ function onStreamAfter (ctx) {
  */
 function getStreamContent (chunks) {
   const toolCalls = []
-  const textParts = []
+  let text = ''
 
   for (const chunk of chunks) {
     if (chunk?.type === 'tool-call') {
       toolCalls.push(chunk)
     } else if (chunk?.type === 'text-delta') {
-      textParts.push(chunk.textDelta)
+      text += chunk.delta ?? chunk.textDelta ?? ''
     }
   }
 
   if (toolCalls.length) return toolCalls
-  const text = textParts.join('')
   return text ? [{ type: 'text', text }] : []
 }
 

package/packages/dd-trace/src/aiguard/messages/anthropic.js

@@ -208,7 +208,14 @@ function walkContentBlocks (blocks) {
 function partsToContent (parts, hasImages) {
   if (!parts.length) return
   if (hasImages) return parts
-  return parts.map(p => p.text).join('\n')
+  let content = ''
+  let isFirstPart = true
+  for (const part of parts) {
+    if (!isFirstPart) content += '\n'
+    content += part.text
+    isFirstPart = false
+  }
+  return content
 }
 
 /**
@@ -269,27 +276,39 @@ function convertServerToolResultContent (content) {
     return content.error_message ? `${content.error_code}: ${content.error_message}` : content.error_code
   }
 
-  const lines = []
+  let lines
   if (Array.isArray(content)) {
     for (const item of content) {
       if (!item || typeof item !== 'object') continue
       // text blocks (MCP / generic), then web-search title + url.
-      if (typeof item.text === 'string') lines.push(item.text)
-      if (typeof item.title === 'string') lines.push(item.title)
-      if (typeof item.url === 'string') lines.push(item.url)
+      if (typeof item.text === 'string') {
+        lines = lines === undefined ? item.text : `${lines}\n${item.text}`
+      }
+      if (typeof item.title === 'string') {
+        lines = lines === undefined ? item.title : `${lines}\n${item.title}`
+      }
+      if (typeof item.url === 'string') {
+        lines = lines === undefined ? item.url : `${lines}\n${item.url}`
+      }
     }
   } else {
-    if (typeof content.stdout === 'string' && content.stdout) lines.push(content.stdout)
-    if (typeof content.stderr === 'string' && content.stderr) lines.push(content.stderr)
-    if (typeof content.content === 'string' && content.content) lines.push(content.content)
+    if (typeof content.stdout === 'string' && content.stdout) lines = content.stdout
+    if (typeof content.stderr === 'string' && content.stderr) {
+      lines = lines === undefined ? content.stderr : `${lines}\n${content.stderr}`
+    }
+    if (typeof content.content === 'string' && content.content) {
+      lines = lines === undefined ? content.content : `${lines}\n${content.content}`
+    }
     if (Array.isArray(content.lines)) {
       for (const line of content.lines) {
-        if (typeof line === 'string') lines.push(line)
+        if (typeof line === 'string') {
+          lines = lines === undefined ? line : `${lines}\n${line}`
+        }
       }
     }
   }
 
-  return lines.join('\n') || '[tool result]'
+  return lines || '[tool result]'
 }
 
 /**

package/packages/dd-trace/src/aiguard/messages/openai.js

@@ -322,7 +322,14 @@ function openAIResponseContentToMessageContent (content) {
 
   if (!parts.length) return
   if (hasImages) return parts
-  return parts.map(part => part.text).join('\n')
+  let textContent = ''
+  let isFirstPart = true
+  for (const part of parts) {
+    if (!isFirstPart) textContent += '\n'
+    textContent += part.text
+    isFirstPart = false
+  }
+  return textContent
 }
 
 /**

package/packages/dd-trace/src/aiguard/messages/vercel-ai.js

@@ -62,19 +62,26 @@ function convertVercelPromptToMessages (prompt) {
         if (hasImages) {
           messages.push({ role: 'user', content: contentParts })
         } else {
-          messages.push({ role: 'user', content: contentParts.map(p => p.text).join('\n') })
+          let content = ''
+          let isFirstPart = true
+          for (const part of contentParts) {
+            if (!isFirstPart) content += '\n'
+            content += part.text
+            isFirstPart = false
+          }
+          messages.push({ role: 'user', content })
         }
         break
       }
 
       case 'assistant': {
-        const textParts = []
+        let text
         const toolCalls = []
         if (!Array.isArray(msg.content)) break
 
         for (const part of msg.content) {
           if (part.type === 'text') {
-            textParts.push(part.text)
+            text = text === undefined ? part.text : `${text}\n${part.text}`
           } else if (part.type === 'tool-call') {
             toolCalls.push({
               id: part.toolCallId,
@@ -88,8 +95,8 @@ function convertVercelPromptToMessages (prompt) {
 
         if (toolCalls.length > 0) {
           messages.push({ role: 'assistant', tool_calls: toolCalls })
-        } else if (textParts.length > 0) {
-          messages.push({ role: 'assistant', content: textParts.join('\n') })
+        } else if (text !== undefined) {
+          messages.push({ role: 'assistant', content: text })
         }
         break
       }
@@ -161,7 +168,14 @@ function buildTextOutputMessages (inputMessages, text) {
 function buildOutputMessages (inputMessages, content) {
   const toolCalls = content.filter(c => c.type === 'tool-call')
   if (toolCalls.length) return buildToolCallOutputMessages(inputMessages, toolCalls)
-  const text = content.filter(c => c.type === 'text').map(c => c.text).join('\n')
+  let text = ''
+  let isFirstTextPart = true
+  for (const part of content) {
+    if (part.type !== 'text') continue
+    if (!isFirstTextPart) text += '\n'
+    text += part.text
+    isFirstTextPart = false
+  }
   if (text) return buildTextOutputMessages(inputMessages, text)
   return []
 }

package/packages/dd-trace/src/aiguard/noop.js

@@ -2,7 +2,15 @@
 
 class NoopAIGuard {
   evaluate (messages, opts) {
-    return Promise.resolve({ action: 'ALLOW', reason: 'AI Guard is not enabled', tags: [], sds: [] })
+    return Promise.resolve({
+      action: 'ALLOW',
+      reason: 'AI Guard is not enabled',
+      tags: [],
+      tagProbabilities: {},
+      sds: [],
+      messages,
+      redactionReplacements: [],
+    })
   }
 }