feat(llmobs): add agent attribution to Java SDK

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

@@ -8,6 +8,7 @@
 import datadog.trace.api.WellKnownTags;
 import datadog.trace.api.llmobs.LLMObs;
 import datadog.trace.api.llmobs.LLMObsContext;
+import datadog.trace.api.llmobs.LLMObsPropagationAccess;
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
 import datadog.trace.api.telemetry.LLMObsMetricCollector;
@@ -48,6 +49,9 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private static final String CONTEXT_VARIABLE_KEYS = "_dd_context_variable_keys";
   private static final String QUERY_VARIABLE_KEYS = "_dd_query_variable_keys";
   private static final String PARENT_ID_TAG_INTERNAL = "parent_id";
+  private static final String PAGENT_SPAN_ID_TAG_INTERNAL =
+      LLMOBS_TAG_PREFIX + "pagent_span_id";
+  private static final String PAGENT_NAME_TAG_INTERNAL = LLMOBS_TAG_PREFIX + "pagent_name";
 
   private static final String SERVICE = LLMOBS_TAG_PREFIX + "service";
   private static final String VERSION = LLMOBS_TAG_PREFIX + "version";
@@ -133,8 +137,72 @@ public DDLLMObsSpan(
       span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
     }
     span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID);
-    // Propagate the effective sessionId to descendant LLMObs spans via the context.
-    scope = LLMObsContext.attach(span.spanContext(), sessionId);
+
+    // Resolve agent attribution (O(1)): identify the nearest agent-kind ancestor.
+    String resolvedPagentSpanId = null;
+    String resolvedPagentName = null;
+
+    if (Tags.LLMOBS_AGENT_SPAN_KIND.equals(kind)) {
+      // This span is itself an agent — it becomes the nearest ancestor for its descendants.
+      resolvedPagentSpanId = String.valueOf(span.getSpanId());
+      resolvedPagentName = agentNameWireSafe(spanName) ? spanName : null;
+    } else {
+      // Inherit from in-process LLMObs parent (set by a parent agent span via context).
+      resolvedPagentSpanId = LLMObsContext.currentParentAgentSpanId();
+      resolvedPagentName = LLMObsContext.currentParentAgentName();
+
+      // Fall back to distributed propagated tags on the root APM span context.
+      if (resolvedPagentSpanId == null) {
+        AgentSpanContext rootCtx = span.getLocalRootSpan().spanContext();
+        if (rootCtx instanceof LLMObsPropagationAccess) {
+          LLMObsPropagationAccess access = (LLMObsPropagationAccess) rootCtx;
+          resolvedPagentSpanId = access.getParentAgentSpanId();
+          resolvedPagentName = access.getParentAgentName();
+        }
+      }
+    }
+
+    // Store pagent values as internal tags so the serializer can emit agent_attribution.
+    if (resolvedPagentSpanId != null) {
+      span.setTag(PAGENT_SPAN_ID_TAG_INTERNAL, resolvedPagentSpanId);
+      if (resolvedPagentName != null) {
+        span.setTag(PAGENT_NAME_TAG_INTERNAL, resolvedPagentName);
+      }
+    }
+
+    // If this span is an agent, stamp the pagent propagation tags for outgoing distributed calls.
+    if (Tags.LLMOBS_AGENT_SPAN_KIND.equals(kind)) {
+      AgentSpanContext rootCtx = span.getLocalRootSpan().spanContext();
+      if (rootCtx instanceof LLMObsPropagationAccess) {
+        LLMObsPropagationAccess access = (LLMObsPropagationAccess) rootCtx;
+        access.setParentAgentSpanId(resolvedPagentSpanId);
+        if (resolvedPagentName != null) {
+          access.setParentAgentName(resolvedPagentName);
+        }
+      }
+    }
+
+    // Propagate the effective sessionId and agent attribution to descendant LLMObs spans.
+    scope = LLMObsContext.attach(span.spanContext(), sessionId, resolvedPagentSpanId, resolvedPagentName);
+  }
+
+  /**
+   * Returns true if the agent name is safe to include in the x-datadog-tags header:
+   * printable ASCII only (0x20–0x7E), no commas (delimiter), no semicolons.
+   * Max 256 UTF-8 bytes. Since the loop rejects all non-ASCII (c > 0x7E), every character that
+   * passes is single-byte in UTF-8, so length() is an exact byte-count proxy.
+   */
+  private static boolean agentNameWireSafe(String name) {
+    if (name == null || name.length() > 256) {
+      return false;
+    }
+    for (int i = 0; i < name.length(); i++) {
+      char c = name.charAt(i);
+      if (c < 0x20 || c > 0x7E || c == ',' || c == ';') {
+        return false;
+      }
+    }
+    return true;
   }
 
   @Override

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

@@ -0,0 +1,174 @@
+package datadog.trace.llmobs.domain;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import datadog.trace.agent.tooling.TracerInstaller;
+import datadog.trace.api.WellKnownTags;
+import datadog.trace.api.llmobs.LLMObsPropagationAccess;
+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 java.lang.reflect.Field;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+class DDLLMObsSpanAgentAttributionTest {
+
+  private static final String PAGENT_SPAN_ID_TAG = "_ml_obs_tag.pagent_span_id";
+  private static final String PAGENT_NAME_TAG = "_ml_obs_tag.pagent_name";
+  private static final Field SPAN_FIELD;
+
+  private static CoreTracer tracer;
+
+  static {
+    try {
+      SPAN_FIELD = DDLLMObsSpan.class.getDeclaredField("span");
+      SPAN_FIELD.setAccessible(true);
+    } catch (ReflectiveOperationException e) {
+      throw new ExceptionInInitializerError(e);
+    }
+  }
+
+  @BeforeAll
+  static void installTracer() {
+    tracer = CoreTracer.builder().build();
+    TracerInstaller.forceInstallGlobalTracer(tracer);
+  }
+
+  @AfterAll
+  static void closeTracer() {
+    TracerInstaller.forceInstallGlobalTracer(null);
+    tracer.close();
+  }
+
+  private static DDLLMObsSpan newSpan(String kind, String name) {
+    WellKnownTags tags =
+        new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java");
+    return new DDLLMObsSpan(kind, name, "test-ml-app", null, "service", tags);
+  }
+
+  private static AgentSpan innerSpan(DDLLMObsSpan llmObsSpan) throws IllegalAccessException {
+    return (AgentSpan) SPAN_FIELD.get(llmObsSpan);
+  }
+
+  @Test
+  void agentSpanStoresOwnIdAndNameAsPagent() throws Exception {
+    DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "my-agent");
+    try {
+      AgentSpan inner = innerSpan(agentSpan);
+      String pagentSpanId = (String) inner.getTag(PAGENT_SPAN_ID_TAG);
+      String pagentName = (String) inner.getTag(PAGENT_NAME_TAG);
+
+      assertEquals(String.valueOf(inner.getSpanId()), pagentSpanId);
+      assertEquals("my-agent", pagentName);
+    } finally {
+      agentSpan.finish();
+    }
+  }
+
+  @Test
+  void agentSpanWithUnsafeNameStoresIdButNullName() throws Exception {
+    // Comma is a separator in x-datadog-tags header — disallowed
+    DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "bad,agent");
+    try {
+      AgentSpan inner = innerSpan(agentSpan);
+      String pagentSpanId = (String) inner.getTag(PAGENT_SPAN_ID_TAG);
+      Object pagentName = inner.getTag(PAGENT_NAME_TAG);
+
+      assertEquals(String.valueOf(inner.getSpanId()), pagentSpanId);
+      assertNull(pagentName);
+    } finally {
+      agentSpan.finish();
+    }
+  }
+
+  @Test
+  void nonAgentChildUnderAgentInheritsAttribution() throws Exception {
+    DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "parent-agent");
+    try {
+      AgentSpan agentInner = innerSpan(agentSpan);
+      String expectedPagentSpanId = String.valueOf(agentInner.getSpanId());
+
+      // Created while agentSpan's ContextScope is active — should inherit attribution
+      DDLLMObsSpan toolSpan = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "child-tool");
+      try {
+        AgentSpan toolInner = innerSpan(toolSpan);
+        assertEquals(expectedPagentSpanId, toolInner.getTag(PAGENT_SPAN_ID_TAG));
+        assertEquals("parent-agent", toolInner.getTag(PAGENT_NAME_TAG));
+      } finally {
+        toolSpan.finish();
+      }
+    } finally {
+      agentSpan.finish();
+    }
+  }
+
+  @Test
+  void transitiveInheritanceAgentToLlmToTool() throws Exception {
+    DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "root-agent");
+    try {
+      AgentSpan agentInner = innerSpan(agentSpan);
+      String expectedPagentSpanId = String.valueOf(agentInner.getSpanId());
+
+      DDLLMObsSpan llmSpan = newSpan(Tags.LLMOBS_LLM_SPAN_KIND, "intermediate-llm");
+      try {
+        DDLLMObsSpan toolSpan = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "leaf-tool");
+        try {
+          AgentSpan toolInner = innerSpan(toolSpan);
+          // Tool must point to the original agent, not the intermediate LLM span
+          assertEquals(expectedPagentSpanId, toolInner.getTag(PAGENT_SPAN_ID_TAG));
+          assertEquals("root-agent", toolInner.getTag(PAGENT_NAME_TAG));
+        } finally {
+          toolSpan.finish();
+        }
+      } finally {
+        llmSpan.finish();
+      }
+    } finally {
+      agentSpan.finish();
+    }
+  }
+
+  @Test
+  void noAgentAncestorProducesNoPagentTags() throws Exception {
+    DDLLMObsSpan workflowSpan = newSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "standalone-workflow");
+    try {
+      AgentSpan inner = innerSpan(workflowSpan);
+      assertNull(inner.getTag(PAGENT_SPAN_ID_TAG));
+      assertNull(inner.getTag(PAGENT_NAME_TAG));
+    } finally {
+      workflowSpan.finish();
+    }
+  }
+
+  @Test
+  void distributedParentPagentValuesAreInherited() throws Exception {
+    // Simulate a distributed parent: an APM root span with pagent propagation tags already set
+    // (e.g. injected by an upstream service during HTTP propagation).
+    AgentSpan rootApmSpan = AgentTracer.get().buildSpan("apm", "http.server.request").start();
+    AgentScope apmScope = AgentTracer.activateSpan(rootApmSpan);
+    try {
+      // Directly stamp the pagent values on the root span context via LLMObsPropagationAccess
+      LLMObsPropagationAccess access = (LLMObsPropagationAccess) rootApmSpan.spanContext();
+      access.setParentAgentSpanId("1234567890abcdef");
+      access.setParentAgentName("upstream-agent");
+
+      // No LLMObs context is active — should fall through to the distributed path
+      DDLLMObsSpan llmSpan = newSpan(Tags.LLMOBS_LLM_SPAN_KIND, "downstream-llm");
+      try {
+        AgentSpan llmInner = innerSpan(llmSpan);
+        assertEquals("1234567890abcdef", llmInner.getTag(PAGENT_SPAN_ID_TAG));
+        assertEquals("upstream-agent", llmInner.getTag(PAGENT_NAME_TAG));
+      } finally {
+        llmSpan.finish();
+      }
+    } finally {
+      apmScope.close();
+      rootApmSpan.finish();
+    }
+  }
+}

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

@@ -21,6 +21,7 @@
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.api.internal.TraceSegment;
+import datadog.trace.api.llmobs.LLMObsPropagationAccess;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.api.sampling.SamplingMechanism;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -61,7 +62,8 @@
  * the associated Span instance
  */
 public class DDSpanContext
-    implements AgentSpanContext, RequestContext, TraceSegment, ProfilerContext {
+    implements AgentSpanContext, RequestContext, TraceSegment, ProfilerContext,
+        LLMObsPropagationAccess {
   private static final Logger log = LoggerFactory.getLogger(DDSpanContext.class);
 
   public static final String PRIORITY_SAMPLING_KEY = "_sampling_priority_v1";
@@ -1492,6 +1494,27 @@ public PropagationTags getPropagationTags() {
     return getRootSpanContextOrThis().propagationTags;
   }
 
+  // LLMObsPropagationAccess implementation — delegates to the root span's propagation tags
+  @Override
+  public String getParentAgentSpanId() {
+    return getPropagationTags().getParentAgentSpanId();
+  }
+
+  @Override
+  public String getParentAgentName() {
+    return getPropagationTags().getParentAgentName();
+  }
+
+  @Override
+  public void setParentAgentSpanId(String value) {
+    getPropagationTags().updateParentAgentSpanId(value);
+  }
+
+  @Override
+  public void setParentAgentName(String value) {
+    getPropagationTags().updateParentAgentName(value);
+  }
+
   /** TraceSegment Implementation */
   @Override
   public void setTagTop(String key, Object value, boolean sanitize) {

dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java

@@ -176,4 +176,16 @@ public HashMap<String, String> createTagMap() {
   }
 
   public abstract void updateAndLockDecisionMaker(PropagationTags source);
+
+  /** Returns the propagated parent agent span ID (_dd.p.llmobs_pagent_span_id), or null. */
+  public abstract String getParentAgentSpanId();
+
+  /** Returns the propagated parent agent name (_dd.p.llmobs_pagent_name), or null. */
+  public abstract String getParentAgentName();
+
+  /** Sets the parent agent span ID for outgoing propagation. Null clears it. */
+  public abstract void updateParentAgentSpanId(String value);
+
+  /** Sets the parent agent name for outgoing propagation. Null clears it. */
+  public abstract void updateParentAgentName(String value);
 }

dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java

@@ -64,6 +64,8 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
     TagValue traceIdTagValue = null;
     int traceSource = 0;
     TagValue orgPropagationMarkerTagValue = null;
+    TagValue parentAgentSpanIdTagValue = null;
+    TagValue parentAgentNameTagValue = null;
     while (tagPos < len) {
       int tagKeyEndsAt =
           validateCharsUntilSeparatorOrEnd(
@@ -102,6 +104,10 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
             traceSource = ProductTraceSource.parseBitfieldHex(tagValue.toString());
           } else if (tagKey.equals(ORG_PROPAGATION_MARKER_TAG)) {
             orgPropagationMarkerTagValue = tagValue;
+          } else if (tagKey.equals(PARENT_AGENT_SPAN_ID_TAG)) {
+            parentAgentSpanIdTagValue = tagValue;
+          } else if (tagKey.equals(PARENT_AGENT_NAME_TAG)) {
+            parentAgentNameTagValue = tagValue;
           } else {
             if (tagPairs == null) {
               // This is roughly the size of a two element linked list but can hold six
@@ -114,12 +120,20 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
       }
       tagPos = tagValueEndsAt + 1;
     }
-    return tagsFactory.createValid(
-        tagPairs,
-        decisionMakerTagValue,
-        traceIdTagValue,
-        traceSource,
-        orgPropagationMarkerTagValue);
+    PropagationTags result =
+        tagsFactory.createValid(
+            tagPairs,
+            decisionMakerTagValue,
+            traceIdTagValue,
+            traceSource,
+            orgPropagationMarkerTagValue);
+    if (parentAgentSpanIdTagValue != null) {
+      result.updateParentAgentSpanId(parentAgentSpanIdTagValue.toString());
+    }
+    if (parentAgentNameTagValue != null) {
+      result.updateParentAgentName(parentAgentNameTagValue.toString());
+    }
+    return result;
   }
 
   @Override