Propagate LLM Observability context across service boundaries #12416
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
ncybul marked this conversation as resolved.
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<String, String> 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")); | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,73 @@ | ||
| 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. | ||
| * | ||
| * <p>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. | ||
| * | ||
| * <p>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 <C> void inject(Context context, C carrier, CarrierSetter<C> 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; | ||
|
ncybul marked this conversation as resolved.
|
||
| } | ||
|
|
||
| spanContext.updateLLMObsContext( | ||
|
ncybul marked this conversation as resolved.
|
||
| LLMObsContext.currentMlApp(), | ||
| LLMObsContext.currentSessionId(), | ||
| LLMObsContext.currentParentAgentSpanId(), | ||
| LLMObsContext.currentParentAgentName(), | ||
| String.valueOf(llmObsContext.getSpanId()), | ||
| LLMObsContext.currentSampleRate(), | ||
| LLMObsContext.currentSamplingDecision()); | ||
| } | ||
|
Comment on lines
+44
to
+64
Contributor
There was a problem hiding this comment. Does this work for concurrent traces? Claude is telling me that it won't. Issue: updateLLMObsContext writes the bundle to getRootSpanContextOrThis().propagationTags — one object shared by the whole local trace — and the tracing propagator serializes it one propagator later, so two concurrent outbound calls under one agent span (parallel tool calls: the normal agentic shape) interleave stage/stage/serialize and ship each other's parent_id/session_id, mis-parenting the downstream service under the wrong agent. volatile prevents a torn bundle but cannot make stage→serialize atomic — the AgentSpanContext.updateLLMObsContext javadoc's atomicity claim covers only the former. |
||
|
|
||
| @Override | ||
| public <C> Context extract(Context context, C carrier, CarrierVisitor<C> 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; | ||
|
ncybul marked this conversation as resolved.
|
||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.