dd-trace 6.15.0 → 6.17.0

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

@@ -9,7 +9,6 @@ const TAGS = require('./tags')
  * Resolves the AI Guard host for a Datadog site.
  *
  * @param {string} site
- * @returns {string}
  */
 function aiGuardHost (site) {
   return site.split('.').length === 2 ? `app.${site}` : site
@@ -60,9 +59,9 @@ class AIGuardClient {
       'DD-AI-GUARD-SOURCE': 'SDK',
       'DD-AI-GUARD-LANGUAGE': 'nodejs',
     }
-    const endpoint = config.experimental.aiguard.endpoint || `https://${aiGuardHost(config.site)}/api/v2/ai-guard`
+    const endpoint = config.aiguard.DD_AI_GUARD_ENDPOINT || `https://${aiGuardHost(config.site)}/api/v2/ai-guard`
     this.#evaluateUrl = `${endpoint}/evaluate`
-    this.#timeout = config.experimental.aiguard.timeout
+    this.#timeout = config.aiguard.DD_AI_GUARD_TIMEOUT
   }
 
   /**

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

@@ -9,6 +9,7 @@ const { keepTrace } = require('../priority_sampler')
 const { extractIp } = require('../plugins/util/ip_extractor')
 const { AI_GUARD } = require('../standalone/product')
 const telemetryMetrics = require('../telemetry/metrics')
+const { truncateString } = require('../util')
 const { normalizeRedactionReplacements, redactMessages } = require('./redaction')
 const TAGS = require('./tags')
 
@@ -96,7 +97,6 @@ function parseEvaluationResponse (body) {
  *
  * @param {boolean} block
  * @param {{ action: string, blockingEnabled: boolean }} evaluation
- * @returns {boolean}
  */
 function shouldBlockEvaluation (block, evaluation) {
   return block && evaluation.blockingEnabled && evaluation.action !== ALLOW
@@ -158,9 +158,9 @@ class EvaluationReporter {
    */
   constructor (config) {
     this.#config = config
-    this.#maxMessagesLength = config.experimental.aiguard.maxMessagesLength
-    this.#maxContentSize = config.experimental.aiguard.maxContentSize
-    this.#redactionEnabled = config.experimental.aiguard.redactionEnabled
+    this.#maxMessagesLength = config.aiguard.DD_AI_GUARD_MAX_MESSAGES_LENGTH
+    this.#maxContentSize = config.aiguard.DD_AI_GUARD_MAX_CONTENT_SIZE
+    this.#redactionEnabled = config.aiguard.DD_AI_GUARD_REDACTION_ENABLED
   }
 
   /**
@@ -211,7 +211,6 @@ class EvaluationReporter {
    *
    * @param {EvaluationReport} report
    * @param {string} errorType
-   * @returns {void}
    */
   fail (report, errorType) {
     report.metaStruct.messages = this.#buildMessagesForMetaStruct(report.messages, report.telemetryTags)
@@ -224,7 +223,6 @@ class EvaluationReporter {
    *
    * @param {EvaluationReport} report
    * @param {EvaluationOutcome} outcome
-   * @returns {void}
    */
   finish (report, outcome) {
     const { result, redaction, shouldBlock } = outcome
@@ -299,14 +297,13 @@ class EvaluationReporter {
    * 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)
+      message.content = truncateString(content, this.#maxContentSize)
       return true
     }
 
@@ -319,7 +316,7 @@ class EvaluationReporter {
       if (typeof text !== 'string') continue
 
       if (text.length > remainingContentSize) {
-        part.text = text.slice(0, remainingContentSize)
+        part.text = truncateString(text, remainingContentSize)
         truncated = true
         remainingContentSize = 0
       } else {
@@ -333,7 +330,6 @@ class EvaluationReporter {
    * Returns whether a message represents a tool call or tool output.
    *
    * @param {Message} message
-   * @returns {boolean}
    */
   #isToolCall (message) {
     return Boolean(message.tool_calls || message.tool_call_id)
@@ -370,7 +366,6 @@ class EvaluationReporter {
    * Adds missing client IP tags to the service entry span.
    *
    * @param {Span} rootSpan
-   * @returns {void}
    */
   #setRootSpanClientIpTags (rootSpan) {
     const currentTags = rootSpan.context().getTags()
@@ -411,7 +406,6 @@ class EvaluationReporter {
    *
    * @param {Span} guardSpan
    * @param {Span} rootSpan
-   * @returns {void}
    */
   #copyServiceEntryTagsToGuardSpan (guardSpan, rootSpan) {
     const rootTags = rootSpan.context().getTags()

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

@@ -17,10 +17,11 @@ function enable (tracer, config) {
 
   try {
     aiguard = new AIGuard(tracer, config)
-    const block = config.experimental?.aiguard?.block !== false
+    const block = config.aiguard.DD_AI_GUARD_BLOCK !== false
+    const analyzeStreamResponses = config.aiguard.DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED
 
     incomingHttpRequestStart.subscribe(onIncomingHttpRequestStart)
-    integrations.enable(aiguard, block)
+    integrations.enable(aiguard, block, analyzeStreamResponses)
 
     isEnabled = true
   } catch (err) {

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

@@ -2,20 +2,20 @@
 
 const { channel } = require('dc-polyfill')
 
-const log = require('../../log')
 const { getMessagesInputMessages, getMessagesOutputMessages } = require('../messages/anthropic')
+const { decode } = require('../messages/utils')
 const { SOURCE_AUTO } = require('../tags')
-const { pushEvaluation } = require('./evaluate')
+const { evaluate } = require('./evaluate')
 
-const messagesBeforeChannel = channel('dd-trace:anthropic:messages:before')
-const messagesAfterChannel = channel('dd-trace:anthropic:messages:after')
+const messagesPrepareChannel = channel('dd-trace:anthropic:messages:prepare')
+const messagesInterceptChannel = channel('dd-trace:anthropic:messages:intercept')
 
 let isEnabled = false
 let aiguard
 let opts
 
 /**
- * Subscribes AI Guard to Anthropic lifecycle channels.
+ * Subscribes AI Guard to the Anthropic interception channels.
  *
  * @param {object} aiguardInstance
  * @param {boolean} block
@@ -26,8 +26,8 @@ function enable (aiguardInstance, block) {
   aiguard = aiguardInstance
   opts = { block, source: SOURCE_AUTO, integration: 'anthropic' }
 
-  messagesBeforeChannel.subscribe(onMessagesBefore)
-  messagesAfterChannel.subscribe(onMessagesAfter)
+  messagesPrepareChannel.subscribe(onMessagesPrepare)
+  messagesInterceptChannel.subscribe(onMessagesIntercept)
 
   isEnabled = true
 }
@@ -35,36 +35,64 @@ function enable (aiguardInstance, block) {
 function disable () {
   if (!isEnabled) return
 
-  messagesBeforeChannel.unsubscribe(onMessagesBefore)
-  messagesAfterChannel.unsubscribe(onMessagesAfter)
+  messagesPrepareChannel.unsubscribe(onMessagesPrepare)
+  messagesInterceptChannel.unsubscribe(onMessagesIntercept)
 
   aiguard = undefined
   opts = undefined
   isEnabled = false
 }
 
-function onMessagesBefore (ctx) {
-  pushEvaluation(ctx, aiguard, getMessagesInputMessages(ctx.args?.[0]), opts)
+/**
+ * Replaces the outgoing options with a JSON snapshot of the evaluated fields, so AI Guard judges
+ * exactly what the SDK serializes and sends, immune to later caller mutation.
+ *
+ * @param {{ arguments: Array<unknown> }} ctx
+ */
+function onMessagesPrepare (ctx) {
+  const options = ctx.arguments[0]
+  if (!options || typeof options !== 'object') return
+
+  const evaluated = { messages: options.messages }
+  if (options.system !== undefined) evaluated.system = options.system
+
+  try {
+    // eslint-disable-next-line unicorn/prefer-structured-clone
+    ctx.arguments[0] = { ...options, ...JSON.parse(JSON.stringify(evaluated)) }
+  } catch {
+    // Unserializable input — judge the caller's options as they are.
+  }
 }
 
-function onMessagesAfter (ctx) {
-  const inputMessages = getMessagesInputMessages(ctx.args?.[0])
+function onMessagesIntercept (ctx) {
+  const inputMessages = getMessagesInputMessages(ctx.arguments?.[0])
   if (!inputMessages?.length) return
 
-  let body = ctx.body
-  if (typeof body === 'string') {
-    try {
-      body = JSON.parse(body)
-    } catch {
-      log.error('AIGuard: unable to decode Anthropic response body')
-      return
-    }
+  // `parse` and `asResponse` are both wrapped, and either may run more than once per call.
+  let inputEvaluation
+  ctx.beforeResult = () => {
+    if (!isEnabled) return
+    inputEvaluation ??= evaluate(ctx, aiguard, [inputMessages], opts)
+    return inputEvaluation
   }
 
-  const outputMessages = getMessagesOutputMessages(body)
-  if (!outputMessages.length) return
-
-  pushEvaluation(ctx, aiguard, [...inputMessages, ...outputMessages], opts)
+  // One model call has one output however many readers observe it: `parse`, `json()`, `text()`
+  // and every `clone()` of the raw response share this callback.
+  let outputEvaluation
+  ctx.onResult = body => {
+    if (!isEnabled) return body
+
+    const outputMessages = decode(
+      () => getMessagesOutputMessages(body),
+      null,
+      'AIGuard: unable to decode Anthropic response body: %s'
+    )
+    if (outputMessages === null) return body
+    if (!outputMessages.length) return body
+
+    outputEvaluation ??= evaluate(ctx, aiguard, [[...inputMessages, ...outputMessages]], opts)
+    return outputEvaluation.then(() => body)
+  }
 }
 
 module.exports = { enable, disable }

package/packages/dd-trace/src/aiguard/integrations/evaluate.js

@@ -3,44 +3,68 @@
 const log = require('../../log')
 
 /**
- * Starts one AI Guard evaluation for a lifecycle ctx.
+ * Runs one evaluation per conversation and rejects if any of them says block.
  *
- * Async evaluations are pushed synchronously during channel publication. If
- * evaluation throws synchronously, the error is handled before publish returns.
+ * Evaluations run concurrently and all of them are awaited, so a slow one cannot hide a block
+ * from a fast one. Any failure that is not a verdict is logged and ignored, because AI Guard
+ * must never break the caller.
  *
- * @param {object} ctx
- * @param {AbortController} ctx.abortController
- * @param {object} [ctx.parentSpan]
- * @param {Array<Promise<void>>} ctx.pending
+ * @param {object} interceptCtx The intercept payload, whose tracing ctx parents the span.
  * @param {object} aiguard
- * @param {Array<object>|undefined} messages
+ * @param {Array<Array<object>>} conversations One complete message list per evaluation.
  * @param {object} opts
+ * @returns {Promise<void>|undefined} `undefined` when there is nothing to evaluate.
  */
-function pushEvaluation (ctx, aiguard, messages, opts) {
-  if (!messages?.length) return
+function evaluate (interceptCtx, aiguard, conversations, opts) {
+  if (!conversations?.length) return
 
-  const evaluateOpts = ctx.parentSpan ? { ...opts, childOf: ctx.parentSpan } : opts
+  const parentSpan = interceptCtx?.tracingContext?.currentStore?.span
+  const evaluateOpts = parentSpan ? { ...opts, childOf: parentSpan } : opts
+  const evaluations = []
 
-  try {
-    ctx.pending.push(aiguard.evaluate(messages, evaluateOpts).catch(handleEvaluationError.bind(null, ctx)))
-  } catch (err) {
-    handleEvaluationError(ctx, err)
+  for (const messages of conversations) {
+    if (messages?.length) evaluations.push(evaluateOne(aiguard, messages, evaluateOpts))
   }
+
+  if (evaluations.length === 0) return
+
+  return Promise.allSettled(evaluations).then(throwFirstBlock)
 }
 
 /**
- * Handles an AI Guard evaluation failure.
+ * Turns a synchronous throw into a rejection, so every outcome is inspected the same way.
  *
- * @param {object} ctx
- * @param {AbortController} ctx.abortController
- * @param {Error} err
+ * @param {object} aiguard
+ * @param {Array<object>} messages
+ * @param {object} opts
+ * @returns {Promise<unknown>}
+ */
+function evaluateOne (aiguard, messages, opts) {
+  try {
+    return aiguard.evaluate(messages, opts)
+  } catch (error) {
+    return Promise.reject(error)
+  }
+}
+
+/**
+ * @param {Array<{ status: string, reason?: Error }>} results
+ * @throws {Error} The first block verdict, if any evaluation returned one.
  */
-function handleEvaluationError (ctx, err) {
-  if (err.name === 'AIGuardAbortError') {
-    ctx.abortController.abort(err)
-  } else {
-    log.error('AIGuard: unexpected error during evaluation: %s', err.message)
+function throwFirstBlock (results) {
+  let block
+
+  for (const result of results) {
+    if (result.status === 'fulfilled') continue
+
+    if (result.reason?.name === 'AIGuardAbortError') {
+      block ??= result.reason
+    } else {
+      log.error('AIGuard: unexpected error during evaluation: %s', result.reason?.message)
+    }
   }
+
+  if (block) throw block
 }
 
-module.exports = { pushEvaluation }
+module.exports = { evaluate }

package/packages/dd-trace/src/aiguard/integrations/index.js

@@ -11,13 +11,14 @@ let isEnabled = false
  *
  * @param {object} aiguard
  * @param {boolean} block
+ * @param {boolean} analyzeStreamResponses
  */
-function enable (aiguard, block) {
+function enable (aiguard, block, analyzeStreamResponses) {
   if (isEnabled) return
 
   anthropic.enable(aiguard, block)
-  openai.enable(aiguard, block)
-  vercelAi.enable(aiguard, block)
+  openai.enable(aiguard, block, analyzeStreamResponses)
+  vercelAi.enable(aiguard, block, analyzeStreamResponses)
 
   isEnabled = true
 }

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

@@ -7,35 +7,37 @@ const {
   getChatCompletionsOutputMessages,
   getResponsesInputMessages,
   getResponsesOutputMessages,
+  getStreamedChatCompletionsOutputMessages,
+  getStreamedResponsesOutputMessages,
 } = require('../messages/openai')
+const { decode } = require('../messages/utils')
 const { SOURCE_AUTO } = require('../tags')
-const { pushEvaluation } = require('./evaluate')
+const { evaluate } = require('./evaluate')
 
-const chatCompletionsBeforeChannel = channel('dd-trace:openai:chat.completions:before')
-const chatCompletionsAfterChannel = channel('dd-trace:openai:chat.completions:after')
-const responsesBeforeChannel = channel('dd-trace:openai:responses:before')
-const responsesAfterChannel = channel('dd-trace:openai:responses:after')
+const chatCompletionsInterceptChannel = channel('dd-trace:openai:chat.completions:intercept')
+const responsesInterceptChannel = channel('dd-trace:openai:responses:intercept')
 
 let isEnabled = false
 let aiguard
 let opts
+let analyzeStreamResponses
 
 /**
- * Subscribes AI Guard to OpenAI lifecycle channels.
+ * Subscribes AI Guard to the OpenAI interception channels.
  *
  * @param {object} aiguardInstance
  * @param {boolean} block
+ * @param {boolean} analyzeStreams
  */
-function enable (aiguardInstance, block) {
+function enable (aiguardInstance, block, analyzeStreams) {
   if (isEnabled) return
 
   aiguard = aiguardInstance
   opts = { block, source: SOURCE_AUTO, integration: 'openai' }
+  analyzeStreamResponses = analyzeStreams
 
-  chatCompletionsBeforeChannel.subscribe(onChatCompletionsBefore)
-  chatCompletionsAfterChannel.subscribe(onChatCompletionsAfter)
-  responsesBeforeChannel.subscribe(onResponsesBefore)
-  responsesAfterChannel.subscribe(onResponsesAfter)
+  chatCompletionsInterceptChannel.subscribe(onChatCompletions)
+  responsesInterceptChannel.subscribe(onResponses)
 
   isEnabled = true
 }
@@ -43,38 +45,139 @@ function enable (aiguardInstance, block) {
 function disable () {
   if (!isEnabled) return
 
-  chatCompletionsBeforeChannel.unsubscribe(onChatCompletionsBefore)
-  chatCompletionsAfterChannel.unsubscribe(onChatCompletionsAfter)
-  responsesBeforeChannel.unsubscribe(onResponsesBefore)
-  responsesAfterChannel.unsubscribe(onResponsesAfter)
+  chatCompletionsInterceptChannel.unsubscribe(onChatCompletions)
+  responsesInterceptChannel.unsubscribe(onResponses)
 
   aiguard = undefined
   opts = undefined
+  analyzeStreamResponses = undefined
   isEnabled = false
 }
 
-function onChatCompletionsBefore (ctx) {
-  pushEvaluation(ctx, aiguard, getChatCompletionsInputMessages(ctx.args?.[0]), opts)
+function onChatCompletions (ctx) {
+  const inputMessages = getChatCompletionsInputMessages(ctx.arguments?.[0])
+  if (!inputMessages?.length) return
+
+  let inputEvaluation
+  ctx.beforeResult = () => {
+    if (!isEnabled) return
+    inputEvaluation ??= evaluate(ctx, aiguard, [inputMessages], opts)
+    return inputEvaluation
+  }
+
+  const isStream = ctx.arguments[0].stream
+  if (isStream && !analyzeStreamResponses) return
+
+  let outputEvaluation
+  ctx.onResult = body => {
+    if (!isEnabled) return body
+
+    if (isStream) {
+      outputEvaluation ??= interceptStream(body, chunks => {
+        return getStreamedChatCompletionsOutputMessages(chunks)
+          .map(message => [...inputMessages, message])
+      }, ctx)
+      return outputEvaluation
+    }
+
+    const conversations = decode(
+      () => getChatCompletionsOutputMessages(body).map(message => [...inputMessages, message]),
+      null,
+      'AIGuard: unable to decode OpenAI response body: %s'
+    )
+    if (!conversations?.length) return body
+
+    outputEvaluation ??= evaluate(ctx, aiguard, conversations, opts)
+    return outputEvaluation.then(() => body)
+  }
 }
 
-function onChatCompletionsAfter (ctx) {
-  const inputMessages = getChatCompletionsInputMessages(ctx.args?.[0])
+function onResponses (ctx) {
+  const inputMessages = getResponsesInputMessages(ctx.arguments?.[0])
   if (!inputMessages?.length) return
-  for (const message of getChatCompletionsOutputMessages(ctx.body)) {
-    pushEvaluation(ctx, aiguard, [...inputMessages, message], opts)
+
+  let inputEvaluation
+  ctx.beforeResult = () => {
+    if (!isEnabled) return
+    inputEvaluation ??= evaluate(ctx, aiguard, [inputMessages], opts)
+    return inputEvaluation
+  }
+
+  const isStream = ctx.arguments[0].stream
+  if (isStream && !analyzeStreamResponses) return
+
+  let outputEvaluation
+  ctx.onResult = body => {
+    if (!isEnabled) return body
+
+    if (isStream) {
+      outputEvaluation ??= interceptStream(body, chunks => {
+        const outputMessages = getStreamedResponsesOutputMessages(chunks)
+        return outputMessages.length ? [[...inputMessages, ...outputMessages]] : []
+      }, ctx)
+      return outputEvaluation
+    }
+
+    const outputMessages = decode(
+      () => getResponsesOutputMessages(body),
+      null,
+      'AIGuard: unable to decode OpenAI response body: %s'
+    )
+    if (!outputMessages?.length) return body
+
+    outputEvaluation ??= evaluate(ctx, aiguard, [[...inputMessages, ...outputMessages]], opts)
+    return outputEvaluation.then(() => body)
   }
 }
 
-function onResponsesBefore (ctx) {
-  pushEvaluation(ctx, aiguard, getResponsesInputMessages(ctx.args?.[0]), opts)
+/**
+ * Uses the OpenAI SDK's stream splitting support to evaluate one branch and return the other.
+ *
+ * @param {object} stream
+ * @param {(chunks: Array<object>) => Array<Array<object>>} getConversations
+ * @param {object} ctx
+ * @returns {object|Promise<object>}
+ */
+function interceptStream (stream, getConversations, ctx) {
+  if (typeof stream?.tee !== 'function') return stream
+
+  let branches
+  try {
+    branches = stream.tee()
+  } catch {
+    return stream
+  }
+
+  const [evaluationStream, resultStream] = branches
+  return drainStream(evaluationStream).then(chunks => {
+    if (!isEnabled) return resultStream
+    const conversations = decode(
+      () => getConversations(chunks),
+      null,
+      'AIGuard: unable to decode the streamed OpenAI response: %s'
+    )
+    if (!conversations?.length) return resultStream
+    return evaluate(ctx, aiguard, conversations, opts).then(() => resultStream)
+  }, () => resultStream)
 }
 
-function onResponsesAfter (ctx) {
-  const inputMessages = getResponsesInputMessages(ctx.args?.[0])
-  if (!inputMessages?.length) return
-  const outputMessages = getResponsesOutputMessages(ctx.body)
-  if (!outputMessages.length) return
-  pushEvaluation(ctx, aiguard, [...inputMessages, ...outputMessages], opts)
+/**
+ * @param {object} stream
+ * @returns {Promise<Array<object>>}
+ */
+function drainStream (stream) {
+  const chunks = []
+  const iterator = stream[Symbol.asyncIterator]()
+
+  function readAll () {
+    return iterator.next().then(({ done, value }) => {
+      if (done) return chunks
+      chunks.push(value)
+      return readAll()
+    })
+  }
+
+  return readAll()
 }
 
 module.exports = { enable, disable }