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 42c6790a0c2..473d118cdc3 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.LLMObsSampler; import datadog.trace.api.llmobs.LLMObsSpan; import datadog.trace.api.llmobs.LLMObsTags; import datadog.trace.api.telemetry.LLMObsMetricCollector; @@ -51,6 +52,8 @@ 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 SAMPLE_RATE_TAG_INTERNAL = "sample_rate"; + private static final String SAMPLING_DECISION_TAG_INTERNAL = "sampling_decision"; 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; @@ -64,6 +67,8 @@ public class DDLLMObsSpan implements LLMObsSpan { private static final Logger LOGGER = LoggerFactory.getLogger(DDLLMObsSpan.class); + private static final LLMObsSampler CONFIGURED_SAMPLER = LLMObsSampler.fromConfig(); + private final AgentSpan span; private final String spanKind; private final String mlApp; @@ -94,6 +99,26 @@ public DDLLMObsSpan( @Nonnull String serviceName, WellKnownTags wellKnownTags, String agentVersion) { + this( + kind, + spanName, + mlApp, + sessionId, + serviceName, + wellKnownTags, + agentVersion, + CONFIGURED_SAMPLER); + } + + DDLLMObsSpan( + @Nonnull String kind, + String spanName, + @Nonnull String mlApp, + String sessionId, + @Nonnull String serviceName, + WellKnownTags wellKnownTags, + String agentVersion, + @Nonnull LLMObsSampler sampler) { if (null == spanName || spanName.isEmpty()) { spanName = kind; @@ -122,12 +147,18 @@ public DDLLMObsSpan( spanKind = kind; this.mlApp = mlApp; span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp); - // Resolve effective parent_id and session_id from the LLMObs context, both gated on - // trace-id consistency. A stale context from a different trace (e.g. async boundary - // leakage) must not contribute either tag. + // Resolve effective parent_id, session_id, agent_version, agent attribution and sampling + // decision from the LLMObs context, all gated on trace-id consistency. A stale context from a + // different trace (e.g. async boundary leakage) must not contribute any of them. Every + // inherited value is read inside the one same-trace branch below, so a newly propagated tag + // cannot ship with a weaker gate of its own. AgentSpanContext parent = LLMObsContext.current(); String parentSpanID = LLMObsContext.ROOT_SPAN_ID; String resolvedAgentVersion = agentVersion; + String sampleRate = null; + String samplingDecision = null; + String resolvedParentAgentSpanId = null; + String resolvedParentAgentName = null; if (null != parent) { if (parent.getTraceId() != span.getTraceId()) { LOGGER.error( @@ -156,9 +187,34 @@ public DDLLMObsSpan( resolvedAgentVersion = inherited; } } + // Inherit the sampling decision from the context if present. + sampleRate = LLMObsContext.currentSampleRate(); + samplingDecision = LLMObsContext.currentSamplingDecision(); + // Inherit agent attribution: the nearest agent-kind ancestor. Overridden just below when + // this span is itself an agent. + resolvedParentAgentSpanId = LLMObsContext.currentParentAgentSpanId(); + resolvedParentAgentName = LLMObsContext.currentParentAgentName(); } } + // An agent span is its own descendants' nearest agent ancestor, replacing anything inherited. + // Use the span name as the initial pagent name; annotateAgentManifest() will update it to the + // manifest name if one is provided later. + if (Tags.LLMOBS_AGENT_SPAN_KIND.equals(kind)) { + resolvedParentAgentSpanId = String.valueOf(span.getSpanId()); + resolvedParentAgentName = spanName; + } + + if (samplingDecision == null || sampleRate == null) { + sampleRate = sampler.formattedRate(); + samplingDecision = + sampler.sample(span.getTraceId().toLong()) + ? LLMObsContext.SAMPLING_DECISION_SAMPLED + : LLMObsContext.SAMPLING_DECISION_DROPPED; + } + span.setTag(LLMOBS_TAG_PREFIX + SAMPLE_RATE_TAG_INTERNAL, sampleRate); + span.setTag(LLMOBS_TAG_PREFIX + SAMPLING_DECISION_TAG_INTERNAL, samplingDecision); + this.hasSessionId = sessionId != null && !sessionId.isEmpty(); if (this.hasSessionId) { span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId); @@ -167,29 +223,6 @@ public DDLLMObsSpan( span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.AGENT_VERSION, resolvedAgentVersion); } span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID); - - // 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); @@ -198,12 +231,15 @@ public DDLLMObsSpan( } } - // Propagate the effective sessionId and agent attribution to descendant LLMObs spans. + // Propagate the effective sessionId, agent_version, sampling decision and agent attribution + // to descendant LLMObs spans via the context. scope = LLMObsContext.attach( span.spanContext(), sessionId, resolvedAgentVersion, + sampleRate, + samplingDecision, resolvedParentAgentSpanId, resolvedParentAgentName); diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanSamplingTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanSamplingTest.java new file mode 100644 index 00000000000..5b411e32edf --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanSamplingTest.java @@ -0,0 +1,134 @@ +package datadog.trace.llmobs.domain; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.trace.agent.tooling.TracerInstaller; +import datadog.trace.api.WellKnownTags; +import datadog.trace.api.llmobs.LLMObsSampler; +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 where the head-based sampling decision is made and where it is inherited: the decision is + * computed once, at the root of an LLMObs trace, and every descendant reports it verbatim. + */ +class DDLLMObsSpanSamplingTest { + private static final String SAMPLE_RATE_TAG = "_ml_obs_tag.sample_rate"; + private static final String SAMPLING_DECISION_TAG = "_ml_obs_tag.sampling_decision"; + + 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 stampsRetainedDecisionAtTheDefaultRate() throws IllegalAccessException { + // The fields are stamped at every rate, including 1.0, matching dd-trace-py. + DDLLMObsSpan llmObsSpan = newSpan(new LLMObsSampler(1.0)); + try { + AgentSpan span = spanOf(llmObsSpan); + assertEquals("1", span.getTag(SAMPLING_DECISION_TAG)); + assertEquals("1", span.getTag(SAMPLE_RATE_TAG)); + } finally { + llmObsSpan.finish(); + } + } + + @Test + void stampsDroppedDecisionOnRoot() throws IllegalAccessException { + DDLLMObsSpan llmObsSpan = newSpan(new LLMObsSampler(0.0)); + try { + AgentSpan span = spanOf(llmObsSpan); + assertEquals("0", span.getTag(SAMPLING_DECISION_TAG)); + assertEquals("0", span.getTag(SAMPLE_RATE_TAG)); + } finally { + llmObsSpan.finish(); + } + } + + @Test + void childInheritsTheRootDecisionInsteadOfRecomputingIt() throws IllegalAccessException { + // The child's sampler drops everything. If the decision were recomputed per span, the child + // would report "0" and the trace would be torn in half at the intake. + DDLLMObsSpan root = newSpan(new LLMObsSampler(1.0)); + try { + AgentSpan rootSpan = spanOf(root); + // Inheritance is gated on the two spans sharing an APM trace, so the root's APM span has to + // be active for the child to be started under it. + try (AgentScope ignored = AgentTracer.activateSpan(rootSpan)) { + DDLLMObsSpan child = newSpan(new LLMObsSampler(0.0)); + try { + AgentSpan childSpan = spanOf(child); + assertEquals("1", childSpan.getTag(SAMPLING_DECISION_TAG)); + assertEquals( + rootSpan.getTag(SAMPLE_RATE_TAG), + childSpan.getTag(SAMPLE_RATE_TAG), + "every span in a trace must report the rate the decision was made at"); + } finally { + child.finish(); + } + } + } finally { + root.finish(); + } + } + + @Test + void childOfADroppedRootStaysDropped() throws IllegalAccessException { + // Symmetric case. Because a decision is stamped at every rate, "the context carries a decision" + // is an unambiguous signal, so a child never mistakes an inherited drop for being a root. + DDLLMObsSpan root = newSpan(new LLMObsSampler(0.0)); + try { + try (AgentScope ignored = AgentTracer.activateSpan(spanOf(root))) { + DDLLMObsSpan child = newSpan(new LLMObsSampler(1.0)); + try { + AgentSpan childSpan = spanOf(child); + assertEquals("0", childSpan.getTag(SAMPLING_DECISION_TAG)); + assertEquals("0", childSpan.getTag(SAMPLE_RATE_TAG)); + } finally { + child.finish(); + } + } + } finally { + root.finish(); + } + } + + private static DDLLMObsSpan newSpan(LLMObsSampler sampler) { + WellKnownTags tags = + new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java"); + return new DDLLMObsSpan( + Tags.LLMOBS_LLM_SPAN_KIND, "span", "ml-app", null, "service", tags, null, sampler); + } + + private static AgentSpan spanOf(DDLLMObsSpan llmObsSpan) throws IllegalAccessException { + return (AgentSpan) SPAN_FIELD.get(llmObsSpan); + } +} 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 8ffa9ca3a5a..ab23b0131c6 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 SAMPLE_RATE = TAG_PREFIX + "sample_rate"; + String SAMPLING_DECISION = TAG_PREFIX + "sampling_decision"; String PAGENT_SPAN_ID = TAG_PREFIX + LLMObsTags.PAGENT_SPAN_ID; String PAGENT_NAME = TAG_PREFIX + LLMObsTags.PAGENT_NAME; 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 aebe9f035d2..796490bdda5 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 @@ -10,6 +10,7 @@ import datadog.trace.api.DDTraceApiInfo; import datadog.trace.api.WellKnownTags; import datadog.trace.api.llmobs.LLMObsContext; +import datadog.trace.api.llmobs.LLMObsSampler; import datadog.trace.api.telemetry.LLMObsMetricCollector; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -44,6 +45,7 @@ public class OpenAiDecorator extends ClientDecorator { private final boolean llmObsEnabled = Config.get().isLlmObsEnabled(); private final WellKnownTags wellKnownTags = Config.get().getWellKnownTags(); + private final LLMObsSampler sampler = LLMObsSampler.fromConfig(); public AgentSpan startSpan(ClientOptions clientOptions) { AgentSpan span = AgentTracer.startSpan(INTEGRATION, SPAN_NAME); @@ -110,10 +112,16 @@ protected void doAfterStart(@Nonnull AgentSpan span) { // Resolve the LLMObs parent context, gated on trace-id consistency: a stale context // from a different trace (e.g. async boundary leakage) must not contribute parent_id, - // session_id, or agent_version to this span. Matches DDLLMObsSpan's manual-span gate. + // session_id, agent_version, agent attribution, or a sampling verdict to this span. + // Matches DDLLMObsSpan's manual-span gate. One flag drives every inherited value, so a + // new propagated tag cannot accidentally ship with a weaker gate of its own. AgentSpanContext parent = LLMObsContext.current(); + boolean inheritable = parent != null && parent.getTraceId().equals(span.getTraceId()); + String parentSpanId = LLMObsContext.ROOT_SPAN_ID; - if (parent != null && parent.getTraceId() == span.getTraceId()) { + String samplingDecision = null; + String sampleRate = null; + if (inheritable) { parentSpanId = String.valueOf(parent.getSpanId()); // Inherit session_id from the active LLMObs parent (e.g. a manual workflow span). @@ -131,13 +139,9 @@ protected void doAfterStart(@Nonnull AgentSpan span) { if (agentVersion != null && !agentVersion.isEmpty()) { span.setTag(CommonTags.AGENT_VERSION, agentVersion); } - } - 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()) { + // Inherit agent attribution: the nearest agent-kind ancestor of this span. The name is + // only meaningful alongside an ID, so it is read inside the ID's branch. String parentAgentSpanId = LLMObsContext.currentParentAgentSpanId(); if (parentAgentSpanId != null) { span.setTag(CommonTags.PAGENT_SPAN_ID, parentAgentSpanId); @@ -146,7 +150,24 @@ protected void doAfterStart(@Nonnull AgentSpan span) { span.setTag(CommonTags.PAGENT_NAME, parentAgentName); } } + + samplingDecision = LLMObsContext.currentSamplingDecision(); + sampleRate = LLMObsContext.currentSampleRate(); + } + span.setTag(CommonTags.PARENT_ID, parentSpanId); + + // Compute the sampling decision if none was inherited (no LLMObs parent), which makes this + // span the root of its own LLMObs trace. Unlike the tags above, this cannot be skipped when + // there is nothing to inherit: an unstamped span is retained at any configured rate. + if (samplingDecision == null || sampleRate == null) { + sampleRate = sampler.formattedRate(); + samplingDecision = + sampler.sample(span.getTraceId().toLong()) + ? LLMObsContext.SAMPLING_DECISION_SAMPLED + : LLMObsContext.SAMPLING_DECISION_DROPPED; } + span.setTag(CommonTags.SAMPLING_DECISION, samplingDecision); + span.setTag(CommonTags.SAMPLE_RATE, sampleRate); } super.doAfterStart(span); } diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy index a16920209dc..0a8553e7930 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ChatCompletionServiceTest.groovy @@ -398,6 +398,8 @@ class ChatCompletionServiceTest extends OpenAiTest { "_ml_obs_metric.cache_read_input_tokens" Long } "_ml_obs_tag.parent_id" "undefined" + "_ml_obs_tag.sampling_decision" "1" + "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String "_ml_obs_tag.service" String "$CommonTags.DDTRACE_VERSION" String diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy index 8ca98ca3677..50529ee02a6 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/CompletionServiceTest.groovy @@ -180,6 +180,8 @@ class CompletionServiceTest extends OpenAiTest { "_ml_obs_metric.output_tokens" Long "_ml_obs_metric.total_tokens" Long "_ml_obs_tag.parent_id" "undefined" + "_ml_obs_tag.sampling_decision" "1" + "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String "_ml_obs_tag.service" String "$CommonTags.DDTRACE_VERSION" String diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy index 0a4f76ff47a..41e341284f3 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/EmbeddingServiceTest.groovy @@ -73,6 +73,8 @@ class EmbeddingServiceTest extends OpenAiTest { } "_ml_obs_tag.output" "[1 embedding(s) returned with size 1536]" "_ml_obs_tag.parent_id" "undefined" + "_ml_obs_tag.sampling_decision" "1" + "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String "_ml_obs_tag.service" String "$CommonTags.DDTRACE_VERSION" String diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy index 43228ce9818..b22389a4d41 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy @@ -475,6 +475,8 @@ class ResponseServiceTest extends OpenAiTest { "_ml_obs_metric.reasoning_output_tokens" Long "_ml_obs_metric.cache_read_input_tokens" Long "_ml_obs_tag.parent_id" "undefined" + "_ml_obs_tag.sampling_decision" "1" + "_ml_obs_tag.sample_rate" "1" "_ml_obs_tag.ml_app" String "$CommonTags.INTEGRATION" "openai" "_ml_obs_tag.service" String diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java index 75caec43331..66d4f6aa285 100644 --- a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java @@ -25,20 +25,16 @@ import org.junit.jupiter.api.Test; /** - * Verifies that auto-instrumented openai.request spans inherit session_id and agent_version from an - * active LLMObs parent context, and that a stale context left over from an unrelated trace does not - * leak either tag onto the span. Forked + @WithConfig used together so the LLMObs system property - * is in place before the agent installs and there's no leakage from prior test state. + * Mock OpenAI backend and request helpers, shared by the LLMObs forked tests in this file. * - *
The mock OpenAI backend returns a minimal 200 response — the test asserts on the span tag set - * by OpenAiDecorator.afterStart(), which runs before the HTTP response is parsed, so the response - * body shape doesn't matter for what's being tested. + *
Subclasses differ only in the {@code @WithConfig} values they declare. One class per
+ * configuration: {@code OpenAiDecorator} reads the LLMObs config once when its {@code DECORATE}
+ * singleton initializes, and {@code forkedTest} forks per test class ({@code forkEvery = 1}).
*/
-@WithConfig(key = "llmobs.enabled", value = "true")
-class LlmObsContextPropagationForkedTest extends AbstractInstrumentationTest {
+abstract class AbstractLlmObsOpenAiForkedTest extends AbstractInstrumentationTest {
- private static HttpServer mockServer;
- private static OpenAIClient openAiClient;
+ protected static HttpServer mockServer;
+ protected static OpenAIClient openAiClient;
@BeforeAll
static void setupMockOpenAi() throws IOException {
@@ -72,6 +68,41 @@ static void tearDownMockOpenAi() {
openAiClient = null;
}
+ protected static ChatCompletionCreateParams buildMinimalChatParams() {
+ return ChatCompletionCreateParams.builder()
+ .model(ChatModel.GPT_4O_MINI)
+ .addSystemMessage("")
+ .addUserMessage("")
+ .build();
+ }
+
+ protected static DDSpan findSpanByOperationName(List Runs at the default sample rate of 1.0. Drop-side coverage lives in {@link
+ * LlmObsZeroSampleRateForkedTest}.
+ *
+ * The mock OpenAI backend returns a minimal 200 response — the test asserts on the span tag set
+ * by OpenAiDecorator.afterStart(), which runs before the HTTP response is parsed, so the response
+ * body shape doesn't matter for what's being tested.
+ */
+@WithConfig(key = "llmobs.enabled", value = "true")
+class LlmObsContextPropagationForkedTest extends AbstractLlmObsOpenAiForkedTest {
+
@Test
void openAiRequestSpanInheritsSessionIdFromActiveContext() throws Exception {
String expectedSessionId = "session-propagation-test-abc";
@@ -138,19 +169,104 @@ void openAiRequestSpanInheritsAgentVersionFromActiveContext() throws Exception {
}
@Test
- void openAiRequestSpanDoesNotInheritSessionIdOrAgentVersionFromStaleCrossTraceContext()
- throws Exception {
+ void openAiRequestSpanInheritsDroppedSamplingDecisionFromActiveContext() throws Exception {
+ AgentSpan parentSpan = AgentTracer.startSpan("test", "parent");
+ try (ContextScope ignored1 = AgentTracer.activateSpan(parentSpan)) {
+ try (ContextScope ignored2 =
+ LLMObsContext.attach(
+ parentSpan.spanContext(),
+ null,
+ null,
+ "0.25",
+ LLMObsContext.SAMPLING_DECISION_DROPPED,
+ null,
+ null)) {
+ try {
+ openAiClient.chat().completions().create(buildMinimalChatParams());
+ } catch (Exception ignored) {
+ }
+ }
+ } finally {
+ parentSpan.finish();
+ }
+
+ writer.waitForTraces(1);
+ DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
+ assertNotNull(openAiSpan, "openai.request span should have been created");
+ assertEquals(
+ LLMObsContext.SAMPLING_DECISION_DROPPED,
+ openAiSpan.getTag("_ml_obs_tag.sampling_decision"));
+ assertEquals("0.25", openAiSpan.getTag("_ml_obs_tag.sample_rate"));
+ }
+
+ @Test
+ void openAiRequestSpanInheritsRetainedSamplingDecisionFromActiveContext() throws Exception {
+ AgentSpan parentSpan = AgentTracer.startSpan("test", "parent");
+ try (ContextScope ignored1 = AgentTracer.activateSpan(parentSpan)) {
+ try (ContextScope ignored2 =
+ LLMObsContext.attach(
+ parentSpan.spanContext(),
+ null,
+ null,
+ "1",
+ LLMObsContext.SAMPLING_DECISION_SAMPLED,
+ null,
+ null)) {
+ try {
+ openAiClient.chat().completions().create(buildMinimalChatParams());
+ } catch (Exception ignored) {
+ }
+ }
+ } finally {
+ parentSpan.finish();
+ }
+
+ writer.waitForTraces(1);
+ DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
+ assertNotNull(openAiSpan, "openai.request span should have been created");
+ assertEquals(
+ LLMObsContext.SAMPLING_DECISION_SAMPLED,
+ openAiSpan.getTag("_ml_obs_tag.sampling_decision"));
+ assertEquals("1", openAiSpan.getTag("_ml_obs_tag.sample_rate"));
+ }
+
+ @Test
+ void openAiRequestSpanComputesItsOwnSamplingDecisionWhenNoLlmObsContext() throws Exception {
+ try {
+ openAiClient.chat().completions().create(buildMinimalChatParams());
+ } catch (Exception ignored) {
+ }
+
+ // No verdict to inherit, so the span is the root of its own LLMObs trace and decides for
+ // itself. The rate of 1.0 retains every trace ID, so the verdict is deterministic without
+ // controlling the trace ID.
+ writer.waitForTraces(1);
+ DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
+ assertNotNull(openAiSpan, "openai.request span should have been created");
+ assertEquals(
+ LLMObsContext.SAMPLING_DECISION_SAMPLED,
+ openAiSpan.getTag("_ml_obs_tag.sampling_decision"));
+ assertEquals("1", openAiSpan.getTag("_ml_obs_tag.sample_rate"));
+ }
+
+ @Test
+ void openAiRequestSpanInheritsNothingFromStaleCrossTraceContext() throws Exception {
// Simulates a stale LLMObsContext leaked across an async boundary: the context is attached,
// but its span is never made the active tracer span, so the openai.request call below starts
// a brand-new trace and the trace-consistency gate in OpenAiDecorator must skip inheritance.
AgentSpan staleParent = AgentTracer.startSpan("test", "stale-parent");
try (ContextScope ignored =
- LLMObsContext.attach(staleParent.spanContext(), "stale-session", "stale-version")) {
+ LLMObsContext.attach(
+ staleParent.spanContext(),
+ "stale-session",
+ "stale-version",
+ "0.25",
+ LLMObsContext.SAMPLING_DECISION_DROPPED,
+ "stale-agent-span-id",
+ "stale-agent")) {
try {
openAiClient.chat().completions().create(buildMinimalChatParams());
} catch (Exception ignored2) {
- // Mock server returns no body — the SDK may throw on parse. The span we care about
- // is already created by the instrumentation advice before this point.
}
} finally {
staleParent.finish();
@@ -159,23 +275,46 @@ void openAiRequestSpanDoesNotInheritSessionIdOrAgentVersionFromStaleCrossTraceCo
writer.waitForTraces(2);
DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
assertNotNull(openAiSpan, "openai.request span should have been created");
+
+ // The stale "0"/"0.25" pair must not leak; the span falls through to deciding for itself at
+ // the configured rate of 1.0 instead.
+ assertEquals(
+ LLMObsContext.SAMPLING_DECISION_SAMPLED,
+ openAiSpan.getTag("_ml_obs_tag.sampling_decision"));
+ assertEquals("1", openAiSpan.getTag("_ml_obs_tag.sample_rate"));
+
+ // The same gate covers parent_id, session_id, agent_version and agent attribution: inheriting
+ // any of them would point this span at a parent in an unrelated trace and file it under an
+ // unrelated session or agent.
+ assertEquals(LLMObsContext.ROOT_SPAN_ID, openAiSpan.getTag("_ml_obs_tag.parent_id"));
assertNull(openAiSpan.getTag("_ml_obs_tag.session_id"));
assertNull(openAiSpan.getTag("_ml_obs_tag.agent_version"));
+ assertNull(openAiSpan.getTag("_ml_obs_tag.pagent_span_id"));
+ assertNull(openAiSpan.getTag("_ml_obs_tag.pagent_name"));
}
+}
- private static ChatCompletionCreateParams buildMinimalChatParams() {
- return ChatCompletionCreateParams.builder()
- .model(ChatModel.GPT_4O_MINI)
- .addSystemMessage("")
- .addUserMessage("")
- .build();
- }
+/**
+ * Verifies that an auto-instrumented openai.request span with no LLMObs parent is stamped as
+ * dropped when the sample rate is 0.
+ */
+@WithConfig(key = "llmobs.enabled", value = "true")
+@WithConfig(key = "llmobs.sample.rate", value = "0")
+class LlmObsZeroSampleRateForkedTest extends AbstractLlmObsOpenAiForkedTest {
- private static DDSpan findSpanByOperationName(List These are skipped while writing rather than removed from the span. {@link
+ * datadog.communication.serialization.msgpack.MsgPackWriter#format} re-invokes {@code map} on the
+ * same span instances after a buffer overflow, so a span mutated on the first pass serializes
+ * differently on the retry — a dropped span would lose its verdict and be re-emitted as retained.
+ */
+ private static final Set> traces, String operationName) {
+ return traces.stream()
+ .flatMap(List::stream)
+ .filter(s -> operationName.equals(s.getOperationName().toString()))
+ .findFirst()
+ .orElse(null);
+ }
+}
+
+/**
+ * Verifies that auto-instrumented openai.request spans inherit session_id, agent_version and the
+ * head-based sampling decision from an active LLMObs parent context, that they compute a sampling
+ * verdict of their own when there is no parent to inherit from, and that a stale context left over
+ * from an unrelated trace leaks none of the three onto the span. Forked + @WithConfig used together
+ * so the LLMObs system property is in place before the agent installs and there's no leakage from
+ * prior test state.
+ *
+ *
> traces, String operationName) {
- return traces.stream()
- .flatMap(List::stream)
- .filter(s -> operationName.equals(s.getOperationName().toString()))
- .findFirst()
- .orElse(null);
+ @Test
+ void parentlessOpenAiRequestSpanIsDroppedAtZeroSampleRate() throws Exception {
+ try {
+ openAiClient.chat().completions().create(buildMinimalChatParams());
+ } catch (Exception ignored) {
+ }
+
+ writer.waitForTraces(1);
+ DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
+ assertNotNull(openAiSpan, "openai.request span should have been created");
+ assertEquals(
+ LLMObsContext.SAMPLING_DECISION_DROPPED,
+ openAiSpan.getTag("_ml_obs_tag.sampling_decision"));
+ assertEquals("0", openAiSpan.getTag("_ml_obs_tag.sample_rate"));
}
}
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
index 20bb7b1740a..eb4cd99ae44 100644
--- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
+++ b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
@@ -195,6 +195,7 @@ public final class ConfigDefaults {
static final boolean DEFAULT_LLM_OBS_ENABLED = false;
static final boolean DEFAULT_LLM_OBS_AGENTLESS_ENABLED = false;
+ static final double DEFAULT_LLM_OBS_SAMPLE_RATE = 1.0;
static final boolean DEFAULT_USM_ENABLED = false;
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/LlmObsConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/LlmObsConfig.java
index 111f8ebc7af..0b2c6c2928c 100644
--- a/dd-trace-api/src/main/java/datadog/trace/api/config/LlmObsConfig.java
+++ b/dd-trace-api/src/main/java/datadog/trace/api/config/LlmObsConfig.java
@@ -12,5 +12,7 @@ public final class LlmObsConfig {
public static final String LLMOBS_AGENTLESS_ENABLED = "llmobs.agentless.enabled";
+ public static final String LLMOBS_SAMPLE_RATE = "llmobs.sample.rate";
+
private LlmObsConfig() {}
}
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 86622291dd7..2b8a85992e1 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
@@ -69,6 +69,9 @@ public class LLMObsSpanMapper implements RemoteMapper {
private static final byte[] APM_TRACE_ID = "apm_trace_id".getBytes(StandardCharsets.UTF_8);
private static final byte[] PARENT_ID = "parent_id".getBytes(StandardCharsets.UTF_8);
private static final byte[] SESSION_ID = "session_id".getBytes(StandardCharsets.UTF_8);
+ private static final byte[] SAMPLE_RATE = "sample_rate".getBytes(StandardCharsets.UTF_8);
+ private static final byte[] SAMPLING_DECISION =
+ "sampling_decision".getBytes(StandardCharsets.UTF_8);
private static final byte[] NAME = "name".getBytes(StandardCharsets.UTF_8);
private static final byte[] DURATION = "duration".getBytes(StandardCharsets.UTF_8);
private static final byte[] START_NS = "start_ns".getBytes(StandardCharsets.UTF_8);
@@ -111,6 +114,35 @@ 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 SAMPLE_RATE_TAG_INTERNAL_FULL = LLMOBS_TAG_PREFIX + "sample_rate";
+ private static final String SAMPLING_DECISION_TAG_INTERNAL_FULL =
+ LLMOBS_TAG_PREFIX + "sampling_decision";
+
+ /**
+ * Fallback pair meaning "retain", used if a span reaches the mapper without a stamped decision.
+ */
+ private static final String SAMPLING_DECISION_SAMPLED = "1";
+
+ private static final String SAMPLE_RATE_ALL = "1";
+
+ /**
+ * Internal tags serialized as dedicated top-level fields, which must therefore not also be
+ * emitted into the {@code tags} array.
+ *
+ *