dd-trace 4.44.0 → 4.46.0

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

@@ -31,6 +31,9 @@ class OpenApiPlugin extends TracingPlugin {
   static get id () { return 'openai' }
   static get operation () { return 'request' }
   static get system () { return 'openai' }
+  static get prefix () {
+    return 'tracing:apm:openai:request'
+  }
 
   constructor (...args) {
     super(...args)
@@ -55,8 +58,10 @@ class OpenApiPlugin extends TracingPlugin {
     super.configure(config)
   }
 
-  start ({ methodName, args, basePath, apiKey }) {
+  bindStart (ctx) {
+    const { methodName, args, basePath, apiKey } = ctx
     const payload = normalizeRequestPayload(methodName, args)
+    const store = storage.getStore() || {}
 
     const span = this.startSpan('openai.request', {
       service: this.config.service,
@@ -87,18 +92,16 @@ class OpenApiPlugin extends TracingPlugin {
         'openai.request.user': payload.user,
         'openai.request.file_id': payload.file_id // deleteFile, retrieveFile, downloadFile
       }
-    })
+    }, false)
 
-    const fullStore = storage.getStore() || {} // certain request body fields are later used for logs
-    const store = Object.create(null)
-    fullStore.openai = store // namespacing these fields
+    const openaiStore = Object.create(null)
 
     const tags = {} // The remaining tags are added one at a time
 
     // createChatCompletion, createCompletion, createImage, createImageEdit, createTranscription, createTranslation
     if (payload.prompt) {
       const prompt = payload.prompt
-      store.prompt = prompt
+      openaiStore.prompt = prompt
       if (typeof prompt === 'string' || (Array.isArray(prompt) && typeof prompt[0] === 'number')) {
         // This is a single prompt, either String or [Number]
         tags['openai.request.prompt'] = normalizeStringOrTokenArray(prompt, true)
@@ -114,7 +117,7 @@ class OpenApiPlugin extends TracingPlugin {
     if (payload.input) {
       const normalized = normalizeStringOrTokenArray(payload.input, false)
       tags['openai.request.input'] = truncateText(normalized)
-      store.input = normalized
+      openaiStore.input = normalized
     }
 
     // createChatCompletion, createCompletion
@@ -141,12 +144,12 @@ class OpenApiPlugin extends TracingPlugin {
       case 'images.edit':
       case 'createImageVariation':
       case 'images.createVariation':
-        commonCreateImageRequestExtraction(tags, payload, store)
+        commonCreateImageRequestExtraction(tags, payload, openaiStore)
         break
 
       case 'createChatCompletion':
       case 'chat.completions.create':
-        createChatCompletionRequestExtraction(tags, payload, store)
+        createChatCompletionRequestExtraction(tags, payload, openaiStore)
         break
 
       case 'createFile':
@@ -160,7 +163,7 @@ class OpenApiPlugin extends TracingPlugin {
       case 'audio.transcriptions.create':
       case 'createTranslation':
       case 'audio.translations.create':
-        commonCreateAudioRequestExtraction(tags, payload, store)
+        commonCreateAudioRequestExtraction(tags, payload, openaiStore)
         break
 
       case 'retrieveModel':
@@ -184,23 +187,32 @@ class OpenApiPlugin extends TracingPlugin {
 
       case 'createEdit':
       case 'edits.create':
-        createEditRequestExtraction(tags, payload, store)
+        createEditRequestExtraction(tags, payload, openaiStore)
         break
     }
 
     span.addTags(tags)
+
+    ctx.currentStore = { ...store, span, openai: openaiStore }
+
+    return ctx.currentStore
   }
 
-  finish (response) {
-    const span = this.activeSpan
+  asyncEnd (ctx) {
+    const { result } = ctx
+    const store = ctx.currentStore
+
+    const span = store?.span
+    if (!span) return
+
     const error = !!span.context()._tags.error
 
     let headers, body, method, path
     if (!error) {
-      headers = response.headers
-      body = response.body
-      method = response.method
-      path = response.path
+      headers = result.headers
+      body = result.data
+      method = result.request.method
+      path = result.request.path
     }
 
     if (!error && headers?.constructor.name === 'Headers') {
@@ -210,10 +222,9 @@ class OpenApiPlugin extends TracingPlugin {
 
     body = coerceResponseBody(body, methodName)
 
-    const fullStore = storage.getStore()
-    const store = fullStore.openai
+    const openaiStore = store.openai
 
-    if (!error && (path.startsWith('https://') || path.startsWith('http://'))) {
+    if (!error && (path?.startsWith('https://') || path?.startsWith('http://'))) {
       // basic checking for if the path was set as a full URL
       // not using a full regex as it will likely be "https://api.openai.com/..."
       path = new URL(path).pathname
@@ -239,11 +250,11 @@ class OpenApiPlugin extends TracingPlugin {
           'openai.response.created_at': body.created_at
         }
 
-    responseDataExtractionByMethod(methodName, tags, body, store)
+    responseDataExtractionByMethod(methodName, tags, body, openaiStore)
     span.addTags(tags)
 
-    super.finish()
-    this.sendLog(methodName, span, tags, store, error)
+    span.finish()
+    this.sendLog(methodName, span, tags, openaiStore, error)
     this.sendMetrics(headers, body, endpoint, span._duration, error, tags)
   }
 
@@ -308,15 +319,15 @@ class OpenApiPlugin extends TracingPlugin {
     }
   }
 
-  sendLog (methodName, span, tags, store, error) {
-    if (!store) return
-    if (!Object.keys(store).length) return
+  sendLog (methodName, span, tags, openaiStore, error) {
+    if (!openaiStore) return
+    if (!Object.keys(openaiStore).length) return
     if (!this.sampler.isSampled()) return
 
     const log = {
       status: error ? 'error' : 'info',
       message: `sampled ${methodName}`,
-      ...store
+      ...openaiStore
     }
 
     this.logger.log(log, span, tags)
@@ -400,21 +411,21 @@ function countTokens (content, model) {
   }
 }
 
-function createEditRequestExtraction (tags, payload, store) {
+function createEditRequestExtraction (tags, payload, openaiStore) {
   const instruction = payload.instruction
   tags['openai.request.instruction'] = instruction
-  store.instruction = instruction
+  openaiStore.instruction = instruction
 }
 
 function retrieveModelRequestExtraction (tags, payload) {
   tags['openai.request.id'] = payload.id
 }
 
-function createChatCompletionRequestExtraction (tags, payload, store) {
+function createChatCompletionRequestExtraction (tags, payload, openaiStore) {
   const messages = payload.messages
   if (!defensiveArrayLength(messages)) return
 
-  store.messages = payload.messages
+  openaiStore.messages = payload.messages
   for (let i = 0; i < payload.messages.length; i++) {
     const message = payload.messages[i]
     tagChatCompletionRequestContent(message.content, i, tags)
@@ -424,20 +435,20 @@ function createChatCompletionRequestExtraction (tags, payload, store) {
   }
 }
 
-function commonCreateImageRequestExtraction (tags, payload, store) {
+function commonCreateImageRequestExtraction (tags, payload, openaiStore) {
   // createImageEdit, createImageVariation
   const img = payload.file || payload.image
   if (img !== null && typeof img === 'object' && img.path) {
     const file = path.basename(img.path)
     tags['openai.request.image'] = file
-    store.file = file
+    openaiStore.file = file
   }
 
   // createImageEdit
   if (payload.mask !== null && typeof payload.mask === 'object' && payload.mask.path) {
     const mask = path.basename(payload.mask.path)
     tags['openai.request.mask'] = mask
-    store.mask = mask
+    openaiStore.mask = mask
   }
 
   tags['openai.request.size'] = payload.size
@@ -445,7 +456,7 @@ function commonCreateImageRequestExtraction (tags, payload, store) {
   tags['openai.request.language'] = payload.language
 }
 
-function responseDataExtractionByMethod (methodName, tags, body, store) {
+function responseDataExtractionByMethod (methodName, tags, body, openaiStore) {
   switch (methodName) {
     case 'createModeration':
     case 'moderations.create':
@@ -458,7 +469,7 @@ function responseDataExtractionByMethod (methodName, tags, body, store) {
     case 'chat.completions.create':
     case 'createEdit':
     case 'edits.create':
-      commonCreateResponseExtraction(tags, body, store, methodName)
+      commonCreateResponseExtraction(tags, body, openaiStore, methodName)
       break
 
     case 'listFiles':
@@ -474,7 +485,7 @@ function responseDataExtractionByMethod (methodName, tags, body, store) {
 
     case 'createEmbedding':
     case 'embeddings.create':
-      createEmbeddingResponseExtraction(tags, body)
+      createEmbeddingResponseExtraction(tags, body, openaiStore)
       break
 
     case 'createFile':
@@ -629,14 +640,14 @@ function deleteFileResponseExtraction (tags, body) {
   tags['openai.response.id'] = body.id
 }
 
-function commonCreateAudioRequestExtraction (tags, body, store) {
+function commonCreateAudioRequestExtraction (tags, body, openaiStore) {
   tags['openai.request.response_format'] = body.response_format
   tags['openai.request.language'] = body.language
 
   if (body.file !== null && typeof body.file === 'object' && body.file.path) {
     const filename = path.basename(body.file.path)
     tags['openai.request.filename'] = filename
-    store.file = filename
+    openaiStore.file = filename
   }
 }
 
@@ -659,8 +670,8 @@ function createRetrieveFileResponseExtraction (tags, body) {
   tags['openai.response.status_details'] = body.status_details
 }
 
-function createEmbeddingResponseExtraction (tags, body) {
-  usageExtraction(tags, body)
+function createEmbeddingResponseExtraction (tags, body, openaiStore) {
+  usageExtraction(tags, body, openaiStore)
 
   if (!body.data) return
 
@@ -694,14 +705,14 @@ function createModerationResponseExtraction (tags, body) {
 }
 
 // createCompletion, createChatCompletion, createEdit
-function commonCreateResponseExtraction (tags, body, store, methodName) {
-  usageExtraction(tags, body, methodName)
+function commonCreateResponseExtraction (tags, body, openaiStore, methodName) {
+  usageExtraction(tags, body, methodName, openaiStore)
 
   if (!body.choices) return
 
   tags['openai.response.choices_count'] = body.choices.length
 
-  store.choices = body.choices
+  openaiStore.choices = body.choices
 
   for (let choiceIdx = 0; choiceIdx < body.choices.length; choiceIdx++) {
     const choice = body.choices[choiceIdx]
@@ -735,7 +746,7 @@ function commonCreateResponseExtraction (tags, body, store, methodName) {
 }
 
 // createCompletion, createChatCompletion, createEdit, createEmbedding
-function usageExtraction (tags, body, methodName) {
+function usageExtraction (tags, body, methodName, openaiStore) {
   let promptTokens = 0
   let completionTokens = 0
   let totalTokens = 0
@@ -743,14 +754,14 @@ function usageExtraction (tags, body, methodName) {
     promptTokens = body.usage.prompt_tokens
     completionTokens = body.usage.completion_tokens
     totalTokens = body.usage.total_tokens
-  } else if (['chat.completions.create', 'completions.create'].includes(methodName)) {
+  } else if (body.model && ['chat.completions.create', 'completions.create'].includes(methodName)) {
     // estimate tokens based on method name for completions and chat completions
     const { model } = body
     let promptEstimated = false
     let completionEstimated = false
 
     // prompt tokens
-    const payload = storage.getStore().openai
+    const payload = openaiStore
     const promptTokensCount = countPromptTokens(methodName, payload, model)
     promptTokens = promptTokensCount.promptTokens
     promptEstimated = promptTokensCount.promptEstimated

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

@@ -14,7 +14,8 @@ const {
   TEST_CONFIGURATION_BROWSER_NAME,
   TEST_IS_NEW,
   TEST_IS_RETRY,
-  TEST_EARLY_FLAKE_ENABLED
+  TEST_EARLY_FLAKE_ENABLED,
+  TELEMETRY_TEST_SESSION
 } = require('../../dd-trace/src/plugins/util/test')
 const { RESOURCE_NAME } = require('../../../ext/tags')
 const { COMPONENT } = require('../../dd-trace/src/constants')
@@ -59,6 +60,7 @@ class PlaywrightPlugin extends CiPlugin {
       this.testSessionSpan.finish()
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'session')
       finishAllTraceSpans(this.testSessionSpan)
+      this.telemetry.count(TELEMETRY_TEST_SESSION, { provider: this.ciProviderName })
       appClosingTelemetry()
       this.tracer._exporter.flush(onDone)
       this.numFailedTests = 0
@@ -160,8 +162,6 @@ class PlaywrightPlugin extends CiPlugin {
         stepSpan.finish(stepStartTime + stepDuration)
       })
 
-      span.finish()
-
       if (testStatus === 'fail') {
         this.numFailedTests++
       }
@@ -169,8 +169,13 @@ class PlaywrightPlugin extends CiPlugin {
       this.telemetry.ciVisEvent(
         TELEMETRY_EVENT_FINISHED,
         'test',
-        { hasCodeOwners: !!span.context()._tags[TEST_CODE_OWNERS] }
+        {
+          hasCodeOwners: !!span.context()._tags[TEST_CODE_OWNERS],
+          isNew,
+          browserDriver: 'playwright'
+        }
       )
+      span.finish()
 
       finishAllTraceSpans(span)
     })

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

@@ -7,9 +7,16 @@ const {
   getTestSuitePath,
   getTestSuiteCommonTags,
   TEST_SOURCE_FILE,
-  TEST_IS_RETRY
+  TEST_IS_RETRY,
+  TEST_CODE_COVERAGE_LINES_PCT,
+  TEST_CODE_OWNERS
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
+const {
+  TELEMETRY_EVENT_CREATED,
+  TELEMETRY_EVENT_FINISHED,
+  TELEMETRY_TEST_SESSION
+} = require('../../dd-trace/src/ci-visibility/telemetry')
 
 // Milliseconds that we subtract from the error test duration
 // so that they do not overlap with the following test
@@ -64,6 +71,9 @@ class VitestPlugin extends CiPlugin {
       const span = store?.span
 
       if (span) {
+        this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', {
+          hasCodeowners: !!span.context()._tags[TEST_CODE_OWNERS]
+        })
         span.setTag(TEST_STATUS, 'pass')
         span.finish(this.taskToFinishTime.get(task))
         finishAllTraceSpans(span)
@@ -75,6 +85,9 @@ class VitestPlugin extends CiPlugin {
       const span = store?.span
 
       if (span) {
+        this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', {
+          hasCodeowners: !!span.context()._tags[TEST_CODE_OWNERS]
+        })
         span.setTag(TEST_STATUS, 'fail')
 
         if (error) {
@@ -91,7 +104,7 @@ class VitestPlugin extends CiPlugin {
 
     this.addSub('ci:vitest:test:skip', ({ testName, testSuiteAbsolutePath }) => {
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      this.startTestSpan(
+      const testSpan = this.startTestSpan(
         testName,
         testSuite,
         this.testSuiteSpan,
@@ -99,7 +112,11 @@ class VitestPlugin extends CiPlugin {
           [TEST_SOURCE_FILE]: testSuite,
           [TEST_STATUS]: 'skip'
         }
-      ).finish()
+      )
+      this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', {
+        hasCodeowners: !!testSpan.context()._tags[TEST_CODE_OWNERS]
+      })
+      testSpan.finish()
     })
 
     this.addSub('ci:vitest:test-suite:start', ({ testSuiteAbsolutePath, frameworkVersion }) => {
@@ -124,6 +141,7 @@ class VitestPlugin extends CiPlugin {
           ...testSuiteMetadata
         }
       })
+      this.telemetry.ciVisEvent(TELEMETRY_EVENT_CREATED, 'suite')
       const store = storage.getStore()
       this.enter(testSuiteSpan, store)
       this.testSuiteSpan = testSuiteSpan
@@ -137,6 +155,7 @@ class VitestPlugin extends CiPlugin {
         span.finish()
         finishAllTraceSpans(span)
       }
+      this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
       // TODO: too frequent flush - find for method in worker to decrease frequency
       this.tracer._exporter.flush(onFinish)
     })
@@ -150,16 +169,23 @@ class VitestPlugin extends CiPlugin {
       }
     })
 
-    this.addSub('ci:vitest:session:finish', ({ status, onFinish, error }) => {
+    this.addSub('ci:vitest:session:finish', ({ status, onFinish, error, testCodeCoverageLinesTotal }) => {
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
       if (error) {
         this.testModuleSpan.setTag('error', error)
         this.testSessionSpan.setTag('error', error)
       }
+      if (testCodeCoverageLinesTotal) {
+        this.testModuleSpan.setTag(TEST_CODE_COVERAGE_LINES_PCT, testCodeCoverageLinesTotal)
+        this.testSessionSpan.setTag(TEST_CODE_COVERAGE_LINES_PCT, testCodeCoverageLinesTotal)
+      }
       this.testModuleSpan.finish()
+      this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'module')
       this.testSessionSpan.finish()
+      this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'session')
       finishAllTraceSpans(this.testSessionSpan)
+      this.telemetry.count(TELEMETRY_TEST_SESSION, { provider: this.ciProviderName })
       this.tracer._exporter.flush(onFinish)
     })
   }

package/packages/datadog-shimmer/src/shimmer.js

@@ -1,5 +1,7 @@
 'use strict'
 
+const log = require('../../dd-trace/src/log')
+
 // Use a weak map to avoid polluting the wrapped function/method.
 const unwrappers = new WeakMap()
 
@@ -18,9 +20,12 @@ function copyProperties (original, wrapped) {
   }
 }
 
-function wrapFn (original, delegate) {
-  assertFunction(delegate)
-  assertNotClass(original) // TODO: support constructors of native classes
+function wrapFunction (original, wrapper) {
+  if (typeof original === 'function') assertNotClass(original)
+  // TODO This needs to be re-done so that this and wrapMethod are distinct.
+  const target = { func: original }
+  wrapMethod(target, 'func', wrapper, typeof original !== 'function')
+  let delegate = target.func
 
   const shim = function shim () {
     return delegate.apply(this, arguments)
@@ -30,17 +35,144 @@ function wrapFn (original, delegate) {
     delegate = original
   })
 
-  copyProperties(original, shim)
+  if (typeof original === 'function') copyProperties(original, shim)
 
   return shim
 }
 
-function wrapMethod (target, name, wrapper) {
-  assertMethod(target, name)
-  assertFunction(wrapper)
+const wrapFn = function (original, delegate) {
+  throw new Error('calling `wrap()` with 2 args is deprecated. Use wrapFunction instead.')
+}
+
+// This is only used in safe mode. It's a simple state machine to track if the
+// original method was called and if it returned. We need this to determine if
+// an error was thrown by the original method, or by us. We'll use one of these
+// per call to a wrapped method.
+class CallState {
+  constructor () {
+    this.called = false
+    this.completed = false
+    this.retVal = undefined
+  }
+
+  startCall () {
+    this.called = true
+  }
+
+  endCall (retVal) {
+    this.completed = true
+    this.retVal = retVal
+  }
+}
+
+function isPromise (obj) {
+  return obj && typeof obj === 'object' && typeof obj.then === 'function'
+}
+
+let safeMode = !!process.env.DD_INEJCTION_ENABLED
+function setSafe (value) {
+  safeMode = value
+}
+
+function wrapMethod (target, name, wrapper, noAssert) {
+  if (!noAssert) {
+    assertMethod(target, name)
+    assertFunction(wrapper)
+  }
 
   const original = target[name]
-  const wrapped = wrapper(original)
+  let wrapped
+
+  if (safeMode && original) {
+    // In this mode, we make a best-effort attempt to handle errors that are thrown
+    // by us, rather than wrapped code. With such errors, we log them, and then attempt
+    // to return the result as if no wrapping was done at all.
+    //
+    // Caveats:
+    //   * If the original function is called in a later iteration of the event loop,
+    //     and we throw _then_, then it won't be caught by this. In practice, we always call
+    //     the original function synchronously, so this is not a problem.
+    //   * While async errors are dealt with here, errors in callbacks are not. This
+    //     is because we don't necessarily know _for sure_ that any function arguments
+    //     are wrapped by us. We could wrap them all anyway and just make that assumption,
+    //     or just assume that the last argument is always a callback set by us if it's a
+    //     function, but those don't seem like things we can rely on. We could add a
+    //     `shimmer.markCallbackAsWrapped()` function that's a no-op outside safe-mode,
+    //     but that means modifying every instrumentation. Even then, the complexity of
+    //     this code increases because then we'd need to effectively do the reverse of
+    //     what we're doing for synchronous functions. This is a TODO.
+
+    // We're going to hold on to current callState in this variable in this scope,
+    // which is fine because any time we reference it, we're referencing it synchronously.
+    // We'll use it in the our wrapper (which, again, is called syncrhonously), and in the
+    // errorHandler, which will already have been bound to this callState.
+    let currentCallState
+
+    // Rather than calling the original function directly from the shim wrapper, we wrap
+    // it again so that we can track if it was called and if it returned. This is because
+    // we need to know if an error was thrown by the original function, or by us.
+    // We could do this inside the `wrapper` function defined below, which would simplify
+    // managing the callState, but then we'd be calling `wrapper` on each invocation, so
+    // instead we do it here, once.
+    const innerWrapped = wrapper(function (...args) {
+      // We need to stash the callState here because of recursion.
+      const callState = currentCallState
+      callState.startCall()
+      const retVal = original.apply(this, args)
+      if (isPromise(retVal)) {
+        retVal.then(callState.endCall.bind(callState))
+      } else {
+        callState.endCall(retVal)
+      }
+      return retVal
+    })
+
+    // This is the crux of what we're doing in safe mode. It handles errors
+    // that _we_ cause, by logging them, and transparently providing results
+    // as if no wrapping was done at all. That means detecting (via callState)
+    // whether the function has already run or not, and if it has, returning
+    // the result, and otherwise calling the original function unwrapped.
+    const handleError = function (args, callState, e) {
+      if (callState.completed) {
+        // error was thrown after original function returned/resolved, so
+        // it was us. log it.
+        log.error(e)
+        // original ran and returned something. return it.
+        return callState.retVal
+      }
+
+      if (!callState.called) {
+        // error was thrown before original function was called, so
+        // it was us. log it.
+        log.error(e)
+        // original never ran. call it unwrapped.
+        return original.apply(this, args)
+      }
+
+      // error was thrown during original function execution, so
+      // it was them. throw.
+      throw e
+    }
+
+    // The wrapped function is the one that will be called by the user.
+    // It calls our version of the original function, which manages the
+    // callState. That way when we use the errorHandler, it can tell where
+    // the error originated.
+    wrapped = function (...args) {
+      currentCallState = new CallState()
+      const errorHandler = handleError.bind(this, args, currentCallState)
+
+      try {
+        const retVal = innerWrapped.apply(this, args)
+        return isPromise(retVal) ? retVal.catch(errorHandler) : retVal
+      } catch (e) {
+        return errorHandler(e)
+      }
+    }
+  } else {
+    // In non-safe mode, we just wrap the original function directly.
+    wrapped = wrapper(original)
+  }
   const descriptor = Object.getOwnPropertyDescriptor(target, name)
 
   const attributes = {
@@ -48,7 +180,7 @@ function wrapMethod (target, name, wrapper) {
     ...descriptor
   }
 
-  copyProperties(original, wrapped)
+  if (typeof original === 'function') copyProperties(original, wrapped)
 
   if (descriptor) {
     unwrappers.set(wrapped, () => Object.defineProperty(target, name, descriptor))
@@ -156,7 +288,9 @@ function assertNotClass (target) {
 
 module.exports = {
   wrap,
+  wrapFunction,
   massWrap,
   unwrap,
-  massUnwrap
+  massUnwrap,
+  setSafe
 }