diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java index 4a901f9b370..94b5d836064 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java @@ -24,6 +24,10 @@ public static void forEachProperty(AgentPropagation.KeyClassifier classifier, St if (acceptJsonProperty(classifier, json, "x-datadog-trace-id")) { acceptJsonProperty(classifier, json, "x-datadog-parent-id"); acceptJsonProperty(classifier, json, "x-datadog-sampling-priority"); + // Propagation tags travel in x-datadog-tags. Without this the whole _dd.p.* set is + // silently dropped at a messaging boundary — including _dd.p.tid, which truncates a + // 128-bit trace id to 64 bits downstream, and the _dd.p.llmobs_* attribution tags. + acceptJsonProperty(classifier, json, "x-datadog-tags"); } if (Config.get().isDataStreamsEnabled()) { acceptJsonProperty(classifier, json, "dd-pathway-ctx-base64"); diff --git a/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParserTest.java b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParserTest.java new file mode 100644 index 00000000000..472f06d5052 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParserTest.java @@ -0,0 +1,31 @@ +package datadog.trace.bootstrap.instrumentation.messaging; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.LinkedHashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; + +/** + * Covers the one behaviour this branch adds to the {@code _datadog} message attribute parser shared + * by the AWS messaging instrumentations (SQS, SNS, EventBridge, Step Functions): {@code + * x-datadog-tags} is forwarded to the extractor. + */ +class DatadogAttributeParserTest { + + @Test + void forwardsPropagationTags() { + Map collected = new LinkedHashMap<>(); + DatadogAttributeParser.forEachProperty( + (key, value) -> { + collected.put(key, value); + return true; + }, + "{\"x-datadog-trace-id\":\"1234567890\"," + + "\"x-datadog-parent-id\":\"9876543210\"," + + "\"x-datadog-sampling-priority\":\"1\"," + + "\"x-datadog-tags\":\"_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000\"}"); + + assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags")); + } +} diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagator.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagator.java new file mode 100644 index 00000000000..cad55e790d8 --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagator.java @@ -0,0 +1,71 @@ +package datadog.trace.llmobs; + +import datadog.context.Context; +import datadog.context.propagation.CarrierSetter; +import datadog.context.propagation.CarrierVisitor; +import datadog.context.propagation.Propagator; +import datadog.trace.api.llmobs.LLMObsContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; + +/** + * Stages the LLM Observability propagation tags onto the span context being injected, so that every + * boundary already covered by automatic instrumentation carries LLMObs context without the + * application having to propagate it by hand. + * + *

This propagator writes nothing to the carrier itself. It runs ahead of the tracing propagator + * (see {@code AgentPropagation.LLMOBS_CONCERN}) and only populates the {@code _dd.p.llmobs_*} + * fields on the span context; the tracing propagator then serializes them into {@code + * x-datadog-tags} / {@code tracestate} along with every other propagation tag. + * + *

Values are resolved from the ambient {@link LLMObsContext} at injection time rather than being + * written once when a span starts, and every injection rewrites the whole set — falling back to + * whatever arrived on the inbound headers when no LLMObs context applies. That way the innermost + * active LLMObs span always wins, leaving an LLMObs scope stops contributing its tags, and a + * service that opens no LLMObs span of its own still forwards its caller's context — all without + * any save/restore bookkeeping. + */ +public class LLMObsContextPropagator implements Propagator { + + @Override + public void inject(Context context, C carrier, CarrierSetter setter) { + AgentSpan span = AgentSpan.fromContext(context); + if (span == null) { + return; + } + AgentSpanContext spanContext = span.spanContext(); + if (spanContext == null) { + return; + } + + // Gate on trace-id consistency, the same way DDLLMObsSpan gates parent_id/session_id + // inheritance. An LLMObs context leaked across an async boundary must not tag an outbound + // request that belongs to an unrelated trace. + AgentSpanContext llmObsContext = LLMObsContext.current(); + if (llmObsContext == null || !llmObsContext.getTraceId().equals(spanContext.getTraceId())) { + // Reset rather than return. These tags are staged on the root span context's propagation + // tags, which the whole local trace shares, so anything an earlier injection wrote would + // otherwise ride along on this one too — shipping a session and an agent attribution that + // are no longer active. Reset restores the extracted values instead of clearing outright: + // the same object also holds what came in on the wire, and a pass-through service must keep + // forwarding its caller's context. + spanContext.resetLLMObsContext(); + return; + } + + spanContext.updateLLMObsContext( + LLMObsContext.currentMlApp(), + LLMObsContext.currentSessionId(), + LLMObsContext.currentParentAgentSpanId(), + LLMObsContext.currentParentAgentName(), + String.valueOf(llmObsContext.getSpanId())); + } + + @Override + public Context extract(Context context, C carrier, CarrierVisitor visitor) { + // Nothing to do: the tracing propagator's codecs already parse the _dd.p.llmobs_* tags back + // into the extracted context's propagation tags, and DDLLMObsSpan reads them from there when + // no in-process LLMObs parent applies. + return context; + } +} diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java index 864cf27eb2c..b6ded19962f 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java @@ -1,6 +1,7 @@ package datadog.trace.llmobs; import datadog.communication.ddagent.SharedCommunicationObjects; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.WellKnownTags; import datadog.trace.api.llmobs.LLMObs; @@ -8,6 +9,7 @@ import datadog.trace.api.llmobs.LLMObsSpan; import datadog.trace.api.llmobs.LLMObsTags; import datadog.trace.api.telemetry.LLMObsMetricCollector; +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.llmobs.domain.DDLLMObsSpan; import datadog.trace.llmobs.domain.LLMObsEval; @@ -46,11 +48,15 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) { String mlApp = config.getLlmObsMlApp(); WellKnownTags wellKnownTags = config.getWellKnownTags(); - LLMObsInternal.setSpanFactory(new LLMObsManualSpanFactory(mlApp, wellKnownTags)); + // The span factory deliberately gets no default ml_app: DDLLMObsSpan applies it last, after + // in-process and propagated values have had their chance. + LLMObsInternal.setSpanFactory(new LLMObsManualSpanFactory(wellKnownTags)); LLMObsInternal.setEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config)); LLMObsInternal.setFeedbackProcessor(new LLMObsCustomFeedbackProcessor(mlApp, sco, config)); + + Propagators.register(AgentPropagation.LLMOBS_CONCERN, new LLMObsContextPropagator()); } private static class LLMObsCustomFeedbackProcessor implements LLMObs.LLMObsFeedbackProcessor { @@ -219,12 +225,10 @@ public void SubmitEvaluation( private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory { - private final String defaultMLApp; private final String serviceName; private final WellKnownTags wellKnownTags; - public LLMObsManualSpanFactory(String defaultMLApp, WellKnownTags wellKnownTags) { - this.defaultMLApp = defaultMLApp; + public LLMObsManualSpanFactory(WellKnownTags wellKnownTags) { this.serviceName = wellKnownTags.getService().toString(); this.wellKnownTags = wellKnownTags; } @@ -239,12 +243,7 @@ public LLMObsSpan startLLMSpan( DDLLMObsSpan span = new DDLLMObsSpan( - Tags.LLMOBS_LLM_SPAN_KIND, - spanName, - getMLApp(mlApp), - sessionId, - serviceName, - wellKnownTags); + Tags.LLMOBS_LLM_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags); if (modelName == null || modelName.isEmpty()) { modelName = CUSTOM_MODEL_VAL; @@ -273,7 +272,7 @@ public LLMObsSpan startAgentSpan( return new DDLLMObsSpan( Tags.LLMOBS_AGENT_SPAN_KIND, spanName, - getMLApp(mlApp), + mlApp, sessionId, serviceName, wellKnownTags, @@ -284,36 +283,21 @@ public LLMObsSpan startAgentSpan( public LLMObsSpan startToolSpan( String spanName, @Nullable String mlApp, @Nullable String sessionId) { return new DDLLMObsSpan( - Tags.LLMOBS_TOOL_SPAN_KIND, - spanName, - getMLApp(mlApp), - sessionId, - serviceName, - wellKnownTags); + Tags.LLMOBS_TOOL_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags); } @Override public LLMObsSpan startTaskSpan( String spanName, @Nullable String mlApp, @Nullable String sessionId) { return new DDLLMObsSpan( - Tags.LLMOBS_TASK_SPAN_KIND, - spanName, - getMLApp(mlApp), - sessionId, - serviceName, - wellKnownTags); + Tags.LLMOBS_TASK_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags); } @Override public LLMObsSpan startWorkflowSpan( String spanName, @Nullable String mlApp, @Nullable String sessionId) { return new DDLLMObsSpan( - Tags.LLMOBS_WORKFLOW_SPAN_KIND, - spanName, - getMLApp(mlApp), - sessionId, - serviceName, - wellKnownTags); + Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags); } @Override @@ -330,7 +314,7 @@ public LLMObsSpan startEmbeddingSpan( new DDLLMObsSpan( Tags.LLMOBS_EMBEDDING_SPAN_KIND, spanName, - getMLApp(mlApp), + mlApp, sessionId, serviceName, wellKnownTags); @@ -342,19 +326,7 @@ public LLMObsSpan startEmbeddingSpan( public LLMObsSpan startRetrievalSpan( String spanName, @Nullable String mlApp, @Nullable String sessionId) { return new DDLLMObsSpan( - Tags.LLMOBS_RETRIEVAL_SPAN_KIND, - spanName, - getMLApp(mlApp), - sessionId, - serviceName, - wellKnownTags); - } - - private String getMLApp(String mlApp) { - if (mlApp == null || mlApp.isEmpty()) { - return defaultMLApp; - } - return mlApp; + Tags.LLMOBS_RETRIEVAL_SPAN_KIND, spanName, mlApp, sessionId, serviceName, wellKnownTags); } } } 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 473d118cdc3..1c9c4b23219 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 @@ -24,6 +24,7 @@ import java.util.List; import java.util.Map; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -84,7 +85,7 @@ public class DDLLMObsSpan implements LLMObsSpan { public DDLLMObsSpan( @Nonnull String kind, String spanName, - @Nonnull String mlApp, + @Nullable String mlApp, String sessionId, @Nonnull String serviceName, WellKnownTags wellKnownTags) { @@ -94,7 +95,7 @@ public DDLLMObsSpan( public DDLLMObsSpan( @Nonnull String kind, String spanName, - @Nonnull String mlApp, + @Nullable String mlApp, String sessionId, @Nonnull String serviceName, WellKnownTags wellKnownTags, @@ -113,7 +114,7 @@ public DDLLMObsSpan( DDLLMObsSpan( @Nonnull String kind, String spanName, - @Nonnull String mlApp, + @Nullable String mlApp, String sessionId, @Nonnull String serviceName, WellKnownTags wellKnownTags, @@ -145,22 +146,22 @@ public DDLLMObsSpan( span.setTag(SPAN_KIND, kind); spanKind = kind; - this.mlApp = mlApp; - span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp); - // 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. + // 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 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() != span.getTraceId()) { + 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(), @@ -168,7 +169,17 @@ public DDLLMObsSpan( span.getTraceId(), span.getSpanId()); } else { + inheritedInProcess = true; parentSpanID = String.valueOf(parent.getSpanId()); + // 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. @@ -197,6 +208,37 @@ public DDLLMObsSpan( } } + if (!inheritedInProcess) { + // No usable in-process LLMObs parent, so fall back to what arrived from another service on + // the span context's propagation tags. + String propagatedParentId = asString(span.spanContext().getLLMObsParentId()); + if (propagatedParentId != null) { + parentSpanID = propagatedParentId; + } + if (resolvedMlApp == null || resolvedMlApp.isEmpty()) { + resolvedMlApp = asString(span.spanContext().getLLMObsMlApp()); + } + if (sessionId == null || sessionId.isEmpty()) { + sessionId = asString(span.spanContext().getLLMObsSessionId()); + } + resolvedParentAgentSpanId = asString(span.spanContext().getLLMObsParentAgentSpanId()); + if (resolvedParentAgentSpanId != null) { + resolvedParentAgentName = asString(span.spanContext().getLLMObsParentAgentName()); + } + } + + // The service default goes last, once the explicit, in-process and propagated values have all + // had their chance — applying it any earlier is indistinguishable from the caller naming the + // service explicitly, which is what would silently discard an upstream ml_app. Config already + // resolves this to DD_LLMOBS_ML_APP or DD_SERVICE, so the full order is + // explicit > in-process parent > propagated > DD_LLMOBS_ML_APP > DD_SERVICE, matching + // dd-trace-py's documented precedence. + if (resolvedMlApp == null || resolvedMlApp.isEmpty()) { + resolvedMlApp = Config.get().getLlmObsMlApp(); + } + this.mlApp = resolvedMlApp; + span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, resolvedMlApp); + // 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. @@ -231,11 +273,12 @@ public DDLLMObsSpan( } } - // Propagate the effective sessionId, agent_version, sampling decision and agent attribution - // to descendant LLMObs spans via the context. + // Propagate the effective mlApp, sessionId, agent_version, sampling decision and agent + // attribution to descendant LLMObs spans via the context. scope = LLMObsContext.attach( span.spanContext(), + resolvedMlApp, sessionId, resolvedAgentVersion, sampleRate, @@ -717,4 +760,9 @@ public DDTraceId getTraceId() { public long getSpanId() { return span.getSpanId(); } + + /** Narrow a propagated tag value to a non-empty String, or null. */ + private static String asString(CharSequence value) { + return value == null || value.length() == 0 ? null : value.toString(); + } } diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagatorTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagatorTest.java new file mode 100644 index 00000000000..3093c835fd1 --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagatorTest.java @@ -0,0 +1,316 @@ +package datadog.trace.llmobs; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.Context; +import datadog.context.propagation.Propagators; +import datadog.trace.agent.tooling.TracerInstaller; +import datadog.trace.api.WellKnownTags; +import datadog.trace.api.llmobs.LLMObsContext; +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; +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 datadog.trace.llmobs.domain.DDLLMObsSpan; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** + * Covers automatic LLM Observability context propagation. Injecting the active span the way + * auto-instrumentation does must carry the LLMObs context. + */ +class LLMObsContextPropagatorTest { + + private static final String ML_APP_TAG = "_dd.p.llmobs_ml_app"; + private static final String SESSION_ID_TAG = "_dd.p.llmobs_sid"; + private static final String PAGENT_SPAN_ID_TAG = "_dd.p.llmobs_pagent_span_id"; + private static final String PAGENT_NAME_TAG = "_dd.p.llmobs_pagent_name"; + private static final String PARENT_ID_TAG = "_dd.p.llmobs_parent_id"; + + private static CoreTracer tracer; + + @BeforeAll + static void installTracer() { + tracer = CoreTracer.builder().build(); + TracerInstaller.forceInstallGlobalTracer(tracer); + Propagators.register(AgentPropagation.LLMOBS_CONCERN, new LLMObsContextPropagator()); + } + + @AfterAll + static void closeTracer() { + TracerInstaller.forceInstallGlobalTracer(null); + tracer.close(); + } + + private static DDLLMObsSpan newSpan(String kind, String name, String mlApp, String sessionId) { + WellKnownTags tags = + new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java"); + return new DDLLMObsSpan(kind, name, mlApp, sessionId, "service", tags); + } + + private static AgentScope startRootApmScope() { + AgentSpan root = AgentTracer.get().buildSpan("apm", "sqs.produce").start(); + return AgentTracer.activateSpan(root); + } + + /** What an auto-instrumented client does: inject the active span into an outbound carrier. */ + private static Map autoInject(AgentSpan span) { + Map carrier = new HashMap<>(); + Propagators.defaultPropagator().inject(span, carrier, Map::put); + return carrier; + } + + /** An outbound carrier as a producer with an active agent span would have injected it. */ + private static Map producerCarrier(String mlApp, String sessionId) { + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan producer = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "dispatcher", mlApp, sessionId); + try { + return autoInject(AgentTracer.activeSpan()); + } finally { + producer.finish(); + } + } + } + + /** Activates an extracted carrier the way a message handler or request filter does. */ + private static AgentSpan extractSpan(Map carrier) { + Context extracted = + Propagators.defaultPropagator() + .extract(Context.root(), carrier, (c, visitor) -> c.forEach(visitor)); + AgentSpan span = AgentSpan.fromContext(extracted); + assertNotNull(span, "expected trace context to be extracted"); + return span; + } + + /** + * The local server span a consumer's entry-point instrumentation opens for the extracted context. + * This is the span whose propagation tags {@code CoreTracer} takes over from the extracted one, + * which is what puts the wire values and anything locally staged in the same place. + */ + private static AgentScope startLocalChildScope(AgentSpan parent) { + AgentSpan local = + AgentTracer.get().buildSpan("apm", "sqs.consume").asChildOf(parent.spanContext()).start(); + return AgentTracer.activateSpan(local); + } + + @Test + void stagesLlmObsTagsOnInjectionWithoutAnyManualPropagation() { + Map carrier; + String agentSpanId; + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan agent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "planner", "my-ml-app", "sess-1"); + agentSpanId = String.valueOf(agent.getSpanId()); + try { + carrier = autoInject(AgentTracer.activeSpan()); + } finally { + agent.finish(); + } + } + + String tags = carrier.get("x-datadog-tags"); + assertNotNull(tags, "expected x-datadog-tags to be injected"); + assertTrue(tags.contains(ML_APP_TAG + "=my-ml-app"), () -> "ml_app missing from " + tags); + assertTrue(tags.contains(SESSION_ID_TAG + "=sess-1"), () -> "session_id missing from " + tags); + assertTrue( + tags.contains(PAGENT_SPAN_ID_TAG + "=" + agentSpanId), + () -> "pagent_span_id missing from " + tags); + assertTrue( + tags.contains(PAGENT_NAME_TAG + "=planner"), () -> "pagent_name missing from " + tags); + assertTrue( + tags.contains(PARENT_ID_TAG + "=" + agentSpanId), () -> "parent_id missing from " + tags); + } + + @Test + void addsNothingWhenNoLlmObsSpanIsActive() { + Map carrier; + try (AgentScope apmScope = startRootApmScope()) { + carrier = autoInject(apmScope.span()); + } + + String tags = carrier.get("x-datadog-tags"); + assertTrue( + tags == null || !tags.contains("_dd.p.llmobs_"), () -> "unexpected LLMObs tags in " + tags); + } + + /** + * The staged tags live on the root span context's propagation tags, which are shared by + * the whole local trace. Once an injection has written them there, a later injection on the same + * trace has to overwrite them — otherwise it ships a session and an agent attribution that are no + * longer active. + */ + @Test + void doesNotLeakStagedTagsIntoALaterInjectionOnTheSameTrace() { + Map duringScope; + Map afterScope; + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan agent = newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "planner", "my-ml-app", "sess-1"); + try { + duringScope = autoInject(AgentTracer.activeSpan()); + } finally { + agent.finish(); + } + afterScope = autoInject(apmScope.span()); + } + + assertTrue( + duringScope.get("x-datadog-tags").contains(SESSION_ID_TAG), + "precondition: the first injection should have staged the LLMObs tags"); + String tags = afterScope.get("x-datadog-tags"); + assertTrue( + tags == null || !tags.contains("_dd.p.llmobs_"), + () -> "stale LLMObs tags leaked into a later injection: " + tags); + } + + /** + * The full cross-process hop, as an SQS producer/worker pair sees it: the producer injects into + * message attributes, the worker extracts and activates them, and an LLMObs span started by the + * worker inherits ml_app, session and agent attribution without any application-level plumbing. + * The worker names no ml_app of its own, so this also covers precedence — the propagated value + * outranks the worker's service default, keeping one logical application intact across the hop. + */ + @Test + void workerInheritsLlmObsContextAcrossTheBoundary() { + Map messageAttributes; + long producerTraceId; + String producerAgentSpanId; + + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan producer = + newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "dispatcher", "checkout", "sess-42"); + producerTraceId = producer.getTraceId().toLong(); + producerAgentSpanId = String.valueOf(producer.getSpanId()); + try { + messageAttributes = autoInject(AgentTracer.activeSpan()); + } finally { + producer.finish(); + } + } + + assertTrue( + messageAttributes.get("x-datadog-tags").contains(ML_APP_TAG + "=checkout"), + () -> "precondition: ml_app should be on the wire: " + messageAttributes); + + // Worker side: a fresh context, as a message handler would have. + AgentSpan consumeSpan = extractSpan(messageAttributes); + try (AgentScope consumeScope = AgentTracer.get().activateSpan(consumeSpan)) { + // The worker names no ml_app, so its own service default would otherwise apply. + DDLLMObsSpan workerTool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", null, null); + try { + assertEquals(producerTraceId, workerTool.getTraceId().toLong(), "trace should be joined"); + // The span publishes its resolved values to the context for its own descendants, so this + // is what the worker's LLMObs span actually settled on. + assertEquals("checkout", LLMObsContext.currentMlApp()); + assertEquals("sess-42", LLMObsContext.currentSessionId()); + assertEquals(producerAgentSpanId, LLMObsContext.currentParentAgentSpanId()); + assertEquals("dispatcher", LLMObsContext.currentParentAgentName()); + // The worker's LLMObs span parents onto the producer's, rather than starting a second + // root — this is the value DDLLMObsSpan reads for its parent_id. + assertEquals( + producerAgentSpanId, String.valueOf(consumeSpan.spanContext().getLLMObsParentId())); + } finally { + workerTool.finish(); + } + } + } + + /** + * A pass-through service — a proxy, a router, or any hop that opens no LLMObs span of its own — + * must keep forwarding the context it received. The staged and the extracted tags share one + * object, so resetting what this hop staged must restore what arrived rather than clear outright. + * + *

That reset also has to survive ordering: an outbound call injected before the + * service opens an LLMObs span of its own must leave the extracted values intact for the span + * that follows, which is what the second half of this test checks. + */ + @Test + void forwardsExtractedContextWhenNoLlmObsSpanIsActive() { + Map inbound = producerCarrier("checkout", "sess-42"); + assertTrue( + inbound.get("x-datadog-tags").contains(SESSION_ID_TAG + "=sess-42"), + () -> "precondition: session_id should be on the wire: " + inbound); + + Map outbound; + try (AgentScope consumeScope = startLocalChildScope(extractSpan(inbound))) { + // No LLMObs span here at all: this hop only relays the call. + outbound = autoInject(consumeScope.span()); + + // Still inside the same hop, after that injection has already reset the staged tags: an + // LLMObs span opened now must still see what arrived on the wire. + DDLLMObsSpan workerTool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", null, null); + try { + assertEquals("checkout", LLMObsContext.currentMlApp()); + assertEquals("sess-42", LLMObsContext.currentSessionId()); + assertEquals("dispatcher", LLMObsContext.currentParentAgentName()); + } finally { + workerTool.finish(); + } + } + + String tags = outbound.get("x-datadog-tags"); + assertNotNull(tags, "expected x-datadog-tags to be injected"); + for (String tag : + new String[] { + ML_APP_TAG + "=checkout", SESSION_ID_TAG + "=sess-42", PAGENT_NAME_TAG + "=dispatcher" + }) { + assertTrue(tags.contains(tag), () -> tag + " dropped by the pass-through hop: " + tags); + } + assertTrue(tags.contains(PARENT_ID_TAG + "="), () -> "parent_id dropped: " + tags); + assertTrue(tags.contains(PAGENT_SPAN_ID_TAG + "="), () -> "pagent_span_id dropped: " + tags); + } + + @Test + void peerSpanDoesNotInheritAFinishedSpansStagedContext() { + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan dispatcher = + newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "dispatcher", "checkout", "sess-42"); + try { + // Stages the five tags on the root span context's propagation tags, which the whole local + // trace shares. Nothing clears them when the span finishes. + autoInject(AgentTracer.activeSpan()); + } finally { + dispatcher.finish(); + } + + // A second LLMObs span on the same trace, with no LLMObs parent of its own. Whatever is still + // staged belongs to a span that has finished, so it isn't upstream context and must not be + // read as such. + DDLLMObsSpan peer = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "unrelated", "billing", null); + try { + assertEquals("billing", LLMObsContext.currentMlApp()); + assertNull(LLMObsContext.currentSessionId()); + assertNull(LLMObsContext.currentParentAgentSpanId()); + assertNull(LLMObsContext.currentParentAgentName()); + } finally { + peer.finish(); + } + } + } + + @Test + void workerWithoutUpstreamLlmObsContextInheritsNothing() { + Map messageAttributes; + try (AgentScope apmScope = startRootApmScope()) { + messageAttributes = autoInject(apmScope.span()); + } + + AgentSpan consumeSpan = extractSpan(messageAttributes); + try (AgentScope consumeScope = AgentTracer.get().activateSpan(consumeSpan)) { + DDLLMObsSpan workerTool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", "my-ml-app", null); + try { + assertNull(LLMObsContext.currentSessionId()); + assertNull(LLMObsContext.currentParentAgentSpanId()); + assertNull(consumeSpan.spanContext().getLLMObsParentId()); + } finally { + workerTool.finish(); + } + } + } +} diff --git a/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanMlAppTest.java b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanMlAppTest.java new file mode 100644 index 00000000000..cbd69b71c9b --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanMlAppTest.java @@ -0,0 +1,131 @@ +package datadog.trace.llmobs.domain; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.trace.agent.tooling.TracerInstaller; +import datadog.trace.api.Config; +import datadog.trace.api.WellKnownTags; +import datadog.trace.api.llmobs.LLMObsTags; +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 ml_app resolution, which mirrors session_id and agent_version inheritance: an explicit + * value always wins, otherwise the enclosing LLMObs span's value applies, and the service default + * is reached only when nothing else names an application. Applying the default any earlier would + * make "the caller passed nothing" indistinguishable from "the caller named the service", which is + * what would discard an ml_app propagated from another service. + */ +class DDLLMObsSpanMlAppTest { + private static final String ML_APP_TAG = "_ml_obs_tag." + LLMObsTags.ML_APP; + + 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 explicitMlAppTagsTheSpan() { + DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", "research-bot"); + try { + assertEquals("research-bot", spanOf(agent).getTag(ML_APP_TAG)); + } finally { + agent.finish(); + } + } + + @Test + void childSpanInheritsMlAppFromParentContext() { + DDLLMObsSpan agent = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", "research-bot"); + try (AgentScope ignored = AgentTracer.activateSpan(spanOf(agent))) { + DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool1", null); + try { + assertEquals("research-bot", spanOf(child).getTag(ML_APP_TAG)); + } finally { + child.finish(); + } + + // An empty ml_app is not a value, so it inherits the same way a null one does. + DDLLMObsSpan blank = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool2", ""); + try { + assertEquals("research-bot", spanOf(blank).getTag(ML_APP_TAG)); + } finally { + blank.finish(); + } + } finally { + agent.finish(); + } + } + + @Test + void explicitMlAppOverridesAnInheritedOneForItsOwnSubtree() { + DDLLMObsSpan outer = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent1", "research-bot"); + try (AgentScope ignored = AgentTracer.activateSpan(spanOf(outer))) { + DDLLMObsSpan inner = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "agent2", "summarizer"); + try (AgentScope innerScope = AgentTracer.activateSpan(spanOf(inner))) { + DDLLMObsSpan child = llmObsSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "tool1", null); + try { + assertEquals("summarizer", spanOf(inner).getTag(ML_APP_TAG)); + assertEquals("summarizer", spanOf(child).getTag(ML_APP_TAG)); + } finally { + child.finish(); + } + } finally { + inner.finish(); + } + } finally { + outer.finish(); + } + } + + @Test + void fallsBackToTheServiceDefaultWhenNothingNamesAnApplication() { + DDLLMObsSpan span = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "work", null); + try { + assertEquals(Config.get().getLlmObsMlApp(), spanOf(span).getTag(ML_APP_TAG)); + } finally { + span.finish(); + } + } + + private static DDLLMObsSpan llmObsSpan(String kind, String name, String mlApp) { + WellKnownTags tags = + new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java"); + return new DDLLMObsSpan(kind, name, mlApp, null, "service", tags); + } + + private static AgentSpan spanOf(DDLLMObsSpan llmObsSpan) { + try { + return (AgentSpan) SPAN_FIELD.get(llmObsSpan); + } catch (IllegalAccessException error) { + throw new AssertionError(error); + } + } +} 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 66d4f6aa285..df8080ff1db 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 @@ -177,6 +177,7 @@ void openAiRequestSpanInheritsDroppedSamplingDecisionFromActiveContext() throws parentSpan.spanContext(), null, null, + null, "0.25", LLMObsContext.SAMPLING_DECISION_DROPPED, null, @@ -208,6 +209,7 @@ void openAiRequestSpanInheritsRetainedSamplingDecisionFromActiveContext() throws parentSpan.spanContext(), null, null, + null, "1", LLMObsContext.SAMPLING_DECISION_SAMPLED, null, @@ -258,6 +260,7 @@ void openAiRequestSpanInheritsNothingFromStaleCrossTraceContext() throws Excepti try (ContextScope ignored = LLMObsContext.attach( staleParent.spanContext(), + null, "stale-session", "stale-version", "0.25", diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java index adf4cd66156..9729f0956e1 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java @@ -1492,6 +1492,47 @@ public PropagationTags getPropagationTags() { return getRootSpanContextOrThis().propagationTags; } + @Override + public CharSequence getLLMObsMlApp() { + return getPropagationTags().getLLMObsMlApp(); + } + + @Override + public CharSequence getLLMObsSessionId() { + return getPropagationTags().getLLMObsSessionId(); + } + + @Override + public CharSequence getLLMObsParentAgentSpanId() { + return getPropagationTags().getLLMObsParentAgentSpanId(); + } + + @Override + public CharSequence getLLMObsParentAgentName() { + return getPropagationTags().getLLMObsParentAgentName(); + } + + @Override + public CharSequence getLLMObsParentId() { + return getPropagationTags().getLLMObsParentId(); + } + + @Override + public void updateLLMObsContext( + CharSequence mlApp, + CharSequence sessionId, + CharSequence parentAgentSpanId, + CharSequence parentAgentName, + CharSequence parentId) { + getPropagationTags() + .updateLLMObsContext(mlApp, sessionId, parentAgentSpanId, parentAgentName, parentId); + } + + @Override + public void resetLLMObsContext() { + getPropagationTags().resetLLMObsContext(); + } + /** TraceSegment Implementation */ @Override public void setTagTop(String key, Object value, boolean sanitize) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ExtractedContext.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ExtractedContext.java index af503e6a6ed..2e22251f285 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ExtractedContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ExtractedContext.java @@ -117,6 +117,31 @@ public PropagationTags getPropagationTags() { return propagationTags; } + @Override + public CharSequence getLLMObsMlApp() { + return propagationTags.getLLMObsMlApp(); + } + + @Override + public CharSequence getLLMObsSessionId() { + return propagationTags.getLLMObsSessionId(); + } + + @Override + public CharSequence getLLMObsParentAgentSpanId() { + return propagationTags.getLLMObsParentAgentSpanId(); + } + + @Override + public CharSequence getLLMObsParentAgentName() { + return propagationTags.getLLMObsParentAgentName(); + } + + @Override + public CharSequence getLLMObsParentId() { + return propagationTags.getLLMObsParentId(); + } + @Override public String toString() { StringBuilder builder = new StringBuilder("ExtractedContext{"); diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java index 3a0c57a4dd8..0757640eb95 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java @@ -169,6 +169,58 @@ public interface Factory { */ public abstract void updateOrgPropagationMarker(CharSequence opm); + /** + * Returns the LLM Observability {@code ml_app} that arrived on the inbound headers as {@code + * _dd.p.llmobs_ml_app}, or {@code null} if none did. + * + *

These five getters read what was extracted, never what a local injection staged + * over it. The two live in the same object — an extracted context's tags become the local root's + * — but only the extracted half is a statement about the caller. A local LLMObs span's tags stay + * staged until the next injection resets them, so a sibling span opened in that window would + * otherwise read a finished span's attribution as if it had come from upstream. + */ + public abstract CharSequence getLLMObsMlApp(); + + /** + * Returns the LLM Observability {@code session_id} that arrived on the inbound headers as {@code + * _dd.p.llmobs_sid}, or {@code null} if none did. See {@link #getLLMObsMlApp()}. + */ + public abstract CharSequence getLLMObsSessionId(); + + /** + * Returns the span id of the parent LLM Observability agent span that arrived on the inbound + * headers as {@code _dd.p.llmobs_pagent_span_id}, or {@code null} if none did. See {@link + * #getLLMObsMlApp()}. + */ + public abstract CharSequence getLLMObsParentAgentSpanId(); + + /** + * Returns the name of the parent LLM Observability agent span that arrived on the inbound headers + * as {@code _dd.p.llmobs_pagent_name}, or {@code null} if none did. See {@link + * #getLLMObsMlApp()}. + */ + public abstract CharSequence getLLMObsParentAgentName(); + + /** + * Returns the span id of the parent LLM Observability span that arrived on the inbound headers as + * {@code _dd.p.llmobs_parent_id}, or {@code null} if none did. See {@link #getLLMObsMlApp()}. + */ + public abstract CharSequence getLLMObsParentId(); + + /** Sets the whole LLM Observability tag set to propagate with this trace. */ + public abstract void updateLLMObsContext( + CharSequence mlApp, + CharSequence sessionId, + CharSequence parentAgentSpanId, + CharSequence parentAgentName, + CharSequence parentId); + + /** + * Discards anything locally staged by {@link #updateLLMObsContext}, restoring the LLM + * Observability tag set that was extracted from the inbound headers. + */ + public abstract void resetLLMObsContext(); + public HashMap createTagMap() { HashMap result = new HashMap<>(); fillTagMap(result); diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java index 3ac0c7ad712..e8aa3a776d1 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java @@ -64,6 +64,11 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) { TagValue traceIdTagValue = null; int traceSource = 0; TagValue orgPropagationMarkerTagValue = null; + TagValue llmObsMlAppTagValue = null; + TagValue llmObsSessionIdTagValue = null; + TagValue llmObsParentAgentSpanIdTagValue = null; + TagValue llmObsParentAgentNameTagValue = null; + TagValue llmObsParentIdTagValue = null; while (tagPos < len) { int tagKeyEndsAt = validateCharsUntilSeparatorOrEnd( @@ -102,6 +107,16 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) { traceSource = ProductTraceSource.parseBitfieldHex(tagValue.toString()); } else if (tagKey.equals(ORG_PROPAGATION_MARKER_TAG)) { orgPropagationMarkerTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_ML_APP_TAG)) { + llmObsMlAppTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_SESSION_ID_TAG)) { + llmObsSessionIdTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_PAGENT_SPAN_ID_TAG)) { + llmObsParentAgentSpanIdTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_PAGENT_NAME_TAG)) { + llmObsParentAgentNameTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_PARENT_ID_TAG)) { + llmObsParentIdTagValue = tagValue; } else { if (tagPairs == null) { // This is roughly the size of a two element linked list but can hold six @@ -119,7 +134,13 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) { decisionMakerTagValue, traceIdTagValue, traceSource, - orgPropagationMarkerTagValue); + orgPropagationMarkerTagValue, + LLMObsTagValues.of( + llmObsMlAppTagValue, + llmObsSessionIdTagValue, + llmObsParentAgentSpanIdTagValue, + llmObsParentAgentNameTagValue, + llmObsParentIdTagValue)); } @Override diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/LLMObsTagValues.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/LLMObsTagValues.java new file mode 100644 index 00000000000..dac9e1d5a89 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/LLMObsTagValues.java @@ -0,0 +1,72 @@ +package datadog.trace.core.propagation.ptags; + +import java.util.Objects; + +/** + * Bundles the five LLM Observability propagation tag values as a single parameter. + * + *

Never {@code null}: use {@link #EMPTY} to say "no LLM Observability tags", and obtain + * instances through {@link #of} so that the common case — an incoming request carrying none of + * these tags, which is every request in a service not using LLM Observability — reuses {@code + * EMPTY} rather than allocating. + */ +final class LLMObsTagValues { + static final LLMObsTagValues EMPTY = new LLMObsTagValues(null, null, null, null, null); + + final TagValue mlApp; + final TagValue sessionId; + final TagValue parentAgentSpanId; + final TagValue parentAgentName; + final TagValue parentId; + + /** Returns {@link #EMPTY} when every value is {@code null}, otherwise a new bundle. */ + static LLMObsTagValues of( + TagValue mlApp, + TagValue sessionId, + TagValue parentAgentSpanId, + TagValue parentAgentName, + TagValue parentId) { + if (mlApp == null + && sessionId == null + && parentAgentSpanId == null + && parentAgentName == null + && parentId == null) { + return EMPTY; + } + return new LLMObsTagValues(mlApp, sessionId, parentAgentSpanId, parentAgentName, parentId); + } + + private LLMObsTagValues( + TagValue mlApp, + TagValue sessionId, + TagValue parentAgentSpanId, + TagValue parentAgentName, + TagValue parentId) { + this.mlApp = mlApp; + this.sessionId = sessionId; + this.parentAgentSpanId = parentAgentSpanId; + this.parentAgentName = parentAgentName; + this.parentId = parentId; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof LLMObsTagValues)) { + return false; + } + LLMObsTagValues other = (LLMObsTagValues) o; + return Objects.equals(mlApp, other.mlApp) + && Objects.equals(sessionId, other.sessionId) + && Objects.equals(parentAgentSpanId, other.parentAgentSpanId) + && Objects.equals(parentAgentName, other.parentAgentName) + && Objects.equals(parentId, other.parentId); + } + + @Override + public int hashCode() { + return Objects.hash(mlApp, sessionId, parentAgentSpanId, parentAgentName, parentId); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java index e2c0658a1d2..7723f0e795a 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java @@ -23,6 +23,11 @@ abstract class PTagsCodec { protected static final String PROPAGATION_ERROR_MALFORMED_TID = "malformed_tid "; protected static final String PROPAGATION_ERROR_INCONSISTENT_TID = "inconsistent_tid "; protected static final TagKey UPSTREAM_SERVICES_DEPRECATED_TAG = TagKey.from("upstream_services"); + protected static final TagKey LLMOBS_ML_APP_TAG = TagKey.from("llmobs_ml_app"); + protected static final TagKey LLMOBS_SESSION_ID_TAG = TagKey.from("llmobs_sid"); + protected static final TagKey LLMOBS_PAGENT_SPAN_ID_TAG = TagKey.from("llmobs_pagent_span_id"); + protected static final TagKey LLMOBS_PAGENT_NAME_TAG = TagKey.from("llmobs_pagent_name"); + protected static final TagKey LLMOBS_PARENT_ID_TAG = TagKey.from("llmobs_parent_id"); static String headerValue(PTagsCodec codec, PTags ptags) { return headerValue(codec, ptags, null); @@ -65,6 +70,22 @@ static String headerValue(PTagsCodec codec, PTags ptags, CharSequence lastParent codec.appendTag( sb, ORG_PROPAGATION_MARKER_TAG, ptags.getOrgPropagationMarkerTagValue(), size); } + LLMObsTagValues llmObsTags = ptags.getLLMObsTagValues(); + if (llmObsTags.mlApp != null) { + size = codec.appendTag(sb, LLMOBS_ML_APP_TAG, llmObsTags.mlApp, size); + } + if (llmObsTags.sessionId != null) { + size = codec.appendTag(sb, LLMOBS_SESSION_ID_TAG, llmObsTags.sessionId, size); + } + if (llmObsTags.parentAgentSpanId != null) { + size = codec.appendTag(sb, LLMOBS_PAGENT_SPAN_ID_TAG, llmObsTags.parentAgentSpanId, size); + } + if (llmObsTags.parentAgentName != null) { + size = codec.appendTag(sb, LLMOBS_PAGENT_NAME_TAG, llmObsTags.parentAgentName, size); + } + if (llmObsTags.parentId != null) { + size = codec.appendTag(sb, LLMOBS_PARENT_ID_TAG, llmObsTags.parentId, size); + } Iterator it = ptags.getTagPairs().iterator(); while (it.hasNext() && !codec.isTooLarge(sb, size)) { TagElement tagKey = it.next(); @@ -137,6 +158,32 @@ static void fillTagMap(PTags propagationTags, Map tagMap) { .forType(Encoding.DATADOG) .toString()); } + LLMObsTagValues llmObsTags = propagationTags.getLLMObsTagValues(); + if (llmObsTags.mlApp != null) { + tagMap.put( + LLMOBS_ML_APP_TAG.forType(Encoding.DATADOG).toString(), + llmObsTags.mlApp.forType(Encoding.DATADOG).toString()); + } + if (llmObsTags.sessionId != null) { + tagMap.put( + LLMOBS_SESSION_ID_TAG.forType(Encoding.DATADOG).toString(), + llmObsTags.sessionId.forType(Encoding.DATADOG).toString()); + } + if (llmObsTags.parentAgentSpanId != null) { + tagMap.put( + LLMOBS_PAGENT_SPAN_ID_TAG.forType(Encoding.DATADOG).toString(), + llmObsTags.parentAgentSpanId.forType(Encoding.DATADOG).toString()); + } + if (llmObsTags.parentAgentName != null) { + tagMap.put( + LLMOBS_PAGENT_NAME_TAG.forType(Encoding.DATADOG).toString(), + llmObsTags.parentAgentName.forType(Encoding.DATADOG).toString()); + } + if (llmObsTags.parentId != null) { + tagMap.put( + LLMOBS_PARENT_ID_TAG.forType(Encoding.DATADOG).toString(), + llmObsTags.parentId.forType(Encoding.DATADOG).toString()); + } if (propagationTags.getError() != null) { tagMap.put(PROPAGATION_ERROR_TAG_KEY, propagationTags.getError()); } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java index 0b5184d448a..c45ded20957 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java @@ -4,6 +4,11 @@ import static datadog.trace.core.propagation.PropagationTags.HeaderType.W3C; import static datadog.trace.core.propagation.ptags.PTagsCodec.DECISION_MAKER_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.KNUTH_SAMPLING_RATE_TAG; +import static datadog.trace.core.propagation.ptags.PTagsCodec.LLMOBS_ML_APP_TAG; +import static datadog.trace.core.propagation.ptags.PTagsCodec.LLMOBS_PAGENT_NAME_TAG; +import static datadog.trace.core.propagation.ptags.PTagsCodec.LLMOBS_PAGENT_SPAN_ID_TAG; +import static datadog.trace.core.propagation.ptags.PTagsCodec.LLMOBS_PARENT_ID_TAG; +import static datadog.trace.core.propagation.ptags.PTagsCodec.LLMOBS_SESSION_ID_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.ORG_PROPAGATION_MARKER_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.TRACE_ID_TAG; import static datadog.trace.core.propagation.ptags.PTagsCodec.TRACE_SOURCE_TAG; @@ -50,7 +55,7 @@ PTagsCodec getDecoderEncoder(@Nonnull HeaderType headerType) { @Override public final PropagationTags empty() { - return createValid(null, null, null, ProductTraceSource.UNSET, null); + return createValid(null, null, null, ProductTraceSource.UNSET, null, LLMObsTagValues.EMPTY); } @Override @@ -71,14 +76,16 @@ PropagationTags createValid( TagValue decisionMakerTagValue, TagValue traceIdTagValue, int productTraceSource, - TagValue orgPropagationMarkerTagValue) { + TagValue orgPropagationMarkerTagValue, + @Nonnull LLMObsTagValues llmObsTagValues) { return new PTags( this, tagPairs, decisionMakerTagValue, traceIdTagValue, productTraceSource, - orgPropagationMarkerTagValue); + orgPropagationMarkerTagValue, + llmObsTagValues); } PropagationTags createInvalid(String error) { @@ -112,6 +119,25 @@ static class PTags extends PropagationTags { private volatile TagValue orgPropagationMarkerTagValue; + /** + * The LLM Observability propagation tags, held as one immutable bundle. Never {@code null} — + * {@link LLMObsTagValues#EMPTY} means "none". + */ + private volatile LLMObsTagValues llmObsTags; + + /** + * The LLM Observability tags as they arrived on the wire, before anything local staged over + * them. Never {@code null}. + * + *

Kept separate because these tags have two writers: the codec, at construction, and the + * LLMObs propagator, at every injection. Both write {@link #llmObsTags}, and an extracted + * context's {@code PropagationTags} become the local root's, so without a record of what was + * extracted the propagator cannot clear its own staging without also deleting the caller's + * context — silently dropping it at any service that forwards a request without opening an + * LLMObs span of its own. + */ + private final LLMObsTagValues extractedLLMObsTags; + // Static cache for the most-recently-seen rate → TagValue. In steady state a service uses one // rate, so this eliminates the char[] + String allocation on every new PTags instance. // Writes are benign-racy: two threads computing the same rate produce equal TagValues. @@ -158,7 +184,8 @@ static class PTags extends PropagationTags { TagValue decisionMakerTagValue, TagValue traceIdTagValue, int traceSource, - TagValue orgPropagationMarkerTagValue) { + TagValue orgPropagationMarkerTagValue, + @Nonnull LLMObsTagValues llmObsTagValues) { this( factory, tagPairs, @@ -168,7 +195,8 @@ static class PTags extends PropagationTags { PrioritySampling.UNSET, null, null, - orgPropagationMarkerTagValue); + orgPropagationMarkerTagValue, + llmObsTagValues); } PTags( @@ -180,7 +208,8 @@ static class PTags extends PropagationTags { int samplingPriority, CharSequence origin, CharSequence lastParentId, - TagValue orgPropagationMarkerTagValue) { + TagValue orgPropagationMarkerTagValue, + @Nonnull LLMObsTagValues llmObsTagValues) { assert tagPairs == null || tagPairs.size() % 2 == 0; this.factory = factory; this.tagPairs = tagPairs; @@ -191,6 +220,8 @@ static class PTags extends PropagationTags { this.origin = origin; this.lastParentId = lastParentId; this.orgPropagationMarkerTagValue = orgPropagationMarkerTagValue; + this.llmObsTags = llmObsTagValues; + this.extractedLLMObsTags = llmObsTagValues; if (traceIdTagValue != null) { CharSequence traceIdHighOrderBitsHex = traceIdTagValue.forType(TagElement.Encoding.DATADOG); this.traceIdHighOrderBits = @@ -212,7 +243,8 @@ static PTags withError(PTagsFactory factory, String error) { PrioritySampling.UNSET, null, null, - null); + null, + LLMObsTagValues.EMPTY); pTags.error = error; return pTags; } @@ -248,8 +280,7 @@ private void doUpdateTraceSamplingPriority(int samplingPriority, int samplingMec TagValue newDM = TagValue.from("-" + samplingMechanism); if (!newDM.equals(decisionMakerTagValue)) { // This should invalidate any cached w3c and datadog header - clearCachedHeader(DATADOG); - clearCachedHeader(W3C); + clearCachedHeaders(); } decisionMakerTagValue = newDM; } @@ -257,8 +288,7 @@ private void doUpdateTraceSamplingPriority(int samplingPriority, int samplingMec // Drop the decision maker tag if (decisionMakerTagValue != null) { // This should invalidate any cached w3c and datadog header - clearCachedHeader(DATADOG); - clearCachedHeader(W3C); + clearCachedHeaders(); } decisionMakerTagValue = null; } @@ -275,8 +305,7 @@ public void addTraceSource(final int product) { } // Invalidate cached headers (atomic context ensures correctness) - clearCachedHeader(DATADOG); - clearCachedHeader(W3C); + clearCachedHeaders(); // Set the bit for the given product return ProductTraceSource.updateProduct(currentValue, product); @@ -301,8 +330,7 @@ public String getDebugPropagation() { @Override public void updateKnuthSamplingRate(double rate) { if (Double.compare(knuthSamplingRate, rate) != 0) { - clearCachedHeader(DATADOG); - clearCachedHeader(W3C); + clearCachedHeaders(); knuthSamplingRate = rate; if (Double.isNaN(rate)) { knuthSamplingRateTagValue = null; @@ -367,8 +395,7 @@ public CharSequence getOrgPropagationMarker() { public void updateOrgPropagationMarker(CharSequence opm) { TagValue newValue = opm == null ? null : TagValue.from(opm); if (!Objects.equals(this.orgPropagationMarkerTagValue, newValue)) { - clearCachedHeader(DATADOG); - clearCachedHeader(W3C); + clearCachedHeaders(); this.orgPropagationMarkerTagValue = newValue; } } @@ -377,6 +404,112 @@ TagValue getOrgPropagationMarkerTagValue() { return orgPropagationMarkerTagValue; } + @Override + public void updateLLMObsContext( + CharSequence mlApp, + CharSequence sessionId, + CharSequence parentAgentSpanId, + CharSequence parentAgentName, + CharSequence parentId) { + LLMObsTagValues updated = + LLMObsTagValues.of( + toTagValue(mlApp), + toTagValue(sessionId), + toTagValue(parentAgentSpanId), + toTagValue(parentAgentName), + toTagValue(parentId)); + if (!updated.equals(llmObsTags)) { + clearCachedHeaders(); + llmObsTags = updated; + } + } + + @Override + public void resetLLMObsContext() { + if (!extractedLLMObsTags.equals(llmObsTags)) { + clearCachedHeaders(); + llmObsTags = extractedLLMObsTags; + } + } + + @Override + public CharSequence getLLMObsMlApp() { + return decoded(extractedLLMObsTags.mlApp); + } + + @Override + public CharSequence getLLMObsSessionId() { + return decoded(extractedLLMObsTags.sessionId); + } + + @Override + public CharSequence getLLMObsParentAgentSpanId() { + return decoded(extractedLLMObsTags.parentAgentSpanId); + } + + @Override + public CharSequence getLLMObsParentAgentName() { + return decoded(extractedLLMObsTags.parentAgentName); + } + + @Override + public CharSequence getLLMObsParentId() { + return decoded(extractedLLMObsTags.parentId); + } + + /** + * The value as the application wrote it, undoing the {@code tracestate} substitutions when the + * value came in on that carrier. {@link TagValue#toString()} would return it in whichever + * encoding it arrived in, so a {@code ml_app} of {@code a=b} would read back as {@code a~b} + * after a W3C-only hop. Same conversion {@link PTagsCodec#fillTagMap} applies to every other + * {@code _dd.p.*} tag. + */ + private static CharSequence decoded(TagValue value) { + return value == null ? null : value.forType(TagElement.Encoding.DATADOG); + } + + LLMObsTagValues getLLMObsTagValues() { + return llmObsTags; + } + + /** + * Wraps a non-empty value as a {@link TagValue}, or {@code null} if it is empty or cannot be + * represented in {@code x-datadog-tags}. + * + *

Unlike every other {@code _dd.p.*} tag, these values come from the application rather than + * the tracer, so they have to be checked before they reach the wire. A value the receiving + * codec rejects doesn't just lose itself: it fails the whole tagset with {@code decoding_error} + * and takes {@code _dd.p.tid} with it, leaving the two services disagreeing about the upper 64 + * bits of the trace id. Dropping the one tag is the cheaper loss. + */ + private static TagValue toTagValue(CharSequence value) { + if (value == null || value.length() == 0 || !isRepresentable(value)) { + return null; + } + return TagValue.from(value); + } + + /** + * Whether every character survives each carrier the value can travel on: printable ASCII, no + * {@code ,} (the {@code x-datadog-tags} separator), nothing the {@code tracestate} conversion + * rewrites, and no {@code "} or {@code \} — AWS messaging carries these headers in a {@code + * _datadog} JSON attribute that is written and parsed without escaping. + */ + private static boolean isRepresentable(CharSequence value) { + for (int i = 0; i < value.length(); i++) { + char c = value.charAt(i); + if (c < ' ' + || c > '~' + || c == ',' + || c == '"' + || c == '\\' + || !TagValue.survivesW3CRoundTrip(c)) { + return false; + } + } + return true; + } + @Override public int getSamplingPriority() { return samplingPriority; @@ -473,6 +606,15 @@ private void setCachedHeader(HeaderType headerType, String header) { cache[headerType.ordinal()] = header; } + /** + * Invalidate every encoding's cached header, and the memoized x-datadog-tags size with them. + * Use this whenever a change affects both wire formats. + */ + private void clearCachedHeaders() { + clearCachedHeader(DATADOG); + clearCachedHeader(W3C); + } + private void clearCachedHeader(HeaderType headerType) { if (headerType == DATADOG) { invalidateXDatadogTagsSize(); @@ -512,6 +654,19 @@ int getXDatadogTagsSize() { size = PTagsCodec.calcXDatadogTagsSize( size, ORG_PROPAGATION_MARKER_TAG, getOrgPropagationMarkerTagValue()); + LLMObsTagValues currentLLMObsTags = llmObsTags; + size = PTagsCodec.calcXDatadogTagsSize(size, LLMOBS_ML_APP_TAG, currentLLMObsTags.mlApp); + size = + PTagsCodec.calcXDatadogTagsSize( + size, LLMOBS_SESSION_ID_TAG, currentLLMObsTags.sessionId); + size = + PTagsCodec.calcXDatadogTagsSize( + size, LLMOBS_PAGENT_SPAN_ID_TAG, currentLLMObsTags.parentAgentSpanId); + size = + PTagsCodec.calcXDatadogTagsSize( + size, LLMOBS_PAGENT_NAME_TAG, currentLLMObsTags.parentAgentName); + size = + PTagsCodec.calcXDatadogTagsSize(size, LLMOBS_PARENT_ID_TAG, currentLLMObsTags.parentId); int currentProductTraceSource = traceSource; if (currentProductTraceSource != ProductTraceSource.UNSET) { size = @@ -553,8 +708,7 @@ public void updateAndLockDecisionMaker(PropagationTags source) { canChangeDecisionMaker = false; decisionMakerTagValue = ((PTags) source).getDecisionMakerTagValue(); if (decisionMakerTagValue != null) { - clearCachedHeader(DATADOG); - clearCachedHeader(W3C); + clearCachedHeaders(); } } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/TagValue.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/TagValue.java index 8c1d6412fec..a6c5f1bbf0b 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/TagValue.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/TagValue.java @@ -111,6 +111,14 @@ private static char convertW3CtoDD(char c) { return c; } + /** + * Whether {@code c} comes back unchanged from a {@code tracestate} hop. The conversion is lossy + * for some characters, so a value carrying one arrives altered at the next service. + */ + static boolean survivesW3CRoundTrip(char c) { + return convertW3CtoDD(convertDDtoW3C(c)) == c; + } + private final CharSequence[] values = new CharSequence[Encoding.getNumValues()]; private final int source; private final int hash; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/W3CPTagsCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/W3CPTagsCodec.java index c0018544188..8a56ab28a05 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/W3CPTagsCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/W3CPTagsCodec.java @@ -99,6 +99,11 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) { int maxUnknownSize = 0; CharSequence lastParentId = null; TagValue orgPropagationMarkerTagValue = null; + TagValue llmObsMlAppTagValue = null; + TagValue llmObsSessionIdTagValue = null; + TagValue llmObsParentAgentSpanIdTagValue = null; + TagValue llmObsParentAgentNameTagValue = null; + TagValue llmObsParentIdTagValue = null; while (tagPos < ddMemberValueEnd) { tagPos = skipEmptyElements(value, tagPos, ddMemberValueEnd); if (tagPos >= ddMemberValueEnd) { @@ -168,6 +173,16 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) { traceSource = ProductTraceSource.parseBitfieldHex(tagValue.toString()); } else if (tagKey.equals(ORG_PROPAGATION_MARKER_TAG)) { orgPropagationMarkerTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_ML_APP_TAG)) { + llmObsMlAppTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_SESSION_ID_TAG)) { + llmObsSessionIdTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_PAGENT_SPAN_ID_TAG)) { + llmObsParentAgentSpanIdTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_PAGENT_NAME_TAG)) { + llmObsParentAgentNameTagValue = tagValue; + } else if (tagKey.equals(LLMOBS_PARENT_ID_TAG)) { + llmObsParentIdTagValue = tagValue; } else { if (tagPairs == null) { // This is roughly the size of a two element linked list but can hold six @@ -201,7 +216,13 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) { ddMemberValueEnd, maxUnknownSize, lastParentId, - orgPropagationMarkerTagValue); + orgPropagationMarkerTagValue, + LLMObsTagValues.of( + llmObsMlAppTagValue, + llmObsSessionIdTagValue, + llmObsParentAgentSpanIdTagValue, + llmObsParentAgentNameTagValue, + llmObsParentIdTagValue)); } @Override @@ -764,7 +785,8 @@ private static W3CPTags empty( ddMemberValueEnd, 0, null, - null); + null, + LLMObsTagValues.EMPTY); } private static class W3CPTags extends PTags { @@ -799,7 +821,8 @@ public W3CPTags( int ddMemberValueEnd, int maxUnknownSize, CharSequence lastParentId, - TagValue orgPropagationMarkerTagValue) { + TagValue orgPropagationMarkerTagValue, + LLMObsTagValues llmObsTagValues) { super( factory, tagPairs, @@ -809,7 +832,8 @@ public W3CPTags( samplingPriority, origin, lastParentId, - orgPropagationMarkerTagValue); + orgPropagationMarkerTagValue, + llmObsTagValues); this.tracestate = original; this.firstMemberStart = firstMemberStart; this.ddMemberStart = ddMemberStart; diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogPropagationTagsTest.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogPropagationTagsTest.java index 73c6c8af9b9..e9a2a9b2697 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogPropagationTagsTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogPropagationTagsTest.java @@ -15,61 +15,65 @@ import java.util.HashMap; import java.util.Map; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.converter.ConvertWith; +import org.junit.jupiter.params.provider.ValueSource; import org.tabletest.junit.TableTest; class DatadogPropagationTagsTest extends DDJavaSpecification { @TableTest({ - "scenario | headerValue | expectedHeaderValue | tags ", - "null input | | | [:] ", - "empty input | '' | | [:] ", - "valid dm tag short | '_dd.p.dm=934086a686-4' | '_dd.p.dm=934086a686-4' | [_dd.p.dm: '934086a686-4'] ", - "valid dm tag 2-digit | '_dd.p.dm=934086a686-10' | '_dd.p.dm=934086a686-10' | [_dd.p.dm: '934086a686-10'] ", - "valid dm tag 3-digit | '_dd.p.dm=934086a686-102' | '_dd.p.dm=934086a686-102' | [_dd.p.dm: '934086a686-102'] ", - "dm tag minus only | '_dd.p.dm=-1' | '_dd.p.dm=-1' | [_dd.p.dm: '-1'] ", - "dm tag with trailing separator | '_dd.p.dm=-4,' | '_dd.p.dm=-4' | [_dd.p.dm: '-4'] ", - "any p tag | '_dd.p.anytag=value' | '_dd.p.anytag=value' | [_dd.p.anytag: 'value'] ", - "non p tag dropped | '_dd.b.somekey=value' | | [:] ", - "upstream services alone dropped | '_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1' | | [:] ", - "dm and anytag | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | [_dd.p.dm: '934086a686-4', _dd.p.anytag: 'value'] ", - "dm with upstream and anytag | '_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | [_dd.p.dm: '934086a686-4', _dd.p.anytag: 'value'] ", - "ddb keyonly with dm upstream | '_dd.b.keyonly=value,_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | [_dd.p.dm: '934086a686-4', _dd.p.anytag: 'value'] ", - "valid p tag with spaces | '_dd.p.ab=1 2 3' | '_dd.p.ab=1 2 3' | [_dd.p.ab: '1 2 3'] ", - "valid p tag leading trail spc | '_dd.p.ab= 123 ' | '_dd.p.ab= 123 ' | [_dd.p.ab: ' 123 '] ", - "key only error | '_dd.p.keyonly' | | [_dd.propagation_error: 'decoding_error'] ", - "leading comma error | ',_dd.p.dm=Value' | | [_dd.propagation_error: 'decoding_error'] ", - "comma only error | ',' | | [_dd.propagation_error: 'decoding_error'] ", - "ddb keyonly with embedded keyonly | '_dd.b.somekey=value,_dd.p.dm=934086a686-4,_dd.p.keyonly,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", - "embedded keyonly with dm upstream | '_dd.p.keyonly,_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", - "leading comma with dm upstream | ',_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", - "double comma in tagset | '_dd.p.dm=934086a686-4,,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", - "space tag in tagset | '_dd.p.dm=934086a686-4, ,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", - "upstream variant dropped alone | '_dd.p.upstream_services=bmV1dHJvbg==|0|1|0.2253' | | [:] ", - "leading space error | ' _dd.p.ab=123' | | [_dd.propagation_error: 'decoding_error'] ", - "key with space error | '_dd.p.a b=123' | | [_dd.propagation_error: 'decoding_error'] ", - "trailing key space error | '_dd.p.ab =123' | | [_dd.propagation_error: 'decoding_error'] ", - "space inside key error | '_dd.p. ab=123' | | [_dd.propagation_error: 'decoding_error'] ", - "tag with eq value | '_dd.p.a=b=1=2' | '_dd.p.a=b=1=2' | [_dd.p.a: 'b=1=2'] ", - "invalid key non-ascii | '_dd.p.1ö2=value' | | [_dd.propagation_error: 'decoding_error'] ", - "value with equals | '_dd.p.ab=1=2' | '_dd.p.ab=1=2' | [_dd.p.ab: '1=2'] ", - "invalid value non-ascii | '_dd.p.ab=1ô2' | | [_dd.propagation_error: 'decoding_error'] ", - "dm tag upper case | '_dd.p.dm=934086A686-4' | | [_dd.propagation_error: 'decoding_error'] ", - "dm tag too short | '_dd.p.dm=934086a66-4' | | [_dd.propagation_error: 'decoding_error'] ", - "dm tag too long | '_dd.p.dm=934086a6653-4' | | [_dd.propagation_error: 'decoding_error'] ", - "dm tag missing separator | '_dd.p.dm=934086a66534' | | [_dd.propagation_error: 'decoding_error'] ", - "dm tag missing mechanism | '_dd.p.dm=934086a665-' | | [_dd.propagation_error: 'decoding_error'] ", - "dm tag invalid mechanism char | '_dd.p.dm=934086a665-a' | | [_dd.propagation_error: 'decoding_error'] ", - "dm tag mechanism with letter | '_dd.p.dm=934086a665-12b' | | [_dd.propagation_error: 'decoding_error'] ", - "tid empty | '_dd.p.tid=' | | [_dd.propagation_error: 'decoding_error'] ", - "tid length 1 | '_dd.p.tid=1' | | [_dd.propagation_error: 'malformed_tid 1'] ", - "tid length 15 | '_dd.p.tid=111111111111111' | | [_dd.propagation_error: 'malformed_tid 111111111111111'] ", - "tid length 17 | '_dd.p.tid=11111111111111111' | | [_dd.propagation_error: 'malformed_tid 11111111111111111']", - "tid invalid uppercase | '_dd.p.tid=123456789ABCDEF0' | | [_dd.propagation_error: 'malformed_tid 123456789ABCDEF0'] ", - "tid invalid non-hex | '_dd.p.tid=123456789abcdefg' | | [_dd.propagation_error: 'malformed_tid 123456789abcdefg'] ", - "tid invalid negative | '_dd.p.tid=-123456789abcdef' | | [_dd.propagation_error: 'malformed_tid -123456789abcdef'] ", - "ts valid 02 | '_dd.p.ts=02' | '_dd.p.ts=02' | [_dd.p.ts: '02'] ", - "ts zero dropped | '_dd.p.ts=00' | | [:] ", - "ts invalid foo | '_dd.p.ts=foo' | | [_dd.propagation_error: 'decoding_error'] " + "scenario | headerValue | expectedHeaderValue | tags ", + "null input | | | [:] ", + "empty input | '' | | [:] ", + "valid dm tag short | '_dd.p.dm=934086a686-4' | '_dd.p.dm=934086a686-4' | [_dd.p.dm: '934086a686-4'] ", + "valid dm tag 2-digit | '_dd.p.dm=934086a686-10' | '_dd.p.dm=934086a686-10' | [_dd.p.dm: '934086a686-10'] ", + "valid dm tag 3-digit | '_dd.p.dm=934086a686-102' | '_dd.p.dm=934086a686-102' | [_dd.p.dm: '934086a686-102'] ", + "dm tag minus only | '_dd.p.dm=-1' | '_dd.p.dm=-1' | [_dd.p.dm: '-1'] ", + "dm tag with trailing separator | '_dd.p.dm=-4,' | '_dd.p.dm=-4' | [_dd.p.dm: '-4'] ", + "any p tag | '_dd.p.anytag=value' | '_dd.p.anytag=value' | [_dd.p.anytag: 'value'] ", + "non p tag dropped | '_dd.b.somekey=value' | | [:] ", + "upstream services alone dropped | '_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1' | | [:] ", + "dm and anytag | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | [_dd.p.dm: '934086a686-4', _dd.p.anytag: 'value'] ", + "dm with upstream and anytag | '_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | [_dd.p.dm: '934086a686-4', _dd.p.anytag: 'value'] ", + "ddb keyonly with dm upstream | '_dd.b.keyonly=value,_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | '_dd.p.dm=934086a686-4,_dd.p.anytag=value' | [_dd.p.dm: '934086a686-4', _dd.p.anytag: 'value'] ", + "valid p tag with spaces | '_dd.p.ab=1 2 3' | '_dd.p.ab=1 2 3' | [_dd.p.ab: '1 2 3'] ", + "valid p tag leading trail spc | '_dd.p.ab= 123 ' | '_dd.p.ab= 123 ' | [_dd.p.ab: ' 123 '] ", + "key only error | '_dd.p.keyonly' | | [_dd.propagation_error: 'decoding_error'] ", + "leading comma error | ',_dd.p.dm=Value' | | [_dd.propagation_error: 'decoding_error'] ", + "comma only error | ',' | | [_dd.propagation_error: 'decoding_error'] ", + "ddb keyonly with embedded keyonly | '_dd.b.somekey=value,_dd.p.dm=934086a686-4,_dd.p.keyonly,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", + "embedded keyonly with dm upstream | '_dd.p.keyonly,_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", + "leading comma with dm upstream | ',_dd.p.dm=934086a686-4,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", + "double comma in tagset | '_dd.p.dm=934086a686-4,,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", + "space tag in tagset | '_dd.p.dm=934086a686-4, ,_dd.p.upstream_services=bWNudWx0eS13ZWI|0|1|0.1,_dd.p.anytag=value' | | [_dd.propagation_error: 'decoding_error'] ", + "upstream variant dropped alone | '_dd.p.upstream_services=bmV1dHJvbg==|0|1|0.2253' | | [:] ", + "leading space error | ' _dd.p.ab=123' | | [_dd.propagation_error: 'decoding_error'] ", + "key with space error | '_dd.p.a b=123' | | [_dd.propagation_error: 'decoding_error'] ", + "trailing key space error | '_dd.p.ab =123' | | [_dd.propagation_error: 'decoding_error'] ", + "space inside key error | '_dd.p. ab=123' | | [_dd.propagation_error: 'decoding_error'] ", + "tag with eq value | '_dd.p.a=b=1=2' | '_dd.p.a=b=1=2' | [_dd.p.a: 'b=1=2'] ", + "invalid key non-ascii | '_dd.p.1ö2=value' | | [_dd.propagation_error: 'decoding_error'] ", + "value with equals | '_dd.p.ab=1=2' | '_dd.p.ab=1=2' | [_dd.p.ab: '1=2'] ", + "invalid value non-ascii | '_dd.p.ab=1ô2' | | [_dd.propagation_error: 'decoding_error'] ", + "dm tag upper case | '_dd.p.dm=934086A686-4' | | [_dd.propagation_error: 'decoding_error'] ", + "dm tag too short | '_dd.p.dm=934086a66-4' | | [_dd.propagation_error: 'decoding_error'] ", + "dm tag too long | '_dd.p.dm=934086a6653-4' | | [_dd.propagation_error: 'decoding_error'] ", + "dm tag missing separator | '_dd.p.dm=934086a66534' | | [_dd.propagation_error: 'decoding_error'] ", + "dm tag missing mechanism | '_dd.p.dm=934086a665-' | | [_dd.propagation_error: 'decoding_error'] ", + "dm tag invalid mechanism char | '_dd.p.dm=934086a665-a' | | [_dd.propagation_error: 'decoding_error'] ", + "dm tag mechanism with letter | '_dd.p.dm=934086a665-12b' | | [_dd.propagation_error: 'decoding_error'] ", + "tid empty | '_dd.p.tid=' | | [_dd.propagation_error: 'decoding_error'] ", + "tid length 1 | '_dd.p.tid=1' | | [_dd.propagation_error: 'malformed_tid 1'] ", + "tid length 15 | '_dd.p.tid=111111111111111' | | [_dd.propagation_error: 'malformed_tid 111111111111111'] ", + "tid length 17 | '_dd.p.tid=11111111111111111' | | [_dd.propagation_error: 'malformed_tid 11111111111111111'] ", + "tid invalid uppercase | '_dd.p.tid=123456789ABCDEF0' | | [_dd.propagation_error: 'malformed_tid 123456789ABCDEF0'] ", + "tid invalid non-hex | '_dd.p.tid=123456789abcdefg' | | [_dd.propagation_error: 'malformed_tid 123456789abcdefg'] ", + "tid invalid negative | '_dd.p.tid=-123456789abcdef' | | [_dd.propagation_error: 'malformed_tid -123456789abcdef'] ", + "llmobs ml_app and sid | '_dd.p.llmobs_ml_app=my-ml-app,_dd.p.llmobs_sid=sess-1' | '_dd.p.llmobs_ml_app=my-ml-app,_dd.p.llmobs_sid=sess-1' | [_dd.p.llmobs_ml_app: 'my-ml-app', _dd.p.llmobs_sid: 'sess-1'] ", + "llmobs parent ids | '_dd.p.llmobs_pagent_span_id=9876543210,_dd.p.llmobs_parent_id=1122334455' | '_dd.p.llmobs_pagent_span_id=9876543210,_dd.p.llmobs_parent_id=1122334455' | [_dd.p.llmobs_pagent_span_id: '9876543210', _dd.p.llmobs_parent_id: '1122334455']", + "ts valid 02 | '_dd.p.ts=02' | '_dd.p.ts=02' | [_dd.p.ts: '02'] ", + "ts zero dropped | '_dd.p.ts=00' | | [:] ", + "ts invalid foo | '_dd.p.ts=foo' | | [_dd.propagation_error: 'decoding_error'] " }) void createPropagationTagsFromHeaderValue( String headerValue, String expectedHeaderValue, Map tags) { @@ -80,10 +84,11 @@ void createPropagationTagsFromHeaderValue( } @TableTest({ - "scenario | headerValue | expectedHeaderValue | tags ", - "single dm tag | '_dd.p.dm=934086a686-4' | 'dd=t.dm:934086a686-4' | [_dd.p.dm: '934086a686-4'] ", - "dm and f tag | '_dd.p.dm=934086a686-4,_dd.p.f=w00t==' | 'dd=t.dm:934086a686-4;t.f:w00t~~' | [_dd.p.dm: '934086a686-4', _dd.p.f: 'w00t==']", - "dm and appsec tag | '_dd.p.dm=934086a686-4,_dd.p.appsec=1' | 'dd=t.dm:934086a686-4;t.appsec:1' | [_dd.p.dm: '934086a686-4', _dd.p.appsec: '1']" + "scenario | headerValue | expectedHeaderValue | tags ", + "single dm tag | '_dd.p.dm=934086a686-4' | 'dd=t.dm:934086a686-4' | [_dd.p.dm: '934086a686-4'] ", + "dm and f tag | '_dd.p.dm=934086a686-4,_dd.p.f=w00t==' | 'dd=t.dm:934086a686-4;t.f:w00t~~' | [_dd.p.dm: '934086a686-4', _dd.p.f: 'w00t=='] ", + "dm and appsec tag | '_dd.p.dm=934086a686-4,_dd.p.appsec=1' | 'dd=t.dm:934086a686-4;t.appsec:1' | [_dd.p.dm: '934086a686-4', _dd.p.appsec: '1'] ", + "llmobs tags | '_dd.p.llmobs_ml_app=my-ml-app,_dd.p.llmobs_sid=sess-1' | 'dd=t.llmobs_ml_app:my-ml-app;t.llmobs_sid:sess-1' | [_dd.p.llmobs_ml_app: 'my-ml-app', _dd.p.llmobs_sid: 'sess-1']" }) void datadogPropagationTagsShouldTranslateToW3cTags( String headerValue, String expectedHeaderValue, Map tags) { @@ -155,6 +160,64 @@ void updatePropagationTagsTraceSourcePropagation( assertEquals(tags, propagationTags.createTagMap()); } + @TableTest({ + "scenario | mlApp | sessionId | pagentSpanId | pagentName | parentId | expectedHeaderValue | tags ", + "all five propagate | 'my-ml-app' | 'sess-1' | '9876543210' | 'planner' | '1122334455' | '_dd.p.llmobs_ml_app=my-ml-app,_dd.p.llmobs_sid=sess-1,_dd.p.llmobs_pagent_span_id=9876543210,_dd.p.llmobs_pagent_name=planner,_dd.p.llmobs_parent_id=1122334455' | [_dd.p.llmobs_ml_app: 'my-ml-app', _dd.p.llmobs_sid: 'sess-1', _dd.p.llmobs_pagent_span_id: '9876543210', _dd.p.llmobs_pagent_name: 'planner', _dd.p.llmobs_parent_id: '1122334455']", + "comma in ml_app dropped | 'planner,west' | 'sess-1' | | | | '_dd.p.llmobs_sid=sess-1' | [_dd.p.llmobs_sid: 'sess-1'] ", + "comma in agent name dropped | 'my-ml-app' | | | 'east,west' | | '_dd.p.llmobs_ml_app=my-ml-app' | [_dd.p.llmobs_ml_app: 'my-ml-app'] ", + "non-ascii ml_app dropped | 'プランナー' | 'sess-1' | | | | '_dd.p.llmobs_sid=sess-1' | [_dd.p.llmobs_sid: 'sess-1'] ", + "w3c-lossy chars dropped | 'app=v1' | 'sess 1' | | 'beta;2~x' | | '_dd.p.llmobs_ml_app=app=v1,_dd.p.llmobs_sid=sess 1' | [_dd.p.llmobs_ml_app: 'app=v1', _dd.p.llmobs_sid: 'sess 1'] ", + "json-unsafe chars dropped | 'a\\\"b' | 'sess-1' | | | | '_dd.p.llmobs_sid=sess-1' | [_dd.p.llmobs_sid: 'sess-1'] ", + "empty and null ignored | '' | | '9876543210' | '' | | '_dd.p.llmobs_pagent_span_id=9876543210' | [_dd.p.llmobs_pagent_span_id: '9876543210'] " + }) + void updatePropagationTagsLLMObsContext( + String mlApp, + String sessionId, + String pagentSpanId, + String pagentName, + String parentId, + String expectedHeaderValue, + Map tags) { + PropagationTags propagationTags = factory().fromHeaderValue(DATADOG, ""); + + propagationTags.updateLLMObsContext(mlApp, sessionId, pagentSpanId, pagentName, parentId); + + assertEquals(expectedHeaderValue, propagationTags.headerValue(DATADOG)); + assertEquals(tags, propagationTags.createTagMap()); + } + + /** + * These values are unrepresentable for the same reason as the table's comma and non-ASCII cases, + * but a literal control character can't be written into a TableTest row. + */ + @ParameterizedTest + @ValueSource(strings = {"planner\nwest", "planner\twest", "planner\u007fwest"}) + void updatePropagationTagsLLMObsContextRejectsControlCharacters(String mlApp) { + PropagationTags propagationTags = factory().fromHeaderValue(DATADOG, ""); + + propagationTags.updateLLMObsContext(mlApp, "sess-1", null, null, null); + + assertEquals("_dd.p.llmobs_sid=sess-1", propagationTags.headerValue(DATADOG)); + } + + @Test + void llmObsValueRejectionPreservesTraceIdHighOrderBits() { + // An unchecked ',' fails the whole tagset with decoding_error at the next hop, taking + // _dd.p.tid with it and leaving the two services disagreeing about the trace id. + PropagationTags propagationTags = + factory().fromHeaderValue(DATADOG, "_dd.p.tid=1234567890abcdef"); + + propagationTags.updateLLMObsContext("planner,west", "sess-1", null, null, null); + PropagationTags reparsed = + factory().fromHeaderValue(DATADOG, propagationTags.headerValue(DATADOG)); + + assertEquals(0x1234567890abcdefL, reparsed.getTraceIdHighOrderBits()); + Map expected = new HashMap<>(); + expected.put("_dd.p.tid", "1234567890abcdef"); + expected.put("_dd.p.llmobs_sid", "sess-1"); + assertEquals(expected, reparsed.createTagMap()); + } + @Test void extractionLimitExceeded() { String tags = "_dd.p.anytag=value"; diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3CPropagationTagsTest.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3CPropagationTagsTest.java index 491b8fb32ab..1a719b17007 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3CPropagationTagsTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/W3CPropagationTagsTest.java @@ -22,6 +22,7 @@ import java.util.Map; import java.util.stream.IntStream; import java.util.stream.Stream; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.converter.ConvertWith; import org.junit.jupiter.params.provider.Arguments; @@ -396,6 +397,20 @@ void propagationTagsShouldBeUpdatedByProductTraceSourcePropagation( assertEquals(tags, propagationTags.createTagMap()); } + @Test + void llmObsGettersDecodeTheTracestateSubstitutions() { + PropagationTags propagationTags = + factory() + .fromHeaderValue( + W3C, "dd=t.llmobs_ml_app:app~v1;t.llmobs_sid:sess~1;t.llmobs_pagent_name:planner"); + + // '=' travels as '~' in tracestate; the getters have to undo that, the way createTagMap does. + assertEquals("app=v1", propagationTags.getLLMObsMlApp().toString()); + assertEquals("sess=1", propagationTags.getLLMObsSessionId().toString()); + assertEquals("planner", propagationTags.getLLMObsParentAgentName().toString()); + assertNull(propagationTags.getLLMObsParentId()); + } + private static String buildHeader(int memberCount) { StringBuilder sb = new StringBuilder(); for (int i = 1; i <= memberCount; i++) { 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 1e768846b10..72cfa3c7bec 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 @@ -19,6 +19,7 @@ private LLMObsContext() { } private static final ContextKey CONTEXT_KEY = ContextKey.named("llmobs_span"); + private static final ContextKey ML_APP_KEY = ContextKey.named("llmobs_ml_app"); private static final ContextKey SESSION_ID_KEY = ContextKey.named("llmobs_session_id"); private static final ContextKey AGENT_VERSION_KEY = ContextKey.named("llmobs_agent_version"); @@ -67,14 +68,17 @@ public static ContextScope attach(AgentSpanContext ctx, String sessionId, String } /** - * 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. + * Attach an LLMObs span context, propagating an ml_app, 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. + * exactly once: four independent mechanisms (application, session, sampling, attribution) share + * one context, so they cannot be attached by separate calls without nesting redundant scopes. + * + *

ml_app is stored here so that distributed propagation can read the innermost active LLMObs + * span's ml_app when injecting, without needing a reference to the span itself. * *

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 @@ -102,6 +106,7 @@ public static ContextScope attach(AgentSpanContext ctx, String sessionId, String */ public static ContextScope attach( AgentSpanContext ctx, + String mlApp, String sessionId, String agentVersion, String sampleRate, @@ -111,6 +116,7 @@ public static ContextScope attach( String decision = emptyToNull(samplingDecision); return Context.current() .with(CONTEXT_KEY, ctx) + .with(ML_APP_KEY, emptyToNull(mlApp)) .with(SESSION_ID_KEY, emptyToNull(sessionId)) .with(AGENT_VERSION_KEY, emptyToNull(agentVersion)) .with(SAMPLING_DECISION_KEY, decision) @@ -124,6 +130,11 @@ public static AgentSpanContext current() { return Context.current().get(CONTEXT_KEY); } + /** Return the ml_app of the innermost active LLMObs span, or null if none is active. */ + public static String currentMlApp() { + return Context.current().get(ML_APP_KEY); + } + /** * Return the session_id propagated from an enclosing LLMObs span, or null if no parent set one. */ diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java index 46a01b3f70f..0cf43469604 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java @@ -23,6 +23,11 @@ public final class AgentPropagation { // TODO DSM propagator should run after the other propagators as it stores the pathway context // TODO into the span context for now. Remove priority after the migration is complete. public static final Concern DSM_CONCERN = withPriority("data-stream-monitoring", 110); + // LLM Observability contributes no headers of its own: it stages the _dd.p.llmobs_* propagation + // tags onto the span context, which the tracing propagator then serializes into x-datadog-tags / + // tracestate. Composite injection runs in reverse priority order, so this must sort after + // TRACING_CONCERN to actually inject before it. + public static final Concern LLMOBS_CONCERN = withPriority("llm-observability", 115); private AgentPropagation() {} diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpanContext.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpanContext.java index 1dba9438168..907bc3be20e 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpanContext.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpanContext.java @@ -55,6 +55,72 @@ default void mergePathwayContext(PathwayContext pathwayContext) {} default void setIntegrationName(CharSequence componentName) {} + /** + * Gets the LLM Observability {@code ml_app} that arrived on the inbound headers, or {@code null} + * if none did or this context implementation doesn't have propagation-tags access. + * + *

These five getters describe the caller, so they report only what was extracted — never what + * a local injection staged onto the same tags for an outbound call. + */ + default CharSequence getLLMObsMlApp() { + return null; + } + + /** + * Gets the LLM Observability {@code session_id} that arrived on the inbound headers, or {@code + * null} if none did or this context implementation doesn't have propagation-tags access. + */ + default CharSequence getLLMObsSessionId() { + return null; + } + + /** + * Gets the span id of the parent LLM Observability agent span that arrived on the inbound + * headers, or {@code null} if none did or this context implementation doesn't have + * propagation-tags access. + */ + default CharSequence getLLMObsParentAgentSpanId() { + return null; + } + + /** + * Gets the name of the parent LLM Observability agent span that arrived on the inbound headers, + * or {@code null} if none did or this context implementation doesn't have propagation-tags + * access. + */ + default CharSequence getLLMObsParentAgentName() { + return null; + } + + /** + * Gets the span id of the parent LLM Observability span that arrived on the inbound headers, or + * {@code null} if none did or this context implementation doesn't have propagation-tags access. + */ + default CharSequence getLLMObsParentId() { + return null; + } + + /** + * Sets the whole LLM Observability tag set to propagate with this trace, replacing any set + * previously staged. Taken together rather than one tag at a time so the update is atomic: a + * concurrent reader never serializes a header mixing values from two different contexts. No-op by + * default. + */ + default void updateLLMObsContext( + CharSequence mlApp, + CharSequence sessionId, + CharSequence parentAgentSpanId, + CharSequence parentAgentName, + CharSequence parentId) {} + + /** + * Discards anything locally staged by {@link #updateLLMObsContext}, restoring the LLM + * Observability tag set that arrived on the inbound headers. Distinct from staging an empty set: + * a service that forwards a request without opening an LLMObs span of its own must keep passing + * the caller's context along. No-op by default. + */ + default void resetLLMObsContext() {} + /** * Gets whether the span context used is part of the local trace or from another service * 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 20f91b51a8f..240cc6aa327 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 @@ -177,7 +177,7 @@ void attachWithSamplingDecisionStoresDecisionAndRate() { AgentSpanContext ctx = mock(AgentSpanContext.class); try (ContextScope scope = LLMObsContext.attach( - ctx, null, null, "0.25", LLMObsContext.SAMPLING_DECISION_DROPPED, null, null)) { + ctx, null, null, null, "0.25", LLMObsContext.SAMPLING_DECISION_DROPPED, null, null)) { assertEquals( LLMObsContext.SAMPLING_DECISION_DROPPED, LLMObsContext.currentSamplingDecision()); assertEquals("0.25", LLMObsContext.currentSampleRate()); @@ -190,7 +190,8 @@ void attachWithSamplingDecisionStoresDecisionAndRate() { 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)) { + try (ContextScope scope = + LLMObsContext.attach(ctx, null, null, null, "0.25", null, null, null)) { assertNull(LLMObsContext.currentSamplingDecision()); assertNull(LLMObsContext.currentSampleRate()); } @@ -202,7 +203,7 @@ void childScopeInheritsParentSamplingDecision() { AgentSpanContext child = mock(AgentSpanContext.class); try (ContextScope parentScope = LLMObsContext.attach( - parent, null, null, "1", LLMObsContext.SAMPLING_DECISION_SAMPLED, null, null)) { + parent, null, null, null, "1", LLMObsContext.SAMPLING_DECISION_SAMPLED, null, null)) { try (ContextScope childScope = LLMObsContext.attach(child)) { assertEquals(child, LLMObsContext.current()); assertEquals( @@ -230,6 +231,7 @@ void fullAttachStoresAllFields() { try (ContextScope scope = LLMObsContext.attach( ctx, + "my-app", "session-1", "v2", "0.5", @@ -237,6 +239,7 @@ void fullAttachStoresAllFields() { "span-99", "my-agent")) { assertEquals(ctx, LLMObsContext.current()); + assertEquals("my-app", LLMObsContext.currentMlApp()); assertEquals("session-1", LLMObsContext.currentSessionId()); assertEquals("v2", LLMObsContext.currentAgentVersion()); assertEquals("0.5", LLMObsContext.currentSampleRate()); @@ -246,6 +249,7 @@ void fullAttachStoresAllFields() { assertEquals("my-agent", LLMObsContext.currentParentAgentName()); } assertNull(LLMObsContext.current()); + assertNull(LLMObsContext.currentMlApp()); assertNull(LLMObsContext.currentSessionId()); assertNull(LLMObsContext.currentAgentVersion()); assertNull(LLMObsContext.currentSampleRate()); @@ -257,7 +261,7 @@ void fullAttachStoresAllFields() { @Test void fullAttachWithNullSessionIdIgnoresSessionId() { AgentSpanContext ctx = mock(AgentSpanContext.class); - try (ContextScope scope = LLMObsContext.attach(ctx, null, null, null, null, null, null)) { + try (ContextScope scope = LLMObsContext.attach(ctx, null, null, null, null, null, null, null)) { assertNull(LLMObsContext.currentSessionId()); assertNull(LLMObsContext.currentAgentVersion()); assertNull(LLMObsContext.currentParentAgentSpanId()); @@ -268,7 +272,8 @@ void fullAttachWithNullSessionIdIgnoresSessionId() { @Test void fullAttachWithEmptySessionIdIgnoresSessionId() { 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)) { + assertNull(LLMObsContext.currentMlApp()); assertNull(LLMObsContext.currentSessionId()); assertNull(LLMObsContext.currentAgentVersion()); } @@ -281,12 +286,12 @@ void fullAttachNullPagentClearsStaleValuesFromOuterScope() { AgentSpanContext outer = mock(AgentSpanContext.class); AgentSpanContext inner = mock(AgentSpanContext.class); try (ContextScope outerScope = - LLMObsContext.attach(outer, "s", "v1", null, null, "agent-span-id", "outer-agent")) { + LLMObsContext.attach(outer, null, "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, null, null)) { + LLMObsContext.attach(inner, null, null, null, null, null, null, null)) { assertNull(LLMObsContext.currentParentAgentSpanId()); assertNull(LLMObsContext.currentParentAgentName()); } @@ -302,9 +307,10 @@ void fullAttachInnerAgentOverridesOuterAgentForDescendants() { AgentSpanContext outer = mock(AgentSpanContext.class); AgentSpanContext inner = mock(AgentSpanContext.class); try (ContextScope outerScope = - LLMObsContext.attach(outer, null, null, null, null, "outer-span-id", "outer-agent")) { + LLMObsContext.attach(outer, null, null, null, null, null, "outer-span-id", "outer-agent")) { try (ContextScope innerScope = - LLMObsContext.attach(inner, null, null, null, null, "inner-span-id", "inner-agent")) { + LLMObsContext.attach( + inner, null, null, null, null, null, "inner-span-id", "inner-agent")) { assertEquals("inner-span-id", LLMObsContext.currentParentAgentSpanId()); assertEquals("inner-agent", LLMObsContext.currentParentAgentName()); } @@ -320,9 +326,9 @@ void fullAttachNullPagentNameClearsNameButNotSpanId() { AgentSpanContext outer = mock(AgentSpanContext.class); AgentSpanContext inner = mock(AgentSpanContext.class); try (ContextScope outerScope = - LLMObsContext.attach(outer, null, null, null, null, "outer-span-id", "outer-agent")) { + LLMObsContext.attach(outer, null, null, null, null, null, "outer-span-id", "outer-agent")) { try (ContextScope innerScope = - LLMObsContext.attach(inner, null, null, null, null, "inner-span-id", null)) { + LLMObsContext.attach(inner, null, null, null, null, null, "inner-span-id", null)) { assertEquals("inner-span-id", LLMObsContext.currentParentAgentSpanId()); assertNull(LLMObsContext.currentParentAgentName()); } @@ -330,13 +336,14 @@ void fullAttachNullPagentNameClearsNameButNotSpanId() { } @Test - void attachPropagatesAllFourMechanismsTogether() { + void attachPropagatesAllMechanismsTogether() { AgentSpanContext parent = mock(AgentSpanContext.class); AgentSpanContext child = mock(AgentSpanContext.class); - // All four propagation mechanisms coexist on one context and are inherited together. + // Every propagation mechanism coexists on one context and is inherited together. try (ContextScope parentScope = LLMObsContext.attach( parent, + "app-abc", "session-abc", "v7", "0.5", @@ -344,6 +351,7 @@ void attachPropagatesAllFourMechanismsTogether() { "agent-span-7", "agent-seven")) { try (ContextScope childScope = LLMObsContext.attach(child)) { + assertEquals("app-abc", LLMObsContext.currentMlApp()); assertEquals("session-abc", LLMObsContext.currentSessionId()); assertEquals("v7", LLMObsContext.currentAgentVersion()); assertEquals(