dd-trace 5.94.0 → 5.96.0

package/packages/datadog-instrumentations/src/helpers/ai-messages.js

@@ -0,0 +1,182 @@
+'use strict'
+
+/**
+ * Converts a LanguageModelV2FilePart with an image mediaType to an AI guard style image_url content part.
+ *
+ * @param {{type: 'file', data: URL|string|Uint8Array, mediaType: string}} part
+ * @returns {{type: 'image_url', image_url: {url: string}}|undefined}
+ */
+function convertFilePartToImageUrl (part) {
+  const { data, mediaType } = part
+
+  if (data instanceof URL) {
+    return { type: 'image_url', image_url: { url: data.toString() } }
+  }
+
+  if (typeof data === 'string') {
+    if (data.startsWith('http') || data.startsWith('data:')) {
+      return { type: 'image_url', image_url: { url: data } }
+    }
+    return { type: 'image_url', image_url: { url: `data:${mediaType};base64,${data}` } }
+  }
+
+  if (data instanceof Uint8Array) {
+    return { type: 'image_url', image_url: { url: `data:${mediaType};base64,${Buffer.from(data).toString('base64')}` } }
+  }
+}
+
+/**
+ * Converts a LanguageModelV2Prompt to the AI guard style message format.
+ *
+ * Vercel AI v2 prompt entries use content arrays with typed parts (e.g. { type: 'text', text },
+ * { type: 'file', data, mediaType }). This function converts them to AI guard style messages.
+ * When file parts with image media types are present, the content is an array of text and
+ * image_url parts; otherwise it is a plain string.
+ *
+ * @param {Array<{role: string, content: string|Array<{type: string}>}>} prompt
+ * @returns {Array<{role: string, content?: string|Array<{type: string}>, tool_calls?: Array, tool_call_id?: string}>}
+ */
+function convertVercelPromptToMessages (prompt) {
+  if (!Array.isArray(prompt)) return []
+
+  const messages = []
+  for (const msg of prompt) {
+    switch (msg.role) {
+      case 'system':
+        messages.push({ role: 'system', content: typeof msg.content === 'string' ? msg.content : '' })
+        break
+
+      case 'user': {
+        if (!Array.isArray(msg.content)) break
+
+        const contentParts = []
+        for (const part of msg.content) {
+          if (part.type === 'text') {
+            contentParts.push({ type: 'text', text: part.text })
+          } else if (part.type === 'file' && part.mediaType?.startsWith('image/')) {
+            const converted = convertFilePartToImageUrl(part)
+            if (converted) contentParts.push(converted)
+          }
+        }
+
+        if (contentParts.length === 0) break
+
+        const hasImages = contentParts.some(p => p.type === 'image_url')
+        if (hasImages) {
+          messages.push({ role: 'user', content: contentParts })
+        } else {
+          messages.push({ role: 'user', content: contentParts.map(p => p.text).join('\n') })
+        }
+        break
+      }
+
+      case 'assistant': {
+        const textParts = []
+        const toolCalls = []
+        if (!Array.isArray(msg.content)) break
+
+        for (const part of msg.content) {
+          if (part.type === 'text') {
+            textParts.push(part.text)
+          } else if (part.type === 'tool-call') {
+            const args = part.args ?? part.input
+            toolCalls.push({
+              id: part.toolCallId,
+              function: {
+                name: part.toolName,
+                arguments: typeof args === 'string' ? args : JSON.stringify(args),
+              },
+            })
+          }
+        }
+
+        if (toolCalls.length > 0) {
+          messages.push({ role: 'assistant', tool_calls: toolCalls })
+        } else if (textParts.length > 0) {
+          messages.push({ role: 'assistant', content: textParts.join('\n') })
+        }
+        break
+      }
+
+      case 'tool': {
+        if (!Array.isArray(msg.content)) break
+
+        for (const part of msg.content) {
+          if (part.type === 'tool-result') {
+            const result = part.result ?? part.output
+            messages.push({
+              role: 'tool',
+              tool_call_id: part.toolCallId,
+              content: typeof result === 'string' ? result : JSON.stringify(result),
+            })
+          }
+        }
+        break
+      }
+    }
+  }
+  return messages
+}
+
+/**
+ * Converts LLM output tool calls to AI guard style message format.
+ *
+ * @param {Array<object>} inputMessages - The input messages already in AI guard style format
+ * @param {Array<{toolCallId: string, toolName: string, args?: unknown, input?: unknown}>} toolCalls
+ * @returns {Array<object>}
+ */
+function buildToolCallOutputMessages (inputMessages, toolCalls) {
+  return [
+    ...inputMessages,
+    {
+      role: 'assistant',
+      tool_calls: toolCalls.map(tc => {
+        const args = tc.args ?? tc.input
+        return {
+          id: tc.toolCallId,
+          function: {
+            name: tc.toolName,
+            arguments: typeof args === 'string' ? args : JSON.stringify(args),
+          },
+        }
+      }),
+    },
+  ]
+}
+
+/**
+ * Builds OpenAI-style output messages for the assistant's text response.
+ *
+ * @param {Array<object>} inputMessages - The input messages already in AI guard style format
+ * @param {string} text - The assistant's text response
+ * @returns {Array<object>}
+ */
+function buildTextOutputMessages (inputMessages, text) {
+  return [
+    ...inputMessages,
+    { role: 'assistant', content: text },
+  ]
+}
+
+/**
+ * Parses a Vercel AI content array and dispatches to the appropriate output message builder.
+ *
+ * @param {Array<object>} inputMessages - The input messages already in AI guard style format
+ * @param {Array<{type: string}>} content - Vercel AI content array from doGenerate/doStream result
+ * @returns {Array<object>}
+ */
+function buildOutputMessages (inputMessages, content) {
+  const toolCalls = content.filter(c => c.type === 'tool-call')
+  const text = content.filter(c => c.type === 'text').map(c => c.text).join('\n')
+  if (toolCalls.length) return buildToolCallOutputMessages(inputMessages, toolCalls)
+  if (text) return buildTextOutputMessages(inputMessages, text)
+  return inputMessages
+}
+
+module.exports = {
+  convertVercelPromptToMessages,
+  convertFilePartToImageUrl,
+  buildToolCallOutputMessages,
+  buildTextOutputMessages,
+  buildOutputMessages,
+}

package/packages/datadog-instrumentations/src/helpers/rewriter/{orchestrion/compiler.js → compiler.js}

@@ -1,6 +1,6 @@
 'use strict'
 
-const log = require('../../../../../dd-trace/src/log')
+const log = require('../../../../dd-trace/src/log')
 
 // eslint-disable-next-line camelcase, no-undef
 const runtimeRequire = typeof __webpack_require__ === 'function' ? __non_webpack_require__ : require
@@ -25,7 +25,7 @@ const compiler = {
       log.error(e)
 
       // Fallback for when OXC is not available.
-      const meriyah = require('../../../../../../vendor/dist/meriyah')
+      const meriyah = require('../../../../../vendor/dist/meriyah')
 
       compiler.parse = (sourceText, { range, sourceType } = {}) => {
         return meriyah.parse(sourceText.toString(), {
@@ -39,16 +39,8 @@ const compiler = {
     return compiler.parse(sourceText, options)
   },
 
-  generate: (...args) => {
-    const astring = require('../../../../../../vendor/dist/astring')
-
-    compiler.generate = astring.generate
-
-    return compiler.generate(...args)
-  },
-
   traverse: (ast, query, visitor) => {
-    const esquery = require('../../../../../../vendor/dist/esquery').default
+    const esquery = require('../../../../../vendor/dist/esquery')
 
     compiler.traverse = (ast, query, visitor) => {
       return esquery.traverse(ast, esquery.parse(query), visitor)
@@ -58,7 +50,7 @@ const compiler = {
   },
 
   query: (ast, query) => {
-    const esquery = require('../../../../../../vendor/dist/esquery').default
+    const esquery = require('../../../../../vendor/dist/esquery')
 
     compiler.query = esquery.query
 
@@ -68,7 +60,6 @@ const compiler = {
 
 module.exports = {
   parse: (...args) => compiler.parse(...args),
-  generate: (...args) => compiler.generate(...args),
   traverse: (...args) => compiler.traverse(...args),
   query: (...args) => compiler.query(...args),
 }

package/packages/datadog-instrumentations/src/helpers/rewriter/index.js

@@ -3,13 +3,27 @@
 const { readFileSync } = require('fs')
 const { join } = require('path')
 const log = require('../../../../dd-trace/src/log')
+const { create } = require('../../../../../vendor/dist/@apm-js-collab/code-transformer')
+const { traceAsyncIterator, traceIterator } = require('./transforms')
 const instrumentations = require('./instrumentations')
-const { create } = require('./orchestrion')
+
+let dcPolyfill
+
+try {
+  dcPolyfill = require.resolve('dc-polyfill').replaceAll('\\', '/')
+} catch {
+  // The `dc-polyfill` module is unavailable for some reason (like bundling).
+  // Let's just keep the default of using `diagnostics-channel` as a fallback
+  // which works for most Node versions.
+}
 
 /** @type {Record<string, string>} map of module base name to version */
 const moduleVersions = {}
 const disabled = new Set()
-const matcher = create(instrumentations, 'dc-polyfill')
+const matcher = create(instrumentations, dcPolyfill)
+
+matcher.addTransform('traceIterator', traceIterator)
+matcher.addTransform('traceAsyncIterator', traceAsyncIterator)
 
 function rewrite (content, filename, format) {
   if (!content) return content

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/ai.js

@@ -75,6 +75,31 @@ module.exports = [
     },
     channelName: 'selectTelemetryAttributes',
   },
+  // resolveLanguageModel called by all LLM entry points, its result is the resolved model instance.
+  {
+    module: {
+      name: 'ai',
+      versionRange: '>=6.0.0',
+      filePath: 'dist/index.js',
+    },
+    functionQuery: {
+      functionName: 'resolveLanguageModel',
+      kind: 'Sync',
+    },
+    channelName: 'resolveLanguageModel',
+  },
+  {
+    module: {
+      name: 'ai',
+      versionRange: '>=6.0.0',
+      filePath: 'dist/index.mjs',
+    },
+    functionQuery: {
+      functionName: 'resolveLanguageModel',
+      kind: 'Sync',
+    },
+    channelName: 'resolveLanguageModel',
+  },
   // tool
   {
     module: {

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/langgraph.js

@@ -10,9 +10,9 @@ module.exports = [
     functionQuery: {
       methodName: 'stream',
       className: 'Pregel',
-      kind: 'AsyncIterator',
     },
     channelName: 'Pregel_stream',
+    transform: 'traceAsyncIterator',
   },
   {
     module: {
@@ -23,8 +23,8 @@ module.exports = [
     functionQuery: {
       methodName: 'stream',
       className: 'Pregel',
-      kind: 'AsyncIterator',
     },
     channelName: 'Pregel_stream',
+    transform: 'traceAsyncIterator',
   },
 ]

package/packages/datadog-instrumentations/src/helpers/rewriter/{orchestrion/transforms.js → transforms.js}

@@ -1,5 +1,7 @@
 'use strict'
 
+// TODO: Move traceIterator to Orchestrion.
+
 const { parse, query, traverse } = require('./compiler')
 
 const tracingChannelPredicate = (node) => (
@@ -36,10 +38,7 @@ const transforms = module.exports = {
   },
 
   traceAsyncIterator: traceAny,
-  traceCallback: traceAny,
   traceIterator: traceAny,
-  tracePromise: traceAny,
-  traceSync: traceAny,
 }
 
 function traceAny (state, node, _parent, ancestry) {
@@ -117,33 +116,10 @@ function traceInstanceMethod (state, node, program) {
 }
 
 function wrap (state, node, program) {
-  const { channelName, operator } = state
+  const { operator } = state
 
   if (operator === 'traceAsyncIterator') return wrapIterator(state, node, program)
-  if (operator === 'traceCallback') return wrapCallback(state, node)
   if (operator === 'traceIterator') return wrapIterator(state, node, program)
-
-  const async = operator === 'tracePromise' ? 'async' : ''
-  const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
-  const wrapper = parse(`
-    function wrapper () {
-      const __apm$traced = ${async} () => {
-        const __apm$wrapped = () => {};
-        return __apm$wrapped.apply(this, arguments);
-      };
-      if (!${channelVariable}.hasSubscribers) return __apm$traced();
-      return ${channelVariable}.${operator}(__apm$traced, {
-        arguments,
-        self: this,
-        moduleVersion: "1.0.0"
-      });
-    }
-  `).body[0].body // Extract only block statement of function body.
-
-  // Replace the right-hand side assignment of `const __apm$wrapped = () => {}`.
-  query(wrapper, '[id.name=__apm$wrapped]')[0].init = node
-
-  return wrapper
 }
 
 function wrapSuper (_state, node) {
@@ -195,68 +171,6 @@ function wrapSuper (_state, node) {
   }
 }
 
-function wrapCallback (state, node) {
-  const { channelName, functionQuery: { callbackIndex = -1 } } = state
-  const channelVariable = 'tr_ch_apm$' + channelName.replaceAll(':', '_')
-  const wrapper = parse(`
-    function wrapper () {
-      const __apm$cb = Array.prototype.at.call(arguments, ${callbackIndex});
-      const __apm$ctx = {
-        arguments,
-        self: this,
-        moduleVersion: "1.0.0"
-      };
-      const __apm$traced = () => {
-        const __apm$wrapped = () => {};
-        return __apm$wrapped.apply(this, arguments);
-      };
-
-      if (!${channelVariable}.start.hasSubscribers) return __apm$traced();
-
-      function __apm$wrappedCb(err, res) {
-        if (err) {
-          __apm$ctx.error = err;
-          ${channelVariable}.error.publish(__apm$ctx);
-        } else {
-          __apm$ctx.result = res;
-        }
-
-        ${channelVariable}.asyncStart.runStores(__apm$ctx, () => {
-          try {
-            if (__apm$cb) {
-              return __apm$cb.apply(this, arguments);
-            }
-          } finally {
-            ${channelVariable}.asyncEnd.publish(__apm$ctx);
-          }
-        });
-      }
-
-      if (typeof __apm$cb !== 'function') {
-        return __apm$traced();
-      }
-      Array.prototype.splice.call(arguments, ${callbackIndex}, 1, __apm$wrappedCb);
-
-      return ${channelVariable}.start.runStores(__apm$ctx, () => {
-        try {
-          return __apm$traced();
-        } catch (err) {
-          __apm$ctx.error = err;
-          ${channelVariable}.error.publish(__apm$ctx);
-          throw err;
-        } finally {
-          ${channelVariable}.end.publish(__apm$ctx);
-        }
-      });
-    }
-  `).body[0].body // Extract only block statement of function body.
-
-  // Replace the right-hand side assignment of `const __apm$wrapped = () => {}`.
-  query(wrapper, '[id.name=__apm$wrapped]')[0].init = node
-
-  return wrapper
-}
-
 function wrapIterator (state, node, program) {
   const { channelName, operator } = state
   const baseChannel = channelName.replaceAll(':', '_')

package/packages/datadog-instrumentations/src/mocha/utils.js

@@ -384,6 +384,11 @@ function getOnTestEndHandler (config) {
       })
     } else if (ctx) { // if there is an afterEach to run, let's store the finalStatus for getOnHookEndHandler
       ctx.finalStatus = finalStatus
+      ctx.hasFailedAllRetries = hasFailedAllRetries
+      ctx.attemptToFixPassed = attemptToFixPassed
+      ctx.attemptToFixFailed = attemptToFixFailed
+      ctx.isAttemptToFixRetry = isAttemptToFixRetry
+      ctx.isAtrRetry = isAtrRetry
     }
   }
 }
@@ -404,6 +409,11 @@ function getOnHookEndHandler () {
             status,
             hasBeenRetried: isMochaRetry(test),
             isLastRetry: getIsLastRetry(test),
+            hasFailedAllRetries: ctx.hasFailedAllRetries,
+            attemptToFixPassed: ctx.attemptToFixPassed,
+            attemptToFixFailed: ctx.attemptToFixFailed,
+            isAttemptToFixRetry: ctx.isAttemptToFixRetry,
+            isAtrRetry: ctx.isAtrRetry,
             ...ctx.currentStore,
             finalStatus: ctx.finalStatus,
           })

package/packages/datadog-plugin-dd-trace-api/src/index.js

@@ -3,14 +3,10 @@
 const Plugin = require('../../dd-trace/src/plugins/plugin')
 const telemetryMetrics = require('../../dd-trace/src/telemetry/metrics')
 const apiMetrics = telemetryMetrics.manager.namespace('tracers')
-const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
 
 // api ==> here
 const objectMap = new WeakMap()
 
-const injectionEnabledTag =
-  `injection_enabled:${getValueFromEnvSources('DD_INJECTION_ENABLED') ? 'yes' : 'no'}`
-
 module.exports = class DdTraceApiPlugin extends Plugin {
   static id = 'dd-trace-api'
 
@@ -18,6 +14,7 @@ module.exports = class DdTraceApiPlugin extends Plugin {
     super(...args)
 
     const tracer = this._tracer
+    const injectionEnabledTag = `injection_enabled:${this._tracerConfig.injectionEnabled ? 'yes' : 'no'}`
 
     this.addSub('datadog-api:v1:tracerinit', ({ proxy }) => {
       const proxyVal = proxy()

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

@@ -0,0 +1,64 @@
+'use strict'
+
+const { channel } = require('dc-polyfill')
+const log = require('../log')
+const AIGuard = require('./sdk')
+
+const aiguardChannel = channel('dd-trace:ai:aiguard')
+
+let isEnabled = false
+let aiguard
+let block
+
+function enable (tracer, config) {
+  if (isEnabled) return
+
+  try {
+    aiguard = new AIGuard(tracer, config)
+    block = config.experimental?.aiguard?.block !== false
+
+    aiguardChannel.subscribe(onEvaluate)
+
+    isEnabled = true
+  } catch (err) {
+    log.error('AIGuard: unexpected error during initialization: %s', err.message)
+    disable()
+  }
+}
+
+function disable () {
+  if (!isEnabled) return
+
+  aiguardChannel.unsubscribe(onEvaluate)
+
+  aiguard = undefined
+  isEnabled = false
+  block = false
+}
+
+/**
+ * Handles channel messages with pre-converted messages.
+ *
+ * @param {{messages: Array<object>, resolve: Function, reject: Function}} ctx
+ */
+function onEvaluate (ctx) {
+  if (!ctx.messages?.length) {
+    ctx.resolve()
+    return
+  }
+
+  aiguard.evaluate(ctx.messages, { block })
+    .then(() => {
+      ctx.resolve()
+    })
+    .catch(err => {
+      if (err.name === 'AIGuardAbortError') {
+        ctx.reject(err)
+      } else {
+        log.error('AIGuard: unexpected error during evaluation: %s', err.message)
+        ctx.resolve()
+      }
+    })
+}
+
+module.exports = { enable, disable }

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

@@ -5,10 +5,9 @@
 const os = require('os')
 const {
   getEnvironmentVariable,
-  getEnvironmentVariables,
   getValueFromEnvSources,
 } = require('./config/helper')
-const { getIsAzureFunction, getIsFlexConsumptionAzureFunction } = require('./serverless')
+const { getIsAzureFunction } = require('./serverless')
 
 function extractSubscriptionID (ownerName) {
   if (ownerName !== undefined) {
@@ -46,32 +45,33 @@ function trimObject (obj) {
 }
 
 function buildMetadata () {
-  const {
-    COMPUTERNAME,
-    FUNCTIONS_EXTENSION_VERSION,
-    FUNCTIONS_WORKER_RUNTIME,
-    FUNCTIONS_WORKER_RUNTIME_VERSION,
-    WEBSITE_INSTANCE_ID,
-    WEBSITE_OWNER_NAME,
-    WEBSITE_OS,
-    WEBSITE_RESOURCE_GROUP,
-    WEBSITE_SITE_NAME,
-  } = getEnvironmentVariables()
+  const COMPUTERNAME = getEnvironmentVariable('COMPUTERNAME')
+  const FUNCTIONS_EXTENSION_VERSION = getEnvironmentVariable('FUNCTIONS_EXTENSION_VERSION')
+  const FUNCTIONS_WORKER_RUNTIME = getEnvironmentVariable('FUNCTIONS_WORKER_RUNTIME')
+  const FUNCTIONS_WORKER_RUNTIME_VERSION = getEnvironmentVariable('FUNCTIONS_WORKER_RUNTIME_VERSION')
+  const WEBSITE_INSTANCE_ID = getEnvironmentVariable('WEBSITE_INSTANCE_ID')
+  const WEBSITE_OWNER_NAME = getEnvironmentVariable('WEBSITE_OWNER_NAME')
+  const WEBSITE_OS = getEnvironmentVariable('WEBSITE_OS')
+  const WEBSITE_RESOURCE_GROUP = getEnvironmentVariable('WEBSITE_RESOURCE_GROUP')
+  const WEBSITE_SITE_NAME = getEnvironmentVariable('WEBSITE_SITE_NAME')
+  const WEBSITE_SKU = getEnvironmentVariable('WEBSITE_SKU')
 
   const DD_AZURE_RESOURCE_GROUP = getValueFromEnvSources('DD_AZURE_RESOURCE_GROUP')
+  const isAzureFunction = FUNCTIONS_EXTENSION_VERSION !== undefined && FUNCTIONS_WORKER_RUNTIME !== undefined
+  const isFlexConsumptionAzureFunction = isAzureFunction && WEBSITE_SKU === 'FlexConsumption'
 
   const subscriptionID = extractSubscriptionID(WEBSITE_OWNER_NAME)
 
   const siteName = WEBSITE_SITE_NAME
 
-  const [siteKind, siteType] = getIsAzureFunction()
+  const [siteKind, siteType] = isAzureFunction
     ? ['functionapp', 'function']
     : ['app', 'app']
 
   // Azure Functions on Flex Consumption plans require the `DD_AZURE_RESOURCE_GROUP` env var.
   // If this logic ever changes, update the logic in `libdatadog`, `serverless-components/src/datadog-trace-agent`,
   // and the serverless compat layers accordingly.
-  const resourceGroup = getIsFlexConsumptionAzureFunction()
+  const resourceGroup = isFlexConsumptionAzureFunction
     ? (DD_AZURE_RESOURCE_GROUP ?? WEBSITE_RESOURCE_GROUP ?? extractResourceGroup(WEBSITE_OWNER_NAME))
     : (WEBSITE_RESOURCE_GROUP ?? extractResourceGroup(WEBSITE_OWNER_NAME))