fix(llmobs): llm observability traces have custom trace IDs (#9460)

* separate llmobs trace ids

* update exportSpan

* tests

* add tagger test

* additional fixups

* update tests from rebase

Author: sabrenner

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

@@ -16,6 +16,8 @@ const {
   SAMPLING_DECISION,
   PROPAGATED_SAMPLE_RATE_KEY,
   PROPAGATED_SAMPLING_DECISION_KEY,
+  TRACE_ID,
+  PROPAGATED_TRACE_ID_KEY,
 } = require('./constants/tags')
 const { storage } = require('./storage')
 const telemetry = require('./telemetry')
@@ -25,6 +27,7 @@ const LLMObsTagger = require('./tagger')
 const LLMObsSpanWriter = require('./writers/spans')
 const { setAgentStrategy } = require('./writers/util')
 const { INCOMPATIBLE_INITIALIZATION } = require('./constants/text')
+const { llmObsTraceIdToWire } = require('./util')
 
 const spanFinishCh = channel('dd-trace:span:finish')
 const evalMetricAppendCh = channel('llmobs:eval-metric:append')
@@ -137,8 +140,12 @@ function handleLLMObsInjection ({ carrier }) {
     mlObsSpanTags?.[SESSION_ID] ??
     parentContext?._trace?.tags?.[SESSION_ID_TRACE_DEFAULT_KEY] ??
     parentContext?._trace?.tags?.[PROPAGATED_SESSION_ID_KEY]
+  const llmobsTraceId = mlObsSpanTags?.[TRACE_ID]
+  const propagatedTraceId = llmobsTraceId === undefined
+    ? parentContext?._trace?.tags?.[PROPAGATED_TRACE_ID_KEY]
+    : llmObsTraceIdToWire(llmobsTraceId)
 
-  if (!parentId && !mlApp && samplingDecision == null && !sessionId) return
+  if (!parentId && !mlApp && samplingDecision == null && !sessionId && !propagatedTraceId) return
 
   // `_injectTags` only writes `x-datadog-tags` when the trace has `_dd.p.*`
   // tags, so it may be undefined here — coalesce before appending.
@@ -149,6 +156,7 @@ function handleLLMObsInjection ({ carrier }) {
   if (sessionId) tags += `${tags ? ',' : ''}${PROPAGATED_SESSION_ID_KEY}=${sessionId}`
   if (sampleRate != null) tags += `${tags ? ',' : ''}${PROPAGATED_SAMPLE_RATE_KEY}=${sampleRate}`
   if (samplingDecision != null) tags += `${tags ? ',' : ''}${PROPAGATED_SAMPLING_DECISION_KEY}=${samplingDecision}`
+  if (propagatedTraceId != null) tags += `${tags ? ',' : ''}${PROPAGATED_TRACE_ID_KEY}=${propagatedTraceId}`
   if (tags !== existing) carrier['x-datadog-tags'] = tags
 }
 

packages/dd-trace/src/llmobs/sdk.js

@@ -11,6 +11,7 @@ const {
   SPAN_KIND,
   OUTPUT_VALUE,
   INPUT_VALUE,
+  TRACE_ID,
 } = require('./constants/tags')
 const {
   getFunctionArguments,
@@ -345,7 +346,7 @@ class LLMObs extends NoopLLMObs {
     }
     try {
       return {
-        traceId: span.context().toTraceId(true),
+        traceId: LLMObsTagger.tagMap.get(span)[TRACE_ID],
         spanId: span.context().toSpanId(),
       }
     } catch {

packages/dd-trace/src/llmobs/span_processor.js

@@ -35,6 +35,7 @@ const {
   LLMOBS_SUBMITTED_TAG_KEY,
   SAMPLE_RATE,
   SAMPLING_DECISION,
+  TRACE_ID,
 } = require('./constants/tags')
 const { UNSERIALIZABLE_VALUE_TEXT } = require('./constants/text')
 const telemetry = require('./telemetry')
@@ -236,8 +237,11 @@ class LLMObsSpanProcessor {
       meta.input.prompt = prompt
     }
 
+    const apmTraceId = span.context().toTraceId(true)
+    const llmobsTraceId = mlObsTags[TRACE_ID] ?? apmTraceId
+
     const llmObsSpanEvent = {
-      trace_id: span.context().toTraceId(true),
+      trace_id: llmobsTraceId,
       span_id: span.context().toSpanId(),
       parent_id: parentId,
       name,
@@ -249,9 +253,10 @@ class LLMObsSpanProcessor {
       metrics,
       _dd: {
         span_id: span.context().toSpanId(),
-        trace_id: span.context().toTraceId(true),
+        trace_id: apmTraceId,
         sample_rate: mlObsTags[SAMPLE_RATE],
         sampling_decision: mlObsTags[SAMPLING_DECISION],
+        apm_trace_id: apmTraceId,
       },
     }
 

packages/dd-trace/src/llmobs/tagger.js

@@ -51,9 +51,18 @@ const {
   SAMPLING_DECISION_DROPPED,
   PROPAGATED_SAMPLE_RATE_KEY,
   PROPAGATED_SAMPLING_DECISION_KEY,
+  TRACE_ID,
+  PROPAGATED_TRACE_ID_KEY,
 } = require('./constants/tags')
 const { storage } = require('./storage')
-const { findGenAIAncestorSpanId, validateCostTags, writeBridgeTags, validateToolDefinitions } = require('./util')
+const {
+  findGenAIAncestorSpanId,
+  validateCostTags,
+  writeBridgeTags,
+  validateToolDefinitions,
+  generateLlmObsTraceId,
+  normalizeLlmObsTraceId,
+} = require('./util')
 
 // global registry of LLMObs spans
 // maps LLMObs spans to their annotations
@@ -126,12 +135,20 @@ class LLMObsTagger {
 
     this._register(span)
 
+    const traceTags = span.context()._trace.tags
+
+    const llmobsTraceId =
+      registry.get(parent)?.[TRACE_ID] ??
+      normalizeLlmObsTraceId(traceTags[PROPAGATED_TRACE_ID_KEY]) ??
+      generateLlmObsTraceId(span._startTime)
+    this._setTag(span, TRACE_ID, llmobsTraceId)
+
     // When the registering span sits below an OTel `gen_ai.*` ancestor, use
     // that ancestor as the parent_id fallback and suppress the bridge
     // parent_id tag so the indexer doesn't invert the trace.
     const genAIAncestorSpanId = findGenAIAncestorSpanId(span)
 
-    writeBridgeTags(span, { includeParentId: genAIAncestorSpanId === null })
+    writeBridgeTags(span, { includeParentId: genAIAncestorSpanId === null, llmobsTraceId })
 
     this._setTag(span, ML_APP, spanMlApp)
 
@@ -141,7 +158,6 @@ class LLMObsTagger {
     if (modelName) this.tagModelName(span, modelName)
     if (modelProvider) this._setTag(span, MODEL_PROVIDER, modelProvider)
 
-    const traceTags = span.context()._trace.tags
     sessionId = sessionId ||
       registry.get(parent)?.[SESSION_ID] ||
       traceTags[SESSION_ID_TRACE_DEFAULT_KEY] ||

packages/dd-trace/src/llmobs/util.js

@@ -1,12 +1,17 @@
 'use strict'
 
+const id = require('../id')
 const log = require('../log')
 const {
   LLMOBS_PARENT_ID_BRIDGE_KEY,
   LLMOBS_TRACE_ID_BRIDGE_KEY,
   SPAN_KINDS,
 } = require('./constants/tags')
 
+const DECIMAL_TRACE_ID_REGEX = /^\d+$/
+const HEX_TRACE_ID_REGEX = /^[0-9a-f]{32}$/i
+const MAX_UINT_64 = (1n << 64n) - 1n
+
 // LLM I/O is overwhelmingly ASCII (English prompts and code). Walk once
 // looking for the first non-ASCII char; if there is none, hand the input
 // straight back. Otherwise pick up the slow path from the byte that needed
@@ -312,12 +317,12 @@ function safeJsonParse (value, fallback) {
 // LLMObs root and hoists the gen_ai ancestors under it, inverting the trace.
 /**
  * @param {import('../opentracing/span')} span
- * @param {{ includeParentId?: boolean }} [opts]
+ * @param {{ includeParentId?: boolean, llmobsTraceId?: string }} [opts]
  */
-function writeBridgeTags (span, { includeParentId = true } = {}) {
+function writeBridgeTags (span, { includeParentId = true, llmobsTraceId } = {}) {
   const traceTags = span?.context?.()._trace?.tags
   if (!traceTags || traceTags[LLMOBS_TRACE_ID_BRIDGE_KEY]) return
-  traceTags[LLMOBS_TRACE_ID_BRIDGE_KEY] = span.context().toTraceId(true)
+  traceTags[LLMOBS_TRACE_ID_BRIDGE_KEY] = llmobsTraceId ?? span.context().toTraceId(true)
   if (includeParentId) {
     traceTags[LLMOBS_PARENT_ID_BRIDGE_KEY] = span.context().toSpanId()
   }
@@ -362,6 +367,51 @@ function findGenAIAncestorSpanId (span) {
   return null
 }
 
+/**
+ * Generate a 128-bit LLMObs trace ID with the span start time encoded in its high bits.
+ * @param {number} startTime
+ * @returns {string}
+ */
+function generateLlmObsTraceId (startTime) {
+  const identifier = id()
+  const traceIdHigh = Math.floor(startTime / 1000)
+    .toString(16)
+    .padStart(8, '0')
+    .padEnd(16, '0')
+
+  return identifier.toTraceIdHex(traceIdHigh).padStart(32, '0')
+}
+
+/**
+ * Convert an internally stored hexadecimal LLMObs trace ID to its distributed wire representation.
+ * @param {string | undefined} traceId
+ * @returns {string | undefined}
+ */
+function llmObsTraceIdToWire (traceId) {
+  if (!traceId) return
+  if (!HEX_TRACE_ID_REGEX.test(traceId)) return traceId
+
+  return BigInt(`0x${traceId}`).toString(10)
+}
+
+/**
+ * Normalize a distributed LLMObs trace ID to the representation expected by LLMObs span events.
+ * @param {string | undefined} traceId
+ * @returns {string | undefined}
+ */
+function normalizeLlmObsTraceId (traceId) {
+  if (!traceId) return
+
+  if (HEX_TRACE_ID_REGEX.test(traceId) && (traceId[0] === '0' || !DECIMAL_TRACE_ID_REGEX.test(traceId))) {
+    return traceId
+  }
+
+  if (!DECIMAL_TRACE_ID_REGEX.test(traceId)) return traceId
+
+  const identifier = BigInt(traceId)
+  return identifier > MAX_UINT_64 ? identifier.toString(16).padStart(32, '0') : traceId
+}
+
 // Maps an audio `format` (e.g. "wav", "mp3") to a MIME type. Defaults to `audio/wav` when the
 // format is missing. Provider-specific overrides (e.g. OpenAI's mp3 -> audio/mpeg) are passed in
 // via `mimeTypeLookup` so this stays provider-agnostic. A non-string `format` is treated as missing
@@ -396,6 +446,9 @@ module.exports = {
   audioMimeTypeFromFormat,
   encodeUnicode,
   findGenAIAncestorSpanId,
+  generateLlmObsTraceId,
+  llmObsTraceIdToWire,
+  normalizeLlmObsTraceId,
   formatAudioPart,
   validateCostTags,
   validateKind,

packages/dd-trace/test/llmobs/index.spec.js

@@ -10,7 +10,13 @@ const sinon = require('sinon')
 const { DD_MAJOR } = require('../../../../version')
 const { INCOMPATIBLE_INITIALIZATION } = require('../../src/llmobs/constants/text')
 const LLMObsTagger = require('../../src/llmobs/tagger')
-const { SAMPLE_RATE, SAMPLING_DECISION, SESSION_ID } = require('../../src/llmobs/constants/tags')
+const {
+  PROPAGATED_TRACE_ID_KEY,
+  SAMPLE_RATE,
+  SAMPLING_DECISION,
+  SESSION_ID,
+  TRACE_ID,
+} = require('../../src/llmobs/constants/tags')
 const { getConfigFresh } = require('../helpers/config')
 const { removeDestroyHandler } = require('./util')
 
@@ -189,6 +195,51 @@ describe('module', () => {
       )
     })
 
+    it('converts the local LLMObs trace id to decimal for propagation', () => {
+      llmobsModule.enable({ llmobs: { mlApp: 'test', agentlessEnabled: false } })
+      store.span = {
+        context () {
+          return {
+            _trace: { tags: {} },
+            toSpanId () { return 'parent-id' },
+          }
+        },
+      }
+      LLMObsTagger.tagMap.set(store.span, {
+        [TRACE_ID]: '6a5f76e7000000001973227978d8110b',
+      })
+
+      const carrier = { 'x-datadog-tags': '' }
+      injectCh.publish({ carrier })
+
+      assert.strictEqual(
+        carrier['x-datadog-tags'],
+        // eslint-disable-next-line @stylistic/max-len
+        '_dd.p.llmobs_parent_id=parent-id,_dd.p.llmobs_ml_app=test,_dd.p.llmobs_trace_id=141393847380800662846519802803680448779'
+      )
+    })
+
+    it('forwards an extracted LLMObs trace id without reinterpreting it', () => {
+      llmobsModule.enable({ llmobs: { mlApp: 'test', agentlessEnabled: false } })
+      const wireTraceId = '12345678901234567890123456789012'
+      store.span = {
+        context () {
+          return {
+            _trace: { tags: { [PROPAGATED_TRACE_ID_KEY]: wireTraceId } },
+            toSpanId () { return 'parent-id' },
+          }
+        },
+      }
+
+      const carrier = { 'x-datadog-tags': '' }
+      injectCh.publish({ carrier })
+
+      assert.strictEqual(
+        carrier['x-datadog-tags'],
+        `_dd.p.llmobs_parent_id=parent-id,_dd.p.llmobs_ml_app=test,_dd.p.llmobs_trace_id=${wireTraceId}`
+      )
+    })
+
     it('does not inject LLMObs parent ID info when there is no parent LLMObs span', () => {
       llmobsModule.enable({ llmobs: { mlApp: 'test', agentlessEnabled: false } })