Add agent_version tag with subtree propagation to LLM Observability spans

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

@@ -261,13 +261,23 @@ public LLMObsSpan startLLMSpan(
     @Override
     public LLMObsSpan startAgentSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId) {
+      return startAgentSpan(spanName, mlApp, sessionId, null);
+    }
+
+    @Override
+    public LLMObsSpan startAgentSpan(
+        String spanName,
+        @Nullable String mlApp,
+        @Nullable String sessionId,
+        @Nullable String version) {
       return new DDLLMObsSpan(
           Tags.LLMOBS_AGENT_SPAN_KIND,
           spanName,
           getMLApp(mlApp),
           sessionId,
           serviceName,
-          wellKnownTags);
+          wellKnownTags,
+          version);
     }
 
     @Override

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

@@ -63,6 +63,7 @@ public class DDLLMObsSpan implements LLMObsSpan {
   private final String mlApp;
   private final ContextScope scope;
   private final boolean hasSessionId;
+  private final boolean hasAgentVersion;
 
   private boolean finished = false;
 
@@ -73,6 +74,17 @@ public DDLLMObsSpan(
       String sessionId,
       @Nonnull String serviceName,
       WellKnownTags wellKnownTags) {
+    this(kind, spanName, mlApp, sessionId, serviceName, wellKnownTags, null);
+  }
+
+  public DDLLMObsSpan(
+      @Nonnull String kind,
+      String spanName,
+      @Nonnull String mlApp,
+      String sessionId,
+      @Nonnull String serviceName,
+      WellKnownTags wellKnownTags,
+      String agentVersion) {
 
     if (null == spanName || spanName.isEmpty()) {
       spanName = kind;
@@ -125,16 +137,30 @@ public DDLLMObsSpan(
             sessionId = inherited;
           }
         }
+        // Inherit agent_version from the enclosing agent span, if this span doesn't set its own.
+        // An explicit value always wins, so a nested agent's own version overrides an ancestor's
+        // for its own subtree, matching session_id's explicit-wins semantics.
+        if (agentVersion == null || agentVersion.isEmpty()) {
+          String inherited = LLMObsContext.currentAgentVersion();
+          if (inherited != null && !inherited.isEmpty()) {
+            agentVersion = inherited;
+          }
+        }
       }
     }
 
     this.hasSessionId = sessionId != null && !sessionId.isEmpty();
     if (this.hasSessionId) {
       span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
     }
+    this.hasAgentVersion = agentVersion != null && !agentVersion.isEmpty();
+    if (this.hasAgentVersion) {
+      span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.AGENT_VERSION, agentVersion);
+    }
     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);
+    // Propagate the effective sessionId and agent_version to descendant LLMObs spans via the
+    // context.
+    scope = LLMObsContext.attach(span.spanContext(), sessionId, agentVersion);
   }
 
   @Override

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

@@ -0,0 +1,175 @@
+package datadog.trace.llmobs.domain;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import datadog.trace.agent.tooling.TracerInstaller;
+import datadog.trace.api.WellKnownTags;
+import datadog.trace.api.llmobs.LLMObsTags;
+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;
+
+/**
+ * Covers agent_version propagation to a versioned agent's subtree, mirroring the session_id
+ * inheritance behavior in {@code DDLLMObsSpanTest}, but with the additional requirement that an
+ * explicit version always wins for its own subtree (so a nested agent's own version overrides an
+ * ancestor's).
+ */
+class DDLLMObsSpanAgentVersionTest {
+  private static final String AGENT_VERSION_TAG = "_ml_obs_tag." + LLMObsTags.AGENT_VERSION;
+
+  private static final Field SPAN_FIELD;
+
+  private static CoreTracer tracer;
+
+  static {
+    try {
+      SPAN_FIELD = DDLLMObsSpan.class.getDeclaredField("span");
+      SPAN_FIELD.setAccessible(true);
+    } catch (ReflectiveOperationException error) {
+      throw new ExceptionInInitializerError(error);
+    }
+  }
+
+  @BeforeAll
+  static void installTracer() {
+    tracer = CoreTracer.builder().build();
+    TracerInstaller.forceInstallGlobalTracer(tracer);
+  }
+
+  @AfterAll
+  static void closeTracer() {
+    TracerInstaller.forceInstallGlobalTracer(null);
+    tracer.close();
+  }
+
+  @Test
+  void agentSpanWithExplicitVersionTagsItself() {
+    DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", "v3");
+    try {
+      assertEquals("v3", spanOf(agent).getTag(AGENT_VERSION_TAG));
+    } finally {
+      agent.finish();
+    }
+  }
+
+  @Test
+  void childSpanInheritsAgentVersionFromParentContext() {
+    DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", "v3");
+    try (AgentScope ignored = AgentTracer.activateSpan(spanOf(agent))) {
+      DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool1", null);
+      try {
+        assertEquals("v3", spanOf(child).getTag(AGENT_VERSION_TAG));
+      } finally {
+        child.finish();
+      }
+    } finally {
+      agent.finish();
+    }
+  }
+
+  @Test
+  void grandchildTransitivelyInheritsAgentVersionThroughIntermediateSpan() {
+    DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", "v3");
+    try (AgentScope agentScope = AgentTracer.activateSpan(spanOf(agent))) {
+      DDLLMObsSpan workflow = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "workflow1", null);
+      try (AgentScope workflowScope = AgentTracer.activateSpan(spanOf(workflow))) {
+        DDLLMObsSpan grandchild = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "llm1", null);
+        try {
+          assertEquals("v3", spanOf(grandchild).getTag(AGENT_VERSION_TAG));
+        } finally {
+          grandchild.finish();
+        }
+      } finally {
+        workflow.finish();
+      }
+    } finally {
+      agent.finish();
+    }
+  }
+
+  @Test
+  void nestedAgentWithOwnVersionOverridesForItsOwnSubtree() {
+    DDLLMObsSpan outerAgent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "outer-agent", "v1");
+    try (AgentScope outerScope = AgentTracer.activateSpan(spanOf(outerAgent))) {
+      DDLLMObsSpan innerAgent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "inner-agent", "v2");
+      try (AgentScope innerScope = AgentTracer.activateSpan(spanOf(innerAgent))) {
+        assertEquals("v2", spanOf(innerAgent).getTag(AGENT_VERSION_TAG));
+
+        DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "inner-tool", null);
+        try {
+          assertEquals(
+              "v2",
+              spanOf(child).getTag(AGENT_VERSION_TAG),
+              "child of the nested agent must inherit the nested agent's own version, not the outer one");
+        } finally {
+          child.finish();
+        }
+      } finally {
+        innerAgent.finish();
+      }
+    } finally {
+      outerAgent.finish();
+    }
+  }
+
+  @Test
+  void noVersionSetAnywhereMeansNoTagOnAnySpanInTheSubtree() {
+    DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", null);
+    try (AgentScope agentScope = AgentTracer.activateSpan(spanOf(agent))) {
+      DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool1", null);
+      try {
+        assertNull(spanOf(agent).getTag(AGENT_VERSION_TAG));
+        assertNull(spanOf(child).getTag(AGENT_VERSION_TAG));
+      } finally {
+        child.finish();
+      }
+    } finally {
+      agent.finish();
+    }
+  }
+
+  @Test
+  void childDoesNotInheritAgentVersionWhenStaleContextIsFromADifferentTrace() {
+    // Simulates a stale LLMObsContext (e.g. leaked across an async boundary): the parent's
+    // context is attached, but its AgentScope is deliberately NOT activated, so the next span
+    // started begins a fresh trace and the trace-consistency gate must skip inheritance.
+    DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "stale-agent", "stale-v1");
+    try {
+      DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool1", null);
+      try {
+        assertNotEquals(
+            spanOf(agent).getTraceId(),
+            spanOf(child).getTraceId(),
+            "sanity: traces must differ for this scenario to be meaningful");
+        assertNull(spanOf(child).getTag(AGENT_VERSION_TAG));
+      } finally {
+        child.finish();
+      }
+    } finally {
+      agent.finish();
+    }
+  }
+
+  private static DDLLMObsSpan llmObsSpan(String kind, String name, String agentVersion) {
+    WellKnownTags tags =
+        new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java");
+    return new DDLLMObsSpan(kind, name, "test-ml-app", null, "service", tags, agentVersion);
+  }
+
+  private static AgentSpan spanOf(DDLLMObsSpan llmObsSpan) {
+    try {
+      return (AgentSpan) SPAN_FIELD.get(llmObsSpan);
+    } catch (IllegalAccessException error) {
+      throw new AssertionError(error);
+    }
+  }
+}

dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java

@@ -40,6 +40,29 @@ public static LLMObsSpan startAgentSpan(
     return SPAN_FACTORY.startAgentSpan(spanName, mlApp, sessionId);
   }
 
+  /**
+   * Starts an agent span, optionally tagging it with a version.
+   *
+   * <p>The version is set as an {@code agent_version} tag on this span and propagated to every
+   * descendant LLMObs span started under it (LLM, tool, workflow, etc.), so that the agent's entire
+   * execution subtree can be filtered or grouped by version. A nested agent span that passes its
+   * own {@code version} overrides the inherited value for its own subtree.
+   *
+   * <p>Setting {@code agent_version} directly via {@link LLMObsSpan#setTag} on a span other than
+   * the one returned here does not propagate to that span's descendants — only this method triggers
+   * propagation.
+   *
+   * @param version the version of this agent, or {@code null}/empty to leave it untagged
+   */
+  public static LLMObsSpan startAgentSpan(
+      String spanName,
+      @Nullable String mlApp,
+      @Nullable String sessionId,
+      @Nullable String version) {
+
+    return SPAN_FACTORY.startAgentSpan(spanName, mlApp, sessionId, version);
+  }
+
   public static LLMObsSpan startToolSpan(
       String spanName, @Nullable String mlApp, @Nullable String sessionId) {
 
@@ -170,6 +193,14 @@ LLMObsSpan startLLMSpan(
 
     LLMObsSpan startAgentSpan(String spanName, @Nullable String mlApp, @Nullable String sessionId);
 
+    default LLMObsSpan startAgentSpan(
+        String spanName,
+        @Nullable String mlApp,
+        @Nullable String sessionId,
+        @Nullable String version) {
+      return startAgentSpan(spanName, mlApp, sessionId);
+    }
+
     LLMObsSpan startToolSpan(String spanName, @Nullable String mlApp, @Nullable String sessionId);
 
     LLMObsSpan startTaskSpan(String spanName, @Nullable String mlApp, @Nullable String sessionId);

dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsTags.java

@@ -4,6 +4,7 @@
 public class LLMObsTags {
   public static final String ML_APP = "ml_app";
   public static final String SESSION_ID = "session_id";
+  public static final String AGENT_VERSION = "agent_version";
 
   // meta
   public static final String METADATA = "metadata";

dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpanFactory.java

@@ -21,6 +21,14 @@ public LLMObsSpan startAgentSpan(
     return NoOpLLMObsSpan.INSTANCE;
   }
 
+  public LLMObsSpan startAgentSpan(
+      String spanName,
+      @Nullable String mlApp,
+      @Nullable String sessionId,
+      @Nullable String version) {
+    return NoOpLLMObsSpan.INSTANCE;
+  }
+
   public LLMObsSpan startToolSpan(
       String spanName, @Nullable String mlApp, @Nullable String sessionId) {
     return NoOpLLMObsSpan.INSTANCE;

dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java

@@ -259,6 +259,34 @@ void testLLMObsSpanMapperSerializesPromptWithoutInputMessages() throws Exception
     tracer.close();
   }
 
+  @Test
+  void testLLMObsSpanMapperSerializesAgentVersionAsAGenericTag() throws Exception {
+    LLMObsSpanMapper mapper = new LLMObsSpanMapper();
+    CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build();
+
+    AgentSpan toolSpan =
+        tracer
+            .buildSpan("datadog", "get_weather")
+            .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_TOOL_SPAN_KIND)
+            .withTag("_ml_obs_tag.agent_version", "v3")
+            .start();
+    toolSpan.setSpanType(InternalSpanTypes.LLMOBS);
+    toolSpan.finish();
+
+    Map<String, Object> spanData = serializeSingleSpan(mapper, toolSpan);
+
+    // Unlike session_id, agent_version needs no top-level field or meta{} remapping — it flows
+    // generically through tags[], the same way ml_app does.
+    assertFalse(spanData.containsKey("agent_version"));
+    Map<String, Object> meta = (Map<String, Object>) spanData.get("meta");
+    assertFalse(meta.containsKey("agent_version"));
+
+    List<String> tags = (List<String>) spanData.get("tags");
+    assertTrue(tags.contains("agent_version:v3"));
+
+    tracer.close();
+  }
+
   @Test
   void testLLMObsSpanMapperPreservesNestedPromptInputCompatibility() throws Exception {
     LLMObsSpanMapper mapper = new LLMObsSpanMapper();