diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java index b5283b73dad..33d1379e101 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java +++ b/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 diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentAttributionTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentAttributionTest.java new file mode 100644 index 00000000000..7b9df554eb7 --- /dev/null +++ b/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(); + } + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java index adf4cd66156..d3ed13ce634 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java +++ b/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) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java index 3a0c57a4dd8..57c865f81ef 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java @@ -176,4 +176,16 @@ public HashMap 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); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java index ec7a9a05feb..bda03e0f9f6 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java +++ b/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 diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java index e2c0658a1d2..d125ab18aa4 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java @@ -20,6 +20,8 @@ abstract class PTagsCodec { protected static final TagKey DEBUG_TAG = TagKey.from("debug"); protected static final TagKey KNUTH_SAMPLING_RATE_TAG = TagKey.from("ksr"); protected static final TagKey ORG_PROPAGATION_MARKER_TAG = TagKey.from("opm"); + protected static final TagKey PARENT_AGENT_SPAN_ID_TAG = TagKey.from("llmobs_pagent_span_id"); + protected static final TagKey PARENT_AGENT_NAME_TAG = TagKey.from("llmobs_pagent_name"); protected static final String PROPAGATION_ERROR_MALFORMED_TID = "malformed_tid "; protected static final String PROPAGATION_ERROR_INCONSISTENT_TID = "inconsistent_tid "; protected static final TagKey UPSTREAM_SERVICES_DEPRECATED_TAG = TagKey.from("upstream_services"); @@ -65,6 +67,15 @@ static String headerValue(PTagsCodec codec, PTags ptags, CharSequence lastParent codec.appendTag( sb, ORG_PROPAGATION_MARKER_TAG, ptags.getOrgPropagationMarkerTagValue(), size); } + if (ptags.getParentAgentSpanIdTagValue() != null) { + size = + codec.appendTag( + sb, PARENT_AGENT_SPAN_ID_TAG, ptags.getParentAgentSpanIdTagValue(), size); + } + if (ptags.getParentAgentNameTagValue() != null) { + size = + codec.appendTag(sb, PARENT_AGENT_NAME_TAG, ptags.getParentAgentNameTagValue(), size); + } Iterator it = ptags.getTagPairs().iterator(); while (it.hasNext() && !codec.isTooLarge(sb, size)) { TagElement tagKey = it.next(); @@ -129,6 +140,16 @@ static void fillTagMap(PTags propagationTags, Map tagMap) { ORG_PROPAGATION_MARKER_TAG.forType(Encoding.DATADOG).toString(), propagationTags.getOrgPropagationMarkerTagValue().forType(Encoding.DATADOG).toString()); } + if (propagationTags.getParentAgentSpanIdTagValue() != null) { + tagMap.put( + PARENT_AGENT_SPAN_ID_TAG.forType(Encoding.DATADOG).toString(), + propagationTags.getParentAgentSpanIdTagValue().forType(Encoding.DATADOG).toString()); + } + if (propagationTags.getParentAgentNameTagValue() != null) { + tagMap.put( + PARENT_AGENT_NAME_TAG.forType(Encoding.DATADOG).toString(), + propagationTags.getParentAgentNameTagValue().forType(Encoding.DATADOG).toString()); + } if (propagationTags.getTraceIdHighOrderBitsHexTagValue() != null) { tagMap.put( TRACE_ID_TAG.forType(Encoding.DATADOG).toString(), diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java index 0b5184d448a..44ca921b468 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java @@ -5,6 +5,8 @@ import static datadog.trace.core.propagation.ptags.PTagsCodec.DECISION_MAKER_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.KNUTH_SAMPLING_RATE_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.ORG_PROPAGATION_MARKER_TAG; +import static datadog.trace.core.propagation.ptags.PTagsCodec.PARENT_AGENT_NAME_TAG; +import static datadog.trace.core.propagation.ptags.PTagsCodec.PARENT_AGENT_SPAN_ID_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.TRACE_ID_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.TRACE_SOURCE_TAG; @@ -112,6 +114,9 @@ static class PTags extends PropagationTags { private volatile TagValue orgPropagationMarkerTagValue; + private volatile TagValue parentAgentSpanIdTagValue; + private volatile TagValue parentAgentNameTagValue; + // Static cache for the most-recently-seen rate → TagValue. In steady state a service uses one // rate, so this eliminates the char[] + String allocation on every new PTags instance. // Writes are benign-racy: two threads computing the same rate produce equal TagValues. @@ -377,6 +382,44 @@ TagValue getOrgPropagationMarkerTagValue() { return orgPropagationMarkerTagValue; } + @Override + public String getParentAgentSpanId() { + TagValue v = parentAgentSpanIdTagValue; + return v == null ? null : v.forType(TagElement.Encoding.DATADOG).toString(); + } + + @Override + public String getParentAgentName() { + TagValue v = parentAgentNameTagValue; + return v == null ? null : v.forType(TagElement.Encoding.DATADOG).toString(); + } + + @Override + public void updateParentAgentSpanId(String value) { + TagValue newValue = value == null ? null : TagValue.from(value); + if (!Objects.equals(this.parentAgentSpanIdTagValue, newValue)) { + clearCachedHeader(DATADOG); + this.parentAgentSpanIdTagValue = newValue; + } + } + + @Override + public void updateParentAgentName(String value) { + TagValue newValue = value == null ? null : TagValue.from(value); + if (!Objects.equals(this.parentAgentNameTagValue, newValue)) { + clearCachedHeader(DATADOG); + this.parentAgentNameTagValue = newValue; + } + } + + TagValue getParentAgentSpanIdTagValue() { + return parentAgentSpanIdTagValue; + } + + TagValue getParentAgentNameTagValue() { + return parentAgentNameTagValue; + } + @Override public int getSamplingPriority() { return samplingPriority; @@ -520,6 +563,11 @@ int getXDatadogTagsSize() { TRACE_SOURCE_TAG, TagValue.from(ProductTraceSource.getBitfieldHex(currentProductTraceSource))); } + size = + PTagsCodec.calcXDatadogTagsSize( + size, PARENT_AGENT_SPAN_ID_TAG, parentAgentSpanIdTagValue); + size = + PTagsCodec.calcXDatadogTagsSize(size, PARENT_AGENT_NAME_TAG, parentAgentNameTagValue); xDatadogTagsSize = size; } return size; diff --git a/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java index 9dd769ee9eb..332e1e94b6e 100644 --- a/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java +++ b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java @@ -68,6 +68,10 @@ public class LLMObsSpanMapper implements RemoteMapper { private static final byte[] ERROR_STACK = "stack".getBytes(StandardCharsets.UTF_8); private static final byte[] META = "meta".getBytes(StandardCharsets.UTF_8); + private static final byte[] AGENT_ATTRIBUTION = + "agent_attribution".getBytes(StandardCharsets.UTF_8); + private static final byte[] PAGENT_NAME = "pagent_name".getBytes(StandardCharsets.UTF_8); + private static final byte[] PAGENT_SPAN_ID = "pagent_span_id".getBytes(StandardCharsets.UTF_8); private static final byte[] METADATA = "metadata".getBytes(StandardCharsets.UTF_8); private static final byte[] PROMPT = "prompt".getBytes(StandardCharsets.UTF_8); private static final byte[] SPAN_KIND = "span.kind".getBytes(StandardCharsets.UTF_8); @@ -94,6 +98,9 @@ public class LLMObsSpanMapper implements RemoteMapper { private static final String PARENT_ID_TAG_INTERNAL_FULL = LLMOBS_TAG_PREFIX + "parent_id"; private static final String SESSION_ID_TAG_INTERNAL_FULL = LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID; + private static final String PAGENT_SPAN_ID_TAG_INTERNAL_FULL = + LLMOBS_TAG_PREFIX + "pagent_span_id"; + private static final String PAGENT_NAME_TAG_INTERNAL_FULL = LLMOBS_TAG_PREFIX + "pagent_name"; private final MetaWriter metaWriter = new MetaWriter(); private final int size; @@ -265,7 +272,9 @@ private static final class MetaWriter implements MetadataConsumer { LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_PROVIDER, LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_VERSION, LLMOBS_TAG_PREFIX + LLMObsTags.TOOL_DEFINITIONS, - LLMOBS_TAG_PREFIX + LLMObsTags.METADATA))); + LLMOBS_TAG_PREFIX + LLMObsTags.METADATA, + PAGENT_SPAN_ID_TAG_INTERNAL_FULL, + PAGENT_NAME_TAG_INTERNAL_FULL))); MetaWriter withWritable(Writable writable, Map errorInfo) { this.writable = writable; @@ -308,6 +317,9 @@ public void accept(Metadata metadata) { String inputPromptTag = LLMOBS_TAG_PREFIX + INPUT_PROMPT; boolean hasInput = tagsToRemapToMeta.containsKey(inputTag); boolean hasInputPrompt = tagsToRemapToMeta.containsKey(inputPromptTag); + boolean hasAgentAttribution = tagsToRemapToMeta.containsKey(PAGENT_SPAN_ID_TAG_INTERNAL_FULL); + boolean hasAgentAttributionName = + tagsToRemapToMeta.containsKey(PAGENT_NAME_TAG_INTERNAL_FULL); Object inputPrompt = null; if (hasInputPrompt) { if (spanKind.equals(Tags.LLMOBS_LLM_SPAN_KIND)) { @@ -342,12 +354,15 @@ public void accept(Metadata metadata) { } // write meta (11) + // pagent_name is always emitted inside agent_attribution (never standalone), so subtract 1 + // whenever it is in the map regardless of whether pagent_span_id is also present. int metaSize = tagsToRemapToMeta.size() - (hasInputPrompt ? 1 : 0) + (inputPrompt != null && !hasInput ? 1 : 0) + 1 - + (null != errorInfo && !errorInfo.isEmpty() ? 1 : 0); + + (null != errorInfo && !errorInfo.isEmpty() ? 1 : 0) + - (hasAgentAttributionName ? 1 : 0); writable.writeUTF8(META); writable.startMap(metaSize); writable.writeUTF8(SPAN_KIND); @@ -378,7 +393,24 @@ public void accept(Metadata metadata) { for (Map.Entry tag : tagsToRemapToMeta.entrySet()) { String key = tag.getKey().substring(LLMOBS_TAG_PREFIX.length()); Object val = tag.getValue(); - if (key.equals(INPUT) || key.equals(OUTPUT)) { + if (key.equals("pagent_name")) { + // Emitted inside the agent_attribution block below; skip standalone entry. + continue; + } else if (key.equals("pagent_span_id")) { + // Emit the structured agent_attribution map. + writable.writeUTF8(AGENT_ATTRIBUTION); + writable.startMap(2); + writable.writeUTF8(PAGENT_NAME); + Object nameVal = tagsToRemapToMeta.get(PAGENT_NAME_TAG_INTERNAL_FULL); + if (nameVal instanceof String) { + writable.writeString((String) nameVal, null); + } else { + writable.writeNull(); + } + writable.writeUTF8(PAGENT_SPAN_ID); + writable.writeObject(val, null); + continue; + } else if (key.equals(INPUT) || key.equals(OUTPUT)) { boolean isDocumentIO = (spanKind.equals(Tags.LLMOBS_EMBEDDING_SPAN_KIND) && key.equals(INPUT)) || (spanKind.equals(Tags.LLMOBS_RETRIEVAL_SPAN_KIND) && key.equals(OUTPUT)); diff --git a/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java b/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java index f683e64ab91..2a1500b01bf 100644 --- a/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.java @@ -546,6 +546,81 @@ void testLLMObsSpanMapperSerializesDocumentIO() throws Exception { tracer.close(); } + @Test + void testAgentAttributionEmittedWithBothFields() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + AgentSpan agentSpan = + tracer + .buildSpan("datadog", "my.agent") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_AGENT_SPAN_KIND) + .withTag("_ml_obs_tag.pagent_span_id", "abc123") + .withTag("_ml_obs_tag.pagent_name", "my-orchestrator") + .start(); + agentSpan.setSpanType(InternalSpanTypes.LLMOBS); + agentSpan.finish(); + + Map spanData = serializeSingleSpan(mapper, agentSpan); + Map meta = (Map) spanData.get("meta"); + + assertTrue(meta.containsKey("agent_attribution")); + Map attribution = (Map) meta.get("agent_attribution"); + assertEquals("abc123", attribution.get("pagent_span_id")); + assertEquals("my-orchestrator", attribution.get("pagent_name")); + + tracer.close(); + } + + @Test + void testAgentAttributionEmitsExplicitNullNameWhenAbsent() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + // Only pagent_span_id is set — pagent_name tag is absent + AgentSpan agentSpan = + tracer + .buildSpan("datadog", "my.agent") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_AGENT_SPAN_KIND) + .withTag("_ml_obs_tag.pagent_span_id", "abc123") + .start(); + agentSpan.setSpanType(InternalSpanTypes.LLMOBS); + agentSpan.finish(); + + Map spanData = serializeSingleSpan(mapper, agentSpan); + Map meta = (Map) spanData.get("meta"); + + assertTrue(meta.containsKey("agent_attribution")); + Map attribution = (Map) meta.get("agent_attribution"); + assertEquals("abc123", attribution.get("pagent_span_id")); + // pagent_name key must be present with an explicit null (not absent) + assertTrue(attribution.containsKey("pagent_name")); + assertNull(attribution.get("pagent_name")); + + tracer.close(); + } + + @Test + void testNoAgentAttributionBlockWhenPagentSpanIdAbsent() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + AgentSpan llmSpan = + tracer + .buildSpan("datadog", "openai.chat") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .start(); + llmSpan.setSpanType(InternalSpanTypes.LLMOBS); + llmSpan.finish(); + + Map spanData = serializeSingleSpan(mapper, llmSpan); + Map meta = (Map) spanData.get("meta"); + + assertFalse(meta.containsKey("agent_attribution")); + + tracer.close(); + } + private static byte[] writeTo(datadog.trace.common.writer.Payload payload) throws IOException { ByteArrayOutputStream channel = new ByteArrayOutputStream(); payload.writeTo( diff --git a/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsContext.java b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsContext.java index 83df56553b8..0814be943fb 100644 --- a/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsContext.java +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsContext.java @@ -14,9 +14,12 @@ private LLMObsContext() { private static final ContextKey CONTEXT_KEY = ContextKey.named("llmobs_span"); private static final ContextKey SESSION_ID_KEY = ContextKey.named("llmobs_session_id"); + private static final ContextKey PAGENT_SPAN_ID_KEY = + ContextKey.named("llmobs_pagent_span_id"); + private static final ContextKey PAGENT_NAME_KEY = ContextKey.named("llmobs_pagent_name"); public static ContextScope attach(AgentSpanContext ctx) { - return attach(ctx, null); + return attach(ctx, null, null, null); } /** @@ -25,10 +28,26 @@ public static ContextScope attach(AgentSpanContext ctx) { * not specify their own sessionId will inherit it via {@link #currentSessionId()}. */ public static ContextScope attach(AgentSpanContext ctx, String sessionId) { + return attach(ctx, sessionId, null, null); + } + + /** + * Attach an LLMObs span context, propagating session_id and agent attribution to descendant + * LLMObs spans. pagentSpanId identifies the nearest agent-kind ancestor; pagentName is its name + * (may be null if it failed wire-safety validation). + */ + public static ContextScope attach( + AgentSpanContext ctx, String sessionId, String pagentSpanId, String pagentName) { Context updated = Context.current().with(CONTEXT_KEY, ctx); if (sessionId != null && !sessionId.isEmpty()) { updated = updated.with(SESSION_ID_KEY, sessionId); } + if (pagentSpanId != null && !pagentSpanId.isEmpty()) { + updated = updated.with(PAGENT_SPAN_ID_KEY, pagentSpanId); + if (pagentName != null) { + updated = updated.with(PAGENT_NAME_KEY, pagentName); + } + } return updated.attach(); } @@ -42,4 +61,18 @@ public static AgentSpanContext current() { public static String currentSessionId() { return Context.current().get(SESSION_ID_KEY); } + + /** + * Return the parent agent span ID propagated from an enclosing agent-kind LLMObs span, or null. + */ + public static String currentParentAgentSpanId() { + return Context.current().get(PAGENT_SPAN_ID_KEY); + } + + /** + * Return the parent agent name propagated from an enclosing agent-kind LLMObs span, or null. + */ + public static String currentParentAgentName() { + return Context.current().get(PAGENT_NAME_KEY); + } } diff --git a/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsPropagationAccess.java b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsPropagationAccess.java new file mode 100644 index 00000000000..e6c53a296a0 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsPropagationAccess.java @@ -0,0 +1,21 @@ +package datadog.trace.api.llmobs; + +/** + * Bridge interface allowing the LLMObs span (in agent-llmobs) to read and write agent attribution + * propagation tags on the underlying APM span context (in dd-trace-core) without a direct module + * dependency. Implemented by DDSpanContext. + */ +public interface LLMObsPropagationAccess { + + /** Returns the propagated parent agent span ID, or null if not set. */ + String getParentAgentSpanId(); + + /** Returns the propagated parent agent name, or null if not set. */ + String getParentAgentName(); + + /** Sets the parent agent span ID to propagate on outgoing requests. */ + void setParentAgentSpanId(String value); + + /** Sets the parent agent name to propagate on outgoing requests. Null clears it. */ + void setParentAgentName(String value); +}