dd-trace 5.132.0 → 5.133.0

package/packages/dd-trace/src/debugger/devtools_client/status.js

@@ -129,6 +129,7 @@ function statusPayload (probeId, probeVersion, status) {
   return {
     ddsource,
     service,
+    timestamp: Date.now(),
     debugger: {
       diagnostics: { probeId, runtimeId, probeVersion, status },
     },

package/packages/dd-trace/src/debugger/guardrail-metrics.js

@@ -9,10 +9,10 @@
  * atomics, and the main thread periodically drains the counters into instrumentation telemetry metrics.
  *
  * The metric names, tags and reason codes are defined by the "Debugger Observability for GA" specification and are
- * shared with the other tracers. All metrics are counts in the `live_debugger` telemetry namespace.
+ * shared with the other tracers. All metrics are counts in the `debugger` telemetry namespace.
  */
 
-const TELEMETRY_NAMESPACE = 'live_debugger'
+const TELEMETRY_NAMESPACE = 'debugger'
 
 const EVENT_TYPE = Object.freeze({
   SNAPSHOT: 0,

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

@@ -533,7 +533,7 @@ class DogStatsDClient {
       for (const [key, value] of Object.entries(config.tags)) {
         // Skip runtime-id unless enabled as cardinality may be too high
         if (typeof value === 'string' && (key !== 'runtime-id' || config.runtimeMetricsRuntimeId)) {
-          // https://docs.datadoghq.com/tagging/#defining-tags
+          // https://docs.datadoghq.com/getting_started/tagging/#define-tags
           const valueStripped = value.replaceAll(/[^a-z0-9_:./-]/ig, '_')
 
           tags.push(`${key}:${valueStripped}`)

package/packages/dd-trace/src/encode/tags-processors.js

@@ -2,7 +2,7 @@
 
 // From agent truncators: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/agent/truncator.go
 
-// Values from: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/traceutil/truncate.go#L22-L27
+// Values from: https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/traceutil/normalize/truncate.go
 // MAX_RESOURCE_NAME_LENGTH the maximum length a span resource can have
 const MAX_RESOURCE_NAME_LENGTH = 5000
 // MAX_META_KEY_LENGTH the maximum length of metadata key
@@ -15,7 +15,7 @@ const MAX_META_VALUE_LENGTH_TEST_OPTIMIZATION = 5000
 const MAX_METRIC_KEY_LENGTH = MAX_META_KEY_LENGTH
 
 // From agent normalizer:
-// https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/traceutil/normalize.go
+// https://github.com/DataDog/datadog-agent/blob/main/pkg/trace/traceutil/normalize/normalize.go
 // DEFAULT_SPAN_NAME is the default name we assign a span if it's missing and we have no reasonable fallback
 const DEFAULT_SPAN_NAME = 'unnamed_operation'
 // DEFAULT_SERVICE_NAME is the default name we assign a service if it's missing and we have no reasonable fallback

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

@@ -86,11 +86,7 @@ function request (data, options, callback) {
   const timeout = options.timeout || 2000
   const isSecure = options.protocol === 'https:'
   const client = isSecure ? https : http
-  let dataArray = data
-
-  if (!Array.isArray(data)) {
-    dataArray = [data]
-  }
+  const dataArray = Array.isArray(data) ? data : [data]
   const contentLength = byteLength(dataArray)
   options.headers['Content-Length'] = contentLength
 

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

@@ -6,6 +6,8 @@ const UINT_MAX = 4_294_967_296
 
 const data = new Uint8Array(8 * 8192)
 const zeroId = new Uint8Array(8)
+// Reused for hex encoding: toString('hex') copies out, so the bytes need not outlive the call.
+const hexScratch = Buffer.alloc(16)
 
 let batch = 0
 
@@ -35,7 +37,7 @@ class Identifier {
    */
   toString (radix = 16) {
     if (radix === 16) {
-      this.#stringHex ??= Buffer.from(this.#buffer).toString('hex')
+      this.#stringHex ??= toHexString(this.#buffer)
       return this.#stringHex
     }
     if (radix === 10) {
@@ -190,6 +192,17 @@ function toNumberString (buffer, radix) {
   return str
 }
 
+// Same string as Buffer.from(buffer).toString('hex'), without allocating a Buffer per call.
+/**
+ * @param {number[] | Uint8Array} buffer
+ */
+function toHexString (buffer) {
+  const length = buffer.length
+  if (length > hexScratch.length) return Buffer.from(buffer).toString('hex')
+  for (let i = 0; i < length; i++) hexScratch[i] = buffer[i]
+  return hexScratch.toString('hex', 0, length)
+}
+
 // Simple pseudo-random 64-bit ID generator.
 /**
  * @returns {number[] | Uint8Array}

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

@@ -1,20 +1,24 @@
 'use strict'
 
-const log = require('../log')
+const { types: { isProxy } } = require('node:util')
+
 const { channel } = require('../../../datadog-instrumentations/src/helpers/instrument')
 const { ERROR_MESSAGE, ERROR_TYPE } = require('../constants')
+const log = require('../log')
 const { ImpendingTimeout } = require('./runtime/errors')
 const { extractContext } = require('./context')
 
+const HANDLER_STREAMING = Symbol.for('aws.lambda.runtime.handler.streaming')
+const HANDLER_HIGHWATERMARK = Symbol.for('aws.lambda.runtime.handler.streaming.highWaterMark')
+const BEFORE_EXIT = Symbol.for('aws.lambda.beforeExit')
 const timeoutChannel = channel('apm:aws:lambda:timeout')
+const monitoredContexts = new WeakSet()
 // Always crash the flushes when a message is received
 // from this channel.
 timeoutChannel.subscribe(_ => {
   crashFlush()
 })
 
-let __lambdaTimeout
-
 /**
  * Publishes to the `apm:aws:lambda:timeout` channel when
  * the AWS Lambda run time is about to end.
@@ -23,11 +27,15 @@ let __lambdaTimeout
  */
 function checkTimeout (context) {
   const remainingTimeInMillis = context.getRemainingTimeInMillis()
-  const apmFlushDeadline = global._ddtrace._tracer._config.DD_APM_FLUSH_DEADLINE_MILLISECONDS
+  const apmFlushDeadline = global._ddtrace?._tracer?._config?.DD_APM_FLUSH_DEADLINE_MILLISECONDS
+  if (!Number.isFinite(remainingTimeInMillis) || !Number.isFinite(apmFlushDeadline)) return
 
-  __lambdaTimeout = setTimeout(() => {
+  const timeout = setTimeout(() => {
     timeoutChannel.publish()
-  }, remainingTimeInMillis - apmFlushDeadline)
+  }, Math.max(0, remainingTimeInMillis - apmFlushDeadline))
+  // eslint-disable-next-line eslint-rules/eslint-timer-unref -- AWS Lambda uses Node.js timers.
+  timeout.unref()
+  return timeout
 }
 
 /**
@@ -58,26 +66,178 @@ function crashFlush () {
 }
 
 /**
- * Patches your AWS Lambda handler function to add some tracing support.
+ * Monitors an AWS Lambda handler for impending timeouts without creating a span.
  *
  * @param {Function} lambdaHandler a Lambda handler function.
+ * @param {boolean} normalized whether datadog-lambda-js owns completion selection.
  */
-exports.datadog = function datadog (lambdaHandler) {
-  return (...args) => {
+exports.datadog = function datadog (lambdaHandler, normalized = false) {
+  function monitored (...args) {
     const context = extractContext(args)
+    if (!context || monitoredContexts.has(context)) return Reflect.apply(lambdaHandler, this, args)
 
-    if (context) {
-      checkTimeout(context)
+    const timeout = checkTimeout(context)
+    if (timeout === undefined) return Reflect.apply(lambdaHandler, this, args)
+
+    let completed = false
+    let callbackPending = false
+    /** @type {Array<[object, string | symbol, Function, Function]>} */
+    const completionMethods = []
+    const clear = () => {
+      if (completed) return
+      completed = true
+      clearTimeout(timeout)
+      for (const [target, name, original, wrapped] of completionMethods) {
+        try {
+          if (target[name] === wrapped) Reflect.set(target, name, original)
+        } catch (error) {
+          log.error('Error restoring Lambda completion %s: %s', name, error)
+        }
+      }
+    }
+    const clearOnHandlerCompletion = normalized
+      ? clear
+      : () => {
+          // Once a native callback selects event-loop drain, later handler signals are ignored by RIC.
+          if (!callbackPending) clear()
+        }
+
+    const contextIndex = args.indexOf(context)
+    const callbackObservable = contextIndex !== 2 && typeof args[2] === 'function'
+    if (!normalized) {
+      for (const name of ['done', 'succeed', 'fail']) {
+        wrapCompletion(context, name, clearOnHandlerCompletion, completionMethods)
+      }
     }
+    // Legacy Lambda runtimes complete idle invocations through this callback instead of the handler's callback.
+    wrapCompletion(globalThis, BEFORE_EXIT, clear, completionMethods)
 
-    const result = lambdaHandler.apply(this, args)
-    if (result && typeof result.then === 'function') {
-      return result.then((res) => {
-        clearTimeout(__lambdaTimeout)
-        return res
-      })
+    if (callbackObservable) {
+      const callback = args[2]
+      const beforeExitObservable = completionMethods.at(-1)?.[1] === BEFORE_EXIT
+      let callbackInvoked = false
+      args[2] = function (...callbackArgs) {
+        if (completed || callbackInvoked) return Reflect.apply(callback, this, callbackArgs)
+        callbackInvoked = true
+        let beforeExit
+        try {
+          if (beforeExitObservable && callbackArgs[0] == null && context.callbackWaitsForEmptyEventLoop) {
+            beforeExit = getDataProperty(globalThis, BEFORE_EXIT)
+          }
+        } catch (error) {
+          log.error('Error reading Lambda callback completion: %s', error)
+        }
+        if (typeof beforeExit !== 'function') clear()
+        try {
+          const result = Reflect.apply(callback, this, callbackArgs)
+          if (!completed) {
+            // A successful legacy callback installs a new listener while waiting for event-loop drain.
+            const replacement = getDataProperty(globalThis, BEFORE_EXIT)
+            if (typeof replacement !== 'function' || replacement === beforeExit ||
+                !wrapCompletion(globalThis, BEFORE_EXIT, clear, completionMethods)) clear()
+            else callbackPending = true
+          }
+          return result
+        } catch (error) {
+          clear()
+          throw error
+        }
+      }
     }
-    clearTimeout(__lambdaTimeout)
-    return result
+
+    try {
+      // The outer completion adapter owns this context during synchronous dispatch to nested wrappers.
+      monitoredContexts.add(context)
+      const result = Reflect.apply(lambdaHandler, this, args)
+      if (result && typeof result.then === 'function') {
+        return result.then(value => {
+          clearOnHandlerCompletion()
+          return value
+        }, error => {
+          clearOnHandlerCompletion()
+          throw error
+        })
+      }
+      // Legacy RIC supplies a callback; newer runtimes complete from the return value.
+      if (!callbackObservable) clear()
+      return result
+    } catch (error) {
+      clearOnHandlerCompletion()
+      throw error
+    } finally {
+      monitoredContexts.delete(context)
+    }
+  }
+
+  if (normalized) return monitored
+
+  // datadog-lambda-js uses arity and the runtime's streaming metadata to choose how to invoke the handler.
+  Object.defineProperty(monitored, 'length', { value: getDataProperty(lambdaHandler, 'length') ?? 0 })
+  for (const name of [HANDLER_STREAMING, HANDLER_HIGHWATERMARK, '_ddWrapped']) {
+    const value = getDataProperty(lambdaHandler, name)
+    if (value !== undefined) monitored[name] = value
+  }
+  if (!isProxy(lambdaHandler)) {
+    for (const name of Reflect.ownKeys(lambdaHandler)) {
+      if (name === 'arguments' || name === 'caller') continue
+      const descriptor = Object.getOwnPropertyDescriptor(lambdaHandler, name)
+      if (!Object.hasOwn(descriptor, 'value')) continue
+      const existing = Object.getOwnPropertyDescriptor(monitored, name)
+      if (existing?.configurable === false) {
+        // Intrinsic function slots such as prototype cannot change their descriptor shape.
+        if (existing.writable) {
+          Object.defineProperty(monitored, name, { value: descriptor.value, writable: descriptor.writable })
+        }
+      } else {
+        Object.defineProperty(monitored, name, descriptor)
+      }
+    }
+  }
+  return monitored
+}
+
+/**
+ * @param {object} target
+ * @param {string | symbol} name
+ * @param {() => void} clear
+ * @param {Array<[object, string | symbol, Function, Function]>} completionMethods
+ */
+function wrapCompletion (target, name, clear, completionMethods) {
+  try {
+    const original = target[name]
+    if (typeof original !== 'function') return true
+    const wrapped = onCompletion(original, clear)
+    if (!Reflect.set(target, name, wrapped) || target[name] !== wrapped) return false
+    completionMethods.push([target, name, original, wrapped])
+    return true
+  } catch (error) {
+    log.error('Error wrapping Lambda completion %s: %s', name, error)
+    return false
+  }
+}
+
+/**
+ * @param {Function} callback
+ * @param {() => void} clear
+ */
+function onCompletion (callback, clear) {
+  return function (...args) {
+    clear()
+    return Reflect.apply(callback, this, args)
+  }
+}
+
+/**
+ * Reads metadata without executing application getters or Proxy traps.
+ *
+ * @param {unknown} target
+ * @param {string | symbol} name
+ */
+function getDataProperty (target, name) {
+  if (target === null || (typeof target !== 'object' && typeof target !== 'function')) return
+  while (target !== null && !isProxy(target)) {
+    const descriptor = Object.getOwnPropertyDescriptor(target, name)
+    if (descriptor) return descriptor.value
+    target = Object.getPrototypeOf(target)
   }
 }

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

@@ -41,7 +41,6 @@ if (!getValueFromEnvSources('DD_TRACE_DISABLED_INSTRUMENTATIONS')?.split(',').in
 
       return moduleExports
     })
-    return
   }
 
   const moduleToPatch = 'datadog-lambda-js'
@@ -62,4 +61,19 @@ if (!getValueFromEnvSources('DD_TRACE_DISABLED_INSTRUMENTATIONS')?.split(',').in
 
     return moduleExports
   })
+
+  // The published datadog-lambda-js ESM entrypoint imports this helper before initializing the tracer.
+  // Revisit cached copies through the registered hook without loading datadog-lambda-js early.
+  const helperSuffixes = [
+    path.join('datadog-lambda-js', 'dist', 'utils', 'handler.js'),
+    path.join('datadog-lambda-js', 'utils', 'handler.js'),
+  ]
+  for (const filename of Object.keys(require.cache)) {
+    if (!helperSuffixes.some(suffix => filename.endsWith(path.sep + suffix))) continue
+    try {
+      require(filename)
+    } catch (error) {
+      log.error('Error executing cached lambda hook for datadog-lambda-js: %s', error)
+    }
+  }
 }

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

@@ -1,6 +1,7 @@
 'use strict'
 
 const path = require('path')
+const { types: { isModuleNamespaceObject } } = require('node:util')
 
 const { datadog } = require('../handler')
 const { addHook } = require('../../../../datadog-instrumentations/src/helpers/instrument')
@@ -23,22 +24,48 @@ function patchDatadogLambdaHandler (datadogHandler) {
   return userHandler => datadogHandler(datadog(userHandler))
 }
 
+/**
+ * Wraps datadog-lambda-js's promisifiedHandler factory, which wraps the customer's handler.
+ * The timeout monitor surrounds that adapter to observe the completion Promise selected by datadog-lambda-js.
+ * Legacy RIC completes returned Promises through context.succeed without waiting for event-loop drain.
+ *
+ * @param {object} moduleExports
+ */
+function patchPromisifiedHandlerModule (moduleExports) {
+  shimmer.wrap(moduleExports, 'promisifiedHandler', promisify => {
+    return handler => datadog(promisify(handler), true)
+  })
+  return moduleExports
+}
+
 /** @param {string} handlerPath */
 function patchLambdaModule (handlerPath) {
   return lambdaModule => {
-    shimmer.wrap(lambdaModule, handlerPath, patchLambdaHandler)
-    return lambdaModule
+    const descriptor = Object.getOwnPropertyDescriptor(lambdaModule, handlerPath)
+    if (typeof descriptor?.value === 'function') {
+      if (!descriptor.configurable && !descriptor.writable) return lambdaModule
+      // The monitor copies Lambda metadata without invoking application accessors or Proxy traps.
+      // Generic shimmer wrapping would copy the original properties back onto that safe wrapper.
+      descriptor.value = datadog(descriptor.value)
+      // Native ESM namespaces reject replacement even when their export reports writable: true.
+      const exports = isModuleNamespaceObject(lambdaModule)
+        ? Object.create(lambdaModule, Object.getOwnPropertyDescriptors(lambdaModule))
+        : lambdaModule
+      Object.defineProperty(exports, handlerPath, descriptor)
+      return exports
+    }
+    return shimmer.wrap(lambdaModule, handlerPath, datadog)
   }
 }
 
-/** @param {Function} lambdaHandler */
-function patchLambdaHandler (lambdaHandler) {
-  return datadog(lambdaHandler)
-}
-
 const lambdaTaskRoot = getEnvironmentVariable('LAMBDA_TASK_ROOT')
 const originalLambdaHandler = getValueFromEnvSources('DD_LAMBDA_HANDLER')
 
+// Observe the completion Promise selected by datadog-lambda-js inside the invocation span's async scope.
+// The npm package keeps dist/; Lambda layer builds place its contents directly in the package root.
+addHook({ name: 'datadog-lambda-js', file: 'dist/utils/handler.js' }, patchPromisifiedHandlerModule)
+addHook({ name: 'datadog-lambda-js', file: 'utils/handler.js' }, patchPromisifiedHandlerModule)
+
 if (originalLambdaHandler === undefined) {
   addHook({ name: 'datadog-lambda-js' }, patchDatadogLambdaModule)
 } else {

package/packages/dd-trace/src/llmobs/audio-utils.js

@@ -106,8 +106,9 @@ function fitsInlineAudioBudget (byteLength, maxBytes = LLMOBS_AUDIO_INLINE_MAX_B
  * @returns {AudioPart | undefined}
  */
 function formatAudioPartWithGuard (audioBytes, mimeType, maxBytes = LLMOBS_AUDIO_INLINE_MAX_BYTES) {
-  if (!audioBytes?.length || !isRenderableAudioMime(mimeType)) return
-  if (!fitsInlineAudioBudget(audioBytes.length, maxBytes)) return
+  if (!audioBytes?.length || !isRenderableAudioMime(mimeType) || !fitsInlineAudioBudget(audioBytes.length, maxBytes)) {
+    return
+  }
 
   return formatAudioPart(audioBytes, mimeType)
 }

package/packages/dd-trace/src/llmobs/constants/tags.js

@@ -7,7 +7,11 @@ const CACHE_READ_INPUT_TOKENS_METRIC_KEY = 'cache_read_input_tokens'
 const CACHE_WRITE_INPUT_TOKENS_METRIC_KEY = 'cache_write_input_tokens'
 const CACHE_WRITE_5M_INPUT_TOKENS_METRIC_KEY = 'ephemeral_5m_input_tokens'
 const CACHE_WRITE_1H_INPUT_TOKENS_METRIC_KEY = 'ephemeral_1h_input_tokens'
+const CACHE_WRITE_30M_INPUT_TOKENS_METRIC_KEY = 'ephemeral_30m_input_tokens'
+const CACHE_WRITE_12H_INPUT_TOKENS_METRIC_KEY = 'ephemeral_12h_input_tokens'
 const REASONING_OUTPUT_TOKENS_METRIC_KEY = 'reasoning_output_tokens'
+const WEB_SEARCH_COUNT_METRIC_KEY = 'web_search_count'
+const STORAGE_SEARCH_COUNT_METRIC_KEY = 'storage_search_count'
 
 module.exports = {
   SPAN_KINDS: ['llm', 'agent', 'workflow', 'task', 'tool', 'embedding', 'retrieval', 'experiment'],
@@ -72,7 +76,11 @@ module.exports = {
   CACHE_WRITE_INPUT_TOKENS_METRIC_KEY,
   CACHE_WRITE_5M_INPUT_TOKENS_METRIC_KEY,
   CACHE_WRITE_1H_INPUT_TOKENS_METRIC_KEY,
+  CACHE_WRITE_30M_INPUT_TOKENS_METRIC_KEY,
+  CACHE_WRITE_12H_INPUT_TOKENS_METRIC_KEY,
   REASONING_OUTPUT_TOKENS_METRIC_KEY,
+  WEB_SEARCH_COUNT_METRIC_KEY,
+  STORAGE_SEARCH_COUNT_METRIC_KEY,
 
   // integrations build metric objects with these camelCase spellings. Null prototype: a metric
   // named after an `Object.prototype` member must not resolve to an inherited property.
@@ -84,7 +92,11 @@ module.exports = {
     cacheWriteTokens: CACHE_WRITE_INPUT_TOKENS_METRIC_KEY,
     cacheWrite5mTokens: CACHE_WRITE_5M_INPUT_TOKENS_METRIC_KEY,
     cacheWrite1hTokens: CACHE_WRITE_1H_INPUT_TOKENS_METRIC_KEY,
+    cacheWrite30mTokens: CACHE_WRITE_30M_INPUT_TOKENS_METRIC_KEY,
+    cacheWrite12hTokens: CACHE_WRITE_12H_INPUT_TOKENS_METRIC_KEY,
     reasoningOutputTokens: REASONING_OUTPUT_TOKENS_METRIC_KEY,
+    webSearchCount: WEB_SEARCH_COUNT_METRIC_KEY,
+    storageSearchCount: STORAGE_SEARCH_COUNT_METRIC_KEY,
   }),
 
   DROPPED_IO_COLLECTION_ERROR: 'dropped_io',

package/packages/dd-trace/src/llmobs/experiments/experiment.js

@@ -253,7 +253,7 @@ class Experiment {
     this.#external = options.external === true
     if (!this.#external) {
       if (!options.dataset) throw new Error('Experiment dataset is required')
-      if (typeof options.task !== 'function') throw new Error('Experiment task is required')
+      if (typeof options.task !== 'function') throw new TypeError('Experiment task is required')
     }
 
     this.#client = client

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

@@ -239,11 +239,19 @@ function mergeChunkUsage (usage, chunk) {
   const mergedUsage = usage ?? { input_tokens: 0, output_tokens: 0 }
   mergedUsage.output_tokens = chunk.usage.output_tokens
 
+  // `message_delta` counts are cumulative; input grows past the `message_start` count when server
+  // tools (e.g. web search) feed results back into the model
+  const inputTokens = chunk.usage.input_tokens
+  if (inputTokens != null) mergedUsage.input_tokens = inputTokens
+
   const cacheCreationTokens = chunk.usage.cache_creation_input_tokens
   const cacheReadTokens = chunk.usage.cache_read_input_tokens
   if (cacheCreationTokens) mergedUsage.cache_creation_input_tokens = cacheCreationTokens
   if (cacheReadTokens) mergedUsage.cache_read_input_tokens = cacheReadTokens
 
+  const serverToolUse = chunk.usage.server_tool_use
+  if (serverToolUse) mergedUsage.server_tool_use = serverToolUse
+
   return mergedUsage
 }
 
@@ -280,6 +288,13 @@ function extractUsage (result) {
     metrics.cacheWrite1hTokens = 0
   }
 
+  // Read the usage counter, not server_tool_use blocks: dynamic filtering nests searches under code
+  // execution, and `response_inclusion: "excluded"` drops the blocks while the searches are still billed.
+  // Failed searches aren't billed and are already excluded. web_fetch_requests is skipped because web
+  // fetch has no per-call price.
+  const webSearchRequests = usage.server_tool_use?.web_search_requests
+  if (webSearchRequests > 0) metrics.webSearchCount = webSearchRequests
+
   return metrics
 }
 

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

@@ -84,8 +84,7 @@ function determineRole (candidate, parts = []) {
  * @param {string} role
  */
 function normalizeRole (role) {
-  if (role === ROLES.MODEL) return ROLES.ASSISTANT
-  if (role === ROLES.ASSISTANT) return ROLES.ASSISTANT
+  if (role === ROLES.MODEL || role === ROLES.ASSISTANT) return ROLES.ASSISTANT
   if (role === ROLES.USER) return ROLES.USER
   if (role === ROLES.REASONING) return ROLES.REASONING
   return ROLES.USER // default

package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js

@@ -62,11 +62,9 @@ class LangChainLLMObsChatModelHandler extends LangChainLLMObsHandler {
    */
   getTokenUsage (results) {
     const tokens = this.checkTokenUsageChatOrLLMResult(results)
-    if (tokens.totalTokens > 0) return tokens
-
     // providers that report usage on each generated message instead of on `llmOutput`; counts are
     // summed per run so a run split across generations is totalled once
-    if (!results.generations) return tokens
+    if (tokens.totalTokens > 0 || !results.generations) return tokens
 
     const tokensPerRunId = {}
     for (const messageSet of results.generations) {

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

@@ -76,7 +76,7 @@ class NextStreamLLMObsPlugin extends LLMObsPlugin {
     const { streamInputs: inputs, chunks } = streamData
     const input = inputs == null ? undefined : formatIO(inputs)
     const lastChunk = chunks.length > 0 ? chunks.at(-1) : undefined
-    const output = !hasError && lastChunk != null ? formatIO(lastChunk) : undefined
+    const output = hasError || lastChunk == null ? undefined : formatIO(lastChunk)
 
     this._tagger.tagTextIO(span, input, output)
 

package/packages/dd-trace/src/llmobs/plugins/nova-sonic.js

@@ -0,0 +1,92 @@
+'use strict'
+
+const { pcm16ToWav } = require('../audio-codec')
+const { fitsInlineAudioBudget, formatAudioPartWithGuard } = require('../audio-utils')
+const { AUDIO_FALLBACK, LLMOBS_AUDIO_INLINE_MAX_BYTES } = require('../constants/audio')
+const { storage } = require('../storage')
+const LLMObsPlugin = require('./base')
+
+class SonicLLMObsPlugin extends LLMObsPlugin {
+  static id = 'bedrockruntime_sonic_llmobs'
+  static integration = 'bedrock'
+  static system = 'aws.bedrock'
+  static prefix = 'tracing:apm:aws:bedrockruntime:sonic:span'
+  static emitsGenAiApmTags = false
+
+  constructor (...args) {
+    super(...args)
+    this.addSub('dd-trace:aws:bedrockruntime:sonic:capture-context', context => {
+      if (!this._llmobsEnabled) return
+      context.enabled = true
+      // Capture once at invocation, even when absent or already finished. Scoped replay restores
+      // the consumer's context on every exit, including exceptions, without repairing IDs later.
+      const store = storage.getStore()
+      const previous = context.runInContext ?? (fn => fn())
+      context.runInContext = fn => previous(() => storage.run(store, fn))
+    })
+  }
+
+  /** @param {object} ctx */
+  getLLMObsSpanRegisterOptions (ctx) {
+    const options = {
+      kind: ctx.kind,
+      name: ctx.name,
+      sessionId: ctx.descriptor.sessionId,
+    }
+    if (ctx.kind === 'llm') {
+      options.modelName = ctx.model
+      options.modelProvider = 'amazon'
+    }
+    return options
+  }
+
+  /** @param {object} ctx */
+  start (ctx) {
+    super.start(ctx)
+    // Replay already has the complete descriptor. Annotate before any end subscriber can finish
+    // the span, including when disabling/re-enabling LLMObs changes subscriber registration order.
+    if (this._llmobsEnabledFor(ctx)) this.setLLMObsTags(ctx)
+  }
+
+  /** @param {object} ctx */
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+    const { turn, history, metadata } = ctx.descriptor
+    const text = turn.finalText || turn.speculativeText
+    if (ctx.kind !== 'llm') {
+      if (ctx.name === 'nova sonic audio turn') this._tagger.tagTextIO(span, turn.userText, text)
+      return
+    }
+
+    const input = { role: 'user', content: turn.userText || (turn.inputAudioObserved ? AUDIO_FALLBACK : '') }
+    const output = { role: 'assistant', content: text || (turn.outputAudioObserved ? AUDIO_FALLBACK : '') }
+    if (turn.tools.length) output.toolCalls = turn.tools
+    const inputs = [...history, input, ...turn.toolResults]
+    const outputs = [output]
+    // Reserve the serialized text/tools and metadata before spending one budget across both WAVs.
+    let budget = Math.max(0, LLMOBS_AUDIO_INLINE_MAX_BYTES -
+      Buffer.byteLength(JSON.stringify([inputs, outputs, metadata])) - 256)
+    for (const [message, pcm, rate] of [
+      [input, turn.inputPcm, turn.inputRate],
+      [output, turn.output.pcm(), turn.output.rate],
+    ]) {
+      if (!pcm?.length || !rate) continue
+      const part = fitsInlineAudioBudget(pcm.length + 44, budget)
+        ? formatAudioPartWithGuard(pcm16ToWav(pcm, rate), 'audio/wav', budget)
+        : undefined
+      if (part) {
+        message.audioParts = [part]
+        budget -= part.content.length
+      } else if (message === output) {
+        metadata.output_audio_omitted_reason = 'payload_limit'
+      }
+    }
+    if (turn.output.omittedReason) metadata.output_audio_omitted_reason = turn.output.omittedReason
+    this._tagger.tagLLMIO(span, inputs, outputs)
+    this._tagger.tagMetadata(span, metadata)
+    this._tagger.tagMetrics(span, turn.metrics)
+  }
+}
+
+module.exports = SonicLLMObsPlugin