Propagate LLM Observability context across service boundaries

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

@@ -0,0 +1,62 @@
+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());
+  }
+
+  @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,23 @@ 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.
+      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 +255,7 @@ public DDLLMObsSpan(
     scope =
         LLMObsContext.attach(
             span.spanContext(),
+            mlApp,
             sessionId,
             resolvedAgentVersion,
             sampleRate,
@@ -717,4 +737,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();
+  }
 }

dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagatorTest.java

@@ -0,0 +1,197 @@
+package datadog.trace.llmobs;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import datadog.context.Context;
+import datadog.context.propagation.Propagators;
+import datadog.trace.agent.tooling.TracerInstaller;
+import datadog.trace.api.WellKnownTags;
+import datadog.trace.api.llmobs.LLMObsContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.core.CoreTracer;
+import datadog.trace.llmobs.domain.DDLLMObsSpan;
+import java.util.HashMap;
+import java.util.Map;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Covers automatic LLM Observability context propagation: no LLMObs propagation API is called
+ * anywhere in these tests. Injecting the active span the way auto-instrumentation does — an HTTP
+ * client, or the SQS interceptor writing message attributes — must carry the LLMObs context.
+ *
+ * <p>The carrier is a plain {@code Map<String, String>}, which is the shape both an HTTP header map
+ * and the SQS {@code _datadog} message attribute reduce to at the propagator boundary.
+ */
+class LLMObsContextPropagatorTest {
+
+  private static final String ML_APP_TAG = "_dd.p.llmobs_ml_app";
+  private static final String SESSION_ID_TAG = "_dd.p.llmobs_sid";
+  private static final String PAGENT_SPAN_ID_TAG = "_dd.p.llmobs_pagent_span_id";
+  private static final String PAGENT_NAME_TAG = "_dd.p.llmobs_pagent_name";
+
+  private static CoreTracer tracer;
+
+  @BeforeAll
+  static void installTracer() {
+    tracer = CoreTracer.builder().build();
+    TracerInstaller.forceInstallGlobalTracer(tracer);
+    Propagators.register(AgentPropagation.LLMOBS_CONCERN, new LLMObsContextPropagator());
+  }
+
+  @AfterAll
+  static void closeTracer() {
+    TracerInstaller.forceInstallGlobalTracer(null);
+    tracer.close();
+  }
+
+  private static DDLLMObsSpan newSpan(String kind, String name, String mlApp, String sessionId) {
+    WellKnownTags tags =
+        new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java");
+    return new DDLLMObsSpan(kind, name, mlApp, sessionId, "service", tags);
+  }
+
+  private static AgentScope startRootApmScope() {
+    AgentSpan root = AgentTracer.get().buildSpan("apm", "sqs.produce").start();
+    return AgentTracer.activateSpan(root);
+  }
+
+  /** What an auto-instrumented client does: inject the active span into an outbound carrier. */
+  private static Map<String, String> autoInject(AgentSpan span) {
+    Map<String, String> carrier = new HashMap<>();
+    Propagators.defaultPropagator().inject(span, carrier, Map::put);
+    return carrier;
+  }
+
+  @Test
+  void stagesLlmObsTagsOnInjectionWithoutAnyManualPropagation() {
+    Map<String, String> carrier;
+    String agentSpanId;
+    try (AgentScope apmScope = startRootApmScope()) {
+      DDLLMObsSpan agent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "planner", "my-ml-app", "sess-1");
+      agentSpanId = String.valueOf(agent.getSpanId());
+      try {
+        carrier = autoInject(AgentTracer.activeSpan());
+      } finally {
+        agent.finish();
+      }
+    }
+
+    String tags = carrier.get("x-datadog-tags");
+    assertNotNull(tags, "expected x-datadog-tags to be injected");
+    assertTrue(tags.contains(ML_APP_TAG + "=my-ml-app"), () -> "ml_app missing from " + tags);
+    assertTrue(tags.contains(SESSION_ID_TAG + "=sess-1"), () -> "session_id missing from " + tags);
+    assertTrue(
+        tags.contains(PAGENT_SPAN_ID_TAG + "=" + agentSpanId),
+        () -> "pagent_span_id missing from " + tags);
+    assertTrue(
+        tags.contains(PAGENT_NAME_TAG + "=planner"), () -> "pagent_name missing from " + tags);
+  }
+
+  @Test
+  void addsNothingWhenNoLlmObsSpanIsActive() {
+    Map<String, String> carrier;
+    try (AgentScope apmScope = startRootApmScope()) {
+      carrier = autoInject(apmScope.span());
+    }
+
+    String tags = carrier.get("x-datadog-tags");
+    assertTrue(
+        tags == null || !tags.contains("_dd.p.llmobs_"), () -> "unexpected LLMObs tags in " + tags);
+  }
+
+  @Test
+  void stopsContributingTagsOnceTheLlmObsScopeIsClosed() {
+    Map<String, String> carrier;
+    try (AgentScope apmScope = startRootApmScope()) {
+      newSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "work", "my-ml-app", "sess-1").finish();
+      // The LLMObs span has finished; a later outbound call on the same APM trace must not be
+      // tagged with a session that is no longer active.
+      carrier = autoInject(apmScope.span());
+    }
+
+    String tags = carrier.get("x-datadog-tags");
+    assertTrue(
+        tags == null || !tags.contains(SESSION_ID_TAG),
+        () -> "session_id leaked after scope close: " + tags);
+  }
+
+  /**
+   * The full cross-process hop, as an SQS producer/worker pair sees it: the producer injects into
+   * message attributes, the worker extracts and activates them, and an LLMObs span started by the
+   * worker inherits the session and agent attribution without any application-level plumbing.
+   */
+  @Test
+  void workerInheritsSessionAndAgentAttributionAcrossTheBoundary() {
+    Map<String, String> messageAttributes;
+    long producerTraceId;
+    String producerAgentSpanId;
+
+    try (AgentScope apmScope = startRootApmScope()) {
+      DDLLMObsSpan producer =
+          newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "dispatcher", "my-ml-app", "sess-42");
+      producerTraceId = producer.getTraceId().toLong();
+      producerAgentSpanId = String.valueOf(producer.getSpanId());
+      try {
+        messageAttributes = autoInject(AgentTracer.activeSpan());
+      } finally {
+        producer.finish();
+      }
+    }
+
+    // Worker side: a fresh context, as a message handler would have.
+    Context extracted =
+        Propagators.defaultPropagator()
+            .extract(
+                Context.root(), messageAttributes, (carrier, visitor) -> carrier.forEach(visitor));
+    AgentSpan consumeSpan = AgentSpan.fromContext(extracted);
+    assertNotNull(consumeSpan, "expected trace context to be extracted");
+
+    try (AgentScope consumeScope = AgentTracer.get().activateSpan(consumeSpan)) {
+      DDLLMObsSpan workerTool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", "my-ml-app", null);
+      try {
+        assertEquals(producerTraceId, workerTool.getTraceId().toLong(), "trace should be joined");
+        // The span publishes its resolved values to the context for its own descendants, so this
+        // is what the worker's LLMObs span actually settled on.
+        assertEquals("sess-42", LLMObsContext.currentSessionId());
+        assertEquals(producerAgentSpanId, LLMObsContext.currentParentAgentSpanId());
+        assertEquals("dispatcher", LLMObsContext.currentParentAgentName());
+      } finally {
+        workerTool.finish();
+      }
+    }
+  }
+
+  @Test
+  void workerWithoutUpstreamLlmObsContextInheritsNothing() {
+    Map<String, String> messageAttributes;
+    try (AgentScope apmScope = startRootApmScope()) {
+      messageAttributes = autoInject(apmScope.span());
+    }
+
+    Context extracted =
+        Propagators.defaultPropagator()
+            .extract(
+                Context.root(), messageAttributes, (carrier, visitor) -> carrier.forEach(visitor));
+    AgentSpan consumeSpan = AgentSpan.fromContext(extracted);
+    assertNotNull(consumeSpan, "expected trace context to be extracted");
+
+    try (AgentScope consumeScope = AgentTracer.get().activateSpan(consumeSpan)) {
+      DDLLMObsSpan workerTool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", "my-ml-app", null);
+      try {
+        assertNull(LLMObsContext.currentSessionId());
+        assertNull(LLMObsContext.currentParentAgentSpanId());
+      } finally {
+        workerTool.finish();
+      }
+    }
+  }
+}

dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java

@@ -1492,6 +1492,46 @@ public PropagationTags getPropagationTags() {
     return getRootSpanContextOrThis().propagationTags;
   }
 
+  @Override
+  public CharSequence getLLMObsMlApp() {
+    return getPropagationTags().getLLMObsMlApp();
+  }
+
+  @Override
+  public void updateLLMObsMlApp(CharSequence mlApp) {
+    getPropagationTags().updateLLMObsMlApp(mlApp);
+  }
+
+  @Override
+  public CharSequence getLLMObsSessionId() {
+    return getPropagationTags().getLLMObsSessionId();
+  }
+
+  @Override
+  public void updateLLMObsSessionId(CharSequence sessionId) {
+    getPropagationTags().updateLLMObsSessionId(sessionId);
+  }
+
+  @Override
+  public CharSequence getLLMObsParentAgentSpanId() {
+    return getPropagationTags().getLLMObsParentAgentSpanId();
+  }
+
+  @Override
+  public void updateLLMObsParentAgentSpanId(CharSequence parentAgentSpanId) {
+    getPropagationTags().updateLLMObsParentAgentSpanId(parentAgentSpanId);
+  }
+
+  @Override
+  public CharSequence getLLMObsParentAgentName() {
+    return getPropagationTags().getLLMObsParentAgentName();
+  }
+
+  @Override
+  public void updateLLMObsParentAgentName(CharSequence parentAgentName) {
+    getPropagationTags().updateLLMObsParentAgentName(parentAgentName);
+  }
+
   /** TraceSegment Implementation */
   @Override
   public void setTagTop(String key, Object value, boolean sanitize) {