dd-trace 6.14.0 → 6.16.0

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)
+  vercelAi.enable(aiguard, block, analyzeStreamResponses)
 
   isEnabled = true
 }

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

@@ -8,20 +8,19 @@ const {
   getResponsesInputMessages,
   getResponsesOutputMessages,
 } = 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
 
 /**
- * Subscribes AI Guard to OpenAI lifecycle channels.
+ * Subscribes AI Guard to the OpenAI interception channels.
  *
  * @param {object} aiguardInstance
  * @param {boolean} block
@@ -32,10 +31,8 @@ function enable (aiguardInstance, block) {
   aiguard = aiguardInstance
   opts = { block, source: SOURCE_AUTO, integration: 'openai' }
 
-  chatCompletionsBeforeChannel.subscribe(onChatCompletionsBefore)
-  chatCompletionsAfterChannel.subscribe(onChatCompletionsAfter)
-  responsesBeforeChannel.subscribe(onResponsesBefore)
-  responsesAfterChannel.subscribe(onResponsesAfter)
+  chatCompletionsInterceptChannel.subscribe(onChatCompletions)
+  responsesInterceptChannel.subscribe(onResponses)
 
   isEnabled = true
 }
@@ -43,38 +40,70 @@ 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
   isEnabled = false
 }
 
-function onChatCompletionsBefore (ctx) {
-  pushEvaluation(ctx, aiguard, getChatCompletionsInputMessages(ctx.args?.[0]), opts)
-}
-
-function onChatCompletionsAfter (ctx) {
-  const inputMessages = getChatCompletionsInputMessages(ctx.args?.[0])
+function onChatCompletions (ctx) {
+  const inputMessages = getChatCompletionsInputMessages(ctx.arguments?.[0])
   if (!inputMessages?.length) return
-  for (const message of getChatCompletionsOutputMessages(ctx.body)) {
-    pushEvaluation(ctx, aiguard, [...inputMessages, message], opts)
+
+  // `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
   }
-}
 
-function onResponsesBefore (ctx) {
-  pushEvaluation(ctx, aiguard, getResponsesInputMessages(ctx.args?.[0]), opts)
+  // One model call has one output however many readers observe it.
+  let outputEvaluation
+  ctx.onResult = body => {
+    if (!isEnabled) return body
+
+    const conversations = decode(
+      () => getChatCompletionsOutputMessages(body).map(message => [...inputMessages, message]),
+      null,
+      'AIGuard: unable to decode OpenAI response body: %s'
+    )
+    if (conversations === null) return body
+    if (conversations.length === 0) return body
+
+    outputEvaluation ??= evaluate(ctx, aiguard, conversations, opts)
+    return outputEvaluation.then(() => body)
+  }
 }
 
-function onResponsesAfter (ctx) {
-  const inputMessages = getResponsesInputMessages(ctx.args?.[0])
+function onResponses (ctx) {
+  const inputMessages = getResponsesInputMessages(ctx.arguments?.[0])
   if (!inputMessages?.length) return
-  const outputMessages = getResponsesOutputMessages(ctx.body)
-  if (!outputMessages.length) return
-  pushEvaluation(ctx, aiguard, [...inputMessages, ...outputMessages], opts)
+
+  let inputEvaluation
+  ctx.beforeResult = () => {
+    if (!isEnabled) return
+    inputEvaluation ??= evaluate(ctx, aiguard, [inputMessages], opts)
+    return inputEvaluation
+  }
+
+  let outputEvaluation
+  ctx.onResult = body => {
+    if (!isEnabled) return body
+
+    const outputMessages = decode(
+      () => getResponsesOutputMessages(body),
+      null,
+      'AIGuard: unable to decode OpenAI 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/vercel-ai.js

@@ -2,35 +2,37 @@
 
 const { channel } = require('dc-polyfill')
 
-const { buildOutputMessages, convertVercelPromptToMessages } = require('../messages/vercel-ai')
+const {
+  buildOutputMessages,
+  convertVercelPromptToMessages,
+  getStreamedContent,
+} = require('../messages/vercel-ai')
+const { decode } = require('../messages/utils')
 const { SOURCE_AUTO } = require('../tags')
-const { pushEvaluation } = require('./evaluate')
+const { evaluate } = require('./evaluate')
 
-const doGenerateBeforeChannel = channel('dd-trace:vercel-ai:doGenerate:before')
-const doGenerateAfterChannel = channel('dd-trace:vercel-ai:doGenerate:after')
-const doStreamBeforeChannel = channel('dd-trace:vercel-ai:doStream:before')
-const doStreamAfterChannel = channel('dd-trace:vercel-ai:doStream:after')
+const modelInterceptChannel = channel('dd-trace:vercel-ai:model:intercept')
 
 let isEnabled = false
 let aiguard
 let opts
+let analyzeStreamResponses
 
 /**
- * Subscribes AI Guard to Vercel AI lifecycle channels.
+ * Subscribes AI Guard to the Vercel AI model call channel.
  *
  * @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: 'ai' }
+  analyzeStreamResponses = analyzeStreams
 
-  doGenerateBeforeChannel.subscribe(onBefore)
-  doGenerateAfterChannel.subscribe(onGenerateAfter)
-  doStreamBeforeChannel.subscribe(onBefore)
-  doStreamAfterChannel.subscribe(onStreamAfter)
+  modelInterceptChannel.subscribe(onModelIntercept)
 
   isEnabled = true
 }
@@ -38,54 +40,126 @@ function enable (aiguardInstance, block) {
 function disable () {
   if (!isEnabled) return
 
-  doGenerateBeforeChannel.unsubscribe(onBefore)
-  doGenerateAfterChannel.unsubscribe(onGenerateAfter)
-  doStreamBeforeChannel.unsubscribe(onBefore)
-  doStreamAfterChannel.unsubscribe(onStreamAfter)
+  modelInterceptChannel.unsubscribe(onModelIntercept)
 
   aiguard = undefined
   opts = undefined
+  analyzeStreamResponses = undefined
   isEnabled = false
 }
 
-function onBefore (ctx) {
-  pushEvaluation(ctx, aiguard, convertVercelPromptToMessages(ctx.prompt), opts)
-}
+function onModelIntercept (ctx) {
+  const inputMessages = convertVercelPromptToMessages(ctx.arguments?.[0]?.prompt)
+  if (!inputMessages.length) return
+
+  // Called exactly once per model call by the instrumentation, so no memoization needed.
+  ctx.beforeResult = () => {
+    if (!isEnabled) return
+    return evaluate(ctx, aiguard, [inputMessages], opts)
+  }
 
-function onGenerateAfter (ctx) {
-  const inputMessages = convertVercelPromptToMessages(ctx.prompt)
-  if (!inputMessages.length || !ctx.result?.content?.length) return
+  if (ctx.method === 'doStream' && !analyzeStreamResponses) return
 
-  pushEvaluation(ctx, aiguard, buildOutputMessages(inputMessages, ctx.result.content), opts)
+  ctx.onResult = ctx.method === 'doStream'
+    ? result => {
+      if (!isEnabled) return result
+      return interceptStreamedResult(ctx, result, inputMessages)
+    }
+    : result => {
+      if (!isEnabled) return result
+
+      const outputMessages = decode(
+        () => buildOutputMessages(inputMessages, result?.content ?? []),
+        null,
+        'AIGuard: unable to decode the model result: %s'
+      )
+      if (outputMessages === null) return result
+      if (!outputMessages.length) return result
+
+      return evaluate(ctx, aiguard, [outputMessages], opts).then(() => result)
+    }
 }
 
-function onStreamAfter (ctx) {
-  const inputMessages = convertVercelPromptToMessages(ctx.prompt)
-  if (!inputMessages.length || !ctx.chunks?.length) return
+/**
+ * Judging a streamed response requires the whole output, so the stream is drained and replaced
+ * with a replay of the collected chunks. That trades streaming latency for output coverage.
+ *
+ * @param {object} ctx
+ * @param {object} result
+ * @param {Array<object>} inputMessages
+ * @returns {object|Promise<object>}
+ */
+function interceptStreamedResult (ctx, result, inputMessages) {
+  if (!result?.stream) return result
+
+  let drained
+  try {
+    drained = drainStream(result.stream)
+  } catch {
+    // Nothing was consumed, so the caller still has the stream they handed us.
+    return result
+  }
+
+  return drained.then(({ chunks, error }) => {
+    // The original stream is spent from here on, so the caller must get the replay either way.
+    const replayed = { ...result, stream: replayChunks(chunks, error) }
 
-  pushEvaluation(ctx, aiguard, buildOutputMessages(inputMessages, getStreamContent(ctx.chunks)), opts)
+    // A stream that failed part-way has no complete output to judge; the replay carries the error.
+    if (error) return replayed
+
+    const outputMessages = decode(
+      () => buildOutputMessages(inputMessages, getStreamedContent(chunks)),
+      null,
+      'AIGuard: unable to decode the streamed model result: %s'
+    )
+    if (outputMessages === null) return replayed
+    if (!outputMessages.length) return replayed
+
+    return evaluate(ctx, aiguard, [outputMessages], opts).then(() => replayed)
+  })
 }
 
 /**
- * Converts Vercel stream chunks into the content shape used by doGenerate results.
+ * Never rejects: a stream that fails part-way is already consumed, so its chunks and its error
+ * are both returned for the replay to hand back to the caller.
  *
- * @param {Array<object>} chunks
- * @returns {Array<object>}
+ * @param {ReadableStream} stream
+ * @returns {Promise<{ chunks: Array<object>, error?: unknown }>}
  */
-function getStreamContent (chunks) {
-  const toolCalls = []
-  let text = ''
-
-  for (const chunk of chunks) {
-    if (chunk?.type === 'tool-call') {
-      toolCalls.push(chunk)
-    } else if (chunk?.type === 'text-delta') {
-      text += chunk.delta ?? chunk.textDelta ?? ''
-    }
+function drainStream (stream) {
+  const chunks = []
+  const reader = stream.getReader()
+
+  function readAll () {
+    return reader.read().then(({ done, value }) => {
+      if (done) return { chunks }
+      chunks.push(value)
+      return readAll()
+    }, error => ({ chunks, error }))
   }
 
-  if (toolCalls.length) return toolCalls
-  return text ? [{ type: 'text', text }] : []
+  return readAll()
+}
+
+/**
+ * @param {Array<object>} chunks
+ * @param {unknown} [error]
+ * @returns {ReadableStream}
+ */
+function replayChunks (chunks, error) {
+  // eslint-disable-next-line n/no-unsupported-features/node-builtins
+  return new ReadableStream({
+    start (controller) {
+      for (const chunk of chunks) {
+        controller.enqueue(chunk)
+      }
+      if (error) {
+        controller.error(error)
+      } else {
+        controller.close()
+      }
+    },
+  })
 }
 
 module.exports = { enable, disable }

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

@@ -1,5 +1,7 @@
 'use strict'
 
+const log = require('../../log')
+
 const FILE_FALLBACK = '[file]'
 const IMAGE_FALLBACK = '[image]'
 
@@ -9,7 +11,13 @@ const IMAGE_FALLBACK = '[image]'
  */
 function stringifyIfNeeded (value) {
   if (value == null) return value
-  return typeof value === 'string' ? value : JSON.stringify(value)
+  if (typeof value === 'string') return value
+
+  try {
+    return JSON.stringify(value)
+  } catch (error) {
+    log.debug('AIGuard: dropping an unserializable message field: %s', error.message)
+  }
 }
 
 /**
@@ -20,9 +28,29 @@ function stringifyOrEmpty (value) {
   return stringifyIfNeeded(value) ?? ''
 }
 
+/**
+ * Runs `attempt`, returning its value. On failure, logs and returns `fallback` instead of
+ * throwing — this runs inside the caller's promise chain, so a bad payload must not break them.
+ *
+ * @template T
+ * @param {() => T} attempt
+ * @param {T} fallback
+ * @param {string} logMessage printf-style, with a single %s for the error message
+ * @returns {T}
+ */
+function decode (attempt, fallback, logMessage) {
+  try {
+    return attempt()
+  } catch (error) {
+    log.error(logMessage, error.message)
+    return fallback
+  }
+}
+
 module.exports = {
   FILE_FALLBACK,
   IMAGE_FALLBACK,
   stringifyIfNeeded,
   stringifyOrEmpty,
+  decode,
 }

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

@@ -180,7 +180,32 @@ function buildOutputMessages (inputMessages, content) {
   return []
 }
 
+/**
+ * Converts Vercel stream chunks into the content shape a doGenerate result carries.
+ *
+ * @param {Array<object>} chunks
+ * @returns {Array<object>}
+ */
+function getStreamedContent (chunks) {
+  const toolCalls = []
+  let text = ''
+
+  for (const chunk of chunks) {
+    if (chunk?.type === 'tool-call') {
+      toolCalls.push(chunk)
+    } else if (chunk?.type === 'text-delta') {
+      // AI SDK 6 renamed `textDelta` to `delta`; both shapes are supported. Coalesced with `??`
+      // because a missing delta must contribute nothing, not "undefined".
+      text += chunk.delta ?? chunk.textDelta ?? ''
+    }
+  }
+
+  if (toolCalls.length) return toolCalls
+  return text ? [{ type: 'text', text }] : []
+}
+
 module.exports = {
+  getStreamedContent,
   convertVercelPromptToMessages,
   convertFilePartToImageUrl,
   buildToolCallOutputMessages,

package/packages/dd-trace/src/appsec/iast/vulnerabilities-formatter/evidence-redaction/sensitive-analyzers/command-sensitive-analyzer.js

@@ -2,7 +2,9 @@
 
 const log = require('../../../../../log')
 
-const COMMAND_PATTERN = String.raw`^(?:\s*(?:sudo|doas)\s+)?\b\S+\b\s(.*)`
+// `\S+\s` accepts command tokens that end in punctuation. The optional prefix uses horizontal
+// whitespace so multiline matching cannot start inside a line-terminator run.
+const COMMAND_PATTERN = String.raw`^(?:[ \t]*(?:sudo|doas)[ \t]+)?\S+\s([\s\S]*)`
 const pattern = new RegExp(COMMAND_PATTERN, 'gmi')
 
 module.exports = function extractSensitiveRanges (evidence) {