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,31 @@
+package datadog.trace.bootstrap.instrumentation.messaging;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Covers the one behaviour this branch adds to the {@code _datadog} message attribute parser shared
+ * by the AWS messaging instrumentations (SQS, SNS, EventBridge, Step Functions): {@code
+ * x-datadog-tags} is forwarded to the extractor.
+ */
+class DatadogAttributeParserTest {
+
+  @Test
+  void forwardsPropagationTags() {
+    Map<String, String> collected = new LinkedHashMap<>();
+    DatadogAttributeParser.forEachProperty(
+        (key, value) -> {
+          collected.put(key, value);
+          return true;
+        },
+        "{\"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\"}");
+
+    assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags"));
+  }
+}

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

@@ -0,0 +1,71 @@
+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 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.
+ *
+ * <p>Values are resolved from the ambient {@link LLMObsContext} at injection time rather than being
+ * written once when a span starts, and every injection rewrites the whole set — falling back to
+ * whatever arrived on the inbound headers when no LLMObs context applies. That way the innermost
+ * active LLMObs span always wins, leaving an LLMObs scope stops contributing its tags, and a
+ * service that opens no LLMObs span of its own still forwards its caller's context — all 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().equals(spanContext.getTraceId())) {
+      // Reset rather than return. These tags are staged on the root span context's propagation
+      // tags, which the whole local trace shares, so anything an earlier injection wrote would
+      // otherwise ride along on this one too — shipping a session and an agent attribution that
+      // are no longer active. Reset restores the extracted values instead of clearing outright:
+      // the same object also holds what came in on the wire, and a pass-through service must keep
+      // forwarding its caller's context.
+      spanContext.resetLLMObsContext();
+      return;
+    }
+
+    spanContext.updateLLMObsContext(
+        LLMObsContext.currentMlApp(),
+        LLMObsContext.currentSessionId(),
+        LLMObsContext.currentParentAgentSpanId(),
+        LLMObsContext.currentParentAgentName(),
+        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;
@@ -46,11 +48,15 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
 
     String mlApp = config.getLlmObsMlApp();
     WellKnownTags wellKnownTags = config.getWellKnownTags();
-    LLMObsInternal.setSpanFactory(new LLMObsManualSpanFactory(mlApp, wellKnownTags));
+    // The span factory deliberately gets no default ml_app: DDLLMObsSpan applies it last, after
+    // in-process and propagated values have had their chance.
+    LLMObsInternal.setSpanFactory(new LLMObsManualSpanFactory(wellKnownTags));
 
     LLMObsInternal.setEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config));
 
     LLMObsInternal.setFeedbackProcessor(new LLMObsCustomFeedbackProcessor(mlApp, sco, config));
+
+    Propagators.register(AgentPropagation.LLMOBS_CONCERN, new LLMObsContextPropagator());
   }
 
   private static class LLMObsCustomFeedbackProcessor implements LLMObs.LLMObsFeedbackProcessor {
@@ -219,12 +225,10 @@ public void SubmitEvaluation(
 
   private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory {
 
-    private final String defaultMLApp;
     private final String serviceName;
     private final WellKnownTags wellKnownTags;
 
-    public LLMObsManualSpanFactory(String defaultMLApp, WellKnownTags wellKnownTags) {
-      this.defaultMLApp = defaultMLApp;
+    public LLMObsManualSpanFactory(WellKnownTags wellKnownTags) {
       this.serviceName = wellKnownTags.getService().toString();
       this.wellKnownTags = wellKnownTags;
     }
@@ -239,12 +243,7 @@ public LLMObsSpan startLLMSpan(
 
       DDLLMObsSpan span =
           new DDLLMObsSpan(
-              Tags.LLMOBS_LLM_SPAN_KIND,
-              spanName,
-              getMLApp(mlApp),
-              sessionId,
-              serviceName,
-              wellKnownTags);
+              Tags.LLMOBS_LLM_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags);
 
       if (modelName == null || modelName.isEmpty()) {
         modelName = CUSTOM_MODEL_VAL;
@@ -273,7 +272,7 @@ public LLMObsSpan startAgentSpan(
       return new DDLLMObsSpan(
           Tags.LLMOBS_AGENT_SPAN_KIND,
           spanName,
-          getMLApp(mlApp),
+          mlApp,
           sessionId,
           serviceName,
           wellKnownTags,
@@ -284,36 +283,21 @@ public LLMObsSpan startAgentSpan(
     public LLMObsSpan startToolSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_TOOL_SPAN_KIND,
-          spanName,
-          getMLApp(mlApp),
-          sessionId,
-          serviceName,
-          wellKnownTags);
+          Tags.LLMOBS_TOOL_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags);
     }
 
     @Override
     public LLMObsSpan startTaskSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_TASK_SPAN_KIND,
-          spanName,
-          getMLApp(mlApp),
-          sessionId,
-          serviceName,
-          wellKnownTags);
+          Tags.LLMOBS_TASK_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags);
     }
 
     @Override
     public LLMObsSpan startWorkflowSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_WORKFLOW_SPAN_KIND,
-          spanName,
-          getMLApp(mlApp),
-          sessionId,
-          serviceName,
-          wellKnownTags);
+          Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags);
     }
 
     @Override
@@ -330,7 +314,7 @@ public LLMObsSpan startEmbeddingSpan(
           new DDLLMObsSpan(
               Tags.LLMOBS_EMBEDDING_SPAN_KIND,
               spanName,
-              getMLApp(mlApp),
+              mlApp,
               sessionId,
               serviceName,
               wellKnownTags);
@@ -342,19 +326,7 @@ public LLMObsSpan startEmbeddingSpan(
     public LLMObsSpan startRetrievalSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
       return new DDLLMObsSpan(
-          Tags.LLMOBS_RETRIEVAL_SPAN_KIND,
-          spanName,
-          getMLApp(mlApp),
-          sessionId,
-          serviceName,
-          wellKnownTags);
-    }
-
-    private String getMLApp(String mlApp) {
-      if (mlApp == null || mlApp.isEmpty()) {
-        return defaultMLApp;
-      }
-      return mlApp;
+          Tags.LLMOBS_RETRIEVAL_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags);
     }
   }
 }