dd-trace 5.101.0 → 5.103.0

package/packages/dd-trace/src/encode/0.5.js

@@ -1,17 +1,18 @@
 'use strict'
 
 const { normalizeSpan } = require('./tags-processors')
-const { AgentEncoder: BaseEncoder } = require('./0.4')
+const { AgentEncoder: BaseEncoder, stringifySpanEvents } = require('./0.4')
 
 const ARRAY_OF_TWO = 0x92
 const ARRAY_OF_TWELVE = 0x9C
 
 function formatSpan (span) {
   span = normalizeSpan(span)
-  // ensure span events are encoded as tags
+  // v0.5 has no native span_events slot; always serialize as a meta tag.
   if (span.span_events) {
-    span.meta.events = JSON.stringify(span.span_events)
-    delete span.span_events
+    span.meta.events = stringifySpanEvents(span.span_events)
+    // `= undefined` over `delete` to keep the span's hidden class.
+    span.span_events = undefined
   }
   return span
 }
@@ -45,22 +46,58 @@ class AgentEncoder extends BaseEncoder {
       this._encodeId(bytes, span.trace_id)
       this._encodeId(bytes, span.span_id)
       this._encodeId(bytes, span.parent_id)
-      this._encodeLong(bytes, span.start || 0)
-      this._encodeLong(bytes, span.duration || 0)
-      this._encodeInteger(bytes, span.error)
+      this._encodeIntOrFloat(bytes, span.start || 0)
+      this._encodeIntOrFloat(bytes, span.duration || 0)
+      this._encodeIntOrFloat(bytes, span.error)
       this._encodeMap(bytes, span.meta || {})
       this._encodeMap(bytes, span.metrics || {})
       this._encodeString(bytes, span.type)
     }
   }
 
+  // Override the inherited 0.4 `_encodeMap` so the v0.5 wire emits each numeric
+  // value via `_encodeIntOrFloat` (compact unsigned/signed int when integer,
+  // float64 otherwise) instead of always float64. The 0.4 base method stays on
+  // float64 because the CI-visibility encoders inherit it and target a
+  // different intake.
+  _encodeMap (bytes, value) {
+    const offset = bytes.length
+    bytes.reserve(5)
+    bytes.buffer[offset] = 0xDF
+
+    let count = 0
+    for (const key of Object.keys(value)) {
+      const entryValue = value[key]
+      if (typeof entryValue === 'string') {
+        this._encodeString(bytes, key)
+        this._encodeString(bytes, entryValue)
+        count++
+      } else if (typeof entryValue === 'number') {
+        this._encodeString(bytes, key)
+        this._encodeIntOrFloat(bytes, entryValue)
+        count++
+      }
+    }
+
+    const target = bytes.buffer
+    target[offset + 1] = count >>> 24
+    target[offset + 2] = count >>> 16
+    target[offset + 3] = count >>> 8
+    target[offset + 4] = count
+  }
+
   _encodeString (bytes, value = '') {
-    this._cacheString(value)
-    this._encodeInteger(bytes, this._stringMap[value])
+    let index = this._stringMap[value]
+    if (index === undefined) {
+      index = this._stringCount++
+      this._stringMap[value] = index
+      this._stringBytes.write(value)
+    }
+    this._encodeInteger(bytes, index)
   }
 
   _cacheString (value) {
-    if (!(value in this._stringMap)) {
+    if (this._stringMap[value] === undefined) {
       this._stringMap[value] = this._stringCount++
       this._stringBytes.write(value)
     }

package/packages/dd-trace/src/encode/agentless-json.js

@@ -85,7 +85,7 @@ function spanToJSON (span) {
 class AgentlessJSONEncoder {
   /**
    * @param {object} writer - Writer instance with a flush() method, called when the buffer exceeds the soft limit
-   * @param {object} [metadata={}] - Shared metadata spread into each trace object (hostname, env, tracerVersion, etc.)
+   * @param {object} [metadata] - Shared metadata spread into each trace object (hostname, env, tracerVersion, etc.)
    */
   constructor (writer, metadata = {}) {
     this._writer = writer

package/packages/dd-trace/src/exporter.js

@@ -7,6 +7,8 @@ const constants = require('./constants')
 
 module.exports = function getExporter (name) {
   switch (name) {
+    case exporters.ELECTRON:
+      return require('./exporters/electron')
     case exporters.LOG:
       return require('./exporters/log')
     case exporters.AGENT:

package/packages/dd-trace/src/exporters/agent/index.js

@@ -50,7 +50,8 @@ class AgentExporter {
       this.#timer = setTimeout(() => {
         this._writer.flush()
         this.#timer = undefined
-      }, flushInterval).unref()
+      }, flushInterval)
+      this.#timer.unref?.()
     }
   }
 

package/packages/dd-trace/src/exporters/agentless/index.js

@@ -18,7 +18,7 @@ class AgentlessExporter {
 
   /**
    * @param {object} config - Configuration object
-   * @param {string} [config.site='datadoghq.com'] - The Datadog site
+   * @param {string} [config.site] - The Datadog site. Defaults to 'datadoghq.com'.
    * @param {string} [config.url] - Override intake URL
    * @param {number} [config.flushInterval] - Batch flush interval in ms
    * @param {string} [config.env] - Environment name
@@ -109,7 +109,8 @@ class AgentlessExporter {
           log.error('Failed to flush traces on timer: %s', err.message)
         }
         this.#timer = undefined
-      }, flushInterval).unref()
+      }, flushInterval)
+      this.#timer.unref?.()
     }
   }
 

package/packages/dd-trace/src/exporters/agentless/writer.js

@@ -19,8 +19,8 @@ class AgentlessWriter extends BaseWriter {
   /**
    * @param {object} options - Writer options
    * @param {URL} [options.url] - The intake URL. If not provided, constructed from site.
-   * @param {string} [options.site='datadoghq.com'] - The Datadog site
-   * @param {object} [options.metadata={}] - Metadata to pass to the encoder (hostname, env, etc.)
+   * @param {string} [options.site] - The Datadog site
+   * @param {object} [options.metadata] - Metadata to pass to the encoder (hostname, env, etc.)
    */
   constructor ({ url, site = 'datadoghq.com', metadata = {} }) {
     super({ url })

package/packages/dd-trace/src/exporters/common/buffering-exporter.js

@@ -39,7 +39,8 @@ class BufferingExporter {
       this[timerKey] = setTimeout(() => {
         writer.flush()
         this[timerKey] = undefined
-      }, flushInterval).unref()
+      }, flushInterval)
+      this[timerKey].unref?.()
     }
   }
 

package/packages/dd-trace/src/exporters/common/request.js

@@ -180,7 +180,7 @@ function request (data, options, callback) {
           // Unref so a pending retry never keeps the host process alive past
           // its natural exit point; long-running apps still retry because the
           // event loop is held open by their own work.
-          setTimeout(attempt, getRetryDelay(options, attemptIndex), attemptIndex + 1).unref()
+          setTimeout(attempt, getRetryDelay(options, attemptIndex), attemptIndex + 1).unref?.()
         } else {
           callback(error)
         }

package/packages/dd-trace/src/exporters/electron/index.js

@@ -0,0 +1,49 @@
+'use strict'
+
+const { channel } = require('dc-polyfill')
+const { truncateSpan, normalizeSpan } = require('../../encode/tags-processors')
+
+const traceChannel = channel('datadog:apm:electron:export')
+
+class ElectronExporter {
+  #timer
+  #traces = []
+
+  constructor (config) {
+    this._config = config
+
+    globalThis[Symbol.for('dd-trace')].beforeExitHandlers.add(this.flush.bind(this))
+  }
+
+  export (spans) {
+    this.#traces.push(spans)
+
+    const { flushInterval } = this._config
+
+    if (flushInterval === 0) {
+      this.flush()
+    } else if (this.#timer === undefined) {
+      this.#timer = setTimeout(() => {
+        this.flush()
+        this.#timer = undefined
+      }, flushInterval)
+      this.#timer.unref?.()
+    }
+  }
+
+  flush (done = () => {}) {
+    clearTimeout(this.#timer)
+    this.#timer = undefined
+
+    const traces = this.#traces.splice(0)
+
+    if (traces.length > 0 && traceChannel.hasSubscribers) {
+      const formattedTraces = traces.map(spans => spans.map(span => normalizeSpan(truncateSpan(span))))
+      traceChannel.publish(formattedTraces)
+    }
+
+    done()
+  }
+}
+
+module.exports = ElectronExporter

package/packages/dd-trace/src/external-logger/src/index.js

@@ -26,7 +26,8 @@ class ExternalLogger {
     }
     this.timer = setInterval(() => {
       this.flush()
-    }, this.interval).unref()
+    }, this.interval)
+    this.timer.unref?.()
 
     tracerLogger.debug(`started log writer to https://${this.intake}${this.endpoint}`)
   }

package/packages/dd-trace/src/git_metadata.js

@@ -12,19 +12,21 @@ const {
   resolveGitHeadSHA,
 } = require('./config/git_properties')
 
-/** @type {{ commitSHA: string | undefined, repositoryUrl: string | undefined } | undefined} */
-let cached
+/**
+ * @typedef {{ commitSHA: string | undefined, repositoryUrl: string | undefined }} GitMetadata
+ * @type {{ enabled?: GitMetadata, disabled?: GitMetadata }}
+ */
+const cache = {}
 
 /**
  * @param {import('./config/config-types').ConfigProperties} config
  */
 function getGitMetadata (config) {
-  if (cached) return cached
-
   if (!config.DD_TRACE_GIT_METADATA_ENABLED) {
-    cached = { commitSHA: undefined, repositoryUrl: undefined }
-    return cached
+    cache.disabled ??= { commitSHA: undefined, repositoryUrl: undefined }
+    return cache.disabled
   }
+  if (cache.enabled) return cache.enabled
 
   let repositoryUrl = removeUserSensitiveInfo(config.DD_GIT_REPOSITORY_URL ?? config.tags[GIT_REPOSITORY_URL])
   let commitSHA = config.DD_GIT_COMMIT_SHA ?? config.tags[GIT_COMMIT_SHA]
@@ -59,8 +61,8 @@ function getGitMetadata (config) {
 
   commitSHA ??= resolveGitHeadSHA(gitFolderPath)
 
-  cached = { commitSHA, repositoryUrl }
-  return cached
+  cache.enabled = { commitSHA, repositoryUrl }
+  return cache.enabled
 }
 
 module.exports = getGitMetadata

package/packages/dd-trace/src/lambda/handler-paths.js

@@ -0,0 +1,52 @@
+/**
+ * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * Modifications copyright 2022 Datadog, Inc.
+ *
+ * Some functions are part of aws-lambda-nodejs-runtime-interface-client
+ * https://github.com/aws/aws-lambda-nodejs-runtime-interface-client/blob/v2.1.0/src/utils/UserFunction.ts
+ */
+'use strict'
+
+const path = require('path')
+
+/**
+ * Example: `'./api/src/index.nested.handler'` → `['./api/src/', 'index.nested.handler']`.
+ *
+ * @param {string} fullHandler
+ */
+function extractModuleRootAndHandler (fullHandler) {
+  const handlerString = path.basename(fullHandler)
+  const moduleRoot = fullHandler.slice(0, Math.max(0, fullHandler.indexOf(handlerString)))
+  return [moduleRoot, handlerString]
+}
+
+/**
+ * Example: `'index.nested.handler'` → `['index', 'nested.handler']`.
+ *
+ * @param {string} handler
+ * @throws {Error} When the handler is not of the form `<module>.<path>`.
+ */
+function extractModuleNameAndHandlerPath (handler) {
+  const match = handler.match(/^([^.]*)\.(.*)$/)
+  if (!match || match.length !== 3) {
+    throw new Error(`Malformed handler name: ${handler}`)
+  }
+  return [match[1], match[2]]
+}
+
+/**
+ * @param {string} lambdaStylePath `LAMBDA_TASK_ROOT` joined with the module root and module name.
+ */
+function getLambdaFilePaths (lambdaStylePath) {
+  return [
+    `${lambdaStylePath}.js`,
+    `${lambdaStylePath}.mjs`,
+    `${lambdaStylePath}.cjs`,
+  ]
+}
+
+module.exports = {
+  extractModuleRootAndHandler,
+  extractModuleNameAndHandlerPath,
+  getLambdaFilePaths,
+}

package/packages/dd-trace/src/lambda/index.js

@@ -1,17 +1,65 @@
 'use strict'
 
-const { getValueFromEnvSources } = require('../config/helper')
-const { registerLambdaHook } = require('./runtime/ritm')
-
-/**
- * It is safe to do it this way, since customers will never be expected to disable
- * this specific instrumentation through the init config object.
- */
-const _DD_TRACE_DISABLED_INSTRUMENTATIONS = getValueFromEnvSources('DD_TRACE_DISABLED_INSTRUMENTATIONS') || ''
-const _disabledInstrumentations = new Set(
-  _DD_TRACE_DISABLED_INSTRUMENTATIONS ? _DD_TRACE_DISABLED_INSTRUMENTATIONS.split(',') : []
-)
-
-if (!_disabledInstrumentations.has('lambda')) {
-  registerLambdaHook()
+const path = require('path')
+
+const log = require('../log')
+const { getEnvironmentVariable, getValueFromEnvSources } = require('../config/helper')
+const Hook = require('../../../datadog-instrumentations/src/helpers/hook')
+const instrumentations = require('../../../datadog-instrumentations/src/helpers/instrumentations')
+const {
+  filename,
+  pathSepExpr,
+} = require('../../../datadog-instrumentations/src/helpers/register')
+const {
+  extractModuleNameAndHandlerPath,
+  extractModuleRootAndHandler,
+  getLambdaFilePaths,
+} = require('./handler-paths')
+
+if (!getValueFromEnvSources('DD_TRACE_DISABLED_INSTRUMENTATIONS')?.split(',').includes('lambda')) {
+  const lambdaTaskRoot = getEnvironmentVariable('LAMBDA_TASK_ROOT')
+  const originalLambdaHandler = getValueFromEnvSources('DD_LAMBDA_HANDLER')
+
+  if (originalLambdaHandler !== undefined && lambdaTaskRoot !== undefined) {
+    const [moduleRoot, moduleAndHandler] = extractModuleRootAndHandler(originalLambdaHandler)
+    const [moduleName] = extractModuleNameAndHandlerPath(moduleAndHandler)
+
+    const lambdaStylePath = path.resolve(lambdaTaskRoot, moduleRoot, moduleName)
+    const lambdaFilePaths = getLambdaFilePaths(lambdaStylePath)
+
+    // TODO: Redo this like any other instrumentation.
+    Hook(lambdaFilePaths, (moduleExports, name, _, moduleVersion) => {
+      require('./runtime/patch')
+
+      for (const { hook } of instrumentations[name]) {
+        try {
+          moduleExports = hook(moduleExports, moduleVersion) ?? moduleExports
+        } catch (error) {
+          log.error('Error executing lambda hook', error)
+        }
+      }
+
+      return moduleExports
+    })
+    return
+  }
+
+  const moduleToPatch = 'datadog-lambda-js'
+  Hook([moduleToPatch], (moduleExports, moduleName, _, moduleVersion) => {
+    moduleName = moduleName.replace(pathSepExpr, '/')
+    require('./runtime/patch')
+
+    for (const { file, hook } of instrumentations[moduleToPatch]) {
+      const fullFilename = filename(moduleToPatch, file)
+      if (moduleName === fullFilename) {
+        try {
+          moduleExports = hook(moduleExports, moduleVersion) ?? moduleExports
+        } catch (error) {
+          log.error('Error executing lambda hook for datadog-lambda-js', error)
+        }
+      }
+    }
+
+    return moduleExports
+  })
 }

package/packages/dd-trace/src/lambda/runtime/patch.js

@@ -6,54 +6,32 @@ const { datadog } = require('../handler')
 const { addHook } = require('../../../../datadog-instrumentations/src/helpers/instrument')
 const shimmer = require('../../../../datadog-shimmer')
 const { getEnvironmentVariable, getValueFromEnvSources } = require('../../config/helper')
-const { _extractModuleNameAndHandlerPath, _extractModuleRootAndHandler, _getLambdaFilePaths } = require('./ritm')
-
-/**
- * Patches a Datadog Lambda module by calling `patchDatadogLambdaHandler`
- * with the handler name `datadog`.
- *
- * @param {object} datadogLambdaModule node module to be patched.
- * @returns a Datadog Lambda module with the `datadog` function from
- * `datadog-lambda-js` patched.
- */
-const patchDatadogLambdaModule = (datadogLambdaModule) => {
+const {
+  extractModuleNameAndHandlerPath,
+  extractModuleRootAndHandler,
+  getLambdaFilePaths,
+} = require('../handler-paths')
+
+/** @param {object} datadogLambdaModule */
+function patchDatadogLambdaModule (datadogLambdaModule) {
   shimmer.wrap(datadogLambdaModule, 'datadog', patchDatadogLambdaHandler)
-
   return datadogLambdaModule
 }
 
-/**
- * Patches a Datadog Lambda handler in order to do
- * Datadog instrumentation by getting the Lambda handler from its
- * arguments.
- *
- * @param {Function} datadogHandler the Datadog Lambda handler to destructure.
- * @returns the datadogHandler with its arguments patched.
- */
+/** @param {Function} datadogHandler */
 function patchDatadogLambdaHandler (datadogHandler) {
-  return (userHandler) => {
-    return datadogHandler(datadog(userHandler))
-  }
+  return userHandler => datadogHandler(datadog(userHandler))
 }
 
-/**
- * Patches a Lambda module on the given handler path.
- *
- * @param {string} handlerPath path of the handler to be patched.
- * @returns a module with the given handler path patched.
- */
-const patchLambdaModule = (handlerPath) => (lambdaModule) => {
-  shimmer.wrap(lambdaModule, handlerPath, patchLambdaHandler)
-
-  return lambdaModule
+/** @param {string} handlerPath */
+function patchLambdaModule (handlerPath) {
+  return lambdaModule => {
+    shimmer.wrap(lambdaModule, handlerPath, patchLambdaHandler)
+    return lambdaModule
+  }
 }
 
-/**
- * Patches a Lambda handler in order to do Datadog instrumentation.
- *
- * @param {Function} lambdaHandler the Lambda handler to be patched.
- * @returns a function which patches the given Lambda handler.
- */
+/** @param {Function} lambdaHandler */
 function patchLambdaHandler (lambdaHandler) {
   return datadog(lambdaHandler)
 }
@@ -62,16 +40,13 @@ const lambdaTaskRoot = getEnvironmentVariable('LAMBDA_TASK_ROOT')
 const originalLambdaHandler = getValueFromEnvSources('DD_LAMBDA_HANDLER')
 
 if (originalLambdaHandler === undefined) {
-  // Instrumentation is done manually.
   addHook({ name: 'datadog-lambda-js' }, patchDatadogLambdaModule)
 } else {
-  const [moduleRoot, moduleAndHandler] = _extractModuleRootAndHandler(originalLambdaHandler)
-  const [_module, handlerPath] = _extractModuleNameAndHandlerPath(moduleAndHandler)
-
-  const lambdaStylePath = path.resolve(lambdaTaskRoot, moduleRoot, _module)
-  const lambdaFilePaths = _getLambdaFilePaths(lambdaStylePath)
+  const [moduleRoot, moduleAndHandler] = extractModuleRootAndHandler(originalLambdaHandler)
+  const [moduleName, handlerPath] = extractModuleNameAndHandlerPath(moduleAndHandler)
 
-  for (const lambdaFilePath of lambdaFilePaths) {
+  const lambdaStylePath = path.resolve(lambdaTaskRoot, moduleRoot, moduleName)
+  for (const lambdaFilePath of getLambdaFilePaths(lambdaStylePath)) {
     addHook({ name: lambdaFilePath }, patchLambdaModule(handlerPath))
   }
 }

package/packages/dd-trace/src/llmobs/index.js

@@ -108,6 +108,10 @@ function disable () {
 // since LLMObs traces can extend between services and be the same trace,
 // we need to propagate the parent id and mlApp.
 function handleLLMObsParentIdInjection ({ carrier }) {
+  // Respect the standard propagator's gate: when trace tag propagation is
+  // disabled, don't write `x-datadog-tags` for LLMObs either.
+  if (globalTracerConfig.DD_TRACE_X_DATADOG_TAGS_MAX_LENGTH === 0) return
+
   const parent = storage.getStore()?.span
   const mlObsSpanTags = LLMObsTagger.tagMap.get(parent)
 
@@ -118,8 +122,15 @@ function handleLLMObsParentIdInjection ({ carrier }) {
     parentContext?._trace?.tags?.[PROPAGATED_ML_APP_KEY] ||
     globalTracerConfig.llmobs.mlApp
 
-  if (parentId) carrier['x-datadog-tags'] += `,${PROPAGATED_PARENT_ID_KEY}=${parentId}`
-  if (mlApp) carrier['x-datadog-tags'] += `,${PROPAGATED_ML_APP_KEY}=${mlApp}`
+  if (!parentId && !mlApp) return
+
+  // `_injectTags` only writes `x-datadog-tags` when the trace has `_dd.p.*`
+  // tags, so it may be undefined here — coalesce before appending.
+  const existing = carrier['x-datadog-tags']
+  let tags = existing || ''
+  if (parentId) tags += `${tags ? ',' : ''}${PROPAGATED_PARENT_ID_KEY}=${parentId}`
+  if (mlApp) tags += `${tags ? ',' : ''}${PROPAGATED_ML_APP_KEY}=${mlApp}`
+  if (tags !== existing) carrier['x-datadog-tags'] = tags
 }
 
 function handleFlush () {

package/packages/dd-trace/src/llmobs/plugins/ai/util.js

@@ -41,8 +41,7 @@ const VERCEL_AI_GENERATION_METADATA_PREFIX = 'ai.settings.'
  */
 function getSpanTags (ctx) {
   const span = ctx.currentStore?.span
-  const carrier = ctx.attributes ?? span?.context()._tags ?? {}
-  return /** @type {SpanTags} */ (carrier)
+  return /** @type {SpanTags} */ (ctx.attributes ?? span?.context()._tags ?? {})
 }
 
 /**

package/packages/dd-trace/src/llmobs/plugins/bedrockruntime.js

@@ -14,7 +14,17 @@ const llmobsStore = storage('llmobs')
 
 const ENABLED_OPERATIONS = new Set(['invokeModel', 'invokeModelWithResponseStream'])
 
-const requestIdsToTokens = {}
+/**
+ * @typedef {{
+ *   inputTokensFromHeaders?: number,
+ *   outputTokensFromHeaders?: number,
+ *   cacheReadTokensFromHeaders?: number,
+ *   cacheWriteTokensFromHeaders?: number,
+ * }} HeaderTokens
+ */
+
+/** @type {Map<string, HeaderTokens>} */
+const pendingTokenHeaders = new Map()
 
 class BedrockRuntimeLLMObsPlugin extends BaseLLMObsPlugin {
   constructor () {
@@ -24,33 +34,39 @@ class BedrockRuntimeLLMObsPlugin extends BaseLLMObsPlugin {
       const { response } = ctx
       const request = response.request
       const operation = request.operation
+
+      // Release the cached headers even for operations the plugin does not tag,
+      // so non-LLM Bedrock calls do not leak entries into pendingTokenHeaders.
+      const tokensFromHeaders = consumeTokenHeaders(response.$metadata?.requestId)
+
       // avoids instrumenting other non supported runtime operations
-      if (!ENABLED_OPERATIONS.has(operation)) {
-        return
-      }
+      if (!ENABLED_OPERATIONS.has(operation)) return
+
       const { modelProvider, modelName } = parseModelId(request.params.modelId)
 
       // avoids instrumenting non llm type
-      if (modelName.includes('embed')) {
-        return
-      }
+      if (modelName.includes('embed')) return
+
       const span = ctx.currentStore?.span
-      this.setLLMObsTags({ ctx, request, span, response, modelProvider, modelName })
+      this.setLLMObsTags({ ctx, request, span, response, modelProvider, modelName, tokensFromHeaders })
     })
 
     this.addSub('apm:aws:response:deserialize:bedrockruntime', ({ headers }) => {
       const requestId = headers['x-amzn-requestid']
+      // No request id means no way to correlate with the :complete: event.
+      if (!requestId) return
+
       const inputTokenCount = headers['x-amzn-bedrock-input-token-count']
       const outputTokenCount = headers['x-amzn-bedrock-output-token-count']
       const cacheReadTokenCount = headers['x-amzn-bedrock-cache-read-input-token-count']
       const cacheWriteTokenCount = headers['x-amzn-bedrock-cache-write-input-token-count']
 
-      requestIdsToTokens[requestId] = {
+      pendingTokenHeaders.set(requestId, {
         inputTokensFromHeaders: inputTokenCount && Number.parseInt(inputTokenCount),
         outputTokensFromHeaders: outputTokenCount && Number.parseInt(outputTokenCount),
         cacheReadTokensFromHeaders: cacheReadTokenCount && Number.parseInt(cacheReadTokenCount),
         cacheWriteTokensFromHeaders: cacheWriteTokenCount && Number.parseInt(cacheWriteTokenCount),
-      }
+      })
     })
 
     this.addSub('apm:aws:response:streamed-chunk:bedrockruntime', ({ ctx, chunk }) => {
@@ -60,7 +76,7 @@ class BedrockRuntimeLLMObsPlugin extends BaseLLMObsPlugin {
     })
   }
 
-  setLLMObsTags ({ ctx, request, span, response, modelProvider, modelName }) {
+  setLLMObsTags ({ ctx, request, span, response, modelProvider, modelName, tokensFromHeaders }) {
     const isStream = request?.operation?.toLowerCase().includes('stream')
     telemetry.incrementLLMObsSpanStartCount({ autoinstrumented: true, integration: 'bedrock' })
 
@@ -97,7 +113,7 @@ class BedrockRuntimeLLMObsPlugin extends BaseLLMObsPlugin {
 
     // add token metrics
     const { inputTokens, outputTokens, totalTokens, cacheReadTokens, cacheWriteTokens } = extractTokens({
-      requestId: response.$metadata.requestId,
+      tokensFromHeaders,
       usage: textAndResponseReason.usage,
     })
     this._tagger.tagMetrics(span, {
@@ -110,14 +126,28 @@ class BedrockRuntimeLLMObsPlugin extends BaseLLMObsPlugin {
   }
 }
 
-function extractTokens ({ requestId, usage }) {
+/**
+ * @param {string | undefined} requestId
+ * @returns {HeaderTokens | undefined}
+ */
+function consumeTokenHeaders (requestId) {
+  const tokens = pendingTokenHeaders.get(requestId)
+  pendingTokenHeaders.delete(requestId)
+  return tokens
+}
+
+/**
+ * Combine response-body usage with header-derived counts, preferring the body.
+ *
+ * @param {{ tokensFromHeaders: HeaderTokens | undefined, usage: Record<string, number | undefined> }} options
+ */
+function extractTokens ({ tokensFromHeaders, usage }) {
   const {
     inputTokensFromHeaders,
     outputTokensFromHeaders,
     cacheReadTokensFromHeaders,
     cacheWriteTokensFromHeaders,
-  } = requestIdsToTokens[requestId] || {}
-  delete requestIdsToTokens[requestId]
+  } = tokensFromHeaders ?? {}
 
   const inputTokens = usage.inputTokens || inputTokensFromHeaders || 0
   const outputTokens = usage.outputTokens || outputTokensFromHeaders || 0