Propagate LLM Observability context across service boundaries

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java

@@ -24,6 +24,10 @@ public static void forEachProperty(AgentPropagation.KeyClassifier classifier, St
       if (acceptJsonProperty(classifier, json, "x-datadog-trace-id")) {
         acceptJsonProperty(classifier, json, "x-datadog-parent-id");
         acceptJsonProperty(classifier, json, "x-datadog-sampling-priority");
+        // Propagation tags travel in x-datadog-tags. Without this the whole _dd.p.* set is
+        // silently dropped at a messaging boundary — including _dd.p.tid, which truncates a
+        // 128-bit trace id to 64 bits downstream, and the _dd.p.llmobs_* attribution tags.
+        acceptJsonProperty(classifier, json, "x-datadog-tags");
       }
       if (Config.get().isDataStreamsEnabled()) {
         acceptJsonProperty(classifier, json, "dd-pathway-ctx-base64");

dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParserTest.java

@@ -0,0 +1,119 @@
+package datadog.trace.bootstrap.instrumentation.messaging;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.ByteBuffer;
+import java.util.Base64;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Covers the {@code _datadog} message attribute parser shared by the AWS messaging instrumentations
+ * (SQS, SNS, EventBridge, Step Functions).
+ */
+class DatadogAttributeParserTest {
+
+  /** What an injected {@code _datadog} attribute looks like on the wire. */
+  private static final String FULL_CONTEXT =
+      "{\"x-datadog-trace-id\":\"1234567890\","
+          + "\"x-datadog-parent-id\":\"9876543210\","
+          + "\"x-datadog-sampling-priority\":\"1\","
+          + "\"x-datadog-tags\":\"_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000\","
+          + "\"traceparent\":\"00-6aa01c5400000000499602d2-000000024cb016ea-01\"}";
+
+  private static Map<String, String> parse(String json) {
+    Map<String, String> collected = new LinkedHashMap<>();
+    DatadogAttributeParser.forEachProperty(
+        (key, value) -> {
+          collected.put(key, value);
+          return true;
+        },
+        json);
+    return collected;
+  }
+
+  @Test
+  void extractsTraceContextAndPropagationTags() {
+    Map<String, String> collected = parse(FULL_CONTEXT);
+
+    assertEquals("1234567890", collected.get("x-datadog-trace-id"));
+    assertEquals("9876543210", collected.get("x-datadog-parent-id"));
+    assertEquals("1", collected.get("x-datadog-sampling-priority"));
+    // Without x-datadog-tags the whole _dd.p.* set is dropped at the messaging boundary: the
+    // 64-bit trace id still joins, but _dd.p.tid is lost so the two services disagree about the
+    // full 128-bit id, and _dd.p.dm is lost with it.
+    assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags"));
+  }
+
+  @Test
+  void extractsPropagationTagsFromByteBufferCarrier() {
+    Map<String, String> collected = new LinkedHashMap<>();
+    DatadogAttributeParser.forEachProperty(
+        (key, value) -> {
+          collected.put(key, value);
+          return true;
+        },
+        ByteBuffer.wrap(FULL_CONTEXT.getBytes(UTF_8)));
+
+    assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags"));
+  }
+
+  @Test
+  void extractsPropagationTagsFromBase64ByteBufferCarrier() {
+    Map<String, String> collected = new LinkedHashMap<>();
+    DatadogAttributeParser.forEachProperty(
+        (key, value) -> {
+          collected.put(key, value);
+          return true;
+        },
+        ByteBuffer.wrap(Base64.getEncoder().encode(FULL_CONTEXT.getBytes(UTF_8))));
+
+    assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags"));
+  }
+
+  @Test
+  void carriesLlmObsPropagationTags() {
+    Map<String, String> collected =
+        parse(
+            "{\"x-datadog-trace-id\":\"1234567890\","
+                + "\"x-datadog-parent-id\":\"9876543210\","
+                + "\"x-datadog-tags\":\"_dd.p.llmobs_ml_app=my-app,_dd.p.llmobs_sid=sess-1,"
+                + "_dd.p.llmobs_parent_id=42\"}");
+
+    String tags = collected.get("x-datadog-tags");
+    assertTrue(tags.contains("_dd.p.llmobs_ml_app=my-app"), tags);
+    assertTrue(tags.contains("_dd.p.llmobs_sid=sess-1"), tags);
+    assertTrue(tags.contains("_dd.p.llmobs_parent_id=42"), tags);
+  }
+
+  @Test
+  void extractsNothingWithoutATraceId() {
+    // Propagation tags on their own describe no trace, so they are not surfaced.
+    Map<String, String> collected =
+        parse("{\"x-datadog-tags\":\"_dd.p.dm=-1\",\"x-datadog-parent-id\":\"9876543210\"}");
+
+    assertTrue(collected.isEmpty(), () -> "expected nothing extracted, got " + collected);
+  }
+
+  @Test
+  void toleratesAMissingTagsProperty() {
+    Map<String, String> collected =
+        parse(
+            "{\"x-datadog-trace-id\":\"1234567890\",\"x-datadog-parent-id\":\"9876543210\","
+                + "\"x-datadog-sampling-priority\":\"1\"}");
+
+    assertEquals("1234567890", collected.get("x-datadog-trace-id"));
+    assertNull(collected.get("x-datadog-tags"));
+  }
+
+  @Test
+  void toleratesMalformedJson() {
+    assertTrue(parse("not json at all").isEmpty());
+    assertTrue(parse("{\"x-datadog-trace-id\":").isEmpty());
+    assertTrue(parse(null).isEmpty());
+  }
+}

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagator.java

@@ -0,0 +1,66 @@
+package datadog.trace.llmobs;
+
+import datadog.context.Context;
+import datadog.context.propagation.CarrierSetter;
+import datadog.context.propagation.CarrierVisitor;
+import datadog.context.propagation.Propagator;
+import datadog.trace.api.llmobs.LLMObsContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+
+/**
+ * Stages the LLM Observability propagation tags onto the span context being injected, so that every
+ * boundary already covered by automatic instrumentation — HTTP, gRPC, SQS, Kafka, ... — carries
+ * LLMObs context without the application having to propagate it by hand.
+ *
+ * <p>This propagator writes nothing to the carrier itself. It runs ahead of the tracing propagator
+ * (see {@code AgentPropagation.LLMOBS_CONCERN}) and only populates the {@code _dd.p.llmobs_*}
+ * fields on the span context; the tracing propagator then serializes them into {@code
+ * x-datadog-tags} / {@code tracestate} along with every other propagation tag. This mirrors
+ * dd-trace-py, where LLMObs subscribes to the generic {@code http.span_inject} hook that {@code
+ * HTTPPropagator.inject} fires on every outbound request, rather than owning a separate wire
+ * format.
+ *
+ * <p>Values are resolved from the ambient {@link LLMObsContext} at injection time rather than being
+ * written once when a span starts. That way the innermost active LLMObs span always wins, and
+ * leaving an LLMObs scope stops contributing its tags without any save/restore bookkeeping.
+ */
+public class LLMObsContextPropagator implements Propagator {
+
+  @Override
+  public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
+    AgentSpan span = AgentSpan.fromContext(context);
+    if (span == null) {
+      return;
+    }
+    AgentSpanContext spanContext = span.spanContext();
+    if (spanContext == null) {
+      return;
+    }
+
+    // Gate on trace-id consistency, the same way DDLLMObsSpan gates parent_id/session_id
+    // inheritance. An LLMObs context leaked across an async boundary must not tag an outbound
+    // request that belongs to an unrelated trace.
+    AgentSpanContext llmObsContext = LLMObsContext.current();
+    if (llmObsContext == null || llmObsContext.getTraceId() != spanContext.getTraceId()) {
+      return;
+    }
+
+    spanContext.updateLLMObsMlApp(LLMObsContext.currentMlApp());
+    spanContext.updateLLMObsSessionId(LLMObsContext.currentSessionId());
+    spanContext.updateLLMObsParentAgentSpanId(LLMObsContext.currentParentAgentSpanId());
+    spanContext.updateLLMObsParentAgentName(LLMObsContext.currentParentAgentName());
+    // The innermost active LLMObs span becomes the downstream span's LLMObs parent, so the
+    // continued trace is a single tree rather than a second root per service. Mirrors
+    // dd-trace-py's _dd.p.llmobs_parent_id.
+    spanContext.updateLLMObsParentId(String.valueOf(llmObsContext.getSpanId()));
+  }
+
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    // Nothing to do: the tracing propagator's codecs already parse the _dd.p.llmobs_* tags back
+    // into the extracted context's propagation tags, and DDLLMObsSpan reads them from there when
+    // no in-process LLMObs parent applies.
+    return context;
+  }
+}

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java

@@ -1,13 +1,15 @@
 package datadog.trace.llmobs;
 
 import datadog.communication.ddagent.SharedCommunicationObjects;
+import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
 import datadog.trace.api.WellKnownTags;
 import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsInternal;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
 import datadog.trace.api.telemetry.LLMObsMetricCollector;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.llmobs.domain.DDLLMObsSpan;
 import datadog.trace.llmobs.domain.LLMObsEval;
@@ -51,6 +53,10 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
     LLMObsInternal.setEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config));
 
     LLMObsInternal.setFeedbackProcessor(new LLMObsCustomFeedbackProcessor(mlApp, sco, config));
+
+    // Carry LLMObs context across every boundary automatic instrumentation already covers, by
+    // staging the _dd.p.llmobs_* tags on each injected span context. See LLMObsContextPropagator.
+    Propagators.register(AgentPropagation.LLMOBS_CONCERN, new LLMObsContextPropagator());
   }
 
   private static class LLMObsCustomFeedbackProcessor implements LLMObs.LLMObsFeedbackProcessor {

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java

@@ -159,6 +159,7 @@ public DDLLMObsSpan(
     String samplingDecision = null;
     String resolvedParentAgentSpanId = null;
     String resolvedParentAgentName = null;
+    boolean inheritedInProcess = false;
     if (null != parent) {
       if (parent.getTraceId() != span.getTraceId()) {
         LOGGER.error(
@@ -168,6 +169,7 @@ public DDLLMObsSpan(
             span.getTraceId(),
             span.getSpanId());
       } else {
+        inheritedInProcess = true;
         parentSpanID = String.valueOf(parent.getSpanId());
         // Inherit session_id from parent context only when it belongs to the same trace.
         // Matches dd-trace-py and dd-trace-js: session_id need only be set on the root
@@ -197,6 +199,30 @@ public DDLLMObsSpan(
       }
     }
 
+    if (!inheritedInProcess) {
+      // No usable in-process LLMObs parent, but this span may still be continuing a trace that
+      // arrived from another service — an SQS worker handling a message, an inbound HTTP request.
+      // The upstream values are on the span context's propagation tags, parsed back out of
+      // x-datadog-tags / tracestate by the tracing propagator.
+      //
+      // This also covers the trace-mismatch branch above: a stale context leaked from an unrelated
+      // trace must not suppress attribution that legitimately arrived over the wire.
+      // Unlike dd-trace-py, a missing parent_id doesn't veto the rest: session and attribution
+      // are inherited independently, so a partially populated upstream still contributes what it
+      // did send.
+      String propagatedParentId = asString(span.spanContext().getLLMObsParentId());
+      if (propagatedParentId != null) {
+        parentSpanID = propagatedParentId;
+      }
+      if (sessionId == null || sessionId.isEmpty()) {
+        sessionId = asString(span.spanContext().getLLMObsSessionId());
+      }
+      resolvedParentAgentSpanId = asString(span.spanContext().getLLMObsParentAgentSpanId());
+      if (resolvedParentAgentSpanId != null) {
+        resolvedParentAgentName = asString(span.spanContext().getLLMObsParentAgentName());
+      }
+    }
+
     // An agent span is its own descendants' nearest agent ancestor, replacing anything inherited.
     // Use the span name as the initial pagent name; annotateAgentManifest() will update it to the
     // manifest name if one is provided later.
@@ -236,6 +262,7 @@ public DDLLMObsSpan(
     scope =
         LLMObsContext.attach(
             span.spanContext(),
+            mlApp,
             sessionId,
             resolvedAgentVersion,
             sampleRate,
@@ -717,4 +744,9 @@ public DDTraceId getTraceId() {
   public long getSpanId() {
     return span.getSpanId();
   }
+
+  /** Narrow a propagated tag value to a non-empty String, or null. */
+  private static String asString(CharSequence value) {
+    return value == null || value.length() == 0 ? null : value.toString();
+  }
 }