dd-trace 5.103.0 → 5.105.0

package/packages/datadog-instrumentations/src/helpers/openai-ai-guard.js

@@ -0,0 +1,269 @@
+'use strict'
+
+const dc = require('dc-polyfill')
+const shimmer = require('../../../datadog-shimmer')
+const {
+  convertOpenAIResponseItemsToMessages,
+  convertOpenAIResponsePromptToMessages,
+  normalizeOpenAIChatMessages,
+} = require('./ai-messages')
+
+// TODO: this channel name is incorrect, instrumentations publish with THEIR name, not with their subscribers names.
+const aiguardChannel = dc.channel('dd-trace:ai:aiguard')
+
+/**
+ * @typedef {object} ResourceHandler
+ * @property {(callArgs: object) => (Array<object>|undefined)} getInputMessages
+ * @property {(body: object) => Array<object>} getOutputMessages
+ * @property {(inputMessages: Array<object>, outputMessages: Array<object>) => Promise<unknown>}
+ *   publishOutputEvaluation
+ */
+
+/**
+ * @typedef {object} Guard
+ * @property {ResourceHandler} handler
+ * @property {Array<object>} inputMessages
+ * @property {() => Promise<void>} getInputEval
+ */
+
+/**
+ * Publishes already-converted AI-style messages to the AI Guard evaluation channel.
+ *
+ * @param {Array<object>} messages - AI-style messages to evaluate.
+ * @returns {Promise<void>}
+ */
+function publishEvaluation (messages) {
+  return new Promise((resolve, reject) => {
+    aiguardChannel.publish({ messages, integration: 'openai', resolve, reject })
+  })
+}
+
+/**
+ * Extracts OpenAI input messages from a `chat.completions.create` call.
+ *
+ * @param {object} callArgs - First argument passed to the wrapped method
+ * @returns {Array<object>|undefined}
+ */
+function getChatCompletionsInputMessages (callArgs) {
+  return normalizeOpenAIChatMessages(callArgs?.messages)
+}
+
+/**
+ * Extracts OpenAI output messages from a `chat.completions.create` parsed body.
+ * Includes any choice whose message carries content (including empty string),
+ * `tool_calls`, a `refusal` field, or the deprecated `function_call` field. GPT-4o
+ * emits `{content: null, refusal: "..."}` on policy refusals, and pre-tool-call
+ * SDK paths still produce `function_call`-only output — AI Guard must still see them.
+ *
+ * @param {object} body - Parsed response body
+ * @returns {Array<object>}
+ */
+function getChatCompletionsOutputMessages (body) {
+  const eligible = []
+  const choices = Array.isArray(body?.choices) ? body.choices : []
+  for (const choice of choices) {
+    const message = choice?.message
+    if (
+      message?.content != null ||
+      message?.tool_calls?.length ||
+      message?.refusal != null ||
+      message?.function_call != null
+    ) {
+      eligible.push(message)
+    }
+  }
+  return normalizeOpenAIChatMessages(eligible) ?? []
+}
+
+/**
+ * Publishes AI Guard After Model evaluation for `chat.completions` output.
+ *
+ * Chat completions may return multiple choices when `n > 1`. Screen every choice
+ * concurrently so any unsafe assistant output rejects `.parse()`, regardless of
+ * which choice the caller ends up using.
+ *
+ * @param {Array<object>} inputMessages
+ * @param {Array<object>} outputMessages - One entry per choice
+ * @returns {Promise<Array<void>>}
+ */
+function publishChatCompletionsOutputEvaluation (inputMessages, outputMessages) {
+  const evals = []
+  for (const message of outputMessages) {
+    evals.push(publishEvaluation([...inputMessages, message]))
+  }
+  return Promise.all(evals)
+}
+
+/**
+ * Extracts OpenAI input messages from a `responses.create` call. The `instructions`
+ * field is treated as a developer prompt — it directly steers model behavior and the
+ * LLMObs OpenAI plugin already surfaces it as one — so AI Guard must screen it too.
+ *
+ * AI Guard `/evaluate` accepts a single leading system/developer message; if the
+ * caller's `input` already begins with one, prepend the `instructions` text to its
+ * content rather than emit a second developer turn.
+ *
+ * @param {object} callArgs - First argument passed to the wrapped method
+ * @returns {Array<object>|undefined}
+ */
+function getResponsesInputMessages (callArgs) {
+  const messages = [
+    ...convertOpenAIResponseItemsToMessages(callArgs?.input, 'user'),
+    ...convertOpenAIResponsePromptToMessages(callArgs?.prompt),
+  ]
+
+  const instructions = typeof callArgs?.instructions === 'string' && callArgs.instructions.length
+    ? callArgs.instructions
+    : null
+  if (!instructions) return messages.length ? messages : undefined
+
+  const first = messages[0]
+  if (first && (first.role === 'developer' || first.role === 'system')) {
+    const merged = { role: 'developer', content: mergeInstructionsWithContent(instructions, first.content) }
+    return [merged, ...messages.slice(1)]
+  }
+  return [{ role: 'developer', content: instructions }, ...messages]
+}
+
+/**
+ * Merges Responses API instructions with an existing leading developer/system content value.
+ *
+ * @param {string} instructions
+ * @param {string|Array<object>|undefined} content
+ * @returns {string|Array<object>}
+ */
+function mergeInstructionsWithContent (instructions, content) {
+  if (Array.isArray(content)) return [{ type: 'text', text: instructions }, ...content]
+  if (typeof content === 'string' && content.length) return `${instructions}\n\n${content}`
+  return instructions
+}
+
+/**
+ * Extracts OpenAI output messages from a `responses.create` parsed body.
+ *
+ * @param {object} body - Parsed response body
+ * @returns {Array<object>}
+ */
+function getResponsesOutputMessages (body) {
+  return convertOpenAIResponseItemsToMessages(body?.output, 'assistant')
+}
+
+/**
+ * Publishes AI Guard After Model evaluation for `responses` output.
+ *
+ * The Responses API returns a single conversation turn whose `output` items form one
+ * coherent message (reasoning steps + final assistant message + tool calls + ...);
+ * they are screened together as a single evaluation.
+ *
+ * @param {Array<object>} inputMessages
+ * @param {Array<object>} outputMessages
+ * @returns {Promise<void>}
+ */
+function publishResponsesOutputEvaluation (inputMessages, outputMessages) {
+  return publishEvaluation([...inputMessages, ...outputMessages])
+}
+
+/**
+ * Per-resource handlers describing how AI Guard reads inputs and screens outputs for
+ * each LLM-prompt-accepting OpenAI endpoint. The keys also serve as the set of
+ * resources eligible for AI Guard evaluation.
+ *
+ * @type {Record<string, ResourceHandler>}
+ */
+const RESOURCE_HANDLERS = {
+  'chat.completions': {
+    getInputMessages: getChatCompletionsInputMessages,
+    getOutputMessages: getChatCompletionsOutputMessages,
+    publishOutputEvaluation: publishChatCompletionsOutputEvaluation,
+  },
+  responses: {
+    getInputMessages: getResponsesInputMessages,
+    getOutputMessages: getResponsesOutputMessages,
+    publishOutputEvaluation: publishResponsesOutputEvaluation,
+  },
+}
+
+/**
+ * Reports whether the AI Guard channel has subscribers. The OpenAI instrumentation
+ * uses this to decide whether to take the AI Guard path at all.
+ *
+ * @returns {boolean}
+ */
+function hasSubscribers () {
+  return aiguardChannel.hasSubscribers
+}
+
+/**
+ * Builds a guard handle when AI Guard is enabled and applicable to this call. The
+ * handle binds the per-resource handler so downstream functions never re-dispatch
+ * on `baseResource`. Returns null when AI Guard does not apply (no subscribers,
+ * non-eligible resource, streaming, or no input messages).
+ *
+ * @param {string} baseResource - e.g. `'chat.completions'` or `'responses'`
+ * @param {object} callArgs - First argument passed to the wrapped OpenAI method
+ * @param {boolean} stream - Whether the caller asked for a streamed response
+ * @returns {Guard|null}
+ */
+function createGuard (baseResource, callArgs, stream) {
+  // Streaming AI Guard support lands in a follow-up PR. For now, provider-level AI
+  // Guard only evaluates non-streaming responses.
+  if (stream || !aiguardChannel.hasSubscribers) return null
+  const handler = RESOURCE_HANDLERS[baseResource]
+  if (!handler) return null
+
+  const inputMessages = handler.getInputMessages(callArgs)
+  if (!inputMessages) return null
+
+  let inputEvalPromise
+  const getInputEval = () => (inputEvalPromise ??= publishEvaluation(inputMessages))
+  return { handler, inputMessages, getInputEval }
+}
+
+/**
+ * Wraps `apiProm.asResponse` so callers that consume the raw `Response` object still
+ * receive the Before Model verdict. After Model evaluation is not performed on this
+ * path because the response body has not been parsed.
+ *
+ * @param {object} apiProm - APIPromise returned from the OpenAI SDK method
+ * @param {Guard} guard
+ */
+function wrapAsResponse (apiProm, guard) {
+  if (typeof apiProm.asResponse !== 'function') return
+  shimmer.wrap(apiProm, 'asResponse', origAsResponse => function (...args) {
+    const responsePromise = origAsResponse.apply(this, args)
+    return Promise.all([guard.getInputEval(), responsePromise]).then(([, response]) => response)
+  })
+}
+
+/**
+ * Gates the parsed-body promise on Before Model evaluation. Resolves to the SDK's
+ * result only once the Before Model verdict is in.
+ *
+ * @param {Promise<unknown>} parsedPromise
+ * @param {Guard} guard
+ * @returns {Promise<unknown>}
+ */
+function gateParse (parsedPromise, guard) {
+  return Promise.all([guard.getInputEval(), parsedPromise]).then(([, result]) => result)
+}
+
+/**
+ * Runs After Model evaluation against the response body.
+ *
+ * @param {Guard} guard
+ * @param {object} body - Parsed OpenAI response body
+ * @returns {Promise<unknown>}
+ */
+function evaluateOutput (guard, body) {
+  const outputMessages = guard.handler.getOutputMessages(body)
+  if (!outputMessages.length) return Promise.resolve()
+  return guard.handler.publishOutputEvaluation(guard.inputMessages, outputMessages)
+}
+
+module.exports = {
+  hasSubscribers,
+  createGuard,
+  wrapAsResponse,
+  gateParse,
+  evaluateOutput,
+}

package/packages/datadog-instrumentations/src/helpers/promise-instrumentor.js

@@ -0,0 +1,42 @@
+'use strict'
+
+const dc = require('dc-polyfill')
+
+const { channel } = require('./instrument')
+
+/**
+ * Shimmer-compatible instrumentor for promise-returning APIs (e.g. `dns.promises.lookup`).
+ * Mirrors `createCallbackInstrumentor`'s channel triplet (`<prefix>:start`, `:finish`, `:error`)
+ * so a plugin subscribing to those channels for the callback variant works for the promise
+ * variant unchanged. `:finish` is the `tracingChannel` `asyncEnd` slot, so it fires after the
+ * promise settles with `ctx.result` set to the resolved value.
+ *
+ * @param {string} prefix
+ * @returns {(buildContext: (thisArg: unknown, args: unknown[]) => object | undefined) =>
+ *   (fn: Function) => Function}
+ */
+function createPromiseInstrumentor (prefix) {
+  const start = channel(prefix + ':start')
+  const finish = channel(prefix + ':finish')
+  const error = channel(prefix + ':error')
+  const tracing = dc.tracingChannel({
+    start,
+    end: channel(prefix + ':end'),
+    asyncStart: channel(prefix + ':asyncStart'),
+    asyncEnd: finish,
+    error,
+  })
+
+  return function instrument (buildContext) {
+    return function wrap (fn) {
+      return function (...args) {
+        if (!start.hasSubscribers) return fn.apply(this, args)
+        const ctx = buildContext(this, args)
+        if (ctx === undefined) return fn.apply(this, args)
+        return tracing.tracePromise(fn, ctx, this, ...args)
+      }
+    }
+  }
+}
+
+module.exports = { createPromiseInstrumentor }

package/packages/datadog-instrumentations/src/helpers/register.js

@@ -133,7 +133,7 @@ for (const name of names) {
         try {
           loadChannel.publish({ name })
 
-          moduleExports = hook(moduleExports, moduleVersion, isIitm) ?? moduleExports
+          moduleExports = hook(moduleExports, moduleVersion, isIitm, { moduleBaseDir, moduleName }) ?? moduleExports
         } catch (error) {
           log.info('Error during ddtrace instrumentation of application, aborting.', error)
           telemetry('error', [

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

@@ -19,9 +19,10 @@ const compiler = {
       // TODO: Figure out ESBuild `createRequire` issue and remove this hack.
       const oxc = runtimeRequire(['oxc', 'parser'].join('-'))
 
-      compiler.parse = (sourceText, options) => {
+      compiler.parse = (sourceText, { range, isModule } = {}) => {
         const { program, errors } = oxc.parseSync('index.js', sourceText, {
-          ...options,
+          range,
+          sourceType: isModule ? 'module' : 'script',
           preserveParens: false,
         })
 

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

@@ -2,15 +2,25 @@
 
 const { readFileSync } = require('fs')
 const { join } = require('path')
+const { pathToFileURL } = require('url')
 const log = require('../../../../dd-trace/src/log')
 const { create } = require('../../../../../vendor/dist/@apm-js-collab/code-transformer')
-const { traceAsyncIterator, traceIterator } = require('./transforms')
+const { waitForAsyncEnd } = require('./transforms')
 const instrumentations = require('./instrumentations')
 
-let dcPolyfill
+// `dc-polyfill` is referenced from injected `require()` (CJS) and `import`
+// (ESM) statements that the transformer splices into the rewritten module.
+// `require()` accepts an absolute filesystem path; the ESM resolver rejects it
+// with `ERR_INVALID_MODULE_SPECIFIER` and needs a `file://` URL instead. We
+// pre-compute both forms here so each matcher hands the transformer a
+// specifier that is valid for the module type it is rewriting.
+let dcPolyfillCjs
+let dcPolyfillEsm
 
 try {
-  dcPolyfill = require.resolve('dc-polyfill').replaceAll('\\', '/')
+  const resolved = require.resolve('dc-polyfill')
+  dcPolyfillCjs = resolved.replaceAll('\\', '/')
+  dcPolyfillEsm = pathToFileURL(resolved).href
 } 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
@@ -20,10 +30,12 @@ try {
 /** @type {Record<string, string>} map of module base name to version */
 const moduleVersions = {}
 const disabled = new Set()
-const matcher = create(instrumentations, dcPolyfill)
+const matcherCjs = create(instrumentations, dcPolyfillCjs)
+const matcherEsm = create(instrumentations, dcPolyfillEsm)
 
-matcher.addTransform('traceIterator', traceIterator)
-matcher.addTransform('traceAsyncIterator', traceAsyncIterator)
+for (const matcher of [matcherCjs, matcherEsm]) {
+  matcher.addTransform('waitForAsyncEnd', waitForAsyncEnd)
+}
 
 function rewrite (content, filename, format) {
   if (!content) return content
@@ -41,6 +53,7 @@ function rewrite (content, filename, format) {
 
   if (disabled.has(moduleName)) return content
 
+  const matcher = moduleType === 'esm' ? matcherEsm : matcherCjs
   const transformer = matcher.getTransformer(moduleName, version, filePath)
 
   if (!transformer) return content

package/packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/azure-cosmos.js

@@ -0,0 +1,50 @@
+'use strict'
+
+module.exports = [{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.1',
+    filePath: 'dist/browser/plugins/Plugin.js',
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async',
+  },
+  channelName: 'executePlugins',
+},
+{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.1',
+    filePath: 'dist/commonjs/plugins/Plugin.js',
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async',
+  },
+  channelName: 'executePlugins',
+},
+{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.1',
+    filePath: 'dist/esm/plugins/Plugin.js',
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async',
+  },
+  channelName: 'executePlugins',
+},
+{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.1',
+    filePath: 'dist/react-native/plugins/Plugin.js',
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async',
+  },
+  channelName: 'executePlugins',
+}]

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

@@ -2,8 +2,10 @@
 
 module.exports = [
   ...require('./ai'),
+  ...require('./azure-cosmos'),
   ...require('./bullmq'),
   ...require('./langchain'),
   ...require('./langgraph'),
   ...require('./modelcontextprotocol-sdk'),
+  ...require('./playwright'),
 ]

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

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

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

@@ -0,0 +1,85 @@
+'use strict'
+
+// Playwright 1.60 bundles several former hook targets into local classes/functions.
+// Keep these rewrites limited to private bundled internals that addHook cannot wrap.
+module.exports = [
+  {
+    module: {
+      name: 'playwright',
+      versionRange: '>=1.60.0',
+      filePath: 'lib/runner/index.js',
+    },
+    functionQuery: {
+      className: 'Dispatcher',
+      methodName: 'run',
+      kind: 'Async',
+    },
+    channelName: 'Dispatcher_run',
+  },
+  {
+    module: {
+      name: 'playwright',
+      versionRange: '>=1.60.0',
+      filePath: 'lib/runner/index.js',
+    },
+    functionQuery: {
+      className: 'Dispatcher',
+      methodName: '_createWorker',
+      kind: 'Sync',
+    },
+    channelName: 'Dispatcher_createWorker',
+  },
+  {
+    module: {
+      name: 'playwright',
+      versionRange: '>=1.60.0',
+      filePath: 'lib/runner/index.js',
+    },
+    functionQuery: {
+      className: 'ProcessHost',
+      methodName: 'startRunner',
+      kind: 'Async',
+    },
+    channelName: 'ProcessHost_startRunner',
+  },
+  {
+    module: {
+      name: 'playwright',
+      versionRange: '>=1.60.0',
+      filePath: 'lib/runner/index.js',
+    },
+    functionQuery: {
+      functionName: 'createRootSuite',
+      kind: 'Async',
+    },
+    channelName: 'createRootSuite',
+  },
+  {
+    module: {
+      name: 'playwright-core',
+      versionRange: '>=1.60.0',
+      filePath: 'lib/coreBundle.js',
+    },
+    astQuery: 'AssignmentExpression[left.name="Page2"] > ClassExpression > ClassBody > ' +
+      'MethodDefinition[kind="method"][key.name="goto"] > FunctionExpression[async], ' +
+      'VariableDeclarator[id.name="Page2"] > ClassExpression > ClassBody > ' +
+      'MethodDefinition[kind="method"][key.name="goto"] > FunctionExpression[async], ' +
+      'ClassDeclaration[id.name="Page2"] > ClassBody > ' +
+      'MethodDefinition[kind="method"][key.name="goto"] > FunctionExpression[async]',
+    functionQuery: {
+      methodName: 'goto',
+      kind: 'Async',
+    },
+    channelName: 'Page_goto',
+  },
+  {
+    module: {
+      name: 'playwright-core',
+      versionRange: '>=1.60.0',
+      filePath: 'lib/coreBundle.js',
+    },
+    astQuery: 'ReturnStatement > CallExpression[callee.object.name="promise"][callee.property.name="then"]',
+    channelName: 'Page_goto',
+    transform: 'waitForAsyncEnd',
+  },
+]