dd-trace 5.66.0 → 5.68.0

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

@@ -73,19 +73,31 @@ function wrapPrettyFactory (prettyFactory) {
   }
 }
 
-addHook({ name: 'pino', versions: ['2 - 3', '4', '>=5 <5.14.0'] }, pino => {
+addHook({ name: 'pino', versions: ['2 - 3', '4'], patchDefault: true }, (pino) => {
   const asJsonSym = (pino.symbols && pino.symbols.asJsonSym) || 'asJson'
 
-  return shimmer.wrapFunction(pino, pino => wrapPino(asJsonSym, wrapAsJson, pino))
+  const wrapped = shimmer.wrapFunction(pino, pino => wrapPino(asJsonSym, wrapAsJson, pino))
+
+  return wrapped
 })
 
-addHook({ name: 'pino', versions: ['>=5.14.0 <6.8.0'] }, pino => {
-  const mixinSym = pino.symbols.mixinSym
+addHook({ name: 'pino', versions: ['>=5 <5.14.0'], patchDefault: true }, (pino) => {
+  const asJsonSym = ((pino.default || pino)?.symbols.asJsonSym) || 'asJson'
+
+  const wrapped = shimmer.wrapFunction(pino, pino => wrapPino(asJsonSym, wrapAsJson, pino))
+
+  return wrapped
+})
+
+addHook({ name: 'pino', versions: ['>=5.14.0 <6.8.0'] }, (pino) => {
+  const mixinSym = (pino.default || pino).symbols.mixinSym
 
-  return shimmer.wrapFunction(pino, pino => wrapPino(mixinSym, wrapMixin, pino))
+  const wrapped = shimmer.wrapFunction(pino, pino => wrapPino(mixinSym, wrapMixin, pino.default || pino))
+
+  return wrapped
 })
 
-addHook({ name: 'pino', versions: ['>=6.8.0'] }, pino => {
+addHook({ name: 'pino', versions: ['>=6.8.0'], patchDefault: false }, (pino, _1, _2, isIitm) => {
   const mixinSym = pino.symbols.mixinSym
 
   const wrapped = shimmer.wrapFunction(pino, pino => wrapPino(mixinSym, wrapMixin, pino))

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

@@ -51,6 +51,7 @@ let remainingTestsByFile = {}
 let isKnownTestsEnabled = false
 let isEarlyFlakeDetectionEnabled = false
 let earlyFlakeDetectionNumRetries = 0
+let isEarlyFlakeDetectionFaulty = false
 let isFlakyTestRetriesEnabled = false
 let flakyTestRetriesCount = 0
 let knownTests = {}
@@ -64,6 +65,10 @@ let quarantinedButNotAttemptToFixFqns = new Set()
 let rootDir = ''
 const MINIMUM_SUPPORTED_VERSION_RANGE_EFD = '>=1.38.0' // TODO: remove this once we drop support for v5
 
+function isValidKnownTests (receivedKnownTests) {
+  return !!receivedKnownTests.playwright
+}
+
 function getTestFullyQualifiedName (test) {
   const fullname = getTestFullname(test)
   return `${test._requireFile} ${fullname}`
@@ -79,8 +84,11 @@ function getTestProperties (test) {
 }
 
 function isNewTest (test) {
+  if (!isValidKnownTests(knownTests)) {
+    return false
+  }
   const testSuite = getTestSuitePath(test._requireFile, rootDir)
-  const testsForSuite = knownTests?.playwright?.[testSuite] || []
+  const testsForSuite = knownTests.playwright[testSuite] || []
 
   return !testsForSuite.includes(getTestFullname(test))
 }
@@ -557,6 +565,11 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
         } else {
           knownTests = receivedKnownTests
         }
+        if (!isValidKnownTests(receivedKnownTests)) {
+          isEarlyFlakeDetectionFaulty = true
+          isEarlyFlakeDetectionEnabled = false
+          isKnownTestsEnabled = false
+        }
       } catch (err) {
         isEarlyFlakeDetectionEnabled = false
         isKnownTestsEnabled = false
@@ -652,6 +665,7 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
     testSessionFinishCh.publish({
       status: STATUS_TO_TEST_STATUS[sessionStatus],
       isEarlyFlakeDetectionEnabled,
+      isEarlyFlakeDetectionFaulty,
       isTestManagementTestsEnabled,
       onDone
     })

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

@@ -7,7 +7,7 @@ const {
 
 const shimmer = require('../../datadog-shimmer')
 
-addHook({ name: 'sequelize', versions: ['>=4'] }, Sequelize => {
+addHook({ name: 'sequelize', versions: ['>=4'], file: ['lib/sequelize.js'] }, Sequelize => {
   const startCh = channel('datadog:sequelize:query:start')
   const finishCh = channel('datadog:sequelize:query:finish')
 

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

@@ -3,6 +3,10 @@
 const { addHook, channel } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 const log = require('../../dd-trace/src/log')
+const {
+  VITEST_WORKER_TRACE_PAYLOAD_CODE,
+  VITEST_WORKER_LOGS_PAYLOAD_CODE
+} = require('../../dd-trace/src/plugins/util/test')
 
 // test hooks
 const testStartCh = channel('ci:vitest:test:start')
@@ -30,6 +34,9 @@ const isEarlyFlakeDetectionFaultyCh = channel('ci:vitest:is-early-flake-detectio
 const testManagementTestsCh = channel('ci:vitest:test-management-tests')
 const impactedTestsCh = channel('ci:vitest:modified-tests')
 
+const workerReportTraceCh = channel('ci:vitest:worker-report:trace')
+const workerReportLogsCh = channel('ci:vitest:worker-report:logs')
+
 const taskToCtx = new WeakMap()
 const taskToStatuses = new WeakMap()
 const newTasks = new WeakSet()
@@ -40,9 +47,26 @@ const modifiedTasks = new WeakSet()
 let isRetryReasonEfd = false
 let isRetryReasonAttemptToFix = false
 const switchedStatuses = new WeakSet()
+const workerProcesses = new WeakSet()
+let isFlakyTestRetriesEnabled = false
+let flakyTestRetriesCount = 0
+let isEarlyFlakeDetectionEnabled = false
+let earlyFlakeDetectionNumRetries = 0
+let isEarlyFlakeDetectionFaulty = false
+let isKnownTestsEnabled = false
+let isTestManagementTestsEnabled = false
+let isImpactedTestsEnabled = false
+let testManagementAttemptToFixRetries = 0
+let isDiEnabled = false
+let testCodeCoverageLinesTotal
+let isSessionStarted = false
 
 const BREAKPOINT_HIT_GRACE_PERIOD_MS = 400
 
+function getTestCommand () {
+  return `vitest ${process.argv.slice(2).join(' ')}`
+}
+
 function waitForHitProbe () {
   return new Promise(resolve => {
     setTimeout(() => {
@@ -51,6 +75,10 @@ function waitForHitProbe () {
   })
 }
 
+function isValidKnownTests (receivedKnownTests) {
+  return !!receivedKnownTests.vitest
+}
+
 function getProvidedContext () {
   try {
     const {
@@ -121,6 +149,10 @@ function getChannelPromise (channelToPublishTo, frameworkVersion) {
   })
 }
 
+function isCliApiPackage (vitestPackage) {
+  return vitestPackage.s?.name === 'startVitest'
+}
+
 function getSessionStatus (state) {
   if (state.getCountOfFailedTests() > 0) {
     return 'fail'
@@ -183,16 +215,6 @@ function getSortWrapper (sort, frameworkVersion) {
     // There isn't any other async function that we seem to be able to hook into
     // So we will use the sort from BaseSequencer. This means that a custom sequencer
     // will not work. This will be a known limitation.
-    let isFlakyTestRetriesEnabled = false
-    let flakyTestRetriesCount = 0
-    let isEarlyFlakeDetectionEnabled = false
-    let earlyFlakeDetectionNumRetries = 0
-    let isEarlyFlakeDetectionFaulty = false
-    let isKnownTestsEnabled = false
-    let isTestManagementTestsEnabled = false
-    let isImpactedTestsEnabled = false
-    let testManagementAttemptToFixRetries = 0
-    let isDiEnabled = false
 
     try {
       const { err, libraryConfig } = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
@@ -235,28 +257,33 @@ function getSortWrapper (sort, frameworkVersion) {
 
         const testFilepaths = await getFilePaths.call(this.ctx)
 
-        isEarlyFlakeDetectionFaultyCh.publish({
-          knownTests: knownTests.vitest || {},
-          testFilepaths,
-          onDone: (isFaulty) => {
-            isEarlyFlakeDetectionFaulty = isFaulty
+        if (isValidKnownTests(knownTests)) {
+          isEarlyFlakeDetectionFaultyCh.publish({
+            knownTests: knownTests.vitest,
+            testFilepaths,
+            onDone: (isFaulty) => {
+              isEarlyFlakeDetectionFaulty = isFaulty
+            }
+          })
+          if (isEarlyFlakeDetectionFaulty) {
+            isEarlyFlakeDetectionEnabled = false
+            log.warn('New test detection is disabled because the number of new tests is too high.')
+          } else {
+            // TODO: use this to pass session and module IDs to the worker, instead of polluting process.env
+            // Note: setting this.ctx.config.provide directly does not work because it's cached
+            try {
+              const workspaceProject = this.ctx.getCoreWorkspaceProject()
+              workspaceProject._provided._ddIsKnownTestsEnabled = isKnownTestsEnabled
+              workspaceProject._provided._ddKnownTests = knownTests
+              workspaceProject._provided._ddIsEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
+              workspaceProject._provided._ddEarlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
+            } catch {
+              log.warn('Could not send known tests to workers so Early Flake Detection will not work.')
+            }
           }
-        })
-        if (isEarlyFlakeDetectionFaulty) {
-          isEarlyFlakeDetectionEnabled = false
-          log.warn('New test detection is disabled because the number of new tests is too high.')
         } else {
-          // TODO: use this to pass session and module IDs to the worker, instead of polluting process.env
-          // Note: setting this.ctx.config.provide directly does not work because it's cached
-          try {
-            const workspaceProject = this.ctx.getCoreWorkspaceProject()
-            workspaceProject._provided._ddIsKnownTestsEnabled = isKnownTestsEnabled
-            workspaceProject._provided._ddKnownTests = knownTests.vitest || {}
-            workspaceProject._provided._ddIsEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
-            workspaceProject._provided._ddEarlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
-          } catch {
-            log.warn('Could not send known tests to workers so Early Flake Detection will not work.')
-          }
+          isEarlyFlakeDetectionFaulty = true
+          isEarlyFlakeDetectionEnabled = false
         }
       }
     }
@@ -303,8 +330,6 @@ function getSortWrapper (sort, frameworkVersion) {
       }
     }
 
-    let testCodeCoverageLinesTotal
-
     if (this.ctx.coverageProvider?.generateCoverage) {
       shimmer.wrap(this.ctx.coverageProvider, 'generateCoverage', generateCoverage => async function () {
         const totalCodeCoverage = await generateCoverage.apply(this, arguments)
@@ -318,48 +343,104 @@ function getSortWrapper (sort, frameworkVersion) {
       })
     }
 
-    shimmer.wrap(this.ctx, 'exit', exit => async function () {
-      let onFinish
+    shimmer.wrap(this.ctx, 'exit', getFinishWrapper)
+    shimmer.wrap(this.ctx, 'close', getFinishWrapper)
 
-      const flushPromise = new Promise(resolve => {
-        onFinish = resolve
-      })
-      const failedSuites = this.state.getFailedFilepaths()
-      let error
-      if (failedSuites.length) {
-        error = new Error(`Test suites failed: ${failedSuites.length}.`)
-      }
+    return sort.apply(this, arguments)
+  }
+}
 
-      testSessionFinishCh.publish({
-        status: getSessionStatus(this.state),
-        testCodeCoverageLinesTotal,
-        error,
-        isEarlyFlakeDetectionEnabled,
-        isEarlyFlakeDetectionFaulty,
-        isTestManagementTestsEnabled,
-        onFinish
-      })
+function getFinishWrapper (exitOrClose) {
+  let isClosed = false
+  return async function () {
+    if (isClosed) { // needed because exit calls close
+      return exitOrClose.apply(this, arguments)
+    }
+    isClosed = true
+    let onFinish
 
-      await flushPromise
+    const flushPromise = new Promise(resolve => {
+      onFinish = resolve
+    })
+    const failedSuites = this.state.getFailedFilepaths()
+    let error
+    if (failedSuites.length) {
+      error = new Error(`Test suites failed: ${failedSuites.length}.`)
+    }
 
-      return exit.apply(this, arguments)
+    testSessionFinishCh.publish({
+      status: getSessionStatus(this.state),
+      testCodeCoverageLinesTotal,
+      error,
+      isEarlyFlakeDetectionEnabled,
+      isEarlyFlakeDetectionFaulty,
+      isTestManagementTestsEnabled,
+      onFinish
     })
 
-    return sort.apply(this, arguments)
+    await flushPromise
+
+    return exitOrClose.apply(this, arguments)
+  }
+}
+
+function getCliOrStartVitestWrapper (frameworkVersion) {
+  return function (oldCliOrStartVitest) {
+    return function () {
+      if (!testSessionStartCh.hasSubscribers || isSessionStarted) {
+        return oldCliOrStartVitest.apply(this, arguments)
+      }
+      isSessionStarted = true
+      testSessionStartCh.publish({ command: getTestCommand(), frameworkVersion })
+      return oldCliOrStartVitest.apply(this, arguments)
+    }
   }
 }
 
 function getCreateCliWrapper (vitestPackage, frameworkVersion) {
-  shimmer.wrap(vitestPackage, 'c', oldCreateCli => function () {
-    if (!testSessionStartCh.hasSubscribers) {
-      return oldCreateCli.apply(this, arguments)
+  shimmer.wrap(vitestPackage, 'c', getCliOrStartVitestWrapper(frameworkVersion))
+
+  return vitestPackage
+}
+
+function threadHandler (thread) {
+  if (workerProcesses.has(thread.process)) {
+    return
+  }
+  workerProcesses.add(thread.process)
+  thread.process.on('message', (message) => {
+    if (message.__tinypool_worker_message__ && message.data) {
+      if (message.interprocessCode === VITEST_WORKER_TRACE_PAYLOAD_CODE) {
+        workerReportTraceCh.publish(message.data)
+      } else if (message.interprocessCode === VITEST_WORKER_LOGS_PAYLOAD_CODE) {
+        workerReportLogsCh.publish(message.data)
+      }
     }
-    const processArgv = process.argv.slice(2).join(' ')
-    testSessionStartCh.publish({ command: `vitest ${processArgv}`, frameworkVersion })
-    return oldCreateCli.apply(this, arguments)
   })
+}
 
-  return vitestPackage
+addHook({
+  name: 'tinypool',
+  versions: ['>=1.0.0'],
+  file: 'dist/index.js'
+}, (TinyPool) => {
+  shimmer.wrap(TinyPool.prototype, 'run', run => async function () {
+    // We have to do this before and after because the threads list gets recycled, that is, the processes are re-created
+    this.threads.forEach(threadHandler)
+    const runResult = await run.apply(this, arguments)
+    this.threads.forEach(threadHandler)
+    return runResult
+  })
+
+  return TinyPool
+})
+
+function getStartVitestWrapper (cliApiPackage, frameworkVersion) {
+  if (!isCliApiPackage(cliApiPackage)) {
+    return cliApiPackage
+  }
+  shimmer.wrap(cliApiPackage, 's', getCliOrStartVitestWrapper(frameworkVersion))
+  return cliApiPackage
 }
 
 addHook({
@@ -735,6 +816,7 @@ addHook({
 })
 
 // Can't specify file because compiled vitest includes hashes in their files
+// Following 3 wrappers are for test session start
 addHook({
   name: 'vitest',
   versions: ['>=1.6.0 <2.0.5'],
@@ -747,6 +829,18 @@ addHook({
   filePattern: 'dist/chunks/cac.*'
 }, getCreateCliWrapper)
 
+addHook({
+  name: 'vitest',
+  versions: ['>=1.6.0 <2.0.5'],
+  filePattern: 'dist/vendor/cli-api.*'
+}, getStartVitestWrapper)
+
+addHook({
+  name: 'vitest',
+  versions: ['>=2.0.5'],
+  filePattern: 'dist/chunks/cli-api.*'
+}, getStartVitestWrapper)
+
 // test suite start and finish
 // only relevant for workers
 addHook({
@@ -857,7 +951,6 @@ addHook({
 
     testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish, ...testSuiteCtx.currentStore })
 
-    // TODO: fix too frequent flushes
     await onFinishPromise
 
     return startTestsResponse

package/packages/datadog-plugin-ai/src/tracing.js

@@ -8,14 +8,14 @@ class VercelAITracingPlugin extends TracingPlugin {
   static prefix = 'tracing:dd-trace:vercel-ai'
 
   bindStart (ctx) {
-    const attributes = ctx.attributes
+    const { attributes, name } = ctx
 
     const model = attributes['ai.model.id']
     const modelProvider = getModelProvider(attributes)
 
-    this.startSpan(ctx.name, {
+    this.startSpan(name, {
       meta: {
-        'resource.name': ctx.name,
+        'resource.name': attributes['resource.name'] ?? name,
         'ai.request.model': model,
         'ai.request.model_provider': modelProvider
       }

package/packages/datadog-plugin-aws-sdk/src/base.js

@@ -4,7 +4,6 @@ const analyticsSampler = require('../../dd-trace/src/analytics_sampler')
 const ClientPlugin = require('../../dd-trace/src/plugins/client')
 const { storage } = require('../../datadog-core')
 const { isTrue } = require('../../dd-trace/src/util')
-const coalesce = require('koalas')
 const { tagsFromRequest, tagsFromResponse } = require('../../dd-trace/src/payload-tagging')
 const { getEnvironmentVariable } = require('../../dd-trace/src/config-helper')
 
@@ -59,6 +58,7 @@ class BaseAwsSdkPlugin extends ClientPlugin {
         'aws.operation': operation,
         'aws.region': awsRegion,
         region: awsRegion,
+        'aws.partition': getPartition(awsRegion),
         aws_service: awsService,
         'aws.service': awsService,
         component: 'aws-sdk'
@@ -115,6 +115,11 @@ class BaseAwsSdkPlugin extends ClientPlugin {
       span.setTag('aws.region', region)
       span.setTag('region', region)
 
+      const partition = getPartition(region)
+      if (partition) {
+        span.setTag('aws.partition', partition)
+      }
+
       if (!this._tracerConfig?._isInServerlessEnvironment()) return
 
       const hostname = getHostname(store, region)
@@ -267,13 +272,10 @@ function normalizeConfig (config, serviceIdentifier) {
   // check if AWS batch propagation or AWS_[SERVICE] batch propagation is enabled via env variable
   const serviceId = serviceIdentifier.toUpperCase()
   const batchPropagationEnabled = isTrue(
-    coalesce(
-      specificConfig.batchPropagationEnabled,
-      getEnvironmentVariable(`DD_TRACE_AWS_SDK_${serviceId}_BATCH_PROPAGATION_ENABLED`),
-      config.batchPropagationEnabled,
-      getEnvironmentVariable('DD_TRACE_AWS_SDK_BATCH_PROPAGATION_ENABLED'),
-      false
-    )
+    specificConfig.batchPropagationEnabled ??
+    getEnvironmentVariable(`DD_TRACE_AWS_SDK_${serviceId}_BATCH_PROPAGATION_ENABLED`) ??
+    config.batchPropagationEnabled ??
+    getEnvironmentVariable('DD_TRACE_AWS_SDK_BATCH_PROPAGATION_ENABLED')
   )
 
   // Merge the specific config back into the main config
@@ -317,4 +319,17 @@ function getHostname (store, region) {
   }
 }
 
+function getPartition (region) {
+  if (!region) return
+
+  let partition = 'aws'
+  if (region.startsWith('cn-')) {
+    partition = 'aws-cn'
+  } else if (region.startsWith('us-gov-')) {
+    partition = 'aws-us-gov'
+  }
+
+  return partition
+}
+
 module.exports = BaseAwsSdkPlugin

package/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js

@@ -3,7 +3,7 @@
 const BaseAwsSdkPlugin = require('../../base')
 const { parseModelId } = require('./utils')
 
-const enabledOperations = new Set(['invokeModel'])
+const enabledOperations = new Set(['invokeModel', 'invokeModelWithResponseStream'])
 
 class BedrockRuntime extends BaseAwsSdkPlugin {
   static id = 'bedrockruntime'
@@ -17,7 +17,7 @@ class BedrockRuntime extends BaseAwsSdkPlugin {
     return super.isEnabled(request)
   }
 
-  generateTags (params, operation, response) {
+  generateTags (params, operation) {
     const { modelProvider, modelName } = parseModelId(params.modelId)
 
     return {

package/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/utils.js

@@ -23,6 +23,100 @@ const PROVIDER = {
   MISTRAL: 'MISTRAL'
 }
 
+/**
+ * Coerce the chunks into a single response body.
+ *
+ * @param {Array<{ chunk: { bytes: Buffer } }>} chunks
+ * @param {string} provider
+ * @returns {Object}
+ */
+function extractTextAndResponseReasonFromStream (chunks, modelProvider, modelName) {
+  const modelProviderUpper = modelProvider.toUpperCase()
+
+  // streaming unsupported for AMAZON embedding models, COHERE embedding models, STABILITY
+  if (
+    (modelProviderUpper === PROVIDER.AMAZON && modelName.includes('embed')) ||
+    (modelProviderUpper === PROVIDER.COHERE && modelName.includes('embed')) ||
+    modelProviderUpper === PROVIDER.STABILITY
+  ) {
+    return {}
+  }
+
+  let message = ''
+  let inputTokens = 0
+  let outputTokens = 0
+  let cacheReadTokens = 0
+  let cacheWriteTokens = 0
+
+  for (const { chunk: { bytes } } of chunks) {
+    const body = JSON.parse(Buffer.from(bytes).toString('utf8'))
+
+    switch (modelProviderUpper) {
+      case PROVIDER.AMAZON: {
+        message += body?.outputText
+
+        inputTokens = body?.inputTextTokenCount
+        outputTokens = body?.totalOutputTextTokenCount
+
+        break
+      }
+      case PROVIDER.AI21: {
+        const content = body?.choices?.[0]?.delta?.content
+        if (content) {
+          message += content
+        }
+
+        break
+      }
+      case PROVIDER.ANTHROPIC: {
+        if (body.completion) {
+          message += body.completion
+        } else if (body.delta?.text) {
+          message += body.delta.text
+        }
+
+        if (body.message?.usage?.input_tokens) inputTokens = body.message.usage.input_tokens
+        if (body.message?.usage?.output_tokens) outputTokens = body.message.usage.output_tokens
+
+        break
+      }
+      case PROVIDER.COHERE: {
+        if (body?.event_type === 'stream-end') {
+          message = body.response?.text
+        }
+
+        break
+      }
+      case PROVIDER.META: {
+        message += body?.generation
+        break
+      }
+      case PROVIDER.MISTRAL: {
+        message += body?.outputs?.[0]?.text
+        break
+      }
+    }
+
+    // by default, it seems newer versions of the AWS SDK include the input/output token counts in the response body
+    const invocationMetrics = body['amazon-bedrock-invocationMetrics']
+    if (invocationMetrics) {
+      inputTokens = invocationMetrics.inputTokenCount
+      outputTokens = invocationMetrics.outputTokenCount
+      cacheReadTokens = invocationMetrics.cacheReadInputTokenCount
+      cacheWriteTokens = invocationMetrics.cacheWriteInputTokenCount
+    }
+  }
+
+  return new Generation({
+    message,
+    role: 'assistant',
+    inputTokens,
+    outputTokens,
+    cacheReadTokens,
+    cacheWriteTokens
+  })
+}
+
 class Generation {
   constructor ({
     message = '',
@@ -30,7 +124,9 @@ class Generation {
     choiceId = '',
     role,
     inputTokens,
-    outputTokens
+    outputTokens,
+    cacheReadTokens,
+    cacheWriteTokens
   } = {}) {
     // stringify message as it could be a single generated message as well as a list of embeddings
     this.message = typeof message === 'string' ? message : JSON.stringify(message) || ''
@@ -39,7 +135,9 @@ class Generation {
     this.role = role
     this.usage = {
       inputTokens,
-      outputTokens
+      outputTokens,
+      cacheReadTokens,
+      cacheWriteTokens
     }
   }
 }
@@ -334,6 +432,7 @@ function extractTextAndResponseReason (response, provider, modelName) {
 module.exports = {
   Generation,
   RequestParams,
+  extractTextAndResponseReasonFromStream,
   parseModelId,
   extractRequestParams,
   extractTextAndResponseReason,

package/packages/datadog-plugin-aws-sdk/src/util.js

@@ -135,7 +135,7 @@ const extractQueueMetadata = queueURL => {
   let partition = 'aws'
   if (region.startsWith('cn-')) {
     partition = 'aws-cn'
-  } else if (region.startsWith('us-gov')) {
+  } else if (region.startsWith('us-gov-')) {
     partition = 'aws-us-gov'
   }