dd-trace 5.80.0 → 5.81.0

package/packages/datadog-instrumentations/src/cucumber.js

@@ -1,5 +1,5 @@
 'use strict'
-const { createCoverageMap } = require('istanbul-lib-coverage')
+const { createCoverageMap } = require('../../../vendor/dist/istanbul-lib-coverage')
 
 const { addHook, channel } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
@@ -43,7 +43,7 @@ const {
   CUCUMBER_WORKER_TRACE_PAYLOAD_CODE,
   getIsFaultyEarlyFlakeDetection
 } = require('../../dd-trace/src/plugins/util/test')
-const satisfies = require('semifies')
+const satisfies = require('../../../vendor/dist/semifies')
 
 const isMarkedAsUnskippable = (pickle) => {
   return pickle.tags.some(tag => tag.name === '@datadog:unskippable')

package/packages/datadog-instrumentations/src/find-my-way.js

@@ -9,13 +9,13 @@ function wrapOn (on) {
   return function onWithTrace (method, path, opts) {
     const index = typeof opts === 'function' ? 2 : 3
     const handler = arguments[index]
-    const wrapper = shimmer.wrapFunction(handler, handler => function (req) {
-      routeChannel.publish({ req, route: path })
-
-      return handler.apply(this, arguments)
-    })
 
     if (typeof handler === 'function') {
+      const wrapper = shimmer.wrapFunction(handler, handler => function (req) {
+        routeChannel.publish({ req, route: path })
+
+        return handler.apply(this, arguments)
+      })
       arguments[index] = wrapper
     }
 
@@ -24,6 +24,7 @@ function wrapOn (on) {
 }
 
 addHook({ name: 'find-my-way', versions: ['>=1'] }, Router => {
+  // No need to wrap the off method as it would not contain the handler function.
   shimmer.wrap(Router.prototype, 'on', wrapOn)
 
   return Router

package/packages/datadog-instrumentations/src/google-genai.js

@@ -0,0 +1,120 @@
+'use strict'
+
+const { addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+const tracingChannel = require('dc-polyfill').tracingChannel
+const channel = require('dc-polyfill').channel
+
+const genaiTracingChannel = tracingChannel('apm:google:genai:request')
+const onStreamedChunkCh = channel('apm:google:genai:request:chunk')
+
+function wrapGenerateContent (method) {
+  return function wrappedGenerateContent (original) {
+    return function (...args) {
+      if (!genaiTracingChannel.start.hasSubscribers) {
+        return original.apply(this, args)
+      }
+
+      const normalizedName = normalizeMethodName(method)
+
+      const ctx = { args, methodName: normalizedName }
+
+      return genaiTracingChannel.start.runStores(ctx, () => {
+        let result
+        try {
+          result = original.apply(this, arguments)
+        } catch (error) {
+          finish(ctx, null, error)
+          throw error
+        } finally {
+          genaiTracingChannel.end.publish(ctx)
+        }
+        return result.then(response => {
+          if (response[Symbol.asyncIterator]) {
+            shimmer.wrap(response, Symbol.asyncIterator, iterator => wrapStreamIterator(iterator, ctx))
+          } else {
+            finish(ctx, response, null)
+          }
+          return response
+        }).catch(error => {
+          finish(ctx, null, error)
+          throw error
+        })
+      })
+    }
+  }
+}
+
+function wrapStreamIterator (iterator, ctx) {
+  return function () {
+    const itr = iterator.apply(this, arguments)
+    shimmer.wrap(itr, 'next', next => function () {
+      return next.apply(this, arguments)
+        .then(res => {
+          const { done, value: chunk } = res
+          onStreamedChunkCh.publish({ ctx, chunk, done })
+
+          if (done) {
+            finish(ctx)
+          }
+
+          return res
+        })
+        .catch(error => {
+          finish(ctx, null, error)
+          throw error
+        })
+    })
+
+    return itr
+  }
+}
+function finish (ctx, result, error) {
+  if (error) {
+    ctx.error = error
+    genaiTracingChannel.error.publish(ctx)
+  }
+
+  // streamed responses are handled and set separately
+  ctx.result ??= result
+
+  genaiTracingChannel.asyncEnd.publish(ctx)
+}
+// Hook the main package entry point
+addHook({
+  name: '@google/genai',
+  versions: ['>=1.19.0']
+}, exports => {
+  // Wrap GoogleGenAI to intercept when it creates Models instances
+  if (!exports.GoogleGenAI) return exports
+
+  shimmer.wrap(exports, 'GoogleGenAI', GoogleGenAI => {
+    return class extends GoogleGenAI {
+      constructor (...args) {
+        super(...args)
+
+        // We are patching the instance instead of the prototype because when it is compiled from
+        // typescript, the models property is not available on the prototype.
+        if (this.models) {
+          if (this.models.generateContent) {
+            shimmer.wrap(this.models, 'generateContent', wrapGenerateContent('generateContent'))
+          }
+          if (this.models.generateContentStream) {
+            shimmer.wrap(this.models, 'generateContentStream', wrapGenerateContent('generateContentStream'))
+          }
+          if (this.models.embedContent) {
+            shimmer.wrap(this.models, 'embedContent', wrapGenerateContent('embedContent'))
+          }
+        }
+      }
+    }
+  })
+  return exports
+})
+
+function normalizeMethodName (methodName) {
+  // Convert camelCase to snake_case and add Models prefix
+  return 'Models.' + methodName
+    .replaceAll(/([a-z0-9])([A-Z])/g, '$1_$2')
+    .toLowerCase()
+}

package/packages/datadog-instrumentations/src/graphql.js

@@ -6,6 +6,8 @@ const {
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
+const ddGlobal = globalThis[Symbol.for('dd-trace')] ??= {}
+
 /** cached objects */
 
 const contexts = new WeakMap()
@@ -360,3 +362,21 @@ addHook({ name: 'graphql', file: 'validation/validate.js', versions: ['>=0.10']
 
   return validate
 })
+
+addHook({ name: 'graphql', file: 'language/printer.js', versions: ['>=0.10'] }, printer => {
+  // HACK: It's possible `graphql` is loaded before `@apollo/gateway` so we
+  //       can't use a channel as the latter plugin would load after the publish
+  //       happened. Not sure how to handle this so for now use a global.
+  ddGlobal.graphql_printer = printer
+  return printer
+})
+
+addHook({ name: 'graphql', file: 'language/visitor.js', versions: ['>=0.10'] }, visitor => {
+  ddGlobal.graphql_visitor = visitor
+  return visitor
+})
+
+addHook({ name: 'graphql', file: 'utilities/index.js', versions: ['>=0.10'] }, utilities => {
+  ddGlobal.graphql_utilities = utilities
+  return utilities
+})

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

@@ -13,6 +13,7 @@ module.exports = {
   '@playwright/test': () => require('../playwright'),
   '@elastic/elasticsearch': () => require('../elasticsearch'),
   '@elastic/transport': () => require('../elasticsearch'),
+  '@google/genai': () => require('../google-genai'),
   '@google-cloud/pubsub': () => require('../google-cloud-pubsub'),
   '@google-cloud/vertexai': () => require('../google-cloud-vertexai'),
   '@graphql-tools/executor': () => require('../graphql'),

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

@@ -2,6 +2,7 @@
 
 const dc = require('dc-polyfill')
 const instrumentations = require('./instrumentations')
+const rewriterInstrumentations = require('./rewriter/instrumentations')
 const { AsyncResource } = require('async_hooks')
 
 const channelMap = {}
@@ -22,6 +23,15 @@ exports.tracingChannel = function (name) {
   return tc
 }
 
+exports.getHooks = function getHooks (names) {
+  names = [names].flat()
+
+  return rewriterInstrumentations
+    .map(inst => inst.module)
+    .filter(({ name }) => names.includes(name))
+    .map(({ name, versionRange, filePath }) => ({ name, versions: [versionRange], file: filePath }))
+}
+
 /**
  * @param {object} args
  * @param {string|string[]} args.name module name

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

@@ -2,7 +2,7 @@
 
 const { channel } = require('dc-polyfill')
 const path = require('path')
-const satisfies = require('semifies')
+const satisfies = require('../../../../vendor/dist/semifies')
 const Hook = require('./hook')
 const requirePackageJson = require('../../../dd-trace/src/require-package-json')
 const log = require('../../../dd-trace/src/log')
@@ -10,6 +10,7 @@ const checkRequireCache = require('./check-require-cache')
 const telemetry = require('../../../dd-trace/src/guardrails/telemetry')
 const { isInServerlessEnvironment } = require('../../../dd-trace/src/serverless')
 const { getEnvironmentVariables } = require('../../../dd-trace/src/config-helper')
+const rewriter = require('./rewriter')
 
 const envs = getEnvironmentVariables()
 
@@ -48,6 +49,10 @@ if (DD_TRACE_DEBUG && DD_TRACE_DEBUG.toLowerCase() !== 'false') {
 const seenCombo = new Set()
 const allInstrumentations = {}
 
+for (const inst of disabledInstrumentations) {
+  rewriter.disable(inst)
+}
+
 // TODO: make this more efficient
 for (const packageName of names) {
   if (disabledInstrumentations.has(packageName)) continue

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

@@ -0,0 +1,27 @@
+'use strict'
+
+let meriyah
+let astring
+let esquery
+
+module.exports = {
+  parse: (...args) => {
+    meriyah ??= require('../../../../../vendor/dist/meriyah')
+
+    return meriyah.parse(...args)
+  },
+
+  generate: (...args) => {
+    astring ??= require('../../../../../vendor/dist/astring')
+
+    return astring.generate(...args)
+  },
+
+  traverse: (ast, query, visitor) => {
+    esquery ??= require('../../../../../vendor/dist/esquery').default
+
+    const selector = esquery.parse(query)
+
+    return esquery.traverse(ast, selector, visitor)
+  },
+}

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

@@ -0,0 +1,152 @@
+'use strict'
+
+/*
+This rewriter is basically a JavaScript version of Orchestrion-JS. The goal is
+not to replace Orchestrion-JS, but rather to make it easier and faster to write
+new integrations in the short-term, especially as many changes to the rewriter
+will be needed as all the patterns we need have not been identified yet. This
+will avoid the back and forth of having to make Rust changes to an external
+library for every integration change or addition that requires something new.
+
+In the meantime, we'll work concurrently on a change to Orchestrion-JS that
+adds an "arbitrary transform" or "plugin" system that can be used from
+JavaScript, in order to enable quick iteration while still using Orchestrion-JS.
+Once that's done we'll use that, so that we can remove this JS approach and
+return to using Orchestrion-JS.
+
+The long term goal is to backport any additional features we add to the JS
+rewriter (or using the plugin system in Orchestrion-JS once we're using that)
+to Orchestrion-JS  once we're confident that the implementation is fairly
+complete and has all features we need.
+
+Here is a list of the additions and changes in this rewriter compared to
+Orchestrion-JS that will need to be backported:
+
+(NOTE: Please keep this list up-to-date whenever new features are added)
+
+- Supports an `astQuery` field to filter AST nodes with an esquery query. This
+  is mostly meant to be used when experimenting or if what needs to be queried
+  is not a function. We'll see over time if something like this is needed to be
+  backported or if it can be replaced by simpler queries.
+- Supports replacing methods of child class instances in the base constructor.
+*/
+
+const { readFileSync } = require('fs')
+const { join } = require('path')
+const semifies = require('../../../../../vendor/dist/semifies')
+const transforms = require('./transforms')
+const { generate, parse, traverse } = require('./compiler')
+const log = require('../../../../dd-trace/src/log')
+const instrumentations = require('./instrumentations')
+const { getEnvironmentVariable } = require('../../../../dd-trace/src/config-helper')
+
+const NODE_OPTIONS = getEnvironmentVariable('NODE_OPTIONS')
+
+const supported = {}
+const disabled = new Set()
+
+// TODO: Source maps without `--enable-source-maps`.
+const enableSourceMaps = NODE_OPTIONS?.includes('--enable-source-maps') ||
+  process.execArgv?.some(arg => arg.includes('--enable-source-maps'))
+
+let SourceMapGenerator
+
+function rewrite (content, filename, format) {
+  if (!content) return content
+
+  try {
+    let ast
+
+    filename = filename.replace('file://', '')
+
+    for (const inst of instrumentations) {
+      const { astQuery, functionQuery = {}, module: { name, versionRange, filePath } } = inst
+      const { kind } = functionQuery
+      const operator = kind === 'Async' ? 'tracePromise' : kind === 'Callback' ? 'traceCallback' : 'traceSync'
+      const transform = transforms[operator]
+
+      if (disabled.has(name)) continue
+      if (!filename.endsWith(`${name}/${filePath}`)) continue
+      if (!transform) continue
+      if (!satisfies(filename, filePath, versionRange)) continue
+
+      ast ??= parse(content.toString(), { loc: true, ranges: true, module: format === 'module' })
+
+      const query = astQuery || fromFunctionQuery(functionQuery)
+      const state = { ...inst, format, functionQuery, operator }
+
+      traverse(ast, query, (...args) => transform(state, ...args))
+    }
+
+    if (ast) {
+      if (!enableSourceMaps) return generate(ast)
+
+      // TODO: Can we use the same version of `source-map` that DI uses?
+      SourceMapGenerator ??= require('../../../../../vendor/dist/@datadog/source-map').SourceMapGenerator
+
+      const sourceMap = new SourceMapGenerator({ file: filename })
+      const code = generate(ast, { sourceMap })
+      const map = Buffer.from(sourceMap.toString()).toString('base64')
+
+      return code + '\n' + `//# sourceMappingURL=data:application/json;base64,${map}`
+    }
+  } catch (e) {
+    log.error(e)
+  }
+
+  return content
+}
+
+function disable (instrumentation) {
+  disabled.add(instrumentation)
+}
+
+function satisfies (filename, filePath, versions) {
+  const [basename] = filename.split(filePath)
+
+  if (supported[basename] === undefined) {
+    try {
+      const pkg = JSON.parse(readFileSync(
+        join(basename, 'package.json'), 'utf8'
+      ))
+
+      supported[basename] = semifies(pkg.version, versions)
+    } catch {
+      supported[basename] = false
+    }
+  }
+
+  return supported[basename]
+}
+
+// TODO: Support index
+function fromFunctionQuery (functionQuery) {
+  const { methodName, functionName, expressionName, className } = functionQuery
+  const queries = []
+
+  if (className) {
+    queries.push(
+      `[id.name="${className}"]`,
+      `[id.name="${className}"] > ClassBody > [key.name="${methodName}"] > [async]`,
+      `[id.name="${className}"] > ClassExpression > ClassBody > [key.name="${methodName}"] > [async]`
+    )
+  } else if (methodName) {
+    queries.push(
+      `ClassBody > [key.name="${methodName}"] > [async]`,
+      `Property[key.name="${methodName}"] > [async]`
+    )
+  }
+
+  if (functionName) {
+    queries.push(`FunctionDeclaration[id.name="${functionName}"][async]`)
+  } else if (expressionName) {
+    queries.push(
+      `FunctionExpression[id.name="${expressionName}"][async]`,
+      `ArrowFunctionExpression[id.name="${expressionName}"][async]`
+    )
+  }
+
+  return queries.join(', ')
+}
+
+module.exports = { rewrite, disable }

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

@@ -0,0 +1,5 @@
+'use strict'
+
+module.exports = [
+  ...require('./langchain')
+]

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

@@ -0,0 +1,237 @@
+'use strict'
+
+module.exports = [
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.cjs'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.cjs'
+    },
+    functionQuery: {
+      methodName: 'batch',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_batch'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/chat_models.cjs'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async',
+      className: 'BaseChatModel'
+    },
+    channelName: 'BaseChatModel_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/llms.cjs'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async'
+    },
+    channelName: 'BaseLLM_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.cjs'
+    },
+    functionQuery: {
+      methodName: 'embedQuery',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedQuery'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.cjs'
+    },
+    functionQuery: {
+      methodName: 'embedDocuments',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedDocuments'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/tools/index.cjs'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'StructuredTool'
+    },
+    channelName: 'Tool_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.cjs'
+    },
+    functionQuery: {
+      methodName: 'similaritySearch',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearch'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.cjs'
+    },
+    functionQuery: {
+      methodName: 'similaritySearchWithScore',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearchWithScore'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.js'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/runnables/base.js'
+    },
+    functionQuery: {
+      methodName: 'batch',
+      kind: 'Async',
+      className: 'RunnableSequence'
+    },
+    channelName: 'RunnableSequence_batch'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/chat_models.js'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async',
+      className: 'BaseChatModel'
+    },
+    channelName: 'BaseChatModel_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/language_models/llms.js'
+    },
+    functionQuery: {
+      methodName: 'generate',
+      kind: 'Async'
+    },
+    channelName: 'BaseLLM_generate'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.js'
+    },
+    functionQuery: {
+      methodName: 'embedQuery',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedQuery'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/embeddings.js'
+    },
+    functionQuery: {
+      methodName: 'embedDocuments',
+      kind: 'Async',
+      className: 'Embeddings'
+    },
+    channelName: 'Embeddings_embedDocuments'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/tools/index.js'
+    },
+    functionQuery: {
+      methodName: 'invoke',
+      kind: 'Async',
+      className: 'StructuredTool'
+    },
+    channelName: 'Tool_invoke'
+  },
+  {
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.js'
+    },
+    functionQuery: {
+      methodName: 'similaritySearch',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearch'
+  },
+  {
+
+    module: {
+      name: '@langchain/core',
+      versionRange: '>=0.1',
+      filePath: 'dist/vectorstores.js'
+    },
+    functionQuery: {
+      methodName: 'similaritySearchWithScore',
+      kind: 'Async',
+      className: 'VectorStore'
+    },
+    channelName: 'VectorStore_similaritySearchWithScore'
+  }
+]

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

@@ -0,0 +1,9 @@
+'use strict'
+
+const Module = require('module')
+const shimmer = require('../../../../datadog-shimmer')
+const { rewrite } = require('./')
+
+shimmer.wrap(Module.prototype, '_compile', compile => function (content, filename, format) {
+  return compile.call(this, rewrite(content, filename, format), filename, format)
+})