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> 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. + * + *

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> 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. + * + *

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 TAGS_WRITTEN_AS_TOP_LEVEL_FIELDS = + Collections.unmodifiableSet( + new HashSet<>( + Arrays.asList( + PARENT_ID_TAG_INTERNAL_FULL, + SAMPLING_DECISION_TAG_INTERNAL_FULL, + SAMPLE_RATE_TAG_INTERNAL_FULL, + SPAN_KIND_TAG_KEY))); + 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 = @@ -175,6 +207,13 @@ public void map(List> trace, Writable writable, boolean re String sessionId = rawSessionId instanceof String ? (String) rawSessionId : null; boolean hasSessionId = sessionId != null && !sessionId.isEmpty(); + // Fallback to set default sampling tags when the fields are absent. + Object rawSamplingDecision = span.getTag(SAMPLING_DECISION_TAG_INTERNAL_FULL); + Object rawSampleRate = span.getTag(SAMPLE_RATE_TAG_INTERNAL_FULL); + boolean stamped = rawSamplingDecision instanceof String && rawSampleRate instanceof String; + String samplingDecision = stamped ? (String) rawSamplingDecision : SAMPLING_DECISION_SAMPLED; + String sampleRate = stamped ? (String) rawSampleRate : SAMPLE_RATE_ALL; + writable.startMap(hasSessionId ? 12 : 11); // 1 writable.writeUTF8(SPAN_ID); @@ -187,7 +226,6 @@ public void map(List> trace, Writable writable, boolean re // 3 writable.writeUTF8(PARENT_ID); writable.writeString(span.getTag(PARENT_ID_TAG_INTERNAL_FULL), null); - span.removeTag(PARENT_ID_TAG_INTERNAL_FULL); // 4 writable.writeUTF8(NAME); @@ -207,13 +245,17 @@ public void map(List> trace, Writable writable, boolean re // 8 writable.writeUTF8(DD); - writable.startMap(3); + writable.startMap(5); writable.writeUTF8(SPAN_ID); writable.writeString(String.valueOf(span.getSpanId()), null); writable.writeUTF8(TRACE_ID); writable.writeString(span.getTraceId().toHexString(), null); writable.writeUTF8(APM_TRACE_ID); writable.writeString(span.getTraceId().toHexString(), null); + writable.writeUTF8(SAMPLING_DECISION); + writable.writeString(samplingDecision, null); + writable.writeUTF8(SAMPLE_RATE); + writable.writeString(sampleRate, null); // 9 — optional top-level session_id field. Required by the LLMObs HTTP intake schema // and by the LLM Trace Explorer's Sessions filter, which keys off this field. @@ -375,6 +417,8 @@ public void accept(Metadata metadata) { String key = tag.getKey(); if (key.equals(SPAN_KIND_TAG_KEY)) { spanKind = String.valueOf(tag.getValue()); + } else if (TAGS_WRITTEN_AS_TOP_LEVEL_FIELDS.contains(key)) { + // Already written as a dedicated field; not counted here so it stays out of the array. } else if (TAGS_FOR_REMAPPING.contains(key)) { tagsToRemapToMeta.put(key, tag.getValue()); } else if (key.startsWith(LLMOBS_METRIC_PREFIX) && tag.getValue() instanceof Number) { @@ -391,9 +435,7 @@ public void accept(Metadata metadata) { } } - if (!spanKind.equals("unknown")) { - metadata.getTags().remove(SPAN_KIND_TAG_KEY); - } else { + if (spanKind.equals("unknown")) { LOGGER.warn("missing span kind"); } @@ -434,7 +476,9 @@ public void accept(Metadata metadata) { for (Map.Entry tag : metadata.getTags().entrySet()) { String key = tag.getKey(); Object value = tag.getValue(); - if (!tagsToRemapToMeta.containsKey(key) && key.startsWith(LLMOBS_TAG_PREFIX)) { + if (!tagsToRemapToMeta.containsKey(key) + && !TAGS_WRITTEN_AS_TOP_LEVEL_FIELDS.contains(key) + && key.startsWith(LLMOBS_TAG_PREFIX)) { writable.writeObject(key.substring(LLMOBS_TAG_PREFIX.length()) + ":" + value, null); } } 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 f2054591432..909d94ed526 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 @@ -1030,6 +1030,118 @@ private static Map serializeSingleSpan(LLMObsSpanMapper mapper, return spans.get(0); } + @Test + void testSamplingFieldsDefaultToRetainWhenTheSpanCarriesNoDecision() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + // Both span producers stamp these tags, so this exercises the mapper's fallback branch. + AgentSpan span = + tracer + .buildSpan("datadog", "chat-completion") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .start(); + span.setSpanType(InternalSpanTypes.LLMOBS); + span.finish(); + + Map spanData = serializeSingleSpan(mapper, span); + Map dd = (Map) spanData.get("_dd"); + assertEquals(5, dd.size()); + assertEquals("1", dd.get("sampling_decision")); + assertEquals("1", dd.get("sample_rate")); + + tracer.close(); + } + + @Test + void testSamplingFieldsAreEmittedAndDoNotLeakIntoTags() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + AgentSpan span = + tracer + .buildSpan("datadog", "chat-completion") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.sampling_decision", "0") + .withTag("_ml_obs_tag.sample_rate", "0.1") + .start(); + span.setSpanType(InternalSpanTypes.LLMOBS); + span.finish(); + + Map spanData = serializeSingleSpan(mapper, span); + Map dd = (Map) spanData.get("_dd"); + assertEquals(5, dd.size()); + assertEquals("0", dd.get("sampling_decision")); + assertEquals("0.1", dd.get("sample_rate")); + + // The mapper writes these as dedicated _dd fields and skips them elsewhere, so they must not + // also appear in tags[]. + List tags = (List) spanData.get("tags"); + assertFalse(tags.stream().anyMatch(tag -> tag.startsWith("sampling_decision:"))); + assertFalse(tags.stream().anyMatch(tag -> tag.startsWith("sample_rate:"))); + assertFalse(tags.stream().anyMatch(tag -> tag.startsWith("parent_id:"))); + assertFalse(tags.stream().anyMatch(tag -> tag.startsWith("span.kind:"))); + + tracer.close(); + } + + @Test + void testTopLevelFieldsSurviveASerializationRetry() throws Exception { + LLMObsSpanMapper mapper = new LLMObsSpanMapper(); + CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build(); + + AgentSpan span = + tracer + .buildSpan("datadog", "chat-completion") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.parent_id", "9876543210") + .withTag("_ml_obs_tag.sampling_decision", "0") + .withTag("_ml_obs_tag.sample_rate", "0.1") + .start(); + span.setSpanType(InternalSpanTypes.LLMOBS); + span.finish(); + List trace = Collections.singletonList((DDSpan) span); + + // On a BufferOverflowException, MsgPackWriter.format discards the partial write and re-invokes + // map(..., retry = true) on the same span instances. Model that contract directly: an attempt + // whose output is thrown away, then the retry. A mapper that consumed tags off the span while + // writing would emit different — and wrong — values the second time: a dropped span would + // come back as retained, and parent_id and span.kind would be lost. + mapWithRetryFlag(mapper, trace, false); + mapper.reset(); + Map retried = mapWithRetryFlag(mapper, trace, true); + + Map dd = (Map) retried.get("_dd"); + assertEquals("0", dd.get("sampling_decision")); + assertEquals("0.1", dd.get("sample_rate")); + assertEquals("9876543210", retried.get("parent_id")); + assertEquals("llm", ((Map) retried.get("meta")).get("span.kind")); + + tracer.close(); + } + + /** + * Invokes {@link LLMObsSpanMapper#map(List, datadog.communication.serialization.Writable, + * boolean)} directly, bypassing {@code MsgPackWriter.format} so the retry flag can be set + * explicitly, and returns the single serialized span. + */ + private static Map mapWithRetryFlag( + LLMObsSpanMapper mapper, List trace, boolean retry) throws IOException { + CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer(); + FlushingBuffer buffer = new FlushingBuffer(16 * 1024, sink); + MsgPackWriter packer = new MsgPackWriter(buffer); + mapper.map(trace, packer, retry); + // format() marks the message complete after a successful map; flush() emits nothing without it. + buffer.mark(); + packer.flush(); + + assertNotNull(sink.captured); + datadog.trace.common.writer.Payload payload = mapper.newPayload(); + payload.withBody(trace.size(), sink.captured); + Map result = objectMapper.readValue(writeTo(payload), Map.class); + return ((List>) result.get("spans")).get(0); + } + static class CapturingByteBufferConsumer implements ByteBufferConsumer { ByteBuffer captured; diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index 18bd4ae31fd..a97f7bda2d0 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -118,6 +118,7 @@ import static datadog.trace.api.ConfigDefaults.DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_ENABLED; import static datadog.trace.api.ConfigDefaults.DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_LIMIT; import static datadog.trace.api.ConfigDefaults.DEFAULT_LLM_OBS_AGENTLESS_ENABLED; +import static datadog.trace.api.ConfigDefaults.DEFAULT_LLM_OBS_SAMPLE_RATE; import static datadog.trace.api.ConfigDefaults.DEFAULT_LOGS_INJECTION_ENABLED; import static datadog.trace.api.ConfigDefaults.DEFAULT_LOGS_OTEL_BATCH_SIZE; import static datadog.trace.api.ConfigDefaults.DEFAULT_LOGS_OTEL_INTERVAL; @@ -478,6 +479,7 @@ import static datadog.trace.api.config.JmxFetchConfig.JMX_TAGS; import static datadog.trace.api.config.LlmObsConfig.LLMOBS_AGENTLESS_ENABLED; import static datadog.trace.api.config.LlmObsConfig.LLMOBS_ML_APP; +import static datadog.trace.api.config.LlmObsConfig.LLMOBS_SAMPLE_RATE; import static datadog.trace.api.config.OtlpConfig.LOGS_OTEL_BATCH_SIZE; import static datadog.trace.api.config.OtlpConfig.LOGS_OTEL_EXPORTER; import static datadog.trace.api.config.OtlpConfig.LOGS_OTEL_INTERVAL; @@ -1157,6 +1159,7 @@ public static String getHostName() { private final boolean llmObsAgentlessEnabled; private final String llmObsAgentlessUrl; private final String llmObsMlApp; + private final double llmObsSampleRate; private final boolean ciVisibilityTraceSanitationEnabled; private final boolean ciVisibilityAgentlessEnabled; @@ -2696,6 +2699,18 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment()) final String tempLlmObsMlApp = configProvider.getString(LLMOBS_ML_APP); llmObsMlApp = tempLlmObsMlApp == null || tempLlmObsMlApp.isEmpty() ? serviceName : tempLlmObsMlApp; + // Fall back to "sample everything" rather than clamping + final double configuredLlmObsSampleRate = + configProvider.getDouble(LLMOBS_SAMPLE_RATE, DEFAULT_LLM_OBS_SAMPLE_RATE); + if (configuredLlmObsSampleRate >= 0.0 && configuredLlmObsSampleRate <= 1.0) { + llmObsSampleRate = configuredLlmObsSampleRate; + } else { + log.warn( + "Invalid value {} for {}: expected a rate between 0.0 and 1.0, falling back to 1.0.", + configuredLlmObsSampleRate, + LLMOBS_SAMPLE_RATE); + llmObsSampleRate = 1.0; + } final String llmObsAgentlessUrlStr = getFinalLLMObsUrl(); URI parsedLLMObsUri = null; @@ -4481,6 +4496,16 @@ public String getLlmObsMlApp() { return llmObsMlApp; } + /** + * The fraction of LLM Observability traces retained by the backend, in {@code [0.0, 1.0]}. + * + *

Independent of APM sampling: the decision made with this rate never affects an APM sampling + * priority, and an APM decision never affects it. + */ + public double getLlmObsSampleRate() { + return llmObsSampleRate; + } + public boolean isCiVisibilityEnabled() { return instrumenterConfig.isCiVisibilityEnabled(); } 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 e925a04f42f..1e768846b10 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 @@ -8,6 +8,12 @@ public final class LLMObsContext { public static final String ROOT_SPAN_ID = "undefined"; + /** Sampling decision value meaning "retain this span". */ + public static final String SAMPLING_DECISION_SAMPLED = "1"; + + /** Sampling decision value meaning "drop this span". */ + public static final String SAMPLING_DECISION_DROPPED = "0"; + private LLMObsContext() { // ~ } @@ -16,6 +22,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 SAMPLE_RATE_KEY = ContextKey.named("llmobs_sample_rate"); + private static final ContextKey SAMPLING_DECISION_KEY = + ContextKey.named("llmobs_sampling_decision"); 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"); @@ -37,7 +46,7 @@ public static ContextScope attach(AgentSpanContext ctx) { public static ContextScope attach(AgentSpanContext ctx, String sessionId) { return Context.current() .with(CONTEXT_KEY, ctx) - .with(SESSION_ID_KEY, sessionId != null && !sessionId.isEmpty() ? sessionId : null) + .with(SESSION_ID_KEY, emptyToNull(sessionId)) .attach(); } @@ -52,42 +61,63 @@ public static ContextScope attach(AgentSpanContext ctx, String sessionId) { public static ContextScope attach(AgentSpanContext ctx, String sessionId, String agentVersion) { return Context.current() .with(CONTEXT_KEY, ctx) - .with(SESSION_ID_KEY, sessionId != null && !sessionId.isEmpty() ? sessionId : null) - .with( - AGENT_VERSION_KEY, - agentVersion != null && !agentVersion.isEmpty() ? agentVersion : null) + .with(SESSION_ID_KEY, emptyToNull(sessionId)) + .with(AGENT_VERSION_KEY, emptyToNull(agentVersion)) .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. + * Attach an LLMObs span context, propagating a session_id, an agent_version, a sampling decision, + * and agent attribution to descendant LLMObs spans. See {@link #attach(AgentSpanContext, String, + * String)} — the same clears-if-null-or-empty semantics apply to every value, so callers are + * expected to pass already-resolved effective values. + * + *

This overload carries every propagated value at once because a span's scope is attached + * exactly once: three independent mechanisms (session, sampling, attribution) share one context, + * so they cannot be attached by separate calls without nesting redundant scopes. + * + *

The sampling decision is computed once at the root of an LLMObs trace and inherited + * unchanged by every descendant, so that a trace is retained or dropped as a whole. Both sampling + * values are carried pre-formatted so that every span in the trace reports byte-identical values. + * The rate travels only alongside a decision — a rate on its own says nothing about whether the + * trace was kept — so a null decision clears both. + * + *

parentAgentSpanId identifies the nearest agent-kind ancestor; parentAgentName is its name + * (may be null). Both pagent keys are always written. 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. + *
+ * + *

In-process only. This context is not serialized into distributed trace + * headers, so each service in a distributed trace decides independently. Because the decision is + * a pure function of the APM trace ID and the configured rate, services configured at the same + * rate agree; services configured at different rates disagree and the trace is retained in part. + * A decision propagated by an upstream dd-trace-py or dd-trace-js service is likewise not read + * here. Closing that gap needs propagated trace tags mirroring the existing {@code _dd.p.ksr}. */ public static ContextScope attach( AgentSpanContext ctx, String sessionId, String agentVersion, + String sampleRate, + String samplingDecision, 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(); + String decision = emptyToNull(samplingDecision); + return Context.current() + .with(CONTEXT_KEY, ctx) + .with(SESSION_ID_KEY, emptyToNull(sessionId)) + .with(AGENT_VERSION_KEY, emptyToNull(agentVersion)) + .with(SAMPLING_DECISION_KEY, decision) + .with(SAMPLE_RATE_KEY, decision == null ? null : emptyToNull(sampleRate)) + .with(PAGENT_SPAN_ID_KEY, parentAgentSpanId) + .with(PAGENT_NAME_KEY, parentAgentName) + .attach(); } public static AgentSpanContext current() { @@ -109,6 +139,23 @@ public static String currentAgentVersion() { return Context.current().get(AGENT_VERSION_KEY); } + /** + * Return the sample rate that produced {@link #currentSamplingDecision()}, or null if no + * enclosing LLMObs span made a sampling decision. + */ + public static String currentSampleRate() { + return Context.current().get(SAMPLE_RATE_KEY); + } + + /** + * Return the sampling decision propagated from an enclosing LLMObs span, or null if none was + * made. A null value identifies the current span as the root of an LLMObs trace, which is the + * only place a decision is computed. + */ + public static String currentSamplingDecision() { + return Context.current().get(SAMPLING_DECISION_KEY); + } + /** * Return the parent agent span ID propagated from an enclosing agent-kind LLMObs span, or null. */ @@ -120,4 +167,8 @@ public static String currentParentAgentSpanId() { public static String currentParentAgentName() { return Context.current().get(PAGENT_NAME_KEY); } + + private static String emptyToNull(String value) { + return value != null && !value.isEmpty() ? value : null; + } } diff --git a/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsSampler.java b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsSampler.java new file mode 100644 index 00000000000..ff92d8368d3 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsSampler.java @@ -0,0 +1,91 @@ +package datadog.trace.api.llmobs; + +import datadog.trace.api.Config; + +/** + * Head-based retention sampler for LLM Observability traces. + * + *

Duplicates the arithmetic in {@code DeterministicSampler} because neither of those modules + * depends on {@code dd-trace-core}; {@code ApiSecurityDownstreamSamplerImpl} does the same. + */ +public final class LLMObsSampler { + + private static final long KNUTH_FACTOR = 1111111111111111111L; + private static final double MAX = Math.pow(2, 64) - 1; + + private final double rate; + private final long threshold; + private final String formattedRate; + + public static LLMObsSampler fromConfig() { + return new LLMObsSampler(Config.get().getLlmObsSampleRate()); + } + + public LLMObsSampler(final double rate) { + // Any rate outside [0.0, 1.0] (including NaN) falls back to 1.0 + this.rate = (rate >= 0.0 && rate <= 1.0) ? rate : 1.0; + this.threshold = cutoff(this.rate); + this.formattedRate = formatRate(this.rate); + } + + /** + * The configured rate, formatted for the wire. Derived from the configured {@code double} rather + * than from a narrowed {@code float}, so the reported rate is the rate actually applied. + */ + public String formattedRate() { + return formattedRate; + } + + /** + * @param samplingId the low-order 64 bits of the APM trace ID. + * @return whether the trace is retained. + */ + public boolean sample(final long samplingId) { + // unsigned 64 bit comparison with cutoff/threshold + return samplingId * KNUTH_FACTOR + Long.MIN_VALUE <= threshold; + } + + private static long cutoff(final double rate) { + if (rate < 0.5) { + return (long) (rate * MAX) + Long.MIN_VALUE; + } + if (rate < 1.0) { + return (long) ((rate * MAX) + Long.MIN_VALUE); + } + return Long.MAX_VALUE; + } + + /** + * Formats a sampling rate with up to 6 decimal digits of precision and no trailing zeros. Mirrors + * {@code format_rate} in dd-trace-py, which stamps {@code sample_rate} through the same helper it + * uses for {@code _dd.p.ksr}, so the two languages report the same rate as the same string. + */ + static String formatRate(final double rate) { + if (rate <= 0.0) { + return "0"; + } + if (rate >= 1.0) { + return "1"; + } + long rounded = Math.round(rate * 1_000_000L); + if (rounded <= 0) { + return "0"; + } + if (rounded >= 1_000_000L) { + return "1"; + } + // Build "0.DDDDDD", then trim trailing zeros. + char[] chars = new char[8]; + chars[0] = '0'; + chars[1] = '.'; + for (int i = 7; i >= 2; i--) { + chars[i] = (char) ('0' + (rounded % 10)); + rounded /= 10; + } + int end = 8; + while (chars[end - 1] == '0') { + end--; + } + return new String(chars, 0, end); + } +} 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 bec825bf887..20f91b51a8f 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 @@ -157,7 +157,62 @@ void childScopeInheritsParentAgentVersion() { } } - // ── 5-arg attach (pagent attribution) ──────────────────────────────────── + @Test + void samplingValuesReturnNullWhenNoContextAttached() { + assertNull(LLMObsContext.currentSamplingDecision()); + assertNull(LLMObsContext.currentSampleRate()); + } + + @Test + void attachWithoutSamplingDecisionLeavesSamplingValuesNull() { + AgentSpanContext ctx = mock(AgentSpanContext.class); + try (ContextScope scope = LLMObsContext.attach(ctx, "session-123")) { + assertNull(LLMObsContext.currentSamplingDecision()); + assertNull(LLMObsContext.currentSampleRate()); + } + } + + @Test + void attachWithSamplingDecisionStoresDecisionAndRate() { + AgentSpanContext ctx = mock(AgentSpanContext.class); + try (ContextScope scope = + LLMObsContext.attach( + ctx, null, null, "0.25", LLMObsContext.SAMPLING_DECISION_DROPPED, null, null)) { + assertEquals( + LLMObsContext.SAMPLING_DECISION_DROPPED, LLMObsContext.currentSamplingDecision()); + assertEquals("0.25", LLMObsContext.currentSampleRate()); + } + assertNull(LLMObsContext.currentSamplingDecision()); + assertNull(LLMObsContext.currentSampleRate()); + } + + @Test + void attachWithNullSamplingDecisionIgnoresSampleRate() { + AgentSpanContext ctx = mock(AgentSpanContext.class); + // The rate is only meaningful alongside a decision, so it is not stored on its own. + try (ContextScope scope = LLMObsContext.attach(ctx, null, null, "0.25", null, null, null)) { + assertNull(LLMObsContext.currentSamplingDecision()); + assertNull(LLMObsContext.currentSampleRate()); + } + } + + @Test + void childScopeInheritsParentSamplingDecision() { + AgentSpanContext parent = mock(AgentSpanContext.class); + AgentSpanContext child = mock(AgentSpanContext.class); + try (ContextScope parentScope = + LLMObsContext.attach( + parent, null, null, "1", LLMObsContext.SAMPLING_DECISION_SAMPLED, null, null)) { + try (ContextScope childScope = LLMObsContext.attach(child)) { + assertEquals(child, LLMObsContext.current()); + assertEquals( + LLMObsContext.SAMPLING_DECISION_SAMPLED, LLMObsContext.currentSamplingDecision()); + assertEquals("1", LLMObsContext.currentSampleRate()); + } + } + } + + // ── full attach (session_id + agent_version + sampling + pagent attribution) ── @Test void currentParentAgentSpanIdReturnsNullWhenNoContextAttached() { @@ -170,26 +225,39 @@ void currentParentAgentNameReturnsNullWhenNoContextAttached() { } @Test - void fiveArgAttachStoresAllFields() { + void fullAttachStoresAllFields() { AgentSpanContext ctx = mock(AgentSpanContext.class); - try (ContextScope scope = LLMObsContext.attach(ctx, "session-1", "v2", "span-99", "my-agent")) { + try (ContextScope scope = + LLMObsContext.attach( + ctx, + "session-1", + "v2", + "0.5", + LLMObsContext.SAMPLING_DECISION_SAMPLED, + "span-99", + "my-agent")) { assertEquals(ctx, LLMObsContext.current()); assertEquals("session-1", LLMObsContext.currentSessionId()); assertEquals("v2", LLMObsContext.currentAgentVersion()); + assertEquals("0.5", LLMObsContext.currentSampleRate()); + assertEquals( + LLMObsContext.SAMPLING_DECISION_SAMPLED, LLMObsContext.currentSamplingDecision()); assertEquals("span-99", LLMObsContext.currentParentAgentSpanId()); assertEquals("my-agent", LLMObsContext.currentParentAgentName()); } assertNull(LLMObsContext.current()); assertNull(LLMObsContext.currentSessionId()); assertNull(LLMObsContext.currentAgentVersion()); + assertNull(LLMObsContext.currentSampleRate()); + assertNull(LLMObsContext.currentSamplingDecision()); assertNull(LLMObsContext.currentParentAgentSpanId()); assertNull(LLMObsContext.currentParentAgentName()); } @Test - void fiveArgAttachWithNullSessionIdIgnoresSessionId() { + void fullAttachWithNullSessionIdIgnoresSessionId() { AgentSpanContext ctx = mock(AgentSpanContext.class); - try (ContextScope scope = LLMObsContext.attach(ctx, null, null, null, null)) { + try (ContextScope scope = LLMObsContext.attach(ctx, null, null, null, null, null, null)) { assertNull(LLMObsContext.currentSessionId()); assertNull(LLMObsContext.currentAgentVersion()); assertNull(LLMObsContext.currentParentAgentSpanId()); @@ -198,26 +266,27 @@ void fiveArgAttachWithNullSessionIdIgnoresSessionId() { } @Test - void fiveArgAttachWithEmptySessionIdIgnoresSessionId() { + void fullAttachWithEmptySessionIdIgnoresSessionId() { AgentSpanContext ctx = mock(AgentSpanContext.class); - try (ContextScope scope = LLMObsContext.attach(ctx, "", "", null, null)) { + try (ContextScope scope = LLMObsContext.attach(ctx, "", "", null, null, null, null)) { assertNull(LLMObsContext.currentSessionId()); assertNull(LLMObsContext.currentAgentVersion()); } } @Test - void fiveArgAttachNullPagentClearsStaleValuesFromOuterScope() { + void fullAttachNullPagentClearsStaleValuesFromOuterScope() { // 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")) { + LLMObsContext.attach(outer, "s", "v1", null, null, "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)) { + try (ContextScope innerScope = + LLMObsContext.attach(inner, null, null, null, null, null, null)) { assertNull(LLMObsContext.currentParentAgentSpanId()); assertNull(LLMObsContext.currentParentAgentName()); } @@ -229,13 +298,13 @@ void fiveArgAttachNullPagentClearsStaleValuesFromOuterScope() { } @Test - void fiveArgAttachInnerAgentOverridesOuterAgentForDescendants() { + void fullAttachInnerAgentOverridesOuterAgentForDescendants() { AgentSpanContext outer = mock(AgentSpanContext.class); AgentSpanContext inner = mock(AgentSpanContext.class); try (ContextScope outerScope = - LLMObsContext.attach(outer, null, null, "outer-span-id", "outer-agent")) { + LLMObsContext.attach(outer, null, null, null, null, "outer-span-id", "outer-agent")) { try (ContextScope innerScope = - LLMObsContext.attach(inner, null, null, "inner-span-id", "inner-agent")) { + LLMObsContext.attach(inner, null, null, null, null, "inner-span-id", "inner-agent")) { assertEquals("inner-span-id", LLMObsContext.currentParentAgentSpanId()); assertEquals("inner-agent", LLMObsContext.currentParentAgentName()); } @@ -245,18 +314,44 @@ void fiveArgAttachInnerAgentOverridesOuterAgentForDescendants() { } @Test - void fiveArgAttachNullPagentNameClearsNameButNotSpanId() { + void fullAttachNullPagentNameClearsNameButNotSpanId() { // 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")) { + LLMObsContext.attach(outer, null, null, null, null, "outer-span-id", "outer-agent")) { try (ContextScope innerScope = - LLMObsContext.attach(inner, null, null, "inner-span-id", null)) { + LLMObsContext.attach(inner, null, null, null, null, "inner-span-id", null)) { assertEquals("inner-span-id", LLMObsContext.currentParentAgentSpanId()); assertNull(LLMObsContext.currentParentAgentName()); } } } + + @Test + void attachPropagatesAllFourMechanismsTogether() { + AgentSpanContext parent = mock(AgentSpanContext.class); + AgentSpanContext child = mock(AgentSpanContext.class); + // All four propagation mechanisms coexist on one context and are inherited together. + try (ContextScope parentScope = + LLMObsContext.attach( + parent, + "session-abc", + "v7", + "0.5", + LLMObsContext.SAMPLING_DECISION_SAMPLED, + "agent-span-7", + "agent-seven")) { + try (ContextScope childScope = LLMObsContext.attach(child)) { + assertEquals("session-abc", LLMObsContext.currentSessionId()); + assertEquals("v7", LLMObsContext.currentAgentVersion()); + assertEquals( + LLMObsContext.SAMPLING_DECISION_SAMPLED, LLMObsContext.currentSamplingDecision()); + assertEquals("0.5", LLMObsContext.currentSampleRate()); + assertEquals("agent-span-7", LLMObsContext.currentParentAgentSpanId()); + assertEquals("agent-seven", LLMObsContext.currentParentAgentName()); + } + } + } } diff --git a/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsSamplerTest.java b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsSamplerTest.java new file mode 100644 index 00000000000..dcf683b31b9 --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsSamplerTest.java @@ -0,0 +1,106 @@ +package datadog.trace.api.llmobs; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.concurrent.ThreadLocalRandom; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +class LLMObsSamplerTest { + + @Test + void keepsEverythingAtRateOne() { + LLMObsSampler sampler = new LLMObsSampler(1.0); + assertEquals("1", sampler.formattedRate()); + for (long id : new long[] {0L, 1L, -1L, Long.MAX_VALUE, Long.MIN_VALUE, 1234567890123L}) { + assertTrue(sampler.sample(id)); + } + } + + @Test + void dropsEverythingAtRateZero() { + LLMObsSampler sampler = new LLMObsSampler(0.0); + assertEquals("0", sampler.formattedRate()); + // 0 is the one sampling id whose Knuth product is the minimum value, so it sits exactly on the + // cutoff at rate 0. Every other id must be dropped. + for (long id : new long[] {1L, -1L, Long.MAX_VALUE, 1234567890123L}) { + assertFalse(sampler.sample(id)); + } + } + + @ParameterizedTest + @ValueSource(doubles = {-1.0, -0.0001, 1.0001, 2.0, Double.NaN}) + void outOfRangeRatesFallBackToKeepingEverything(double rate) { + LLMObsSampler sampler = new LLMObsSampler(rate); + // Including negatives and NaN: a bad rate must never be read as "drop everything". + assertEquals("1", sampler.formattedRate()); + for (long id : new long[] {0L, 1L, -1L, Long.MAX_VALUE, Long.MIN_VALUE, 1234567890123L}) { + assertTrue(sampler.sample(id)); + } + } + + @ParameterizedTest + @CsvSource({ + "1.0, 1", + "0.0, 0", + "0.5, 0.5", + "0.25, 0.25", + "0.1, 0.1", + "0.123456, 0.123456", + // Beyond 6 digits of precision the rate rounds, matching the _dd.p.ksr format exactly. + "0.1234567, 0.123457", + "0.0000001, 0", + "0.999999, 0.999999", + }) + void formatsRateLikeKnuthSamplingRateTag(double rate, String expected) { + assertEquals(expected, LLMObsSampler.formatRate(rate)); + } + + @Test + void isDeterministicForTheSameSamplingId() { + LLMObsSampler sampler = new LLMObsSampler(0.5); + long id = 987654321987654321L; + boolean first = sampler.sample(id); + for (int i = 0; i < 100; i++) { + assertEquals(first, sampler.sample(id)); + } + // Two samplers configured with the same rate must agree, which is what lets independently + // configured services reach the same decision for a distributed trace. + assertEquals(first, new LLMObsSampler(0.5).sample(id)); + } + + @ParameterizedTest + @ValueSource(doubles = {0.1, 0.25, 0.5, 0.75, 0.9}) + void keepsRoughlyTheConfiguredFraction(double rate) { + LLMObsSampler sampler = new LLMObsSampler(rate); + int iterations = 100_000; + int kept = 0; + for (int i = 0; i < iterations; i++) { + if (sampler.sample(ThreadLocalRandom.current().nextLong())) { + kept++; + } + } + double observed = (double) kept / iterations; + assertTrue( + Math.abs(observed - rate) < 0.02, + "expected ~" + rate + " of traces kept but observed " + observed); + } + + @Test + void higherRatesAreASupersetOfLowerRates() { + // A trace kept at 10% must also be kept at 50%: the cutoff only moves outward as the rate + // rises. This is what makes a service configured at a higher rate safe to add to a trace. + LLMObsSampler low = new LLMObsSampler(0.1); + LLMObsSampler high = new LLMObsSampler(0.5); + for (int i = 0; i < 10_000; i++) { + long id = ThreadLocalRandom.current().nextLong(); + if (low.sample(id)) { + assertTrue(high.sample(id), "id " + id + " kept at 0.1 but dropped at 0.5"); + } + } + } +} diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 4db6b59fdda..1892f63130c 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -2332,6 +2332,14 @@ "aliases": [] } ], + "DD_LLMOBS_SAMPLE_RATE": [ + { + "version": "A", + "type": "decimal", + "default": "1.0", + "aliases": [] + } + ], "DD_LOGS_INJECTION": [ { "version": "B",