Propagate LLM Observability context (session, ml_app, agent attribution) across process boundaries

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

@@ -0,0 +1,102 @@
+package datadog.trace.llmobs;
+
+import datadog.context.Context;
+import datadog.context.ContextScope;
+import datadog.context.propagation.Propagators;
+import datadog.trace.api.llmobs.LLMObs;
+import datadog.trace.api.llmobs.LLMObsContext;
+import datadog.trace.api.llmobs.LLMObsSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.llmobs.domain.DDLLMObsSpan;
+import java.io.Closeable;
+import java.util.Map;
+import java.util.Objects;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Explicit, manual distributed tracing propagation for LLM Observability, for boundaries that
+ * automatic instrumentation doesn't cover — e.g. an SQS worker reading its own message attributes.
+ *
+ * <p>Boundaries that <em>are</em> auto-instrumented (HTTP, gRPC, ...) need none of this: {@link
+ * LLMObsContextPropagator} is registered as a propagation concern and stages the same {@code
+ * _dd.p.llmobs_*} tags on every injection. This class is the manual equivalent for carriers no
+ * instrumentation reaches, mirroring dd-trace-py's {@code inject_distributed_headers} / {@code
+ * activate_distributed_headers}.
+ *
+ * <p>The standard APM trace context (trace id, parent id, sampling, {@code x-datadog-tags}, ...) is
+ * injected/extracted via the normal {@link Propagators}. LLMObs-specific values (ml_app,
+ * session_id, agent attribution) are written onto the span's {@link AgentSpanContext} as dedicated
+ * propagation-tags fields before injection, so the same {@link Propagators} call serializes them as
+ * additional {@code _dd.p.llmobs_*} tags — via {@code x-datadog-tags} or {@code tracestate},
+ * whichever the configured propagation style carries — the wire container dd-trace-py/js/go already
+ * use for these tags, so a mixed-language pipeline can still join a trace across this hop.
+ */
+public class DDLLMObsPropagator implements LLMObs.LLMObsPropagator {
+  private static final Logger LOGGER = LoggerFactory.getLogger(DDLLMObsPropagator.class);
+
+  @Override
+  public Map<String, String> injectDistributedHeaders(
+      LLMObsSpan span, Map<String, String> headers) {
+    Objects.requireNonNull(span, "span");
+    Objects.requireNonNull(headers, "headers");
+    if (!(span instanceof DDLLMObsSpan)) {
+      LOGGER.debug(
+          "injectDistributedHeaders requires a span started by the LLM Observability SDK, got {}; ignoring",
+          span.getClass());
+      return headers;
+    }
+    DDLLMObsSpan llmObsSpan = (DDLLMObsSpan) span;
+    AgentSpan agentSpan = llmObsSpan.getAgentSpan();
+
+    // Stage this span's own LLMObs values rather than relying on the ambient context: the caller
+    // may hand us a span that isn't the innermost active one, and may not even be inside its scope.
+    AgentSpanContext spanContext = agentSpan.spanContext();
+    spanContext.updateLLMObsMlApp(llmObsSpan.getMlApp());
+    spanContext.updateLLMObsSessionId(llmObsSpan.getSessionId());
+    spanContext.updateLLMObsParentAgentSpanId(llmObsSpan.getParentAgentSpanId());
+    spanContext.updateLLMObsParentAgentName(llmObsSpan.getParentAgentName());
+
+    Propagators.defaultPropagator().inject(agentSpan, headers, Map::put);
+    return headers;
+  }
+
+  @Override
+  public Closeable activateDistributedHeaders(Map<String, String> headers) {
+    Objects.requireNonNull(headers, "headers");
+
+    Context extracted =
+        Propagators.defaultPropagator()
+            .extract(Context.root(), headers, (carrier, visitor) -> carrier.forEach(visitor));
+    AgentSpan extractedSpan = AgentSpan.fromContext(extracted);
+    if (extractedSpan == null) {
+      LOGGER.debug(
+          "no distributed trace context found in headers; activateDistributedHeaders is a no-op");
+      return () -> {};
+    }
+
+    AgentSpanContext extractedContext = extractedSpan.spanContext();
+    CharSequence mlApp = extractedContext.getLLMObsMlApp();
+    CharSequence sessionId = extractedContext.getLLMObsSessionId();
+    CharSequence pagentSpanId = extractedContext.getLLMObsParentAgentSpanId();
+    CharSequence pagentName = extractedContext.getLLMObsParentAgentName();
+
+    AgentScope apmScope = AgentTracer.get().activateSpan(extractedSpan);
+    ContextScope llmObsScope =
+        LLMObsContext.attach(
+            extractedContext,
+            mlApp == null ? null : mlApp.toString(),
+            sessionId == null ? null : sessionId.toString(),
+            null,
+            pagentSpanId == null ? null : pagentSpanId.toString(),
+            pagentName == null ? null : pagentName.toString());
+
+    return () -> {
+      llmObsScope.close();
+      apmScope.close();
+    };
+  }
+}

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,13 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
     LLMObsInternal.setEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config));
 
     LLMObsInternal.setFeedbackProcessor(new LLMObsCustomFeedbackProcessor(mlApp, sco, config));
+
+    LLMObsInternal.setPropagator(new DDLLMObsPropagator());
+
+    // Automatic propagation: every boundary that injects trace context now carries the LLMObs
+    // propagation tags too, matching dd-trace-py. DDLLMObsPropagator stays as the manual entry
+    // point for carriers no instrumentation covers (e.g. SQS message attributes).
+    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

@@ -68,6 +68,9 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private final String spanKind;
   private final String mlApp;
   private final boolean hasSessionId;
+  private final String sessionId;
+  private final String parentAgentSpanId;
+  private final String parentAgentName;
   private final ContextScope scope;
   // Non-null only for agent-kind spans started without an ambient APM root. Activating the
   // agent's APM span keeps children in the same APM trace so the trace-ID gate passes and
@@ -159,7 +162,18 @@ public DDLLMObsSpan(
       }
     }
 
+    // No in-process LLMObs parent: this span may still be continuing a trace that arrived from
+    // another service, in which case the upstream session_id is on the span context's propagation
+    // tags (parsed back out of x-datadog-tags / tracestate by the tracing propagator).
+    if (null == parent && (sessionId == null || sessionId.isEmpty())) {
+      CharSequence propagated = span.spanContext().getLLMObsSessionId();
+      if (propagated != null && propagated.length() > 0) {
+        sessionId = propagated.toString();
+      }
+    }
+
     this.hasSessionId = sessionId != null && !sessionId.isEmpty();
+    this.sessionId = this.hasSessionId ? sessionId : null;
     if (this.hasSessionId) {
       span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
     }
@@ -187,6 +201,14 @@ public DDLLMObsSpan(
       if (null != parent && parent.getTraceId() == span.getTraceId()) {
         resolvedParentAgentSpanId = LLMObsContext.currentParentAgentSpanId();
         resolvedParentAgentName = LLMObsContext.currentParentAgentName();
+      } else if (null == parent) {
+        // Continuing a distributed trace: attribute to the upstream agent carried on the wire.
+        CharSequence propagatedId = span.spanContext().getLLMObsParentAgentSpanId();
+        if (propagatedId != null && propagatedId.length() > 0) {
+          resolvedParentAgentSpanId = propagatedId.toString();
+          CharSequence propagatedName = span.spanContext().getLLMObsParentAgentName();
+          resolvedParentAgentName = propagatedName == null ? null : propagatedName.toString();
+        }
       }
     }
 
@@ -197,11 +219,14 @@ public DDLLMObsSpan(
         span.setTag(PAGENT_NAME_TAG_INTERNAL, resolvedParentAgentName);
       }
     }
+    this.parentAgentSpanId = resolvedParentAgentSpanId;
+    this.parentAgentName = resolvedParentAgentName;
 
     // Propagate the effective sessionId and agent attribution to descendant LLMObs spans.
     scope =
         LLMObsContext.attach(
             span.spanContext(),
+            mlApp,
             sessionId,
             resolvedAgentVersion,
             resolvedParentAgentSpanId,
@@ -681,4 +706,35 @@ public DDTraceId getTraceId() {
   public long getSpanId() {
     return span.getSpanId();
   }
+
+  /** Internal accessor for the underlying APM span, used by {@code DDLLMObsPropagator}. */
+  public AgentSpan getAgentSpan() {
+    return span;
+  }
+
+  /** Internal accessor for this span's effective ml_app, used by {@code DDLLMObsPropagator}. */
+  public String getMlApp() {
+    return mlApp;
+  }
+
+  /**
+   * Internal accessor for this span's effective session_id (including one inherited from an
+   * enclosing LLMObs span), used by {@code DDLLMObsPropagator}. May be {@code null}.
+   */
+  public String getSessionId() {
+    return sessionId;
+  }
+
+  /**
+   * Internal accessor for this span's effective agent attribution, used by {@code
+   * DDLLMObsPropagator}. May be {@code null}.
+   */
+  public String getParentAgentSpanId() {
+    return parentAgentSpanId;
+  }
+
+  /** See {@link #getParentAgentSpanId()}. May be {@code null}. */
+  public String getParentAgentName() {
+    return parentAgentName;
+  }
 }

dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy

@@ -796,6 +796,9 @@ class DDLLMObsSpanTest  extends DDSpecification{
     def innerSpan = (AgentSpan) test.span
     innerSpan.getTag(LLMOBS_TAG_PREFIX + "team") == "backend"
     innerSpan.getTag(LLMOBS_TAG_PREFIX + "owner") == "ml-platform"
+
+    cleanup:
+    test.finish()
   }
 
   def "agent manifest full annotation sets correct tag"() {