dd-trace 4.38.1 → 4.40.0

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

@@ -15,13 +15,15 @@ const V4_PACKAGE_SHIMS = [
     file: 'resources/chat/completions.js',
     targetClass: 'Completions',
     baseResource: 'chat.completions',
-    methods: ['create']
+    methods: ['create'],
+    streamedResponse: true
   },
   {
     file: 'resources/completions.js',
     targetClass: 'Completions',
     baseResource: 'completions',
-    methods: ['create']
+    methods: ['create'],
+    streamedResponse: true
   },
   {
     file: 'resources/embeddings.js',
@@ -121,7 +123,7 @@ addHook({ name: 'openai', file: 'dist/api.js', versions: ['>=3.0.0 <4'] }, expor
 
       return fn.apply(this, arguments)
         .then((response) => {
-          finishCh.publish({
+          finish({
             headers: response.headers,
             body: response.data,
             path: response.request.path,
@@ -130,10 +132,10 @@ addHook({ name: 'openai', file: 'dist/api.js', versions: ['>=3.0.0 <4'] }, expor
 
           return response
         })
-        .catch((err) => {
-          errorCh.publish({ err })
+        .catch(error => {
+          finish(undefined, error)
 
-          throw err
+          throw error
         })
     })
   }
@@ -141,9 +143,140 @@ addHook({ name: 'openai', file: 'dist/api.js', versions: ['>=3.0.0 <4'] }, expor
   return exports
 })
 
+function addStreamedChunk (content, chunk) {
+  content.usage = chunk.usage // add usage if it was specified to be returned
+  for (const choice of chunk.choices) {
+    const choiceIdx = choice.index
+    const oldChoice = content.choices.find(choice => choice?.index === choiceIdx)
+    if (!oldChoice) {
+      // we don't know which choices arrive in which order
+      content.choices[choiceIdx] = choice
+    } else {
+      if (!oldChoice.finish_reason) {
+        oldChoice.finish_reason = choice.finish_reason
+      }
+
+      // delta exists on chat completions
+      const delta = choice.delta
+
+      if (delta) {
+        const content = delta.content
+        if (content) {
+          if (oldChoice.delta.content) { // we don't want to append to undefined
+            oldChoice.delta.content += content
+          } else {
+            oldChoice.delta.content = content
+          }
+        }
+      } else {
+        const text = choice.text
+        if (text) {
+          if (oldChoice.text) {
+            oldChoice.text += text
+          } else {
+            oldChoice.text = text
+          }
+        }
+      }
+
+      // tools only exist on chat completions
+      const tools = delta && choice.delta.tool_calls
+
+      if (tools) {
+        oldChoice.delta.tool_calls = tools.map((newTool, toolIdx) => {
+          const oldTool = oldChoice.delta.tool_calls[toolIdx]
+
+          if (oldTool) {
+            oldTool.function.arguments += newTool.function.arguments
+          }
+
+          return oldTool
+        })
+      }
+    }
+  }
+}
+
+function convertBufferstoObjects (chunks = []) {
+  return Buffer
+    .concat(chunks) // combine the buffers
+    .toString() // stringify
+    .split(/(?=data:)/) // split on "data:"
+    .map(chunk => chunk.split('\n').join('')) // remove newlines
+    .map(chunk => chunk.substring(6)) // remove 'data: ' from the front
+    .slice(0, -1) // remove the last [DONE] message
+    .map(JSON.parse) // parse all of the returned objects
+}
+
+/**
+ * For streamed responses, we need to accumulate all of the content in
+ * the chunks, and let the combined content be the final response.
+ * This way, spans look the same as when not streamed.
+ */
+function wrapStreamIterator (response, options, n) {
+  let processChunksAsBuffers = false
+  let chunks = []
+  return function (itr) {
+    return function () {
+      const iterator = itr.apply(this, arguments)
+      shimmer.wrap(iterator, 'next', next => function () {
+        return next.apply(this, arguments)
+          .then(res => {
+            const { done, value: chunk } = res
+
+            if (chunk) {
+              chunks.push(chunk)
+              if (chunk instanceof Buffer) {
+                // this operation should be safe
+                // if one chunk is a buffer (versus a plain object), the rest should be as well
+                processChunksAsBuffers = true
+              }
+            }
+
+            if (done) {
+              let body = {}
+              chunks = chunks.filter(chunk => chunk != null) // filter null or undefined values
+
+              if (chunks) {
+                if (processChunksAsBuffers) {
+                  chunks = convertBufferstoObjects(chunks)
+                }
+
+                if (chunks.length) {
+                  // define the initial body having all the content outside of choices from the first chunk
+                  // this will include import data like created, id, model, etc.
+                  body = { ...chunks[0], choices: Array.from({ length: n }) }
+                  // start from the first chunk, and add its choices into the body
+                  for (let i = 0; i < chunks.length; i++) {
+                    addStreamedChunk(body, chunks[i])
+                  }
+                }
+              }
+
+              finish({
+                headers: response.headers,
+                body,
+                path: response.url,
+                method: options.method
+              })
+            }
+
+            return res
+          })
+          .catch(err => {
+            finish(undefined, err)
+
+            throw err
+          })
+      })
+      return iterator
+    }
+  }
+}
+
 for (const shim of V4_PACKAGE_SHIMS) {
-  const { file, targetClass, baseResource, methods } = shim
-  addHook({ name: 'openai', file, versions: shim.versions || ['>=4'] }, exports => {
+  const { file, targetClass, baseResource, methods, versions, streamedResponse } = shim
+  addHook({ name: 'openai', file, versions: versions || ['>=4'] }, exports => {
     const targetPrototype = exports[targetClass].prototype
 
     for (const methodName of methods) {
@@ -152,6 +285,22 @@ for (const shim of V4_PACKAGE_SHIMS) {
           return methodFn.apply(this, arguments)
         }
 
+        // The OpenAI library lets you set `stream: true` on the options arg to any method
+        // However, we only want to handle streamed responses in specific cases
+        // chat.completions and completions
+        const stream = streamedResponse && getOption(arguments, 'stream', false)
+
+        // we need to compute how many prompts we are sending in streamed cases for completions
+        // not applicable for chat completiond
+        let n
+        if (stream) {
+          n = getOption(arguments, 'n', 1)
+          const prompt = getOption(arguments, 'prompt')
+          if (Array.isArray(prompt) && typeof prompt[0] !== 'number') {
+            n *= prompt.length
+          }
+        }
+
         const client = this._client || this.client
 
         startCh.publish({
@@ -170,19 +319,29 @@ for (const shim of V4_PACKAGE_SHIMS) {
             // the original response is wrapped in a promise, so we need to unwrap it
             .then(body => Promise.all([this.responsePromise, body]))
             .then(([{ response, options }, body]) => {
-              finishCh.publish({
-                headers: response.headers,
-                body,
-                path: response.url,
-                method: options.method
-              })
+              if (stream) {
+                if (body.iterator) {
+                  shimmer.wrap(body, 'iterator', wrapStreamIterator(response, options, n))
+                } else {
+                  shimmer.wrap(
+                    body.response.body, Symbol.asyncIterator, wrapStreamIterator(response, options, n)
+                  )
+                }
+              } else {
+                finish({
+                  headers: response.headers,
+                  body,
+                  path: response.url,
+                  method: options.method
+                })
+              }
 
               return body
             })
-            .catch(err => {
-              errorCh.publish({ err })
+            .catch(error => {
+              finish(undefined, error)
 
-              throw err
+              throw error
             })
             .finally(() => {
               // maybe we don't want to unwrap here in case the promise is re-used?
@@ -197,3 +356,15 @@ for (const shim of V4_PACKAGE_SHIMS) {
     return exports
   })
 }
+
+function finish (response, error) {
+  if (error) {
+    errorCh.publish({ error })
+  }
+
+  finishCh.publish(response)
+}
+
+function getOption (args, option, defaultValue) {
+  return args[args.length - 1]?.[option] || defaultValue
+}

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

@@ -301,7 +301,7 @@ function dispatcherRunWrapperNew (run) {
     if (!this._allTests) {
       // Removed in https://github.com/microsoft/playwright/commit/1e52c37b254a441cccf332520f60225a5acc14c7
       // Not available from >=1.44.0
-      this._allTests = testGroups.map(g => g.tests).flat()
+      this._ddAllTests = testGroups.map(g => g.tests).flat()
     }
     remainingTestsByFile = getTestsBySuiteFromTestGroups(arguments[0])
     return run.apply(this, arguments)
@@ -339,8 +339,9 @@ function getTestByTestId (dispatcher, testId) {
   if (dispatcher._testById) {
     return dispatcher._testById.get(testId)?.test
   }
-  if (dispatcher._allTests) {
-    return dispatcher._allTests.find(({ id }) => id === testId)
+  const allTests = dispatcher._allTests || dispatcher._ddAllTests
+  if (allTests) {
+    return allTests.find(({ id }) => id === testId)
   }
 }
 

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

@@ -1,6 +1,6 @@
 'use strict'
 
-const METHODS = require('methods').concat('all')
+const METHODS = require('http').METHODS.map(v => v.toLowerCase()).concat('all')
 const pathToRegExp = require('path-to-regexp')
 const shimmer = require('../../datadog-shimmer')
 const { addHook, channel } = require('./helpers/instrument')

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

@@ -23,6 +23,9 @@ addHook({
 }, (seleniumPackage, seleniumVersion) => {
   // TODO: do not turn this into async. Use promises
   shimmer.wrap(seleniumPackage.WebDriver.prototype, 'get', get => async function () {
+    if (!ciSeleniumDriverGetStartCh.hasSubscribers) {
+      return get.apply(this, arguments)
+    }
     let traceId
     const setTraceId = (inputTraceId) => {
       traceId = inputTraceId
@@ -40,15 +43,20 @@ addHook({
       isRumActive
     })
 
-    await this.manage().addCookie({
-      name: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
-      value: traceId
-    })
+    if (traceId && isRumActive) {
+      await this.manage().addCookie({
+        name: DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME,
+        value: traceId
+      })
+    }
 
     return getResult
   })
 
   shimmer.wrap(seleniumPackage.WebDriver.prototype, 'quit', quit => async function () {
+    if (!ciSeleniumDriverGetStartCh.hasSubscribers) {
+      return quit.apply(this, arguments)
+    }
     const isRumActive = await this.executeScript(RUM_STOP_SESSION_SCRIPT)
 
     if (isRumActive) {
@@ -58,10 +66,9 @@ addHook({
           resolve()
         }, DD_CIVISIBILITY_RUM_FLUSH_WAIT_MILLIS)
       })
+      await this.manage().deleteCookie(DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME)
     }
 
-    await this.manage().deleteCookie(DD_CIVISIBILITY_TEST_EXECUTION_ID_COOKIE_NAME)
-
     return quit.apply(this, arguments)
   })
 

package/packages/datadog-plugin-graphql/src/resolve.js

@@ -80,6 +80,10 @@ class GraphQLResolvePlugin extends TracingPlugin {
     // this will disable resolve subscribers if `config.depth` is set to 0
     super.configure(config.depth === 0 ? false : config)
   }
+
+  finish (finishTime) {
+    this.activeSpan.finish(finishTime)
+  }
 }
 
 // helpers

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

@@ -20,7 +20,14 @@ const {
   removeEfdStringFromTestName,
   TEST_IS_NEW,
   TEST_IS_RETRY,
-  TEST_EARLY_FLAKE_ENABLED
+  TEST_EARLY_FLAKE_ENABLED,
+  TEST_SESSION_ID,
+  TEST_MODULE_ID,
+  TEST_MODULE,
+  TEST_SUITE_ID,
+  TEST_COMMAND,
+  TEST_SUITE,
+  MOCHA_IS_PARALLEL
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const {
@@ -33,6 +40,22 @@ const {
   TELEMETRY_ITR_UNSKIPPABLE,
   TELEMETRY_CODE_COVERAGE_NUM_FILES
 } = require('../../dd-trace/src/ci-visibility/telemetry')
+const id = require('../../dd-trace/src/id')
+const log = require('../../dd-trace/src/log')
+
+function getTestSuiteLevelVisibilityTags (testSuiteSpan) {
+  const testSuiteSpanContext = testSuiteSpan.context()
+  const suiteTags = {
+    [TEST_SUITE_ID]: testSuiteSpanContext.toSpanId(),
+    [TEST_SESSION_ID]: testSuiteSpanContext.toTraceId(),
+    [TEST_COMMAND]: testSuiteSpanContext._tags[TEST_COMMAND],
+    [TEST_MODULE]: 'mocha'
+  }
+  if (testSuiteSpanContext._parentId) {
+    suiteTags[TEST_MODULE_ID] = testSuiteSpanContext._parentId.toString(10)
+  }
+  return suiteTags
+}
 
 class MochaPlugin extends CiPlugin {
   static get id () {
@@ -50,7 +73,8 @@ class MochaPlugin extends CiPlugin {
       if (!this.libraryConfig?.isCodeCoverageEnabled) {
         return
       }
-      const testSuiteSpan = this._testSuites.get(suiteFile)
+      const testSuite = getTestSuitePath(suiteFile, this.sourceRoot)
+      const testSuiteSpan = this._testSuites.get(testSuite)
 
       if (!coverageFiles.length) {
         this.telemetry.count(TELEMETRY_CODE_COVERAGE_EMPTY)
@@ -73,16 +97,20 @@ class MochaPlugin extends CiPlugin {
     })
 
     this.addSub('ci:mocha:test-suite:start', ({
-      testSuite,
+      testSuiteAbsolutePath,
       isUnskippable,
       isForcedToRun,
       itrCorrelationId
     }) => {
-      const store = storage.getStore()
+      // If the test module span is undefined, the plugin has not been initialized correctly and we bail out
+      if (!this.testModuleSpan) {
+        return
+      }
+      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.sourceRoot)
       const testSuiteMetadata = getTestSuiteCommonTags(
         this.command,
         this.frameworkVersion,
-        getTestSuitePath(testSuite, this.sourceRoot),
+        testSuite,
         'mocha'
       )
       if (isUnskippable) {
@@ -109,6 +137,7 @@ class MochaPlugin extends CiPlugin {
       if (itrCorrelationId) {
         testSuiteSpan.setTag(ITR_CORRELATION_ID, itrCorrelationId)
       }
+      const store = storage.getStore()
       this.enter(testSuiteSpan, store)
       this._testSuites.set(testSuite, testSuiteSpan)
     })
@@ -142,6 +171,10 @@ class MochaPlugin extends CiPlugin {
       this.enter(span, store)
     })
 
+    this.addSub('ci:mocha:worker:finish', () => {
+      this.tracer._exporter.flush()
+    })
+
     this.addSub('ci:mocha:test:finish', (status) => {
       const store = storage.getStore()
       const span = store?.span
@@ -194,7 +227,8 @@ class MochaPlugin extends CiPlugin {
       hasForcedToRunSuites,
       hasUnskippableSuites,
       error,
-      isEarlyFlakeDetectionEnabled
+      isEarlyFlakeDetectionEnabled,
+      isParallel
     }) => {
       if (this.testSessionSpan) {
         const { isSuitesSkippingEnabled, isCodeCoverageEnabled } = this.libraryConfig || {}
@@ -206,6 +240,10 @@ class MochaPlugin extends CiPlugin {
           this.testModuleSpan.setTag('error', error)
         }
 
+        if (isParallel) {
+          this.testSessionSpan.setTag(MOCHA_IS_PARALLEL, 'true')
+        }
+
         addIntelligentTestRunnerSpanTags(
           this.testSessionSpan,
           this.testModuleSpan,
@@ -234,6 +272,37 @@ class MochaPlugin extends CiPlugin {
       this.libraryConfig = null
       this.tracer._exporter.flush()
     })
+
+    this.addSub('ci:mocha:worker-report:trace', (traces) => {
+      const formattedTraces = JSON.parse(traces).map(trace =>
+        trace.map(span => {
+          const formattedSpan = {
+            ...span,
+            span_id: id(span.span_id),
+            trace_id: id(span.trace_id),
+            parent_id: id(span.parent_id)
+          }
+          if (formattedSpan.name === 'mocha.test') {
+            const testSuite = span.meta[TEST_SUITE]
+            const testSuiteSpan = this._testSuites.get(testSuite)
+            if (!testSuiteSpan) {
+              log.warn(`Test suite span not found for test span with test suite ${testSuite}`)
+              return formattedSpan
+            }
+            const suiteTags = getTestSuiteLevelVisibilityTags(testSuiteSpan)
+            formattedSpan.meta = {
+              ...formattedSpan.meta,
+              ...suiteTags
+            }
+          }
+          return formattedSpan
+        })
+      )
+
+      formattedTraces.forEach(trace => {
+        this.tracer._exporter.export(trace)
+      })
+    })
   }
 
   startTestSpan (testInfo) {
@@ -242,7 +311,8 @@ class MochaPlugin extends CiPlugin {
       title,
       isNew,
       isEfdRetry,
-      testStartLine
+      testStartLine,
+      isParallel
     } = testInfo
 
     const testName = removeEfdStringFromTestName(testInfo.testName)
@@ -257,8 +327,12 @@ class MochaPlugin extends CiPlugin {
       extraTags[TEST_SOURCE_START] = testStartLine
     }
 
+    if (isParallel) {
+      extraTags[MOCHA_IS_PARALLEL] = 'true'
+    }
+
     const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.sourceRoot)
-    const testSuiteSpan = this._testSuites.get(testSuiteAbsolutePath)
+    const testSuiteSpan = this._testSuites.get(testSuite)
 
     if (this.repositoryRoot !== this.sourceRoot && !!this.repositoryRoot) {
       extraTags[TEST_SOURCE_FILE] = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)

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

@@ -6,7 +6,7 @@ const analyticsSampler = require('../../dd-trace/src/analytics_sampler')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const web = require('../../dd-trace/src/plugins/util/web')
 
-const errorPages = ['/404', '/500', '/_error', '/_not-found']
+const errorPages = ['/404', '/500', '/_error', '/_not-found', '/_not-found/page']
 
 class NextPlugin extends ServerPlugin {
   static get id () {
@@ -120,7 +120,6 @@ class NextPlugin extends ServerPlugin {
       'resource.name': `${req.method} ${page}`.trim(),
       'next.page': page
     })
-
     web.setRoute(req, page)
   }