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..446a06a0c29 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParserTest.java @@ -0,0 +1,119 @@ +package datadog.trace.bootstrap.instrumentation.messaging; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.ByteBuffer; +import java.util.Base64; +import java.util.LinkedHashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; + +/** + * Covers the {@code _datadog} message attribute parser shared by the AWS messaging instrumentations + * (SQS, SNS, EventBridge, Step Functions). + */ +class DatadogAttributeParserTest { + + /** What an injected {@code _datadog} attribute looks like on the wire. */ + private static final String FULL_CONTEXT = + "{\"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\"," + + "\"traceparent\":\"00-6aa01c5400000000499602d2-000000024cb016ea-01\"}"; + + private static Map parse(String json) { + Map collected = new LinkedHashMap<>(); + DatadogAttributeParser.forEachProperty( + (key, value) -> { + collected.put(key, value); + return true; + }, + json); + return collected; + } + + @Test + void extractsTraceContextAndPropagationTags() { + Map collected = parse(FULL_CONTEXT); + + assertEquals("1234567890", collected.get("x-datadog-trace-id")); + assertEquals("9876543210", collected.get("x-datadog-parent-id")); + assertEquals("1", collected.get("x-datadog-sampling-priority")); + // Without x-datadog-tags the whole _dd.p.* set is dropped at the messaging boundary: the + // 64-bit trace id still joins, but _dd.p.tid is lost so the two services disagree about the + // full 128-bit id, and _dd.p.dm is lost with it. + assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags")); + } + + @Test + void extractsPropagationTagsFromByteBufferCarrier() { + Map collected = new LinkedHashMap<>(); + DatadogAttributeParser.forEachProperty( + (key, value) -> { + collected.put(key, value); + return true; + }, + ByteBuffer.wrap(FULL_CONTEXT.getBytes(UTF_8))); + + assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags")); + } + + @Test + void extractsPropagationTagsFromBase64ByteBufferCarrier() { + Map collected = new LinkedHashMap<>(); + DatadogAttributeParser.forEachProperty( + (key, value) -> { + collected.put(key, value); + return true; + }, + ByteBuffer.wrap(Base64.getEncoder().encode(FULL_CONTEXT.getBytes(UTF_8)))); + + assertEquals("_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000", collected.get("x-datadog-tags")); + } + + @Test + void carriesLlmObsPropagationTags() { + Map collected = + parse( + "{\"x-datadog-trace-id\":\"1234567890\"," + + "\"x-datadog-parent-id\":\"9876543210\"," + + "\"x-datadog-tags\":\"_dd.p.llmobs_ml_app=my-app,_dd.p.llmobs_sid=sess-1," + + "_dd.p.llmobs_parent_id=42\"}"); + + String tags = collected.get("x-datadog-tags"); + assertTrue(tags.contains("_dd.p.llmobs_ml_app=my-app"), tags); + assertTrue(tags.contains("_dd.p.llmobs_sid=sess-1"), tags); + assertTrue(tags.contains("_dd.p.llmobs_parent_id=42"), tags); + } + + @Test + void extractsNothingWithoutATraceId() { + // Propagation tags on their own describe no trace, so they are not surfaced. + Map collected = + parse("{\"x-datadog-tags\":\"_dd.p.dm=-1\",\"x-datadog-parent-id\":\"9876543210\"}"); + + assertTrue(collected.isEmpty(), () -> "expected nothing extracted, got " + collected); + } + + @Test + void toleratesAMissingTagsProperty() { + Map collected = + parse( + "{\"x-datadog-trace-id\":\"1234567890\",\"x-datadog-parent-id\":\"9876543210\"," + + "\"x-datadog-sampling-priority\":\"1\"}"); + + assertEquals("1234567890", collected.get("x-datadog-trace-id")); + assertNull(collected.get("x-datadog-tags")); + } + + @Test + void toleratesMalformedJson() { + assertTrue(parse("not json at all").isEmpty()); + assertTrue(parse("{\"x-datadog-trace-id\":").isEmpty()); + assertTrue(parse(null).isEmpty()); + } +} 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..056358b8f90 --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagator.java @@ -0,0 +1,67 @@ +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 — HTTP, gRPC, SQS, Kafka, ... — 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. This mirrors + * dd-trace-py, where LLMObs subscribes to the generic {@code http.span_inject} hook that {@code + * HTTPPropagator.inject} fires on every outbound request, rather than owning a separate wire + * format. + * + *

Values are resolved from the ambient {@link LLMObsContext} at injection time rather than being + * written once when a span starts. That way the innermost active LLMObs span always wins, and + * leaving an LLMObs scope stops contributing its tags 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() != spanContext.getTraceId()) { + return; + } + + // The innermost active LLMObs span becomes the downstream span's LLMObs parent, so the + // continued trace is a single tree rather than a second root per service. Mirrors + // dd-trace-py's _dd.p.llmobs_parent_id. + 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..a819598b14d 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; @@ -51,6 +53,10 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) { LLMObsInternal.setEvalProcessor(new LLMObsCustomEvalProcessor(mlApp, sco, config)); LLMObsInternal.setFeedbackProcessor(new LLMObsCustomFeedbackProcessor(mlApp, sco, config)); + + // Carry LLMObs context across every boundary automatic instrumentation already covers, by + // staging the _dd.p.llmobs_* tags on each injected span context. See LLMObsContextPropagator. + Propagators.register(AgentPropagation.LLMOBS_CONCERN, new LLMObsContextPropagator()); } private static class LLMObsCustomFeedbackProcessor implements LLMObs.LLMObsFeedbackProcessor { 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..ac693f9e79f 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 @@ -159,6 +159,7 @@ public DDLLMObsSpan( String samplingDecision = null; String resolvedParentAgentSpanId = null; String resolvedParentAgentName = null; + boolean inheritedInProcess = false; if (null != parent) { if (parent.getTraceId() != span.getTraceId()) { LOGGER.error( @@ -168,6 +169,7 @@ public DDLLMObsSpan( span.getTraceId(), span.getSpanId()); } else { + inheritedInProcess = true; parentSpanID = String.valueOf(parent.getSpanId()); // 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 @@ -197,6 +199,30 @@ public DDLLMObsSpan( } } + if (!inheritedInProcess) { + // No usable in-process LLMObs parent, but this span may still be continuing a trace that + // arrived from another service — an SQS worker handling a message, an inbound HTTP request. + // The upstream values are on the span context's propagation tags, parsed back out of + // x-datadog-tags / tracestate by the tracing propagator. + // + // This also covers the trace-mismatch branch above: a stale context leaked from an unrelated + // trace must not suppress attribution that legitimately arrived over the wire. + // Unlike dd-trace-py, a missing parent_id doesn't veto the rest: session and attribution + // are inherited independently, so a partially populated upstream still contributes what it + // did send. + String propagatedParentId = asString(span.spanContext().getLLMObsParentId()); + if (propagatedParentId != null) { + parentSpanID = propagatedParentId; + } + if (sessionId == null || sessionId.isEmpty()) { + sessionId = asString(span.spanContext().getLLMObsSessionId()); + } + resolvedParentAgentSpanId = asString(span.spanContext().getLLMObsParentAgentSpanId()); + if (resolvedParentAgentSpanId != null) { + resolvedParentAgentName = asString(span.spanContext().getLLMObsParentAgentName()); + } + } + // 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. @@ -236,6 +262,7 @@ public DDLLMObsSpan( scope = LLMObsContext.attach( span.spanContext(), + mlApp, sessionId, resolvedAgentVersion, sampleRate, @@ -717,4 +744,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..b65270761e3 --- /dev/null +++ b/dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/LLMObsContextPropagatorTest.java @@ -0,0 +1,208 @@ +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: no LLMObs propagation API is called + * anywhere in these tests. Injecting the active span the way auto-instrumentation does — an HTTP + * client, or the SQS interceptor writing message attributes — must carry the LLMObs context. + * + *

The carrier is a plain {@code Map}, which is the shape both an HTTP header map + * and the SQS {@code _datadog} message attribute reduce to at the propagator boundary. + */ +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; + } + + @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); + } + + @Test + void stopsContributingTagsOnceTheLlmObsScopeIsClosed() { + Map carrier; + try (AgentScope apmScope = startRootApmScope()) { + newSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "work", "my-ml-app", "sess-1").finish(); + // The LLMObs span has finished; a later outbound call on the same APM trace must not be + // tagged with a session that is no longer active. + carrier = autoInject(apmScope.span()); + } + + String tags = carrier.get("x-datadog-tags"); + assertTrue( + tags == null || !tags.contains(SESSION_ID_TAG), + () -> "session_id leaked after scope close: " + tags); + assertTrue( + tags == null || !tags.contains(PARENT_ID_TAG), + () -> "parent_id leaked after scope close: " + 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 the session and agent attribution without any application-level plumbing. + */ + @Test + void workerInheritsSessionAndAgentAttributionAcrossTheBoundary() { + Map messageAttributes; + long producerTraceId; + String producerAgentSpanId; + + try (AgentScope apmScope = startRootApmScope()) { + DDLLMObsSpan producer = + newSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "dispatcher", "my-ml-app", "sess-42"); + producerTraceId = producer.getTraceId().toLong(); + producerAgentSpanId = String.valueOf(producer.getSpanId()); + try { + messageAttributes = autoInject(AgentTracer.activeSpan()); + } finally { + producer.finish(); + } + } + + // Worker side: a fresh context, as a message handler would have. + Context extracted = + Propagators.defaultPropagator() + .extract( + Context.root(), messageAttributes, (carrier, visitor) -> carrier.forEach(visitor)); + AgentSpan consumeSpan = AgentSpan.fromContext(extracted); + assertNotNull(consumeSpan, "expected trace context to be extracted"); + + try (AgentScope consumeScope = AgentTracer.get().activateSpan(consumeSpan)) { + DDLLMObsSpan workerTool = newSpan(Tags.LLMOBS_TOOL_SPAN_KIND, "handler", "my-ml-app", 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("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(); + } + } + } + + @Test + void workerWithoutUpstreamLlmObsContextInheritsNothing() { + Map messageAttributes; + try (AgentScope apmScope = startRootApmScope()) { + messageAttributes = autoInject(apmScope.span()); + } + + Context extracted = + Propagators.defaultPropagator() + .extract( + Context.root(), messageAttributes, (carrier, visitor) -> carrier.forEach(visitor)); + AgentSpan consumeSpan = AgentSpan.fromContext(extracted); + assertNotNull(consumeSpan, "expected trace context to be extracted"); + + 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-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java index adf4cd66156..44d9862fba4 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,42 @@ public PropagationTags getPropagationTags() { return getRootSpanContextOrThis().propagationTags; } + @Override + public CharSequence getLLMObsMlApp() { + return getPropagationTags().getLLMObsMlApp(); + } + + @Override + public void updateLLMObsContext( + CharSequence mlApp, + CharSequence sessionId, + CharSequence parentAgentSpanId, + CharSequence parentAgentName, + CharSequence parentId) { + getPropagationTags() + .updateLLMObsContext(mlApp, sessionId, parentAgentSpanId, parentAgentName, parentId); + } + + @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(); + } + /** 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..39d2ceff1e0 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,48 @@ public interface Factory { */ public abstract void updateOrgPropagationMarker(CharSequence opm); + /** + * Returns the LLM Observability {@code ml_app} currently propagated with this trace, encoded as + * {@code _dd.p.llmobs_ml_app}. Returns {@code null} if none is set. + */ + public abstract CharSequence getLLMObsMlApp(); + + /** + * 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. + */ + public abstract void updateLLMObsContext( + CharSequence mlApp, + CharSequence sessionId, + CharSequence parentAgentSpanId, + CharSequence parentAgentName, + CharSequence parentId); + + /** + * Returns the LLM Observability {@code session_id} currently propagated with this trace, encoded + * as {@code _dd.p.llmobs_sid}. Returns {@code null} if none is set. + */ + public abstract CharSequence getLLMObsSessionId(); + + /** + * Returns the span id of the parent LLM Observability agent span currently propagated with this + * trace, encoded as {@code _dd.p.llmobs_pagent_span_id}. Returns {@code null} if none is set. + */ + public abstract CharSequence getLLMObsParentAgentSpanId(); + + /** + * Returns the name of the parent LLM Observability agent span currently propagated with this + * trace, encoded as {@code _dd.p.llmobs_pagent_name}. Returns {@code null} if none is set. + */ + public abstract CharSequence getLLMObsParentAgentName(); + + /** + * Returns the span id of the parent LLM Observability span currently propagated with this trace, + * encoded as {@code _dd.p.llmobs_parent_id}. Returns {@code null} if none is set. + */ + public abstract CharSequence getLLMObsParentId(); + 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..f3dabe3dae2 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/LLMObsTagValues.java @@ -0,0 +1,67 @@ +package datadog.trace.core.propagation.ptags; + +import java.util.Objects; + +/** + * Bundles the five LLM Observability propagation tag values ({@code ml_app}, {@code session_id}, + * parent agent span id, parent agent name, parent span id) extracted from an incoming header, so + * they can be threaded through {@link PTagsFactory.PTags} construction 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; + } + + /** + * Whether {@code other} carries the same five values. Used to skip cache invalidation when an + * injection re-stages tags a span already has; not {@code equals} because these are never used as + * map keys and identity equality is the useful default elsewhere in this package. + */ + boolean sameAs(LLMObsTagValues other) { + return this == other + || (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)); + } +} 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..ec865e79c6f 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,23 @@ static String headerValue(PTagsCodec codec, PTags ptags, CharSequence lastParent codec.appendTag( sb, ORG_PROPAGATION_MARKER_TAG, ptags.getOrgPropagationMarkerTagValue(), size); } + // One snapshot, so a concurrent injection can't have us encode a mix of old and new values. + 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 +159,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..2e10b18a532 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,14 @@ static class PTags extends PropagationTags { private volatile TagValue orgPropagationMarkerTagValue; + /** + * The LLM Observability propagation tags, held as one immutable bundle rather than five fields + * so that an update is a single reference swap. Readers therefore always observe a tag set that + * actually existed, instead of a mix of values from before and after an injection. Never {@code + * null} — {@link LLMObsTagValues#EMPTY} means "none". + */ + private volatile LLMObsTagValues llmObsTags = LLMObsTagValues.EMPTY; + // 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 +173,8 @@ static class PTags extends PropagationTags { TagValue decisionMakerTagValue, TagValue traceIdTagValue, int traceSource, - TagValue orgPropagationMarkerTagValue) { + TagValue orgPropagationMarkerTagValue, + @Nonnull LLMObsTagValues llmObsTagValues) { this( factory, tagPairs, @@ -168,7 +184,8 @@ static class PTags extends PropagationTags { PrioritySampling.UNSET, null, null, - orgPropagationMarkerTagValue); + orgPropagationMarkerTagValue, + llmObsTagValues); } PTags( @@ -180,7 +197,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 +209,7 @@ static class PTags extends PropagationTags { this.origin = origin; this.lastParentId = lastParentId; this.orgPropagationMarkerTagValue = orgPropagationMarkerTagValue; + this.llmObsTags = llmObsTagValues; if (traceIdTagValue != null) { CharSequence traceIdHighOrderBitsHex = traceIdTagValue.forType(TagElement.Encoding.DATADOG); this.traceIdHighOrderBits = @@ -212,7 +231,8 @@ static PTags withError(PTagsFactory factory, String error) { PrioritySampling.UNSET, null, null, - null); + null, + LLMObsTagValues.EMPTY); pTags.error = error; return pTags; } @@ -248,8 +268,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 +276,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 +293,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 +318,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 +383,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 +392,70 @@ 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)); + // Re-injecting the same context onto the same span is the common case; don't invalidate. + if (!updated.sameAs(llmObsTags)) { + clearCachedHeaders(); + llmObsTags = updated; + } + } + + @Override + public CharSequence getLLMObsMlApp() { + return llmObsTags.mlApp; + } + + @Override + public CharSequence getLLMObsSessionId() { + return llmObsTags.sessionId; + } + + @Override + public CharSequence getLLMObsParentAgentSpanId() { + return llmObsTags.parentAgentSpanId; + } + + @Override + public CharSequence getLLMObsParentAgentName() { + return llmObsTags.parentAgentName; + } + + @Override + public CharSequence getLLMObsParentId() { + return llmObsTags.parentId; + } + + LLMObsTagValues getLLMObsTagValues() { + return llmObsTags; + } + + /** + * Wraps a non-empty value as a {@link TagValue}, or {@code null} if empty. No length capping is + * applied here — matching dd-trace-py, which writes these free-form values (ml_app, session_id, + * agent id/name) as-is and relies on the codecs' own overflow handling (dropping the whole + * {@code x-datadog-tags} header on the Datadog codec, or dropping individual overlong tags on + * the W3C codec) rather than a fixed per-field character limit. + */ + private static TagValue toTagValue(CharSequence value) { + if (value == null || value.length() == 0) { + return null; + } + return TagValue.from(value); + } + @Override public int getSamplingPriority() { return samplingPriority; @@ -473,6 +552,16 @@ 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; the single-encoding {@link + * #clearCachedHeader} calls that remain are deliberate. + */ + private void clearCachedHeaders() { + clearCachedHeader(DATADOG); + clearCachedHeader(W3C); + } + private void clearCachedHeader(HeaderType headerType) { if (headerType == DATADOG) { invalidateXDatadogTagsSize(); @@ -512,6 +601,21 @@ int getXDatadogTagsSize() { size = PTagsCodec.calcXDatadogTagsSize( size, ORG_PROPAGATION_MARKER_TAG, getOrgPropagationMarkerTagValue()); + // One snapshot: sizing a mix of old and new values would gate the header on a tag set that + // never existed, and this total is what decides whether x-datadog-tags is emitted at all. + 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 +657,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/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/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..ed3a5bb68a4 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"); @@ -108,9 +109,38 @@ public static ContextScope attach( String samplingDecision, String parentAgentSpanId, String parentAgentName) { + return attach( + ctx, + null, + sessionId, + agentVersion, + sampleRate, + samplingDecision, + parentAgentSpanId, + parentAgentName); + } + + /** + * Attach an LLMObs span context, propagating ml_app alongside everything {@link + * #attach(AgentSpanContext, String, String, String, String, String, String)} carries. + * + *

ml_app is held here — rather than only as a span tag — so that distributed propagation can + * read the innermost active LLMObs span's ml_app when injecting, without needing a reference to + * the span itself. + */ + public static ContextScope attach( + AgentSpanContext ctx, + String mlApp, + String sessionId, + String agentVersion, + String sampleRate, + String samplingDecision, + String parentAgentSpanId, + String parentAgentName) { 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 +154,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..b6141a6418a 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,60 @@ default void mergePathwayContext(PathwayContext pathwayContext) {} default void setIntegrationName(CharSequence componentName) {} + /** + * Gets the LLM Observability {@code ml_app} propagated with this trace, or {@code null} if none + * is set or this context implementation doesn't have propagation-tags access. + */ + default CharSequence getLLMObsMlApp() { + 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) {} + + /** + * Gets the LLM Observability {@code session_id} propagated with this trace, or {@code null} if + * none is set 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 propagated with this trace, or + * {@code null} if none is set 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 propagated with this trace, or {@code + * null} if none is set 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 propagated with this trace, or {@code + * null} if none is set or this context implementation doesn't have propagation-tags access. + */ + default CharSequence getLLMObsParentId() { + return null; + } + /** * Gets whether the span context used is part of the local trace or from another service *