dd-trace 6.14.0 → 6.16.0

package/packages/datadog-esbuild/src/utils.js

@@ -1,13 +1,13 @@
 'use strict'
 
 // The content of this file is copied from the `import-in-the-middle` package with minor modifications (https://www.npmjs.com/package/import-in-the-middle)
-const { pathToFileURL, fileURLToPath } = require('node:url')
 const fs = require('node:fs')
 const path = require('node:path')
-const { NODE_MAJOR, NODE_MINOR } = require('../../../version.js')
+const { fileURLToPath, pathToFileURL } = require('node:url')
 
-const LOAD_OPERATION = 0
-const RESOLVE_OPERATION = 1
+const { createEsmResolver, driveGetExportsGenerator } = require('./resolver')
+
+/** @typedef {ReturnType<typeof createEsmResolver>} EsmResolver */
 
 const getExportsImporting = (url) => import(url).then(Object.keys)
 let getExportsModulePromise
@@ -19,16 +19,33 @@ const loadGetExportsModule = () => {
   return getExportsModulePromise
 }
 
-const getExports = NODE_MAJOR >= 20 || (NODE_MAJOR === 18 && NODE_MINOR >= 19)
-  ? async (srcUrl, context, getSource) => {
-    const mod = await loadGetExportsModule()
-    const exportNames = mod.getExports(srcUrl, context, getSource)
-    if (exportNames?.next) {
-      return driveGetExportsGenerator(exportNames, getSource)
-    }
-    return exportNames
+/**
+ * @param {URL} srcUrl
+ * @param {object} context
+ * @param {(url: URL, context: object) => { source: string, format: string }} getSource
+ * @param {EsmResolver} resolver
+ * @returns {Promise<{
+ *   exportNames: Iterable<string>,
+ *   hasModuleExportsCJSDefault: boolean,
+ *   starReexports?: Array<{ specifier: string, parentURL: string }>
+ * }>}
+ */
+const getExports = async (srcUrl, context, getSource, resolver) => {
+  const mod = await loadGetExportsModule()
+  const exportsGenerator = mod.getExports(srcUrl, context, getSource)
+  /**
+   * @param {string} specifier
+   * @param {{ parentURL: URL }} operationContext
+   * @returns {Promise<{ format: string, url: URL }>}
+   */
+  const resolve = (specifier, operationContext) => resolveModule(specifier, operationContext, resolver)
+  const result = await driveGetExportsGenerator(exportsGenerator, getSource, resolve)
+  return {
+    exportNames: result.exportNames,
+    hasModuleExportsCJSDefault: mod.hasModuleExportsCJSDefault,
+    starReexports: result.starReexports,
   }
-  : getExportsImporting
+}
 
 function isStarExportLine (line) {
   return /^\* from /.test(line)
@@ -58,82 +75,37 @@ function isBareSpecifier (specifier) {
   }
 }
 
-function resolve (specifier, context) {
-  // This comes from an import, that is why import makes preference
-  const conditions = ['import']
-
-  if (specifier.startsWith('file://')) {
-    specifier = fileURLToPath(specifier)
-  }
-
-  const resolved = require.resolve(specifier, {
-    paths: [fileURLToPath(context.parentURL)],
-    // @ts-expect-error - Node.js 22+ unofficially supports a conditions option
-    conditions,
-  })
-
-  return {
-    url: pathToFileURL(resolved),
-    format: isESMFile(resolved) ? 'module' : 'commonjs',
-  }
-}
-
-function getSource (url, { format }) {
-  return {
-    source: fs.readFileSync(fileURLToPath(url), 'utf8'),
-    format,
-  }
-}
-
 /**
- * @typedef {[typeof LOAD_OPERATION, URL, object] | [typeof RESOLVE_OPERATION, string, object]} GetExportsOperation
- */
-
-/**
- * @typedef {{ done: false, value: GetExportsOperation } | { done: true, value: Set<string> }} GetExportsResult
- */
-
-/**
- * @typedef {{
- *   next: (value?: unknown) => GetExportsResult,
- *   throw: (error?: unknown) => GetExportsResult,
- * }} GetExportsGenerator
+ * @param {string} modulePath
+ * @param {string} format
+ * @returns {string}
  */
+function getModuleFormat (modulePath, format) {
+  if (modulePath.endsWith('.mts')) return 'module-typescript'
+  if (modulePath.endsWith('.cts')) return 'commonjs-typescript'
+  if (modulePath.endsWith('.ts')) return `${format}-typescript`
+  return format
+}
 
 /**
- * Drives the generator returned by import-in-the-middle >=3.1.0 export discovery.
+ * Resolves a module with the import conditions used by ESM instrumentation.
  *
- * @param {GetExportsGenerator} exportsGenerator Generator returned by getExports
- * @param {(url: URL, context: object) => { source: string, format: string }} getSource
- * Function that loads module source
- * @returns {Set<string>}
+ * @param {string} specifier
+ * @param {{ parentURL: URL }} context
+ * @param {EsmResolver} resolver
+ * @returns {Promise<{ format: string, url: URL }>}
  */
-function driveGetExportsGenerator (exportsGenerator, getSource) {
-  let next = exportsGenerator.next()
-  while (next.done === false) {
-    let result
-    let error
-    let threw = false
-
-    try {
-      const operation = next.value
-      const operationType = operation[0]
-
-      if (operationType === LOAD_OPERATION) {
-        result = getSource(operation[1], operation[2])
-      } else if (operationType === RESOLVE_OPERATION) {
-        result = resolve(operation[1], operation[2])
-      } else {
-        throw new Error(`Unsupported import-in-the-middle getExports operation: ${operationType}`)
-      }
-    } catch (err) {
-      threw = true
-      error = err
-    }
-
-    next = threw ? exportsGenerator.throw(error) : exportsGenerator.next(result)
+async function resolveModule (specifier, context, resolver) {
+  const url = new URL(await resolver.resolve(specifier, context.parentURL))
+  if (url.protocol === 'node:') return { format: 'builtin', url }
+  if (url.protocol !== 'file:') throw new Error(`Unsupported ESM resolution URL: ${url.href}`)
+
+  const resolved = fileURLToPath(url)
+  if (resolved.endsWith('.node')) {
+    throw new Error(`Unsupported ESM analysis target: ${resolved}`)
   }
-  return next.value
+  const format = isESMFile(resolved) ? 'module' : 'commonjs'
+  return { format, url }
 }
 
 /**
@@ -144,16 +116,72 @@ function driveGetExportsGenerator (exportsGenerator, getSource) {
  * @param {boolean} [moduleData.internal]
  * @param {object} moduleData.context
  * @param {boolean} [moduleData.excludeDefault]
+ * @param {Map<string, string>} [moduleData.moduleSources]
+ * @param {EsmResolver} [moduleData.resolver]
+ * @param {(source: string, options: { loader: 'ts' }) => { code: string }} [moduleData.transform]
+ * @param {Set<string>} [activeModules]
  * @returns {Promise<Map>}
  */
-async function processModule ({ path, internal = false, context, excludeDefault = false }) {
-  let exportNames, srcUrl
+async function processModule (
+  { path, internal = false, context, excludeDefault = false, moduleSources = new Map(), resolver, transform },
+  activeModules
+) {
+  const ownsResolver = resolver === undefined
+  resolver ??= createEsmResolver()
+  try {
+    return await processModuleWithResolver(
+      { path, internal, context, excludeDefault, moduleSources, transform },
+      activeModules,
+      resolver
+    )
+  } finally {
+    if (ownsResolver) await resolver.close()
+  }
+}
+
+/**
+ * @param {object} moduleData
+ * @param {string} moduleData.path
+ * @param {boolean} moduleData.internal
+ * @param {object} moduleData.context
+ * @param {boolean} moduleData.excludeDefault
+ * @param {Map<string, string>} moduleData.moduleSources
+ * @param {(source: string, options: { loader: 'ts' }) => { code: string }} [moduleData.transform]
+ * @param {Set<string>} [activeModules]
+ * @param {EsmResolver} resolver
+ * @returns {Promise<Map>}
+ */
+async function processModuleWithResolver (
+  { path, internal, context, excludeDefault, moduleSources, transform },
+  activeModules,
+  resolver
+) {
+  let moduleExports, srcUrl
   if (internal) {
     // we can not read and parse of internal modules
-    exportNames = await getExportsImporting(path)
+    moduleExports = { exportNames: await getExportsImporting(path) }
   } else {
+    context = { ...context, format: getModuleFormat(path, context.format) }
     srcUrl = pathToFileURL(path)
-    exportNames = await getExports(srcUrl, context, getSource)
+    const readSource = (url) => {
+      const modulePath = fileURLToPath(url)
+      let source = moduleSources.get(modulePath)
+      if (source === undefined) {
+        source = fs.readFileSync(modulePath, 'utf8')
+        moduleSources.set(modulePath, source)
+      }
+      return source
+    }
+    const loadSource = (url, { format }) => {
+      let source = readSource(url)
+      if (transform && (format === 'module-typescript' || format === 'commonjs-typescript')) {
+        source = transform(source, { loader: 'ts' }).code
+        format = format === 'module-typescript' ? 'module' : 'commonjs'
+      }
+      return { format, source }
+    }
+    readSource(srcUrl)
+    moduleExports = await getExports(srcUrl, context, loadSource, resolver)
   }
 
   const starExports = new Set()
@@ -186,43 +214,31 @@ async function processModule ({ path, internal = false, context, excludeDefault
     }
   }
 
-  for (const n of exportNames) {
-    if (n === 'default' && excludeDefault) continue
+  let starReexports = moduleExports.starReexports
+  for (const n of moduleExports.exportNames) {
+    const isDefault = n === 'default' || (
+      n === 'module.exports' &&
+      moduleExports.hasModuleExportsCJSDefault &&
+      (context.format === 'commonjs' || context.format === 'commonjs-typescript')
+    )
+    if (isDefault && excludeDefault) continue
 
     if (isStarExportLine(n)) {
-      // export * from 'wherever'
-      const [, modFile] = n.split('* from ')
-
-      // Relative paths need to be resolved relative to the parent module
-      const newSpecifier = isBareSpecifier(modFile) ? modFile : new URL(modFile, srcUrl).href
-      // We need to call `parentResolve` to resolve bare specifiers to a full
-      // URL. We also need to call `parentResolve` for all sub-modules to get
-      // the `format`. We can't rely on the parents `format` to know if this
-      // sub-module is ESM or CJS!
+      starReexports ??= []
+      starReexports.push({ parentURL: srcUrl.href, specifier: n.slice('* from '.length) })
+      continue
+    }
 
-      const result = resolve(newSpecifier, { parentURL: srcUrl })
+    const variableName = `$dd${Buffer.from(n).toString('hex')}`
+    const objectKey = JSON.stringify(n)
+    const reExportedName = n === 'default' ? n : objectKey
 
-      // eslint-disable-next-line no-await-in-loop
-      const subSetters = await processModule({
-        path: fileURLToPath(result.url),
-        context: { ...context, format: result.format },
-        excludeDefault: true,
-      })
-
-      for (const [name, setter] of subSetters.entries()) {
-        addSetter(name, setter, true)
-      }
-    } else {
-      const variableName = `$${n.replaceAll(/[^a-zA-Z0-9_$]/g, '_')}`
-      const objectKey = JSON.stringify(n)
-      const reExportedName = n === 'default' ? n : objectKey
-
-      addSetter(n, `
+    addSetter(n, `
       let ${variableName}
       try {
         ${variableName} = _[${objectKey}] = namespace[${objectKey}]
-      } catch (err) {
-        if (!(err instanceof ReferenceError)) throw err
+      } catch (error) {
+        if (!(error instanceof ReferenceError)) throw error
       }
       export { ${variableName} as ${reExportedName} }
       set[${objectKey}] = (v) => {
@@ -231,6 +247,34 @@ async function processModule ({ path, internal = false, context, excludeDefault
       }
       get[${objectKey}] = () => ${variableName}
       `)
+  }
+
+  if (starReexports) {
+    for (const { parentURL, specifier } of starReexports) {
+      const baseUrl = new URL(parentURL)
+      const resolvedSpecifier = isBareSpecifier(specifier) ? specifier : new URL(specifier, baseUrl).href
+      // The runtime's import conditions and the declaring module's URL own star-export resolution.
+      // eslint-disable-next-line no-await-in-loop
+      const result = await resolveModule(resolvedSpecifier, { parentURL: baseUrl }, resolver)
+
+      activeModules ??= new Set([srcUrl.href])
+      if (activeModules.has(result.url.href)) continue
+      activeModules.add(result.url.href)
+
+      // eslint-disable-next-line no-await-in-loop
+      const subSetters = await processModuleWithResolver({
+        path: result.format === 'builtin' ? result.url.href : fileURLToPath(result.url),
+        internal: result.format === 'builtin',
+        context: { ...context, format: result.format },
+        excludeDefault: true,
+        moduleSources,
+        transform,
+      }, activeModules, resolver)
+      activeModules.delete(result.url.href)
+
+      for (const [name, setter] of subSetters.entries()) {
+        addSetter(name, setter, true)
+      }
     }
   }
 
@@ -246,8 +290,8 @@ async function processModule ({ path, internal = false, context, excludeDefault
  * @returns {boolean}
  */
 function isESMFile (fullPathToModule, modulePackageJsonPath, packageJson = {}) {
-  if (fullPathToModule.endsWith('.mjs')) return true
-  if (fullPathToModule.endsWith('.cjs')) return false
+  if (fullPathToModule.endsWith('.mjs') || fullPathToModule.endsWith('.mts')) return true
+  if (fullPathToModule.endsWith('.cjs') || fullPathToModule.endsWith('.cts')) return false
 
   const pathParts = fullPathToModule.split(path.sep)
   do {
@@ -271,6 +315,6 @@ function isESMFile (fullPathToModule, modulePackageJsonPath, packageJson = {}) {
 }
 
 module.exports = {
-  processModule,
   isESMFile,
+  processModule,
 }

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

@@ -6,112 +6,48 @@ const { addHook, getHooks } = require('./helpers/instrument')
 
 const vercelAiTracingChannel = tracingChannel('dd-trace:vercel-ai')
 const vercelAiSpanSetAttributesChannel = channel('dd-trace:vercel-ai:span:setAttributes')
-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 tracers = new WeakSet()
-const wrappedModels = new WeakSet()
+// Published synchronously per model call with the native call data. A subscriber may set either
+// callback and this instrumentation applies it, so the wrapping stays here and only the policy
+// lives outside:
+//   beforeResult () => Promise<void>|undefined     holds the result back until it settles
+//   onResult (result) => unknown|Promise<unknown>  inspects or replaces the delivered result
+const modelInterceptChannel = channel('dd-trace:vercel-ai:model:intercept')
 
-/**
- * Publishes a provider-native lifecycle payload to a cancelable lifecycle channel.
- *
- * Subscribers push async work into `pending` synchronously during publication and
- * abort `abortController` with an error before the pushed promise resolves to block.
- *
- * @param {object} lifecycleChannel
- * @param {object} payload
- * @returns {Promise<void>}
- */
-function publishLifecycle (lifecycleChannel, payload) {
-  const abortController = new AbortController()
-  const ctx = { ...payload, abortController, pending: [] }
+const INTERCEPTED_MODEL_METHODS = ['doGenerate', 'doStream']
 
-  lifecycleChannel.publish(ctx)
-
-  return Promise.all(ctx.pending).then(() => {
-    if (abortController.signal.aborted) {
-      throw abortController.signal.reason
-    }
-  })
-}
+const tracers = new WeakSet()
+const wrappedModels = new WeakSet()
 
 /**
- * Wraps a Vercel AI language model's doGenerate and doStream lifecycle methods.
+ * Wraps a language model's call methods so subscribers can interpose on each call.
  *
- * @param {object} model - A Vercel AI language model instance
+ * @param {object} model
  */
-function wrapModelWithLifecycle (model) {
+function wrapModel (model) {
   if (!model || wrappedModels.has(model)) return
   wrappedModels.add(model)
 
-  if (typeof model.doGenerate === 'function') {
-    shimmer.wrap(model, 'doGenerate', function (original) {
-      return function (options) {
-        const originalResult = original.call(this, options)
+  for (const method of INTERCEPTED_MODEL_METHODS) {
+    if (typeof model[method] !== 'function') continue
 
-        if (!doGenerateBeforeChannel.hasSubscribers && !doGenerateAfterChannel.hasSubscribers) return originalResult
-        if (!options.prompt?.length) return originalResult
+    shimmer.wrap(model, method, original => function (...args) {
+      const result = original.apply(this, args)
+      if (!modelInterceptChannel.hasSubscribers) return result
 
-        const beforeEvaluation = doGenerateBeforeChannel.hasSubscribers
-          ? publishLifecycle(doGenerateBeforeChannel, { prompt: options.prompt, options })
-          : Promise.resolve()
+      const interceptCtx = { method, arguments: args }
+      modelInterceptChannel.publish(interceptCtx)
 
-        return Promise.all([beforeEvaluation, originalResult])
-          .then(([, result]) => {
-            if (!doGenerateAfterChannel.hasSubscribers || !result.content?.length) return result
-            const payload = { prompt: options.prompt, options, result }
-            return publishLifecycle(doGenerateAfterChannel, payload).then(() => result)
-          })
-      }
-    })
-  }
+      if (!interceptCtx.beforeResult && !interceptCtx.onResult) return result
 
-  if (typeof model.doStream === 'function') {
-    shimmer.wrap(model, 'doStream', function (original) {
-      return function (options) {
-        const originalResult = original.call(this, options)
-
-        if (!doStreamBeforeChannel.hasSubscribers && !doStreamAfterChannel.hasSubscribers) return originalResult
-        if (!options.prompt?.length) return originalResult
-
-        const beforeEvaluation = doStreamBeforeChannel.hasSubscribers
-          ? publishLifecycle(doStreamBeforeChannel, { prompt: options.prompt, options })
-          : Promise.resolve()
-
-        return Promise.all([beforeEvaluation, originalResult])
-          .then(([, result]) => {
-            if (!doStreamAfterChannel.hasSubscribers) return result
-
-            const chunks = []
-            const reader = result.stream.getReader()
-
-            function readAll () {
-              return reader.read().then(({ done, value }) => {
-                if (done) return
-                chunks.push(value)
-                return readAll()
-              })
-            }
-
-            return readAll().then(() => {
-              return publishLifecycle(doStreamAfterChannel, { prompt: options.prompt, options, chunks })
-                .then(() => {
-                  // eslint-disable-next-line n/no-unsupported-features/node-builtins
-                  const stream = new ReadableStream({
-                    start (controller) {
-                      for (const chunk of chunks) {
-                        controller.enqueue(chunk)
-                      }
-                      controller.close()
-                    },
-                  })
-                  return { ...result, stream }
-                })
-            })
-          })
-      }
+      // A model may return any thenable, so normalize before attaching handlers.
+      const settled = Promise.resolve(result)
+      // The interceptor's rejection wins over an earlier SDK rejection.
+      settled.catch(() => {})
+
+      return Promise.resolve(interceptCtx.beforeResult?.())
+        .then(() => settled)
+        .then(value => (interceptCtx.onResult ? interceptCtx.onResult(value) : value))
     })
   }
 }
@@ -209,7 +145,9 @@ function wrapTracer (tracer) {
   })
 }
 
-for (const hook of getHooks('ai')) {
+let orchestrionSubscribed = false
+
+for (const hook of getHooks('ai').values()) {
   if (hook.file === 'dist/index.js') {
     // if not removed, the below hook will never match correctly
     // however, it is still needed in the orchestrion definition
@@ -217,6 +155,9 @@ for (const hook of getHooks('ai')) {
   }
 
   addHook(hook, exports => {
+    if (orchestrionSubscribed) return exports
+    orchestrionSubscribed = true
+
     const getTracerChannel = tracingChannel('orchestrion:ai:getTracer')
     getTracerChannel.subscribe({
       end (ctx) {
@@ -258,11 +199,14 @@ for (const hook of getHooks('ai')) {
     tracingChannel('orchestrion:ai:resolveLanguageModel').subscribe({
       end (ctx) {
         const model = ctx.arguments[0]
+
+        // The SDK builds a model from a string id, in which case only the resolved instance is
+        // worth wrapping; when the caller passed an instance, that is the one the SDK calls.
         if (typeof model !== 'string' && model !== ctx.result) {
-          wrapModelWithLifecycle(model)
+          wrapModel(model)
           wrappedModels.add(ctx.result)
         } else {
-          wrapModelWithLifecycle(ctx.result)
+          wrapModel(ctx.result)
         }
       },
     })
@@ -315,5 +259,3 @@ addHook({ name: 'ai', versions: ['>=7.0.0'] }, exports => {
 
   return exports
 })
-
-module.exports = { wrapModelWithLifecycle }