dd-trace 5.108.0 → 5.109.0

package/packages/datadog-plugin-bullmq/src/producer.js

@@ -3,6 +3,9 @@
 const log = require('../../dd-trace/src/log')
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
 const { DsmPathwayCodec, getMessageSize } = require('../../dd-trace/src/datastreams')
+const { getFilter } = require('./filter')
+
+const filteredJobs = Symbol('bullmq.filteredJobs')
 
 // Customer-controlled metadata may be malformed JSON. Returning a fresh `{}`
 // on parse failure keeps the publish path alive instead of throwing into
@@ -31,6 +34,19 @@ class BaseBullmqProducerPlugin extends ProducerPlugin {
   }
 
   bindStart (ctx) {
+    let instrument = true
+    if (this.config.producerFilter) {
+      try {
+        instrument = this.shouldInstrument(ctx)
+      } catch (error) {
+        log.error('bullmq: producerFilter threw, filtering is disabled: %s', error.message)
+      }
+    }
+
+    if (!instrument) {
+      return { noop: true }
+    }
+
     const { resource, meta } = this.getSpanData(ctx)
     const span = this.startSpan({
       resource,
@@ -48,6 +64,17 @@ class BaseBullmqProducerPlugin extends ProducerPlugin {
     return ctx.currentStore
   }
 
+  configure (config) {
+    if (typeof config === 'boolean') {
+      return super.configure(config)
+    }
+    return super.configure({ ...config, producerFilter: getFilter(config) })
+  }
+
+  shouldInstrument (ctx) {
+    throw new Error('shouldInstrument must be implemented by subclass')
+  }
+
   getSpanData (ctx) {
     throw new Error('getSpanData must be implemented by subclass')
   }
@@ -87,6 +114,15 @@ class BaseBullmqProducerPlugin extends ProducerPlugin {
 class QueueAddPlugin extends BaseBullmqProducerPlugin {
   static prefix = 'tracing:orchestrion:bullmq:Queue_add'
 
+  shouldInstrument (ctx) {
+    return this.config.producerFilter({
+      name: ctx.arguments?.[0],
+      data: ctx.arguments?.[1],
+      opts: ctx.arguments?.[2],
+      queueName: ctx.self?.name,
+    })
+  }
+
   getSpanData (ctx) {
     const queueName = ctx.self?.name || 'bullmq'
     return {
@@ -128,6 +164,40 @@ class QueueAddPlugin extends BaseBullmqProducerPlugin {
 class QueueAddBulkPlugin extends BaseBullmqProducerPlugin {
   static prefix = 'tracing:orchestrion:bullmq:Queue_addBulk'
 
+  shouldInstrument (ctx) {
+    const jobs = this.#getFilteredJobs(ctx)
+    // Empty bulk calls are publish attempts and should keep producing a span.
+    return jobs === undefined || jobs.length > 0 || ctx.arguments?.[0].length === 0
+  }
+
+  #getFilteredJobs (ctx) {
+    if (Object.hasOwn(ctx, filteredJobs)) return ctx[filteredJobs]
+
+    const jobs = ctx.arguments?.[0]
+    if (!Array.isArray(jobs)) return
+
+    let allowedJobs
+    if (this.config.producerFilter) {
+      const queueName = ctx.self?.name
+      try {
+        allowedJobs = []
+        for (const job of jobs) {
+          if (job && this.config.producerFilter({ name: job.name, data: job.data, opts: job.opts, queueName })) {
+            allowedJobs.push(job)
+          }
+        }
+      } catch (error) {
+        log.error('bullmq: producerFilter threw, filtering is disabled: %s', error.message)
+        allowedJobs = jobs.filter(Boolean)
+      }
+    } else {
+      allowedJobs = jobs
+    }
+
+    ctx[filteredJobs] = allowedJobs
+    return allowedJobs
+  }
+
   operationName () {
     return 'bullmq.addBulk'
   }
@@ -139,14 +209,14 @@ class QueueAddBulkPlugin extends BaseBullmqProducerPlugin {
       resource: queueName,
       meta: {
         'messaging.destination.name': ctx.self?.name,
-        'messaging.batch.message_count': Array.isArray(jobs) ? jobs.length : undefined,
+        'messaging.batch.message_count': jobs?.length,
       },
     }
   }
 
   injectTraceContext (span, ctx) {
-    const jobs = ctx.arguments?.[0]
-    if (!Array.isArray(jobs)) return
+    const jobs = this.#getFilteredJobs(ctx)
+    if (!jobs) return
 
     const cache = []
     for (let i = 0; i < jobs.length; i++) {
@@ -159,7 +229,7 @@ class QueueAddBulkPlugin extends BaseBullmqProducerPlugin {
   }
 
   setProducerCheckpoint (span, ctx) {
-    const jobs = ctx.arguments?.[0] || []
+    const jobs = this.#getFilteredJobs(ctx) || []
     const queueName = ctx.self?.name || 'bullmq'
     const edgeTags = ['direction:out', `topic:${queueName}`, 'type:bullmq']
     const cache = ctx._ddMetadata
@@ -180,6 +250,16 @@ class QueueAddBulkPlugin extends BaseBullmqProducerPlugin {
 class FlowProducerAddPlugin extends BaseBullmqProducerPlugin {
   static prefix = 'tracing:orchestrion:bullmq:FlowProducer_add'
 
+  shouldInstrument (ctx) {
+    const flow = ctx.arguments?.[0]
+    return this.config.producerFilter({
+      name: flow?.name,
+      data: flow?.data,
+      opts: flow?.opts,
+      queueName: flow?.queueName,
+    })
+  }
+
   getSpanData (ctx) {
     const flow = ctx.arguments?.[0]
     const queueName = flow?.queueName || 'bullmq'

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

@@ -5,6 +5,7 @@ const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
 class FsPlugin extends TracingPlugin {
   static id = 'fs'
   static operation = 'operation'
+  static experimental = true
 
   configure (...args) {
     return super.configure(...args)

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

@@ -28,7 +28,7 @@ class RedisPlugin extends CachePlugin {
   }
 
   bindStart (ctx) {
-    const { db, command, args, argsStartIndex, connectionOptions, connectionName } = ctx
+    const { command, args, argsStartIndex, connectionOptions, connectionName } = ctx
 
     const resource = command
     const normalizedCommand = command.toUpperCase()
@@ -55,7 +55,6 @@ class RedisPlugin extends CachePlugin {
       type: this._spanType,
       meta: {
         'db.type': this._spanType,
-        'db.name': db || '0',
         [this.#rawCommandKey]: formatCommand(normalizedCommand, args, argsStartIndex),
         'out.host': connectionOptions.host,
         [CLIENT_PORT_KEY]: connectionOptions.port,

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

@@ -63,6 +63,8 @@ class VitestPlugin extends CiPlugin {
         testSessionId: testSessionSpanContext?.toTraceId(),
         testModuleId: testModuleSpanContext?.toSpanId(),
         testCommand: this.command,
+        repositoryRoot: this.repositoryRoot,
+        codeOwnersEntries: this.codeOwnersEntries,
       })
     })
 
@@ -307,10 +309,11 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addBind('ci:vitest:test-suite:start', (ctx) => {
-      const { testSuiteAbsolutePath, frameworkVersion } = ctx
+      const { codeOwnersEntries, repositoryRoot, testSuiteAbsolutePath, frameworkVersion } = ctx
 
       const testCommand = ctx.testCommand || 'vitest run'
       const { testSessionId, testModuleId } = ctx
+      this._setRepositoryRoot(repositoryRoot, codeOwnersEntries)
       this.command = testCommand
       this.frameworkVersion = frameworkVersion
       const testSessionSpanContext = testSessionId && testModuleId

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

@@ -3,6 +3,5 @@
 const dc = require('dc-polyfill')
 
 module.exports = {
-  aiguardChannel: dc.channel('dd-trace:ai:aiguard'),
   incomingHttpRequestStart: dc.channel('dd-trace:incomingHttpRequestStart'),
 }

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

@@ -1,13 +1,16 @@
 'use strict'
 
 const log = require('../log')
-const { incomingHttpRequestStart, aiguardChannel } = require('./channels')
+const { incomingHttpRequestStart } = require('./channels')
+const openaiIntegration = require('./integrations/openai')
+const vercelAiIntegration = require('./integrations/vercel-ai')
 const AIGuard = require('./sdk')
-const { SOURCE_AUTO, INTEGRATION_NONE } = require('./tags')
 
 let isEnabled = false
 let aiguard
 let block
+let disableOpenAIIntegration
+let disableVercelAiIntegration
 
 function onIncomingHttpRequestStart () {
   // No-op: subscribing ensures the HTTP plugin spreads req onto the store
@@ -21,7 +24,8 @@ function enable (tracer, config) {
     block = config.experimental?.aiguard?.block !== false
 
     incomingHttpRequestStart.subscribe(onIncomingHttpRequestStart)
-    aiguardChannel.subscribe(onEvaluate)
+    disableOpenAIIntegration = openaiIntegration.enable(aiguard, block)
+    disableVercelAiIntegration = vercelAiIntegration.enable(aiguard, block)
 
     isEnabled = true
   } catch (err) {
@@ -34,56 +38,14 @@ function disable () {
   if (!isEnabled) return
 
   incomingHttpRequestStart.unsubscribe(onIncomingHttpRequestStart)
-  aiguardChannel.unsubscribe(onEvaluate)
+  disableOpenAIIntegration?.()
+  disableVercelAiIntegration?.()
 
   aiguard = undefined
   isEnabled = false
   block = false
-}
-
-/**
- * Handles channel messages with pre-converted messages.
- *
- * @param {object} ctx
- * @param {Array<object>} ctx.messages
- * @param {string} [ctx.integration]
- * @param {object} [ctx.parentSpan] - LLM span to parent the `ai_guard` span under.
- * @param {AbortController} ctx.abortController
- * @param {Array<Promise<void>>} ctx.pending - Subscribers push only when they evaluate.
- */
-function onEvaluate (ctx) {
-  // Decline to evaluate empty payloads by not pushing to pending.
-  if (!ctx.messages?.length) {
-    return
-  }
-
-  const opts = {
-    block,
-    source: SOURCE_AUTO,
-    integration: ctx.integration || INTEGRATION_NONE,
-    childOf: ctx.parentSpan,
-  }
-
-  try {
-    ctx.pending.push(aiguard.evaluate(ctx.messages, opts).catch(handleEvaluationError.bind(null, ctx)))
-  } catch (err) {
-    ctx.pending.push(Promise.resolve().then(() => handleEvaluationError(ctx, err)))
-  }
-}
-
-/**
- * Handles an AI Guard evaluation failure.
- *
- * @param {object} ctx
- * @param {AbortController} ctx.abortController
- * @param {Error} err
- */
-function handleEvaluationError (ctx, err) {
-  if (err.name === 'AIGuardAbortError') {
-    ctx.abortController.abort(err)
-  } else {
-    log.error('AIGuard: unexpected error during evaluation: %s', err.message)
-  }
+  disableOpenAIIntegration = undefined
+  disableVercelAiIntegration = undefined
 }
 
 module.exports = { enable, disable }

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

@@ -0,0 +1,46 @@
+'use strict'
+
+const log = require('../../log')
+
+/**
+ * Starts one AI Guard evaluation for a lifecycle ctx.
+ *
+ * Async evaluations are pushed synchronously during channel publication. If
+ * evaluation throws synchronously, the error is handled before publish returns.
+ *
+ * @param {object} ctx
+ * @param {AbortController} ctx.abortController
+ * @param {object} [ctx.parentSpan]
+ * @param {Array<Promise<void>>} ctx.pending
+ * @param {object} aiguard
+ * @param {Array<object>|undefined} messages
+ * @param {object} opts
+ */
+function pushEvaluation (ctx, aiguard, messages, opts) {
+  if (!messages?.length) return
+
+  const evaluateOpts = ctx.parentSpan ? { ...opts, childOf: ctx.parentSpan } : opts
+
+  try {
+    ctx.pending.push(aiguard.evaluate(messages, evaluateOpts).catch(handleEvaluationError.bind(null, ctx)))
+  } catch (err) {
+    handleEvaluationError(ctx, err)
+  }
+}
+
+/**
+ * Handles an AI Guard evaluation failure.
+ *
+ * @param {object} ctx
+ * @param {AbortController} ctx.abortController
+ * @param {Error} err
+ */
+function handleEvaluationError (ctx, err) {
+  if (err.name === 'AIGuardAbortError') {
+    ctx.abortController.abort(err)
+  } else {
+    log.error('AIGuard: unexpected error during evaluation: %s', err.message)
+  }
+}
+
+module.exports = { pushEvaluation }

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

@@ -0,0 +1,66 @@
+'use strict'
+
+const { channel } = require('dc-polyfill')
+
+const {
+  getChatCompletionsInputMessages,
+  getChatCompletionsOutputMessages,
+  getResponsesInputMessages,
+  getResponsesOutputMessages,
+} = require('../messages/openai')
+const { SOURCE_AUTO } = require('../tags')
+const { pushEvaluation } = 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')
+
+/**
+ * Subscribes AI Guard to OpenAI lifecycle channels.
+ *
+ * @param {object} aiguard
+ * @param {boolean} block
+ * @returns {() => void}
+ */
+function enable (aiguard, block) {
+  const opts = { block, source: SOURCE_AUTO, integration: 'openai' }
+
+  function onChatCompletionsBefore (ctx) {
+    pushEvaluation(ctx, aiguard, getChatCompletionsInputMessages(ctx.args?.[0]), opts)
+  }
+
+  function onChatCompletionsAfter (ctx) {
+    const inputMessages = getChatCompletionsInputMessages(ctx.args?.[0])
+    if (!inputMessages?.length) return
+    for (const message of getChatCompletionsOutputMessages(ctx.body)) {
+      pushEvaluation(ctx, aiguard, [...inputMessages, message], opts)
+    }
+  }
+
+  function onResponsesBefore (ctx) {
+    pushEvaluation(ctx, aiguard, getResponsesInputMessages(ctx.args?.[0]), opts)
+  }
+
+  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)
+  }
+
+  chatCompletionsBeforeChannel.subscribe(onChatCompletionsBefore)
+  chatCompletionsAfterChannel.subscribe(onChatCompletionsAfter)
+  responsesBeforeChannel.subscribe(onResponsesBefore)
+  responsesAfterChannel.subscribe(onResponsesAfter)
+
+  return function disable () {
+    chatCompletionsBeforeChannel.unsubscribe(onChatCompletionsBefore)
+    chatCompletionsAfterChannel.unsubscribe(onChatCompletionsAfter)
+    responsesBeforeChannel.unsubscribe(onResponsesBefore)
+    responsesAfterChannel.unsubscribe(onResponsesAfter)
+  }
+}
+
+module.exports = { enable }

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

@@ -0,0 +1,78 @@
+'use strict'
+
+const { channel } = require('dc-polyfill')
+
+const { buildOutputMessages, convertVercelPromptToMessages } = require('../messages/vercel-ai')
+const { SOURCE_AUTO } = require('../tags')
+const { pushEvaluation } = 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')
+
+/**
+ * Subscribes AI Guard to Vercel AI lifecycle channels.
+ *
+ * @param {object} aiguard
+ * @param {boolean} block
+ * @returns {() => void}
+ */
+function enable (aiguard, block) {
+  const opts = { block, source: SOURCE_AUTO, integration: 'ai' }
+
+  function onBefore (ctx) {
+    pushEvaluation(ctx, aiguard, convertVercelPromptToMessages(ctx.prompt), opts)
+  }
+
+  function onGenerateAfter (ctx) {
+    const inputMessages = convertVercelPromptToMessages(ctx.prompt)
+    if (!inputMessages.length || !ctx.result?.content?.length) return
+
+    pushEvaluation(ctx, aiguard, buildOutputMessages(inputMessages, ctx.result.content), opts)
+  }
+
+  function onStreamAfter (ctx) {
+    const inputMessages = convertVercelPromptToMessages(ctx.prompt)
+    if (!inputMessages.length || !ctx.chunks?.length) return
+
+    pushEvaluation(ctx, aiguard, buildOutputMessages(inputMessages, getStreamContent(ctx.chunks)), opts)
+  }
+
+  doGenerateBeforeChannel.subscribe(onBefore)
+  doGenerateAfterChannel.subscribe(onGenerateAfter)
+  doStreamBeforeChannel.subscribe(onBefore)
+  doStreamAfterChannel.subscribe(onStreamAfter)
+
+  return function disable () {
+    doGenerateBeforeChannel.unsubscribe(onBefore)
+    doGenerateAfterChannel.unsubscribe(onGenerateAfter)
+    doStreamBeforeChannel.unsubscribe(onBefore)
+    doStreamAfterChannel.unsubscribe(onStreamAfter)
+  }
+}
+
+/**
+ * Converts Vercel stream chunks into the content shape used by doGenerate results.
+ *
+ * @param {Array<object>} chunks
+ * @returns {Array<object>}
+ */
+function getStreamContent (chunks) {
+  const toolCalls = []
+  const textParts = []
+
+  for (const chunk of chunks) {
+    if (chunk?.type === 'tool-call') {
+      toolCalls.push(chunk)
+    } else if (chunk?.type === 'text-delta') {
+      textParts.push(chunk.textDelta)
+    }
+  }
+
+  if (toolCalls.length) return toolCalls
+  const text = textParts.join('')
+  return text ? [{ type: 'text', text }] : []
+}
+
+module.exports = { enable }