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..2237f4b4da0 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; @@ -63,6 +64,11 @@ public class DDLLMObsSpan implements LLMObsSpan { private final String mlApp; private final ContextScope scope; private final boolean hasSessionId; + // Non-null only for a standalone agent span (no ambient APM root). Activating the underlying + // APM span ensures that subsequently auto-instrumented outgoing calls (HTTP, gRPC, …) are + // created as children of this agent, keeping all LLMObs spans in one APM trace so the + // trace-ID consistency gate in the constructor does not break parent_id / session_id inheritance. + private final AgentScope standaloneApmScope; private boolean finished = false; @@ -135,6 +141,16 @@ public DDLLMObsSpan( 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); + + // When there is no ambient APM root, activate this span so subsequent auto-instrumented + // outgoing calls become children of it, keeping all LLMObs spans within one APM trace. + // Without this, each new DDLLMObsSpan would start a fresh APM root with a different trace ID, + // causing the trace-ID consistency gate above to reject parent_id and session_id inheritance. + if (Tags.LLMOBS_AGENT_SPAN_KIND.equals(kind) && span.getLocalRootSpan() == span) { + standaloneApmScope = AgentTracer.activateSpan(span); + } else { + standaloneApmScope = null; + } } @Override @@ -493,6 +509,9 @@ public void finish() { } span.finish(); scope.close(); + if (standaloneApmScope != null) { + standaloneApmScope.close(); + } finished = true; boolean isRootSpan = span.getLocalRootSpan() == span; LLMObsMetricCollector.get() 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 bef904409a2..60c2398fd3a 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 @@ -228,6 +228,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"() { @@ -510,6 +513,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/DDLLMObsSpanStandaloneApmScopeTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanStandaloneApmScopeTest.java new file mode 100644 index 00000000000..72ba76651ef --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanStandaloneApmScopeTest.java @@ -0,0 +1,87 @@ +package datadog.trace.llmobs.domain; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; + +import datadog.trace.agent.tooling.TracerInstaller; +import datadog.trace.api.WellKnownTags; +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 DDLLMObsSpanStandaloneApmScopeTest { + + private static final Field STANDALONE_APM_SCOPE_FIELD; + + private static CoreTracer tracer; + + static { + try { + STANDALONE_APM_SCOPE_FIELD = DDLLMObsSpan.class.getDeclaredField("standaloneApmScope"); + STANDALONE_APM_SCOPE_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 AgentScope standaloneApmScope(DDLLMObsSpan span) throws IllegalAccessException { + return (AgentScope) STANDALONE_APM_SCOPE_FIELD.get(span); + } + + @Test + void standaloneAgentSpanActivatesApmScopeForOutgoingPropagation() throws Exception { + // A standalone agent span (no ambient APM root) must activate its APM scope so that + // auto-instrumented outgoing calls created within the same workflow become children of this + // agent's APM span, keeping all LLMObs spans in one APM trace. Without this, each new + // DDLLMObsSpan would start a fresh APM root with a different trace ID, causing the + // trace-ID consistency gate to break parent_id and session_id inheritance. + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "standalone-agent"); + try { + assertNotNull(standaloneApmScope(agentSpan)); + } finally { + agentSpan.finish(); + } + } + + @Test + void nonStandaloneAgentSpanDoesNotActivateApmScope() throws Exception { + // When an ambient APM scope already exists, the agent span is NOT the local root, so it must + // NOT activate an additional APM scope — doing so would corrupt the active scope stack and + // cause APM spans created after finish() to be incorrectly parented. + AgentSpan root = AgentTracer.get().buildSpan("apm", "http.server.request").start(); + try (AgentScope rootScope = AgentTracer.activateSpan(root)) { + DDLLMObsSpan agentSpan = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "non-standalone-agent"); + try { + assertNull(standaloneApmScope(agentSpan)); + } finally { + agentSpan.finish(); + } + } finally { + root.finish(); + } + } +}