feat(llmobs): add agent attribution to Java SDK

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

@@ -11,6 +11,7 @@
 import datadog.trace.api.llmobs.LLMObsSpan;
 import datadog.trace.api.llmobs.LLMObsTags;
 import datadog.trace.api.telemetry.LLMObsMetricCollector;
+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;
@@ -50,6 +51,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 + LLMObsTags.PAGENT_SPAN_ID;
+  private static final String PAGENT_NAME_TAG_INTERNAL = LLMOBS_TAG_PREFIX + LLMObsTags.PAGENT_NAME;
 
   private static final String SERVICE = LLMOBS_TAG_PREFIX + "service";
   private static final String VERSION = LLMOBS_TAG_PREFIX + "version";
@@ -63,9 +67,12 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private final AgentSpan span;
   private final String spanKind;
   private final String mlApp;
-  private final ContextScope scope;
   private final boolean hasSessionId;
-  private final boolean hasAgentVersion;
+  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
+  // they inherit agent attribution correctly.
+  private final AgentScope standaloneApmScope;
 
   private boolean finished = false;
 
@@ -156,14 +163,58 @@ public DDLLMObsSpan(
     if (this.hasSessionId) {
       span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
     }
-    this.hasAgentVersion = resolvedAgentVersion != null && !resolvedAgentVersion.isEmpty();
-    if (this.hasAgentVersion) {
+    if (resolvedAgentVersion != null && !resolvedAgentVersion.isEmpty()) {
       span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.AGENT_VERSION, resolvedAgentVersion);
     }
     span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID);
-    // Propagate the effective sessionId and agent_version to descendant LLMObs spans via the
-    // context.
-    scope = LLMObsContext.attach(span.spanContext(), sessionId, resolvedAgentVersion);
+
+    // Resolve agent attribution (O(1)): identify the nearest agent-kind ancestor.
+    String resolvedParentAgentSpanId = null;
+    String resolvedParentAgentName = null;
+
+    if (Tags.LLMOBS_AGENT_SPAN_KIND.equals(kind)) {
+      // This span is itself an agent — it becomes the nearest ancestor for its descendants.
+      // Use the span name as the initial pagent name; annotateAgentManifest() will update it
+      // to the manifest name if one is provided later.
+      resolvedParentAgentSpanId = String.valueOf(span.getSpanId());
+      resolvedParentAgentName = spanName;
+    } else {
+      // Inherit from in-process LLMObs parent only when the context belongs to the same trace.
+      // Matches the gate applied to parent_id and session_id above: a stale LLMObsContext
+      // leaked across an async boundary would otherwise attribute a span to an agent from a
+      // different trace. For standalone agent spans (no ambient APM root), standaloneApmScope
+      // ensures descendants are started under the agent's APM span so this gate passes.
+      if (null != parent && parent.getTraceId() == span.getTraceId()) {
+        resolvedParentAgentSpanId = LLMObsContext.currentParentAgentSpanId();
+        resolvedParentAgentName = LLMObsContext.currentParentAgentName();
+      }
+    }
+
+    // Store pagent values as internal tags so the serializer can emit agent_attribution.
+    if (resolvedParentAgentSpanId != null) {
+      span.setTag(PAGENT_SPAN_ID_TAG_INTERNAL, resolvedParentAgentSpanId);
+      if (resolvedParentAgentName != null) {
+        span.setTag(PAGENT_NAME_TAG_INTERNAL, resolvedParentAgentName);
+      }
+    }
+
+    // Propagate the effective sessionId and agent attribution to descendant LLMObs spans.
+    scope =
+        LLMObsContext.attach(
+            span.spanContext(),
+            sessionId,
+            resolvedAgentVersion,
+            resolvedParentAgentSpanId,
+            resolvedParentAgentName);
+
+    // For standalone agent spans (no ambient APM root), activate the underlying APM span so
+    // that child LLMObs spans share the same trace ID and pass the trace-ID gate. Without
+    // this, children start a fresh APM trace, the gate rejects the agent context, and
+    // agent attribution is silently dropped.
+    standaloneApmScope =
+        Tags.LLMOBS_AGENT_SPAN_KIND.equals(kind) && span.getLocalRootSpan() == span
+            ? AgentTracer.activateSpan(span)
+            : null;
   }
 
   @Override
@@ -340,6 +391,10 @@ public void annotateAgentManifest(LLMObs.AgentManifest manifest) {
     mergeManifest(base, manifest);
     base.put("framework", MANUAL_FRAMEWORK);
     span.setTag(AGENT_MANIFEST, base);
+
+    // Sync pagent name to the manifest name so the serializer emits the manifest name in
+    // agent_attribution. The manifest name takes priority over the span name set at construction.
+    span.setTag(PAGENT_NAME_TAG_INTERNAL, (String) base.get("name"));
   }
 
   private void mergeManifest(Map<String, Object> base, LLMObs.AgentManifest manifest) {
@@ -601,6 +656,9 @@ public void finish() {
       return;
     }
     span.finish();
+    if (standaloneApmScope != null) {
+      standaloneApmScope.close();
+    }
     scope.close();
     finished = true;
     boolean isRootSpan = span.getLocalRootSpan() == span;

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

@@ -229,6 +229,9 @@ class DDLLMObsSpanTest  extends DDSpecification{
     "v1" == tagVersion.toString()
 
     DDTraceApiInfo.VERSION == innerSpan.getTag(LLMOBS_TAG_PREFIX + "ddtrace.version")
+
+    cleanup:
+    test.finish()
   }
 
   def "test llm span string input formatted to messages"() {
@@ -511,6 +514,9 @@ class DDLLMObsSpanTest  extends DDSpecification{
     innerSpan.getTag(INPUT_PROMPT) == null
     innerSpan.getTag(PROMPT_TRACKING_INSTRUMENTATION_METHOD) == null
 
+    cleanup:
+    test.finish()
+
     where:
     spanKind << [
       Tags.LLMOBS_AGENT_SPAN_KIND,