diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagationSource.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagationSource.java index 2ad01fb2b2b..110148e7c55 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagationSource.java +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagationSource.java @@ -3,8 +3,8 @@ import datadog.trace.api.llmobs.LLMObsContext; import datadog.trace.api.llmobs.LLMObsPropagationSource; import datadog.trace.api.llmobs.LLMObsPropagationValues; +import datadog.trace.api.llmobs.LLMObsTraceId; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; -import datadog.trace.llmobs.domain.LLMObsTraceId; import javax.annotation.Nullable; /** 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 bc78b63151c..9f643307257 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 @@ -9,13 +9,13 @@ import datadog.trace.api.llmobs.GenAiApmTags; import datadog.trace.api.llmobs.LLMObs; import datadog.trace.api.llmobs.LLMObsContext; -import datadog.trace.api.llmobs.LLMObsPropagationValues; +import datadog.trace.api.llmobs.LLMObsParent; import datadog.trace.api.llmobs.LLMObsSampler; import datadog.trace.api.llmobs.LLMObsSpan; import datadog.trace.api.llmobs.LLMObsTags; +import datadog.trace.api.llmobs.LLMObsTraceId; import datadog.trace.api.telemetry.LLMObsMetricCollector; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.Tags; import java.util.ArrayList; @@ -152,108 +152,40 @@ public DDLLMObsSpan( span.setTag(SPAN_KIND, kind); spanKind = kind; - // Resolve effective ml_app, 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 resolvedTraceId = null; - String resolvedMlApp = mlApp; - String resolvedAgentVersion = agentVersion; - String sampleRate = null; - String samplingDecision = null; - String resolvedParentAgentSpanId = null; - String resolvedParentAgentName = null; - boolean inheritedInProcess = false; - if (null != parent) { - if (!parent.getTraceId().equals(span.getTraceId())) { - LOGGER.error( - "trace ID mismatch, retrieved parent from context trace_id={}, span_id={}, started span trace_id={}, span_id={}", - parent.getTraceId(), - parent.getSpanId(), - span.getTraceId(), - span.getSpanId()); - } else { - inheritedInProcess = true; - parentSpanID = String.valueOf(parent.getSpanId()); - // Stay on the enclosing span's LLMObs trace. Never re-seeded from the APM trace id here: - // the enclosing span may itself have adopted a trace id from another service, and every - // span below it has to keep reporting that one. - resolvedTraceId = LLMObsContext.currentTraceId(); - // Inherit ml_app from the enclosing LLMObs span, if this span doesn't name its own, so a - // whole agent subtree stays in one application rather than each nested span falling back - // to the service default. - if (resolvedMlApp == null || resolvedMlApp.isEmpty()) { - String inherited = LLMObsContext.currentMlApp(); - if (inherited != null && !inherited.isEmpty()) { - resolvedMlApp = inherited; - } - } - // Inherit session_id from parent context only when it belongs to the same trace. - // Matches dd-trace-py and dd-trace-js: session_id need only be set on the root - // span; descendants inherit transitively via context propagation. - if (sessionId == null || sessionId.isEmpty()) { - String inherited = LLMObsContext.currentSessionId(); - if (inherited != null && !inherited.isEmpty()) { - sessionId = inherited; - } - } - // Inherit agent_version from the enclosing agent span, if this span doesn't set its own. - // An explicit value always wins, so a nested agent's own version overrides an ancestor's - // for its own subtree, matching session_id's explicit-wins semantics. - if (agentVersion == null || agentVersion.isEmpty()) { - String inherited = LLMObsContext.currentAgentVersion(); - if (inherited != null && !inherited.isEmpty()) { - 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(); - } - } - - // No usable in-process LLMObs parent, so fall back to what arrived from another service on - // the span context's propagation tags. Null when nothing did, which leaves every value below - // at the default it already holds. - LLMObsPropagationValues propagated = - inheritedInProcess ? null : span.spanContext().getExtractedLLMObsValues(); - if (propagated != null) { - if (propagated.parentId != null) { - parentSpanID = propagated.parentId; - } - // Adopt the caller's LLMObs trace id so the trace stays whole across the boundary. The APM - // trace id can't stand in for it: an intermediate service that starts a fresh APM trace - // would split the LLMObs trace in two, and dd-trace-py has been propagating this since it - // gained the tag. The value arrives as decimal and is stored as hex. - resolvedTraceId = LLMObsTraceId.fromWire(propagated.traceId); - if (resolvedMlApp == null || resolvedMlApp.isEmpty()) { - resolvedMlApp = propagated.mlApp; - } - if (sessionId == null || sessionId.isEmpty()) { - sessionId = propagated.sessionId; - } - resolvedParentAgentSpanId = propagated.parentAgentSpanId; - if (resolvedParentAgentSpanId != null) { - resolvedParentAgentName = propagated.parentAgentName; - } - // Adopt the upstream sampling verdict so a distributed LLMObs trace is retained or dropped - // as a whole. Re-rolling locally only agrees with the caller when both services happen to - // be configured at the same rate; honouring the propagated decision removes that - // coincidence. Both tags are required together — the rate is meaningless without the - // decision it produced, and a lone decision would be reported against this service's rate, - // which is not the rate that produced it. - if (propagated.sampleRate != null && propagated.samplingDecision != null) { - sampleRate = propagated.sampleRate; - samplingDecision = propagated.samplingDecision; - } - } + // Resolve the LLMObs parent: the in-process context when it is on this trace, else what arrived + // from another service, else none. See LLMObsParent for the rules; what follows layers this + // span's own explicit values on top. + LLMObsParent parent = LLMObsParent.resolve(span); + if (parent.ignoredContext != null) { + LOGGER.error( + "trace ID mismatch, retrieved parent from context trace_id={}, span_id={}, started span trace_id={}, span_id={}", + parent.ignoredContext.getTraceId(), + parent.ignoredContext.getSpanId(), + span.getTraceId(), + span.getSpanId()); + } + String parentSpanID = parent.parentId; + // Stay on the parent's LLMObs trace. Never re-seeded from the APM trace id when there is one: + // the parent may itself have adopted a trace id from another service, and every span below it + // has to keep reporting that one. + String resolvedTraceId = parent.traceId; + // An explicit value wins; an empty one counts as unset. Inheriting ml_app keeps a whole agent + // subtree in one application rather than each nested span falling back to the service + // default. session_id need only be set on the root span, matching dd-trace-py and dd-trace-js. + // A nested agent's own agent_version overrides an ancestor's for its own subtree. + String resolvedMlApp = isEmpty(mlApp) ? parent.mlApp : mlApp; + if (isEmpty(sessionId) && parent.sessionId != null) { + sessionId = parent.sessionId; + } + String resolvedAgentVersion = + isEmpty(agentVersion) && parent.agentVersion != null ? parent.agentVersion : agentVersion; + // Honouring an inherited verdict keeps a distributed LLMObs trace retained or dropped as a + // whole, rather than agreeing with the caller only when both are configured at the same rate. + String sampleRate = parent.sampleRate; + String samplingDecision = parent.samplingDecision; + // The nearest agent-kind ancestor. Overridden just below when this span is itself an agent. + String resolvedParentAgentSpanId = parent.parentAgentSpanId; + String resolvedParentAgentName = parent.parentAgentName; // Root of an LLMObs trace: seed the LLMObs trace id from the APM one, in the shape the other // tracers store it in — hex above 2^64, decimal below. The shape has to match what fromWire @@ -286,14 +218,11 @@ public DDLLMObsSpan( } // Roll a decision only at the true root of an LLMObs trace - if (!inheritedInProcess && propagated == null) { + if (parent.isRoot()) { sampleRate = sampler.formattedRate(); - samplingDecision = - sampler.sample(span.getTraceId().toLong()) - ? LLMObsContext.SAMPLING_DECISION_SAMPLED - : LLMObsContext.SAMPLING_DECISION_DROPPED; + samplingDecision = sampler.decide(span.getTraceId()); } - // Both tags move together for the same reason they are adopted together above: a rate without + // Both tags move together for the same reason LLMObsParent adopts them together: a rate without // the decision it produced says nothing, and a decision on its own would be read against a // rate that did not produce it. if (sampleRate != null && samplingDecision != null) { @@ -838,4 +767,8 @@ public String getLLMObsTraceId() { public long getSpanId() { return span.getSpanId(); } + + private static boolean isEmpty(String value) { + return value == null || value.isEmpty(); + } } diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagationSourceTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagationSourceTest.java index 9ffd37952a8..7a51c5e1190 100644 --- a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagationSourceTest.java +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagationSourceTest.java @@ -424,6 +424,76 @@ void continuationDoesNotRerollWhenTheUpstreamSentNoVerdict() { } } + /** + * An agent span on the receiving side is its own descendants' nearest agent: the upstream + * attribution it received is replaced, for itself and for everything below it, by its own id and + * name. + */ + @Test + void agentSpanReceivingPropagatedAttributionBecomesTheParentAgent() { + Map inbound = producerCarrier("checkout", "sess-42"); + try (ContextScope consumeScope = startLocalChildScope(extractSpan(inbound))) { + DDLLMObsSpan workerAgent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "worker-agent", null, null); + try { + assertEquals( + String.valueOf(workerAgent.getSpanId()), LLMObsContext.currentParentAgentSpanId()); + assertEquals("worker-agent", LLMObsContext.currentParentAgentName()); + assertEquals("sess-42", LLMObsContext.currentSessionId()); + assertEquals("checkout", LLMObsContext.currentMlApp()); + } finally { + workerAgent.finish(); + } + } + } + + /** + * A queue consumer whose previous message left an LLMObs scope open on the thread. The leaked + * context is on another trace, so it is ignored and the span takes what arrived with the current + * message rather than starting a new LLMObs trace. + */ + @Test + void staleContextFromAnotherTraceDoesNotShadowThePropagatedOne() { + Map inbound = producerCarrier("checkout", "sess-42"); + + AgentSpan staleApm = AgentTracer.get().buildSpan("apm", "previous.message").start(); + try (ContextScope stale = LLMObsContext.attach(staleApm.spanContext(), "stale-session")) { + try (ContextScope consumeScope = startLocalChildScope(extractSpan(inbound))) { + DDLLMObsSpan worker = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", null, null); + try { + assertEquals("sess-42", LLMObsContext.currentSessionId()); + assertEquals("checkout", LLMObsContext.currentMlApp()); + assertEquals("dispatcher", LLMObsContext.currentParentAgentName()); + } finally { + worker.finish(); + } + } + } finally { + staleApm.finish(); + } + } + + /** An explicit empty session id counts as unset, so the propagated one still applies. */ + @Test + void emptyExplicitSessionIdInheritsThePropagatedOne() { + Map inbound = producerCarrier("checkout", "sess-42"); + + try (ContextScope consumeScope = startLocalChildScope(extractSpan(inbound))) { + DDLLMObsSpan worker = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", null, ""); + try { + assertEquals("sess-42", LLMObsContext.currentSessionId()); + // And an in-process child naming an empty session inherits the worker's. + DDLLMObsSpan child = newSpan(Tags.LLMOBS_LLM_SPAN_KIND, "generate", null, ""); + try { + assertEquals("sess-42", LLMObsContext.currentSessionId()); + } finally { + child.finish(); + } + } finally { + worker.finish(); + } + } + } + @Test void peerSpanDoesNotInheritAFinishedSpansInjectedContext() { try (ContextScope apmScope = startRootApmScope()) { diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java index 3e36b979c62..9c674b973b2 100644 --- a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanAgentVersionTest.java @@ -140,6 +140,21 @@ void noVersionSetAnywhereMeansNoTagOnAnySpanInTheSubtree() { } } + @Test + void emptyExplicitVersionCountsAsUnsetAndInheritsFromParent() { + DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", "v3"); + try (ContextScope ignored = AgentTracer.activateSpan(spanOf(agent))) { + DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool1", ""); + try { + assertEquals("v3", spanOf(child).getTag(AGENT_VERSION_TAG)); + } finally { + child.finish(); + } + } finally { + agent.finish(); + } + } + @Test void childDoesNotInheritAgentVersionWhenStaleContextIsFromADifferentTrace() throws ReflectiveOperationException { 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 aebd6223b32..9bef3e5a454 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 @@ -11,11 +11,10 @@ import datadog.trace.api.DDTraceApiInfo; import datadog.trace.api.WellKnownTags; import datadog.trace.api.llmobs.GenAiApmTags; -import datadog.trace.api.llmobs.LLMObsContext; +import datadog.trace.api.llmobs.LLMObsParent; 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; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; import datadog.trace.bootstrap.instrumentation.api.Tags; @@ -115,72 +114,39 @@ protected void doAfterStart(@Nonnull AgentSpan span) { span.setTag(CommonTags.SOURCE, "integration"); span.setTag(CommonTags.INTEGRATION, INTEGRATION); - // 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 trace_id, ml_app, - // parent_id, 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; - String samplingDecision = null; - String sampleRate = null; - if (inheritable) { - parentSpanId = String.valueOf(parent.getSpanId()); - - // Stay in the parent's LLMObs trace and application. The parent may have adopted both - // from another service, so the APM trace id and the service default can differ from them. - String traceId = LLMObsContext.currentTraceId(); - if (traceId != null && !traceId.isEmpty()) { - span.setTag(CommonTags.TRACE_ID, traceId); - } - String mlApp = LLMObsContext.currentMlApp(); - if (mlApp != null && !mlApp.isEmpty()) { - span.setTag(CommonTags.ML_APP, mlApp); - } - - // Inherit session_id from the active LLMObs parent (e.g. a manual workflow span). - // Matches dd-trace-py / dd-trace-js, where auto-instrumented LLM spans inherit - // session_id from the workflow root via context propagation. Without this, the - // auto-instrumented openai.request span would not appear under its session in - // the LLM Trace Explorer's Sessions view. - String sessionId = LLMObsContext.currentSessionId(); - if (sessionId != null && !sessionId.isEmpty()) { - span.setTag(CommonTags.SESSION_ID, sessionId); - } - - // Inherit agent_version from the active LLMObs parent. - String agentVersion = LLMObsContext.currentAgentVersion(); - if (agentVersion != null && !agentVersion.isEmpty()) { - span.setTag(CommonTags.AGENT_VERSION, agentVersion); - } - - // 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); - String parentAgentName = LLMObsContext.currentParentAgentName(); - if (parentAgentName != null) { - span.setTag(CommonTags.PAGENT_NAME, parentAgentName); - } + // Resolve the LLMObs parent the same way DDLLMObsSpan does: the in-process context when it + // is on this trace, else what arrived from another service (an HTTP handler or queue + // consumer calling OpenAI without an LLMObs span of its own), else none. Inheriting + // session_id is what files this span under its session in the LLM Trace Explorer. + LLMObsParent parent = LLMObsParent.resolve(span); + span.setTag(CommonTags.PARENT_ID, parent.parentId); + // With no trace id to inherit the tag is left unset, and the mapper seeds it from the APM + // trace id the same way DDLLMObsSpan does. + if (parent.traceId != null) { + span.setTag(CommonTags.TRACE_ID, parent.traceId); + } + if (parent.mlApp != null) { + span.setTag(CommonTags.ML_APP, parent.mlApp); + } + if (parent.sessionId != null) { + span.setTag(CommonTags.SESSION_ID, parent.sessionId); + } + if (parent.agentVersion != null) { + span.setTag(CommonTags.AGENT_VERSION, parent.agentVersion); + } + // The name is only meaningful alongside an id. + if (parent.parentAgentSpanId != null) { + span.setTag(CommonTags.PAGENT_SPAN_ID, parent.parentAgentSpanId); + if (parent.parentAgentName != null) { + span.setTag(CommonTags.PAGENT_NAME, parent.parentAgentName); } - - samplingDecision = LLMObsContext.currentSamplingDecision(); - sampleRate = LLMObsContext.currentSampleRate(); } - span.setTag(CommonTags.PARENT_ID, parentSpanId); - // Roll a decision only when this span is the root of its own LLMObs trace. A parent with no - // decision continued a trace whose caller sent none, so this span leaves both tags unset - // too, as DDLLMObsSpan does; a local roll here could drop the span while its parent is kept. - if (!inheritable) { + String sampleRate = parent.sampleRate; + String samplingDecision = parent.samplingDecision; + if (parent.isRoot()) { sampleRate = sampler.formattedRate(); - samplingDecision = - sampler.sample(span.getTraceId().toLong()) - ? LLMObsContext.SAMPLING_DECISION_SAMPLED - : LLMObsContext.SAMPLING_DECISION_DROPPED; + samplingDecision = sampler.decide(span.getTraceId()); } if (samplingDecision != null && sampleRate != null) { span.setTag(CommonTags.SAMPLING_DECISION, samplingDecision); 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 664fb56073f..76da10afb7c 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 @@ -14,7 +14,9 @@ import com.openai.models.embeddings.EmbeddingModel; import com.sun.net.httpserver.HttpHandler; import com.sun.net.httpserver.HttpServer; +import datadog.context.Context; import datadog.context.ContextScope; +import datadog.context.propagation.Propagators; import datadog.trace.agent.test.AbstractInstrumentationTest; import datadog.trace.api.Config; import datadog.trace.api.llmobs.GenAiApmTags; @@ -25,9 +27,12 @@ import datadog.trace.test.junit.utils.config.WithConfig; import java.io.IOException; import java.io.OutputStream; +import java.math.BigInteger; import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; +import java.util.HashMap; import java.util.List; +import java.util.Map; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -130,9 +135,11 @@ protected static DDSpan findSpanByOperationName(List> traces, Strin * 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. + * from an unrelated trace leaks none of the three onto the span. With no in-process parent, they + * adopt the LLMObs context that arrived on the inbound request's {@code _dd.p.llmobs_*} tags, so a + * service calling OpenAI directly stays in its caller's LLMObs trace. 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}. @@ -410,6 +417,203 @@ void openAiRequestSpanInheritsNothingFromStaleCrossTraceContext() throws Excepti assertNull(openAiSpan.getTag("_ml_obs_tag.pagent_span_id")); assertNull(openAiSpan.getTag("_ml_obs_tag.pagent_name")); } + + @Test + void openAiRequestSpanAdoptsPropagatedContextWithNoLlmObsSpanInTheService() throws Exception { + // An HTTP handler or queue consumer that calls OpenAI directly: the upstream LLMObs context + // arrived on the request headers, and nothing in this service opened an LLMObs span. + try (ServerScope ignored = + startServerScope(inboundHeaders(UPSTREAM_TAGS + "," + SAMPLED_TAGS))) { + openAiClient.chat().completions().create(buildMinimalChatParams()); + } + + writer.waitForTraces(1); + DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request"); + assertNotNull(openAiSpan, "openai.request span should have been created"); + assertEquals(UPSTREAM_LLMOBS_TRACE_ID, openAiSpan.getTag("_ml_obs_tag.trace_id")); + assertEquals(UPSTREAM_PARENT_ID, openAiSpan.getTag("_ml_obs_tag.parent_id")); + assertEquals("upstream-app", openAiSpan.getTag("_ml_obs_tag.ml_app")); + assertEquals("upstream-session", openAiSpan.getTag("_ml_obs_tag.session_id")); + assertEquals(UPSTREAM_AGENT_SPAN_ID, openAiSpan.getTag("_ml_obs_tag.pagent_span_id")); + assertEquals("planner-agent", openAiSpan.getTag("_ml_obs_tag.pagent_name")); + assertEquals( + LLMObsContext.SAMPLING_DECISION_SAMPLED, + openAiSpan.getTag("_ml_obs_tag.sampling_decision")); + assertEquals("0.5", openAiSpan.getTag("_ml_obs_tag.sample_rate")); + } + + @Test + void openAiRequestSpanHonoursAPropagatedDropDecision() throws Exception { + // The caller dropped this LLMObs trace; re-rolling locally at the rate of 1.0 would keep the + // LLM span while the rest of the trace is discarded. + try (ServerScope ignored = + startServerScope( + inboundHeaders( + UPSTREAM_TAGS + + ",_dd.p.llmobs_sr=0.25,_dd.p.llmobs_sd=" + + LLMObsContext.SAMPLING_DECISION_DROPPED))) { + openAiClient.chat().completions().create(buildMinimalChatParams()); + } + + 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 openAiRequestSpanRollsNoDecisionWhenThePropagatedContextCarriesNone() throws Exception { + // The caller continued an LLMObs trace but sent no verdict. Same rule as an in-process parent + // with none: the span must not decide for itself. + try (ServerScope ignored = startServerScope(inboundHeaders(UPSTREAM_TAGS))) { + openAiClient.chat().completions().create(buildMinimalChatParams()); + } + + writer.waitForTraces(1); + DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request"); + assertNotNull(openAiSpan, "openai.request span should have been created"); + assertEquals(UPSTREAM_PARENT_ID, openAiSpan.getTag("_ml_obs_tag.parent_id")); + assertNull(openAiSpan.getTag("_ml_obs_tag.sampling_decision")); + assertNull(openAiSpan.getTag("_ml_obs_tag.sample_rate")); + } + + @Test + void openAiRequestSpanPrefersTheInProcessParentOverThePropagatedContext() throws Exception { + try (ServerScope ignored1 = + startServerScope(inboundHeaders(UPSTREAM_TAGS + "," + SAMPLED_TAGS))) { + AgentSpan localParent = AgentTracer.startSpan("test", "local-workflow"); + try (ContextScope ignored2 = AgentTracer.activateSpan(localParent); + ContextScope ignored3 = + LLMObsContext.attach(localParent.spanContext(), "local-session")) { + openAiClient.chat().completions().create(buildMinimalChatParams()); + } finally { + localParent.finish(); + } + } + + writer.waitForTraces(1); + DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request"); + assertNotNull(openAiSpan, "openai.request span should have been created"); + DDSpan localParent = findSpanByOperationName(writer, "local-workflow"); + assertEquals( + String.valueOf(localParent.getSpanId()), openAiSpan.getTag("_ml_obs_tag.parent_id")); + assertEquals("local-session", openAiSpan.getTag("_ml_obs_tag.session_id")); + } + + @Test + void openAiRequestSpanIgnoresAStaleContextAndUsesThePropagatedOne() throws Exception { + // A queue consumer whose previous message left an LLMObs scope open. The leaked context is on + // another trace, so it must neither contribute values nor shadow what arrived with this one. + AgentSpan staleParent = AgentTracer.startSpan("test", "stale-parent"); + try (ContextScope ignored1 = LLMObsContext.attach(staleParent.spanContext(), "stale-session"); + ServerScope ignored2 = + startServerScope(inboundHeaders(UPSTREAM_TAGS + "," + SAMPLED_TAGS))) { + openAiClient.chat().completions().create(buildMinimalChatParams()); + } finally { + staleParent.finish(); + } + + writer.waitForTraces(2); + DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request"); + assertNotNull(openAiSpan, "openai.request span should have been created"); + assertEquals(UPSTREAM_PARENT_ID, openAiSpan.getTag("_ml_obs_tag.parent_id")); + assertEquals("upstream-session", openAiSpan.getTag("_ml_obs_tag.session_id")); + assertEquals(UPSTREAM_LLMOBS_TRACE_ID, openAiSpan.getTag("_ml_obs_tag.trace_id")); + } + + @Test + void openAiRequestSpanIsAnLlmObsRootWhenTheInboundRequestCarriesNoLlmObsContext() + throws Exception { + // An APM-only caller: the trace is distributed, but there is no LLMObs context to adopt. + try (ServerScope ignored = startServerScope(inboundHeaders(null))) { + openAiClient.chat().completions().create(buildMinimalChatParams()); + } + + writer.waitForTraces(1); + DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request"); + assertNotNull(openAiSpan, "openai.request span should have been created"); + assertEquals(LLMObsContext.ROOT_SPAN_ID, openAiSpan.getTag("_ml_obs_tag.parent_id")); + assertNull(openAiSpan.getTag("_ml_obs_tag.trace_id")); + assertNull(openAiSpan.getTag("_ml_obs_tag.session_id")); + assertEquals(Config.get().getLlmObsMlApp(), openAiSpan.getTag("_ml_obs_tag.ml_app")); + assertEquals( + LLMObsContext.SAMPLING_DECISION_SAMPLED, + openAiSpan.getTag("_ml_obs_tag.sampling_decision")); + } + + /** The upstream LLMObs trace id, deliberately not the APM one, in the hex form it is stored. */ + private static final String UPSTREAM_LLMOBS_TRACE_ID = "6ab41f6e00000000fea4f6b7295bd5e9"; + + private static final String UPSTREAM_PARENT_ID = "12571653173807582101"; + private static final String UPSTREAM_AGENT_SPAN_ID = "12571653173807582101"; + + /** The upstream context as dd-trace-py writes it; the trace id goes out as decimal. */ + private static final String UPSTREAM_TAGS = + "_dd.p.llmobs_trace_id=" + + new BigInteger(UPSTREAM_LLMOBS_TRACE_ID, 16) + + ",_dd.p.llmobs_parent_id=" + + UPSTREAM_PARENT_ID + + ",_dd.p.llmobs_ml_app=upstream-app" + + ",_dd.p.llmobs_sid=upstream-session" + + ",_dd.p.llmobs_pagent_span_id=" + + UPSTREAM_AGENT_SPAN_ID + + ",_dd.p.llmobs_pagent_name=planner-agent"; + + private static final String SAMPLED_TAGS = + "_dd.p.llmobs_sr=0.5,_dd.p.llmobs_sd=" + LLMObsContext.SAMPLING_DECISION_SAMPLED; + + private static Map inboundHeaders(String llmObsTags) { + Map headers = new HashMap<>(); + headers.put("x-datadog-trace-id", "4242"); + headers.put("x-datadog-parent-id", "4343"); + headers.put("x-datadog-sampling-priority", "1"); + if (llmObsTags != null) { + headers.put("x-datadog-tags", llmObsTags); + } + return headers; + } + + /** + * Extracts inbound headers and activates the local span a server or consumer instrumentation + * would open for them. That local span takes over the extracted propagation tags, which is where + * the openai.request span started under it reads the upstream LLMObs context from. + */ + /** Closes the server span's scope and finishes the span, as the server instrumentation would. */ + private static ServerScope startServerScope(Map headers) { + return new ServerScope(headers); + } + + /** + * A nested class rather than a lambda: the test class itself may not declare methods whose + * signatures carry bootstrap types such as {@code AgentSpan} or {@code ContextScope}. + */ + private static final class ServerScope implements AutoCloseable { + private final AgentSpan server; + private final ContextScope scope; + + ServerScope(Map headers) { + Context extracted = + Propagators.defaultPropagator() + .extract(Context.root(), headers, (c, visitor) -> c.forEach(visitor)); + AgentSpan remoteParent = AgentSpan.fromContext(extracted); + assertNotNull(remoteParent, "expected trace context to be extracted"); + server = + AgentTracer.get() + .buildSpan("test", "http.request") + .asChildOf(remoteParent.spanContext()) + .start(); + scope = AgentTracer.activateSpan(server); + } + + @Override + public void close() { + scope.close(); + server.finish(); + } + } } /** diff --git a/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsParent.java b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsParent.java new file mode 100644 index 00000000000..015f7cec1c4 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsParent.java @@ -0,0 +1,193 @@ +package datadog.trace.api.llmobs; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import javax.annotation.Nullable; + +/** + * The LLM Observability parent a span inherits from, resolved once when the span starts. Shared by + * manual spans ({@code DDLLMObsSpan}) and auto-instrumented ones, so that every LLMObs span picks + * its parent the same way. + * + *

There are two candidates, tried in order: + * + *

    + *
  1. The active in-process {@link LLMObsContext}, but only when it belongs to the same trace as + * the span. A context from another trace (leaked across an async boundary, or left open by a + * previously handled message) contributes nothing. + *
  2. The {@code _dd.p.llmobs_*} values that arrived on the inbound request, from {@link + * AgentSpanContext#getExtractedLLMObsValues()}. + *
+ * + *

If neither applies the span is the root of its own LLMObs trace, which {@link #isRoot()} + * reports. One gate decides every inherited value, so a newly propagated value cannot ship with a + * weaker gate of its own. + * + *

Resolution only: nothing is written to the span. Values the caller sets explicitly, the + * service-default {@code ml_app}, an agent span standing in as its own descendants' parent agent, + * and a fresh sampling decision at the root are all the caller's to apply. + * + *

Inherited strings are {@code null} when absent; an empty value reads as absent, as it does on + * the wire. In-process values are otherwise passed through exactly as the context holds them, since + * a span re-attaches them for its own descendants and injections. Propagated values follow the two + * pairing rules: the parent agent name is taken only alongside its id, and the sampling rate only + * together with the decision it produced. + */ +public final class LLMObsParent { + + /** Where the inherited values came from. */ + public enum Source { + /** The active in-process LLMObs context on the same trace. */ + IN_PROCESS, + /** The LLMObs context that arrived on the inbound request. */ + PROPAGATED, + /** Neither: the span is the root of its own LLMObs trace. */ + NONE + } + + private static final LLMObsParent ROOT = + new LLMObsParent( + Source.NONE, + LLMObsContext.ROOT_SPAN_ID, + null, + null, + null, + null, + null, + null, + null, + null, + null); + + public final Source source; + + /** The parent's span id, or {@link LLMObsContext#ROOT_SPAN_ID} when there is none. */ + public final String parentId; + + /** + * The LLMObs trace id in its stored form (hex above 2^64, decimal below), or {@code null} when + * there is none to inherit and the caller seeds it from the APM trace id. + */ + @Nullable public final String traceId; + + @Nullable public final String mlApp; + @Nullable public final String sessionId; + + /** Only ever inherited in-process; it is not propagated across services. */ + @Nullable public final String agentVersion; + + @Nullable public final String parentAgentSpanId; + @Nullable public final String parentAgentName; + @Nullable public final String sampleRate; + @Nullable public final String samplingDecision; + + /** + * The in-process context that was active but ignored because it belongs to another trace, or + * {@code null} if there was none. Exposed so a caller can report it. + */ + @Nullable public final AgentSpanContext ignoredContext; + + private LLMObsParent( + Source source, + String parentId, + @Nullable String traceId, + @Nullable String mlApp, + @Nullable String sessionId, + @Nullable String agentVersion, + @Nullable String parentAgentSpanId, + @Nullable String parentAgentName, + @Nullable String sampleRate, + @Nullable String samplingDecision, + @Nullable AgentSpanContext ignoredContext) { + this.source = source; + this.parentId = parentId; + this.traceId = traceId; + this.mlApp = mlApp; + this.sessionId = sessionId; + this.agentVersion = agentVersion; + this.parentAgentSpanId = parentAgentSpanId; + this.parentAgentName = parentAgentName; + this.sampleRate = sampleRate; + this.samplingDecision = samplingDecision; + this.ignoredContext = ignoredContext; + } + + /** Resolves the LLMObs parent of a span that has just started. */ + public static LLMObsParent resolve(AgentSpan span) { + AgentSpanContext active = LLMObsContext.current(); + AgentSpanContext ignored = null; + if (active != null) { + if (active.getTraceId().equals(span.getTraceId())) { + return new LLMObsParent( + Source.IN_PROCESS, + String.valueOf(active.getSpanId()), + emptyToNull(LLMObsContext.currentTraceId()), + emptyToNull(LLMObsContext.currentMlApp()), + emptyToNull(LLMObsContext.currentSessionId()), + emptyToNull(LLMObsContext.currentAgentVersion()), + emptyToNull(LLMObsContext.currentParentAgentSpanId()), + emptyToNull(LLMObsContext.currentParentAgentName()), + emptyToNull(LLMObsContext.currentSampleRate()), + emptyToNull(LLMObsContext.currentSamplingDecision()), + null); + } + ignored = active; + } + + LLMObsPropagationValues propagated = span.spanContext().getExtractedLLMObsValues(); + if (propagated != null) { + String parentAgentSpanId = emptyToNull(propagated.parentAgentSpanId); + // Adopted together or not at all: a rate means nothing without the decision it produced, and + // a lone decision would be read against this service's rate rather than the one behind it. + boolean hasVerdict = + emptyToNull(propagated.sampleRate) != null + && emptyToNull(propagated.samplingDecision) != null; + return new LLMObsParent( + Source.PROPAGATED, + propagated.parentId != null && !propagated.parentId.isEmpty() + ? propagated.parentId + : LLMObsContext.ROOT_SPAN_ID, + // Adopt the caller's LLMObs trace id so the trace stays whole across the boundary. It + // arrives as decimal and is stored as hex. + LLMObsTraceId.fromWire(propagated.traceId), + emptyToNull(propagated.mlApp), + emptyToNull(propagated.sessionId), + null, + parentAgentSpanId, + parentAgentSpanId != null ? emptyToNull(propagated.parentAgentName) : null, + hasVerdict ? propagated.sampleRate : null, + hasVerdict ? propagated.samplingDecision : null, + ignored); + } + + if (ignored == null) { + return ROOT; + } + return new LLMObsParent( + Source.NONE, + LLMObsContext.ROOT_SPAN_ID, + null, + null, + null, + null, + null, + null, + null, + null, + ignored); + } + + /** + * Whether the span is the root of its own LLMObs trace, and so the only one entitled to roll a + * sampling decision. A parent that carries no decision is not a root: it continued a trace whose + * caller sent none, and a local roll could drop the span while its parent is kept. + */ + public boolean isRoot() { + return source == Source.NONE; + } + + @Nullable + private static String emptyToNull(@Nullable String value) { + return value == null || value.isEmpty() ? null : value; + } +} 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 index ff92d8368d3..7f663086a41 100644 --- a/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsSampler.java +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsSampler.java @@ -1,6 +1,7 @@ package datadog.trace.api.llmobs; import datadog.trace.api.Config; +import datadog.trace.api.DDTraceId; /** * Head-based retention sampler for LLM Observability traces. @@ -45,6 +46,16 @@ public boolean sample(final long samplingId) { return samplingId * KNUTH_FACTOR + Long.MIN_VALUE <= threshold; } + /** + * The sampling decision for the root of an LLMObs trace, in its wire form: {@link + * LLMObsContext#SAMPLING_DECISION_SAMPLED} or {@link LLMObsContext#SAMPLING_DECISION_DROPPED}. + */ + public String decide(final DDTraceId traceId) { + return sample(traceId.toLong()) + ? LLMObsContext.SAMPLING_DECISION_SAMPLED + : LLMObsContext.SAMPLING_DECISION_DROPPED; + } + private static long cutoff(final double rate) { if (rate < 0.5) { return (long) (rate * MAX) + Long.MIN_VALUE; diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsTraceId.java b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsTraceId.java similarity index 98% rename from dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsTraceId.java rename to internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsTraceId.java index 726d2b6576c..cfa760941f4 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsTraceId.java +++ b/internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsTraceId.java @@ -1,4 +1,4 @@ -package datadog.trace.llmobs.domain; +package datadog.trace.api.llmobs; import datadog.trace.api.DDTraceId; import datadog.trace.api.internal.util.LongStringUtils; @@ -40,7 +40,7 @@ private LLMObsTraceId() {} * renders only the low 64 bits — is the right decimal. It is not safe to call outside this * branch, which is the reason to go through here rather than at each call site. */ - static String format(DDTraceId traceId) { + public static String format(DDTraceId traceId) { return traceId.toHighOrderLong() != 0 ? traceId.toHexString() : traceId.toString(); } diff --git a/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsParentTest.java b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsParentTest.java new file mode 100644 index 00000000000..ccd7751d7d5 --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsParentTest.java @@ -0,0 +1,269 @@ +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.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import datadog.context.ContextScope; +import datadog.trace.api.DDTraceId; +import datadog.trace.api.llmobs.LLMObsParent.Source; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import org.junit.jupiter.api.Test; + +class LLMObsParentTest { + + private static final DDTraceId TRACE = DDTraceId.from(1234); + private static final DDTraceId OTHER_TRACE = DDTraceId.from(5678); + + /** 0x6ab41f6e00000000fea4f6b7295bd5e9, as the propagation tag carries it. */ + private static final String WIRE_TRACE_ID = "141833418457532941290148377431948121577"; + + private static final String STORED_TRACE_ID = "6ab41f6e00000000fea4f6b7295bd5e9"; + + @Test + void noContextAndNothingPropagatedIsARoot() { + LLMObsParent parent = LLMObsParent.resolve(span(TRACE, null)); + + assertEquals(Source.NONE, parent.source); + assertTrue(parent.isRoot()); + assertEquals(LLMObsContext.ROOT_SPAN_ID, parent.parentId); + assertNull(parent.traceId); + assertNull(parent.mlApp); + assertNull(parent.sessionId); + assertNull(parent.sampleRate); + assertNull(parent.samplingDecision); + assertNull(parent.ignoredContext); + } + + @Test + void sameTraceInProcessContextWinsOverPropagatedValues() { + AgentSpanContext active = spanContext(TRACE, 42); + try (ContextScope ignored = + LLMObsContext.attach( + active, + "local-trace", + "local-app", + "local-session", + "v1", + "0.5", + LLMObsContext.SAMPLING_DECISION_SAMPLED, + "7", + "local-agent")) { + LLMObsParent parent = LLMObsParent.resolve(span(TRACE, propagated())); + + assertEquals(Source.IN_PROCESS, parent.source); + assertFalse(parent.isRoot()); + assertEquals("42", parent.parentId); + assertEquals("local-trace", parent.traceId); + assertEquals("local-app", parent.mlApp); + assertEquals("local-session", parent.sessionId); + assertEquals("v1", parent.agentVersion); + assertEquals("7", parent.parentAgentSpanId); + assertEquals("local-agent", parent.parentAgentName); + assertEquals("0.5", parent.sampleRate); + assertEquals(LLMObsContext.SAMPLING_DECISION_SAMPLED, parent.samplingDecision); + assertNull(parent.ignoredContext); + } + } + + /** + * In-process values are passed through as the context holds them, partial pairs included: the + * span re-attaches them for its own descendants and injections, which read each one separately. + */ + @Test + void inProcessValuesPassThroughUnpaired() { + try (ContextScope ignored = + LLMObsContext.attach( + spanContext(TRACE, 42), + null, + null, + null, + null, + null, + LLMObsContext.SAMPLING_DECISION_SAMPLED, + null, + "orphan-name")) { + LLMObsParent parent = LLMObsParent.resolve(span(TRACE, null)); + + assertEquals(Source.IN_PROCESS, parent.source); + assertFalse( + parent.isRoot(), "a parent without a full verdict is not a root and must not roll"); + assertNull(parent.sampleRate); + assertEquals(LLMObsContext.SAMPLING_DECISION_SAMPLED, parent.samplingDecision); + assertNull(parent.parentAgentSpanId); + assertEquals("orphan-name", parent.parentAgentName); + } + } + + @Test + void emptyInProcessValuesReadAsAbsent() { + try (ContextScope ignored = + LLMObsContext.attach(spanContext(TRACE, 42), "", "", "", "", "", "", "", "")) { + LLMObsParent parent = LLMObsParent.resolve(span(TRACE, null)); + + assertNull(parent.traceId); + assertNull(parent.mlApp); + assertNull(parent.sessionId); + assertNull(parent.agentVersion); + assertNull(parent.parentAgentSpanId); + assertNull(parent.parentAgentName); + assertNull(parent.sampleRate); + assertNull(parent.samplingDecision); + } + } + + @Test + void staleContextFallsBackToPropagatedValues() { + AgentSpanContext stale = spanContext(OTHER_TRACE, 99); + try (ContextScope ignored = LLMObsContext.attach(stale, "stale-session")) { + LLMObsParent parent = LLMObsParent.resolve(span(TRACE, propagated())); + + assertEquals(Source.PROPAGATED, parent.source); + assertEquals("upstream-session", parent.sessionId); + assertSame(stale, parent.ignoredContext); + } + } + + @Test + void staleContextWithNothingPropagatedIsARootThatReportsTheIgnoredContext() { + AgentSpanContext stale = spanContext(OTHER_TRACE, 99); + try (ContextScope ignored = + LLMObsContext.attach( + stale, + null, + "stale-app", + "stale-session", + "stale-v", + "0.25", + LLMObsContext.SAMPLING_DECISION_DROPPED, + "8", + "stale-agent")) { + LLMObsParent parent = LLMObsParent.resolve(span(TRACE, null)); + + assertTrue(parent.isRoot()); + assertEquals(LLMObsContext.ROOT_SPAN_ID, parent.parentId); + assertNull(parent.mlApp); + assertNull(parent.sessionId); + assertNull(parent.agentVersion); + assertNull(parent.parentAgentSpanId); + assertNull(parent.sampleRate); + assertNull(parent.samplingDecision); + assertSame(stale, parent.ignoredContext); + } + } + + @Test + void propagatedValuesAreAdoptedWithTheTraceIdConvertedFromTheWire() { + LLMObsParent parent = LLMObsParent.resolve(span(TRACE, propagated())); + + assertEquals(Source.PROPAGATED, parent.source); + assertFalse(parent.isRoot()); + assertEquals("12571653173807582101", parent.parentId); + assertEquals(STORED_TRACE_ID, parent.traceId); + assertEquals("upstream-app", parent.mlApp); + assertEquals("upstream-session", parent.sessionId); + assertNull(parent.agentVersion, "agent_version is not propagated"); + assertEquals("555", parent.parentAgentSpanId); + assertEquals("planner-agent", parent.parentAgentName); + assertEquals("0.25", parent.sampleRate); + assertEquals(LLMObsContext.SAMPLING_DECISION_DROPPED, parent.samplingDecision); + } + + /** + * Context that arrived without a parent id or a verdict still continues the caller's trace: the + * span is not a root, so it neither claims a parent of its own nor rolls a decision. + */ + @Test + void propagatedContextWithoutParentIdOrVerdictIsNotARoot() { + LLMObsParent parent = + LLMObsParent.resolve( + span( + TRACE, + new LLMObsPropagationValues( + WIRE_TRACE_ID, "upstream-app", null, null, null, null, null, null))); + + assertEquals(Source.PROPAGATED, parent.source); + assertFalse(parent.isRoot()); + assertEquals(LLMObsContext.ROOT_SPAN_ID, parent.parentId); + assertNull(parent.sampleRate); + assertNull(parent.samplingDecision); + } + + @Test + void propagatedRateWithoutADecisionIsDropped() { + LLMObsParent parent = + LLMObsParent.resolve( + span( + TRACE, + new LLMObsPropagationValues(null, null, null, null, null, "1", "0.5", null))); + + assertNull(parent.sampleRate); + assertNull(parent.samplingDecision); + } + + @Test + void propagatedDecisionWithoutARateIsDropped() { + LLMObsParent parent = + LLMObsParent.resolve( + span( + TRACE, + new LLMObsPropagationValues( + null, + null, + null, + null, + null, + "1", + null, + LLMObsContext.SAMPLING_DECISION_SAMPLED))); + + assertNull(parent.sampleRate); + assertNull(parent.samplingDecision); + } + + @Test + void propagatedAgentNameWithoutAnIdIsDropped() { + LLMObsParent parent = + LLMObsParent.resolve( + span( + TRACE, + new LLMObsPropagationValues( + null, null, null, null, "planner-agent", "1", null, null))); + + assertNull(parent.parentAgentSpanId); + assertNull(parent.parentAgentName); + } + + private static LLMObsPropagationValues propagated() { + return new LLMObsPropagationValues( + WIRE_TRACE_ID, + "upstream-app", + "upstream-session", + "555", + "planner-agent", + "12571653173807582101", + "0.25", + LLMObsContext.SAMPLING_DECISION_DROPPED); + } + + private static AgentSpanContext spanContext(DDTraceId traceId, long spanId) { + AgentSpanContext context = mock(AgentSpanContext.class); + when(context.getTraceId()).thenReturn(traceId); + when(context.getSpanId()).thenReturn(spanId); + return context; + } + + private static AgentSpan span(DDTraceId traceId, LLMObsPropagationValues extracted) { + AgentSpanContext context = mock(AgentSpanContext.class); + when(context.getExtractedLLMObsValues()).thenReturn(extracted); + AgentSpan span = mock(AgentSpan.class); + when(span.getTraceId()).thenReturn(traceId); + when(span.spanContext()).thenReturn(context); + return span; + } +} 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 index dcf683b31b9..7a7a1bb0e47 100644 --- a/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsSamplerTest.java +++ b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsSamplerTest.java @@ -4,6 +4,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; +import datadog.trace.api.DDTraceId; import java.util.concurrent.ThreadLocalRandom; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -103,4 +104,11 @@ void higherRatesAreASupersetOfLowerRates() { } } } + + @Test + void decideRendersTheVerdictInItsWireForm() { + DDTraceId traceId = DDTraceId.from(1234567890123L); + assertEquals(LLMObsContext.SAMPLING_DECISION_SAMPLED, new LLMObsSampler(1.0).decide(traceId)); + assertEquals(LLMObsContext.SAMPLING_DECISION_DROPPED, new LLMObsSampler(0.0).decide(traceId)); + } } diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/LLMObsTraceIdTest.java b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsTraceIdTest.java similarity index 99% rename from dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/LLMObsTraceIdTest.java rename to internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsTraceIdTest.java index 61c6e1c8427..9badf019104 100644 --- a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/LLMObsTraceIdTest.java +++ b/internal-api/src/test/java/datadog/trace/api/llmobs/LLMObsTraceIdTest.java @@ -1,4 +1,4 @@ -package datadog.trace.llmobs.domain; +package datadog.trace.api.llmobs; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull;