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 71344a15444..42c6790a0c2 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 @@ -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 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; diff --git a/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy b/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy index bd5ce69ab57..2f45f9268db 100644 --- a/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy +++ b/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, 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..db45fa7ce33 --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentAttributionTest.java @@ -0,0 +1,283 @@ +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.LLMObs; +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); + } + + /** Starts a root APM span and activates it, so all LLMObs spans created within share a trace. */ + private static AgentScope startRootApmScope() { + AgentSpan root = AgentTracer.get().buildSpan("apm", "http.server.request").start(); + return AgentTracer.activateSpan(root); + } + + @Test + void agentSpanStoresOwnIdAndNameAsPagent() throws Exception { + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "my-agent"); + try { + AgentSpan inner = innerSpan(agentSpan); + String parentAgentSpanId = (String) inner.getTag(PAGENT_SPAN_ID_TAG); + String parentAgentName = (String) inner.getTag(PAGENT_NAME_TAG); + + assertEquals(String.valueOf(inner.getSpanId()), parentAgentSpanId); + assertEquals("my-agent", parentAgentName); + } finally { + agentSpan.finish(); + apmScope.span().finish(); + } + } + } + + @Test + void agentSpanNameIsPreservedRegardlessOfCharacters() throws Exception { + // Agent names are stored as-is — no character restrictions since the only wire consumer + // is the msgpack intake mapper, which accepts any string. + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "router~v2,résumé"); + try { + AgentSpan inner = innerSpan(agentSpan); + assertEquals(String.valueOf(inner.getSpanId()), inner.getTag(PAGENT_SPAN_ID_TAG)); + assertEquals("router~v2,résumé", inner.getTag(PAGENT_NAME_TAG)); + } finally { + agentSpan.finish(); + apmScope.span().finish(); + } + } + } + + @Test + void innerAgentNameOverridesOuterAgentNameForDescendants() throws Exception { + // Descendants of an inner agent must see the inner agent's name, not the outer agent's. + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan outerAgent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "outer-agent"); + try { + DDLLMObsSpan innerAgent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "inner-agent"); + try { + DDLLMObsSpan tool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool"); + try { + AgentSpan toolInner = innerSpan(tool); + AgentSpan innerAgentSpan = innerSpan(innerAgent); + assertEquals( + String.valueOf(innerAgentSpan.getSpanId()), toolInner.getTag(PAGENT_SPAN_ID_TAG)); + assertEquals("inner-agent", toolInner.getTag(PAGENT_NAME_TAG)); + } finally { + tool.finish(); + } + } finally { + innerAgent.finish(); + } + } finally { + outerAgent.finish(); + apmScope.span().finish(); + } + } + } + + @Test + void nonAgentChildUnderAgentInheritsAttribution() throws Exception { + // All LLMObs spans share the same APM trace so the trace-ID consistency gate passes. + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "parent-agent"); + try { + AgentSpan agentInner = innerSpan(agentSpan); + String expectedParentAgentSpanId = 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(expectedParentAgentSpanId, toolInner.getTag(PAGENT_SPAN_ID_TAG)); + assertEquals("parent-agent", toolInner.getTag(PAGENT_NAME_TAG)); + } finally { + toolSpan.finish(); + } + } finally { + agentSpan.finish(); + apmScope.span().finish(); + } + } + } + + @Test + void transitiveInheritanceAgentToLlmToTool() throws Exception { + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "root-agent"); + try { + AgentSpan agentInner = innerSpan(agentSpan); + String expectedParentAgentSpanId = 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(expectedParentAgentSpanId, toolInner.getTag(PAGENT_SPAN_ID_TAG)); + assertEquals("root-agent", toolInner.getTag(PAGENT_NAME_TAG)); + } finally { + toolSpan.finish(); + } + } finally { + llmSpan.finish(); + } + } finally { + agentSpan.finish(); + apmScope.span().finish(); + } + } + } + + @Test + void noAgentAncestorProducesNoPagentTags() throws Exception { + try (AgentScope apmScope = startRootApmScope()) { + 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(); + apmScope.span().finish(); + } + } + } + + @Test + void innerAgentFinishRestoresOuterAgentPropagationTags() throws Exception { + // Outer agent's LLMObsContext is active. Inner agent starts (its own context pushed on top). + // After inner agent finishes its context is popped, restoring outer agent's context. + // A sibling span then sees outer agent's attribution via LLMObsContext. + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan outerAgent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "outer-agent"); + try { + AgentSpan outerInner = innerSpan(outerAgent); + String outerParentAgentSpanId = String.valueOf(outerInner.getSpanId()); + + DDLLMObsSpan innerAgent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "inner-agent"); + innerAgent.finish(); + // After finish(), inner agent's LLMObsContext scope is closed — outer agent's is restored. + + // Sibling span created now sees outer agent's attribution via the restored LLMObsContext. + DDLLMObsSpan siblingTool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "sibling-tool"); + try { + AgentSpan siblingInner = innerSpan(siblingTool); + assertEquals(outerParentAgentSpanId, siblingInner.getTag(PAGENT_SPAN_ID_TAG)); + assertEquals("outer-agent", siblingInner.getTag(PAGENT_NAME_TAG)); + } finally { + siblingTool.finish(); + } + } finally { + outerAgent.finish(); + apmScope.span().finish(); + } + } + } + + @Test + void staleContextInDifferentTraceDoesNotInheritPagent() throws Exception { + // Create an agent span in one APM trace, then create a non-agent span in a different APM + // trace. The stale LLMObsContext from the first trace must not leak pagent attribution. + AgentSpan firstRoot = AgentTracer.get().buildSpan("apm", "http.request.1").start(); + AgentScope firstScope = AgentTracer.activateSpan(firstRoot); + + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "first-trace-agent"); + // agentSpan's LLMObsContext is now active in the current thread + + // Close the first APM scope (but do NOT close agentSpan's scope yet) — simulate a context + // leak where the LLMObsContext outlives the APM scope it was created in. + firstScope.close(); + + // Start a fresh APM root (different trace) while the first trace's LLMObsContext is active + AgentSpan secondRoot = AgentTracer.get().buildSpan("apm", "http.request.2").start(); + AgentScope secondScope = AgentTracer.activateSpan(secondRoot); + try { + DDLLMObsSpan toolInSecondTrace = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "second-trace-tool"); + try { + // The tool span's APM parent is secondRoot (different trace from agentSpan). + // The trace-ID gate must block inheritance from the stale LLMObsContext. + AgentSpan toolInner = innerSpan(toolInSecondTrace); + assertNull(toolInner.getTag(PAGENT_SPAN_ID_TAG)); + assertNull(toolInner.getTag(PAGENT_NAME_TAG)); + } finally { + toolInSecondTrace.finish(); + } + } finally { + secondScope.close(); + secondRoot.finish(); + agentSpan.finish(); + firstRoot.finish(); + } + } + + @Test + void manifestNameOverridesSpanNameForPagent() throws Exception { + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "span-name"); + try { + AgentSpan inner = innerSpan(agentSpan); + assertEquals("span-name", inner.getTag(PAGENT_NAME_TAG)); + + agentSpan.annotateAgentManifest( + LLMObs.AgentManifest.builder().name("manifest-name").build()); + + // The internal tag (used by the serializer for this span's agent_attribution) updates + // to the manifest name. Context-propagated pagent name for children remains the span name. + assertEquals("manifest-name", inner.getTag(PAGENT_NAME_TAG)); + } finally { + agentSpan.finish(); + apmScope.span().finish(); + } + } + } +} diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java index fdfef883615..0f221ac5bd5 100644 --- a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java @@ -27,6 +27,7 @@ class DDLLMObsSpanAgentVersionTest { private static final String AGENT_VERSION_TAG = "_ml_obs_tag." + LLMObsTags.AGENT_VERSION; private static final Field SPAN_FIELD; + private static final Field STANDALONE_APM_SCOPE_FIELD; private static CoreTracer tracer; @@ -34,6 +35,8 @@ class DDLLMObsSpanAgentVersionTest { try { SPAN_FIELD = DDLLMObsSpan.class.getDeclaredField("span"); SPAN_FIELD.setAccessible(true); + STANDALONE_APM_SCOPE_FIELD = DDLLMObsSpan.class.getDeclaredField("standaloneApmScope"); + STANDALONE_APM_SCOPE_FIELD.setAccessible(true); } catch (ReflectiveOperationException error) { throw new ExceptionInInitializerError(error); } @@ -138,11 +141,19 @@ void noVersionSetAnywhereMeansNoTagOnAnySpanInTheSubtree() { } @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. + void childDoesNotInheritAgentVersionWhenStaleContextIsFromADifferentTrace() + throws ReflectiveOperationException { + // Simulates a stale LLMObsContext (e.g. leaked across an async boundary): the agent's + // LLMObs context is still active but its APM scope is closed (as it would be on a different + // thread or after an async handoff). The child starts a fresh APM trace and the + // trace-consistency gate must skip inheritance. DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "stale-agent", "stale-v1"); + // Close the standalone APM scope to simulate async boundary — LLMObs context leaks but + // APM scope does not. + AgentScope apmScope = (AgentScope) STANDALONE_APM_SCOPE_FIELD.get(agent); + if (apmScope != null) { + apmScope.close(); + } try { DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool1", null); try (AgentScope childScope = AgentTracer.activateSpan(spanOf(child))) { diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CommonTags.java b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CommonTags.java index 33c83be9778..8ffa9ca3a5a 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CommonTags.java +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CommonTags.java @@ -32,6 +32,8 @@ interface CommonTags { String PARENT_ID = TAG_PREFIX + "parent_id"; String SESSION_ID = TAG_PREFIX + LLMObsTags.SESSION_ID; String AGENT_VERSION = TAG_PREFIX + LLMObsTags.AGENT_VERSION; + String PAGENT_SPAN_ID = TAG_PREFIX + LLMObsTags.PAGENT_SPAN_ID; + String PAGENT_NAME = TAG_PREFIX + LLMObsTags.PAGENT_NAME; String TOOL_DEFINITIONS = TAG_PREFIX + "tool_definitions"; diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java index a094eff8ab0..aebe9f035d2 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java @@ -133,6 +133,20 @@ protected void doAfterStart(@Nonnull AgentSpan span) { } } span.setTag(CommonTags.PARENT_ID, parentSpanId); + + // Inherit agent attribution only when the LLMObs context belongs to the same trace. + // Mirrors the gate in DDLLMObsSpan: a stale LLMObsContext from a different async trace + // must not stamp its agent ID onto this span. + if (parent != null && parent.getTraceId() == span.getTraceId()) { + String parentAgentSpanId = LLMObsContext.currentParentAgentSpanId(); + if (parentAgentSpanId != null) { + span.setTag(CommonTags.PAGENT_SPAN_ID, parentAgentSpanId); + String parentAgentName = LLMObsContext.currentParentAgentName(); + if (parentAgentName != null) { + span.setTag(CommonTags.PAGENT_NAME, parentAgentName); + } + } + } } super.doAfterStart(span); } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsTags.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsTags.java index 4ae11578d2e..0a7931589b5 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsTags.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObsTags.java @@ -15,4 +15,8 @@ public class LLMObsTags { public static final String MODEL_PROVIDER = "model_provider"; public static final String TOOL_DEFINITIONS = "tool_definitions"; public static final String AGENT_MANIFEST = "agent_manifest"; + + // Agent attribution + public static final String PAGENT_SPAN_ID = "pagent_span_id"; + public static final String PAGENT_NAME = "pagent_name"; } 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 42b7222e7d7..86622291dd7 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 @@ -79,6 +79,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[] AGENT_MANIFEST_KEY = "agent_manifest".getBytes(StandardCharsets.UTF_8); @@ -107,6 +111,10 @@ 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 + LLMObsTags.PAGENT_SPAN_ID; + private static final String PAGENT_NAME_TAG_INTERNAL_FULL = + LLMOBS_TAG_PREFIX + LLMObsTags.PAGENT_NAME; private final MetaWriter metaWriter = new MetaWriter(); private final int size; @@ -348,7 +356,9 @@ private static final class MetaWriter implements MetadataConsumer { LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_VERSION, LLMOBS_TAG_PREFIX + LLMObsTags.TOOL_DEFINITIONS, LLMOBS_TAG_PREFIX + LLMObsTags.METADATA, - LLMOBS_TAG_PREFIX + LLMObsTags.AGENT_MANIFEST))); + LLMOBS_TAG_PREFIX + LLMObsTags.AGENT_MANIFEST, + PAGENT_SPAN_ID_TAG_INTERNAL_FULL, + PAGENT_NAME_TAG_INTERNAL_FULL))); MetaWriter withWritable(Writable writable, Map errorInfo) { this.writable = writable; @@ -391,6 +401,11 @@ public void accept(Metadata metadata) { String inputPromptTag = LLMOBS_TAG_PREFIX + INPUT_PROMPT; boolean hasInput = tagsToRemapToMeta.containsKey(inputTag); boolean hasInputPrompt = tagsToRemapToMeta.containsKey(inputPromptTag); + Object parentAgentSpanIdVal = tagsToRemapToMeta.get(PAGENT_SPAN_ID_TAG_INTERNAL_FULL); + boolean hasAgentAttribution = + parentAgentSpanIdVal instanceof String && !((String) parentAgentSpanIdVal).isEmpty(); + boolean hasAgentAttributionName = + tagsToRemapToMeta.containsKey(PAGENT_NAME_TAG_INTERNAL_FULL); Object inputPrompt = null; if (hasInputPrompt) { if (spanKind.equals(Tags.LLMOBS_LLM_SPAN_KIND)) { @@ -425,12 +440,20 @@ 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. + // When pagent_span_id exists but is invalid (non-string or empty), the whole + // agent_attribution block is skipped; subtract 1 for that entry too. + boolean hasInvalidParentAgentSpanId = + tagsToRemapToMeta.containsKey(PAGENT_SPAN_ID_TAG_INTERNAL_FULL) && !hasAgentAttribution; 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) + - (hasInvalidParentAgentSpanId ? 1 : 0); writable.writeUTF8(META); writable.startMap(metaSize); writable.writeUTF8(SPAN_KIND); @@ -461,7 +484,28 @@ 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")) { + if (!hasAgentAttribution) { + // Value was invalid (non-string or empty); skip — subtracted from metaSize above. + continue; + } + // 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.writeString((String) parentAgentSpanIdVal, 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 425666c3fda..f2054591432 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 @@ -4,6 +4,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -686,6 +687,31 @@ void testLLMObsSpanMapperPreservesStringRetrievalOutput() 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 testLLMObsSpanProcessorCanDropSpan() throws Exception { LLMObs.registerProcessor(span -> "true".equals(span.getTag("drop")) ? null : span); @@ -917,6 +943,55 @@ private static List> serialize(List trace, LLMObsSpa return (List>) result.get("spans"); } + @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 testNoAgentAttributionBlockWhenParentAgentSpanIdAbsent() 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 008baedfcfd..e925a04f42f 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 @@ -16,6 +16,9 @@ private LLMObsContext() { private static final ContextKey SESSION_ID_KEY = ContextKey.named("llmobs_session_id"); private static final ContextKey AGENT_VERSION_KEY = ContextKey.named("llmobs_agent_version"); + 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"); /** * Attach an LLMObs span context, leaving any inherited session_id/agent_version from an enclosing @@ -56,6 +59,37 @@ public static ContextScope attach(AgentSpanContext ctx, String sessionId, String .attach(); } + /** + * Attach an LLMObs span context, propagating session_id, agent_version, and agent attribution to + * descendant LLMObs spans. pagentSpanId identifies the nearest agent-kind ancestor; pagentName is + * its name (may be null). Both pagent keys are always written — null clears stale values from an + * outer scope. + */ + public static ContextScope attach( + AgentSpanContext ctx, + String sessionId, + String agentVersion, + String parentAgentSpanId, + String parentAgentName) { + Context updated = Context.current().with(CONTEXT_KEY, ctx); + if (sessionId != null && !sessionId.isEmpty()) { + updated = updated.with(SESSION_ID_KEY, sessionId); + } + updated = + updated.with( + AGENT_VERSION_KEY, + agentVersion != null && !agentVersion.isEmpty() ? agentVersion : null); + // Always write both pagent keys. Per the Context API contract (Context.java: "Mapping to a + // null value will remove the key-value from the context copy"), with(key, null) clears any + // stale value inherited from an outer scope. This prevents two leakage scenarios: + // 1. An unsafe-named inner agent must not let descendants see the outer agent's name. + // 2. A non-agent span whose trace-ID gate blocked attribution must not let its same-trace + // children pick up a pagent ID that belongs to a different trace. + updated = updated.with(PAGENT_SPAN_ID_KEY, parentAgentSpanId); + updated = updated.with(PAGENT_NAME_KEY, parentAgentName); + return updated.attach(); + } + public static AgentSpanContext current() { return Context.current().get(CONTEXT_KEY); } @@ -74,4 +108,16 @@ public static String currentSessionId() { public static String currentAgentVersion() { return Context.current().get(AGENT_VERSION_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/test/java/datadog/trace/api/llmobs/LLMObsContextTest.java b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsContextTest.java index d60ac03228b..bf4a6cfd5b8 100644 --- a/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsContextTest.java +++ b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsContextTest.java @@ -156,4 +156,108 @@ void childScopeInheritsParentAgentVersion() { } } } + + // ── 5-arg attach (pagent attribution) ──────────────────────────────────── + + @Test + void currentParentAgentSpanIdReturnsNullWhenNoContextAttached() { + assertNull(LLMObsContext.currentParentAgentSpanId()); + } + + @Test + void currentParentAgentNameReturnsNullWhenNoContextAttached() { + assertNull(LLMObsContext.currentParentAgentName()); + } + + @Test + void fiveArgAttachStoresAllFields() { + AgentSpanContext ctx = mock(AgentSpanContext.class); + try (ContextScope scope = + LLMObsContext.attach(ctx, "session-1", "v2", "span-99", "my-agent")) { + assertEquals(ctx, LLMObsContext.current()); + assertEquals("session-1", LLMObsContext.currentSessionId()); + assertEquals("v2", LLMObsContext.currentAgentVersion()); + assertEquals("span-99", LLMObsContext.currentParentAgentSpanId()); + assertEquals("my-agent", LLMObsContext.currentParentAgentName()); + } + assertNull(LLMObsContext.current()); + assertNull(LLMObsContext.currentSessionId()); + assertNull(LLMObsContext.currentAgentVersion()); + assertNull(LLMObsContext.currentParentAgentSpanId()); + assertNull(LLMObsContext.currentParentAgentName()); + } + + @Test + void fiveArgAttachWithNullSessionIdIgnoresSessionId() { + AgentSpanContext ctx = mock(AgentSpanContext.class); + try (ContextScope scope = LLMObsContext.attach(ctx, null, null, null, null)) { + assertNull(LLMObsContext.currentSessionId()); + assertNull(LLMObsContext.currentAgentVersion()); + assertNull(LLMObsContext.currentParentAgentSpanId()); + assertNull(LLMObsContext.currentParentAgentName()); + } + } + + @Test + void fiveArgAttachWithEmptySessionIdIgnoresSessionId() { + AgentSpanContext ctx = mock(AgentSpanContext.class); + try (ContextScope scope = LLMObsContext.attach(ctx, "", "", null, null)) { + assertNull(LLMObsContext.currentSessionId()); + assertNull(LLMObsContext.currentAgentVersion()); + } + } + + @Test + void fiveArgAttachNullPagentClearsStaleValuesFromOuterScope() { + // When an inner (non-agent) span attaches with null pagent keys, the outer agent's + // pagent ID and name must not leak through to that span's descendants. + AgentSpanContext outer = mock(AgentSpanContext.class); + AgentSpanContext inner = mock(AgentSpanContext.class); + try (ContextScope outerScope = + LLMObsContext.attach(outer, "s", "v1", "agent-span-id", "outer-agent")) { + assertEquals("agent-span-id", LLMObsContext.currentParentAgentSpanId()); + assertEquals("outer-agent", LLMObsContext.currentParentAgentName()); + + try (ContextScope innerScope = LLMObsContext.attach(inner, null, null, null, null)) { + assertNull(LLMObsContext.currentParentAgentSpanId()); + assertNull(LLMObsContext.currentParentAgentName()); + } + + // Outer values are restored after inner scope closes. + assertEquals("agent-span-id", LLMObsContext.currentParentAgentSpanId()); + assertEquals("outer-agent", LLMObsContext.currentParentAgentName()); + } + } + + @Test + void fiveArgAttachInnerAgentOverridesOuterAgentForDescendants() { + AgentSpanContext outer = mock(AgentSpanContext.class); + AgentSpanContext inner = mock(AgentSpanContext.class); + try (ContextScope outerScope = + LLMObsContext.attach(outer, null, null, "outer-span-id", "outer-agent")) { + try (ContextScope innerScope = + LLMObsContext.attach(inner, null, null, "inner-span-id", "inner-agent")) { + assertEquals("inner-span-id", LLMObsContext.currentParentAgentSpanId()); + assertEquals("inner-agent", LLMObsContext.currentParentAgentName()); + } + assertEquals("outer-span-id", LLMObsContext.currentParentAgentSpanId()); + assertEquals("outer-agent", LLMObsContext.currentParentAgentName()); + } + } + + @Test + void fiveArgAttachNullPagentNameClearsNameButNotSpanId() { + // An agent with a null name (e.g. manifest not yet set) must not let outer agent's name + // leak into its scope — only the span ID is set. + AgentSpanContext outer = mock(AgentSpanContext.class); + AgentSpanContext inner = mock(AgentSpanContext.class); + try (ContextScope outerScope = + LLMObsContext.attach(outer, null, null, "outer-span-id", "outer-agent")) { + try (ContextScope innerScope = + LLMObsContext.attach(inner, null, null, "inner-span-id", null)) { + assertEquals("inner-span-id", LLMObsContext.currentParentAgentSpanId()); + assertNull(LLMObsContext.currentParentAgentName()); + } + } + } }