Skip to content

Commit 529329c

Browse files
Yun-Kimclaude
andcommitted
Keep auto-instrumented OpenAI spans in their parent's LLMObs trace
An OpenAI span under a manual LLMObs span reported the APM trace id and the service-default ml_app. When the parent had adopted both from another service, the LLM span landed in a different LLMObs trace from its parent. It now inherits the LLMObs trace id and ml_app from the active context. It also rolled its own sampling decision when the parent had none. A parent with no decision continued a trace whose caller sent none, so the span now leaves both sampling tags unset, as DDLLMObsSpan does. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1 parent 6d00d1b commit 529329c

3 files changed

Lines changed: 124 additions & 9 deletions

File tree

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CommonTags.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ interface CommonTags {
2929

3030
String ENV = TAG_PREFIX + "env";
3131
String SERVICE = TAG_PREFIX + "service";
32+
String TRACE_ID = TAG_PREFIX + "trace_id";
3233
String PARENT_ID = TAG_PREFIX + "parent_id";
3334
String SESSION_ID = TAG_PREFIX + LLMObsTags.SESSION_ID;
3435
String AGENT_VERSION = TAG_PREFIX + LLMObsTags.AGENT_VERSION;

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -116,9 +116,9 @@ protected void doAfterStart(@Nonnull AgentSpan span) {
116116
span.setTag(CommonTags.INTEGRATION, INTEGRATION);
117117

118118
// Resolve the LLMObs parent context, gated on trace-id consistency: a stale context
119-
// from a different trace (e.g. async boundary leakage) must not contribute parent_id,
120-
// session_id, agent_version, agent attribution, or a sampling verdict to this span.
121-
// Matches DDLLMObsSpan's manual-span gate. One flag drives every inherited value, so a
119+
// from a different trace (e.g. async boundary leakage) must not contribute trace_id, ml_app,
120+
// parent_id, session_id, agent_version, agent attribution, or a sampling verdict to this
121+
// span. Matches DDLLMObsSpan's manual-span gate. One flag drives every inherited value, so a
122122
// new propagated tag cannot accidentally ship with a weaker gate of its own.
123123
AgentSpanContext parent = LLMObsContext.current();
124124
boolean inheritable = parent != null && parent.getTraceId().equals(span.getTraceId());
@@ -129,6 +129,17 @@ protected void doAfterStart(@Nonnull AgentSpan span) {
129129
if (inheritable) {
130130
parentSpanId = String.valueOf(parent.getSpanId());
131131

132+
// Stay in the parent's LLMObs trace and application. The parent may have adopted both
133+
// from another service, so the APM trace id and the service default can differ from them.
134+
String traceId = LLMObsContext.currentTraceId();
135+
if (traceId != null && !traceId.isEmpty()) {
136+
span.setTag(CommonTags.TRACE_ID, traceId);
137+
}
138+
String mlApp = LLMObsContext.currentMlApp();
139+
if (mlApp != null && !mlApp.isEmpty()) {
140+
span.setTag(CommonTags.ML_APP, mlApp);
141+
}
142+
132143
// Inherit session_id from the active LLMObs parent (e.g. a manual workflow span).
133144
// Matches dd-trace-py / dd-trace-js, where auto-instrumented LLM spans inherit
134145
// session_id from the workflow root via context propagation. Without this, the
@@ -161,18 +172,20 @@ protected void doAfterStart(@Nonnull AgentSpan span) {
161172
}
162173
span.setTag(CommonTags.PARENT_ID, parentSpanId);
163174

164-
// Compute the sampling decision if none was inherited (no LLMObs parent), which makes this
165-
// span the root of its own LLMObs trace. Unlike the tags above, this cannot be skipped when
166-
// there is nothing to inherit: an unstamped span is retained at any configured rate.
167-
if (samplingDecision == null || sampleRate == null) {
175+
// Roll a decision only when this span is the root of its own LLMObs trace. A parent with no
176+
// decision continued a trace whose caller sent none, so this span leaves both tags unset
177+
// too, as DDLLMObsSpan does; a local roll here could drop the span while its parent is kept.
178+
if (!inheritable) {
168179
sampleRate = sampler.formattedRate();
169180
samplingDecision =
170181
sampler.sample(span.getTraceId().toLong())
171182
? LLMObsContext.SAMPLING_DECISION_SAMPLED
172183
: LLMObsContext.SAMPLING_DECISION_DROPPED;
173184
}
174-
span.setTag(CommonTags.SAMPLING_DECISION, samplingDecision);
175-
span.setTag(CommonTags.SAMPLE_RATE, sampleRate);
185+
if (samplingDecision != null && sampleRate != null) {
186+
span.setTag(CommonTags.SAMPLING_DECISION, samplingDecision);
187+
span.setTag(CommonTags.SAMPLE_RATE, sampleRate);
188+
}
176189
}
177190
super.doAfterStart(span);
178191
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/LlmObsContextPropagationForkedTest.java‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.instrumentation.openai_java;
22

33
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
45
import static org.junit.jupiter.api.Assertions.assertNotNull;
56
import static org.junit.jupiter.api.Assertions.assertNull;
67

@@ -15,6 +16,7 @@
1516
import com.sun.net.httpserver.HttpServer;
1617
import datadog.context.ContextScope;
1718
import datadog.trace.agent.test.AbstractInstrumentationTest;
19+
import datadog.trace.api.Config;
1820
import datadog.trace.api.llmobs.GenAiApmTags;
1921
import datadog.trace.api.llmobs.LLMObsContext;
2022
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -267,6 +269,105 @@ void openAiRequestSpanComputesItsOwnSamplingDecisionWhenNoLlmObsContext() throws
267269
assertEquals("1", openAiSpan.getTag("_ml_obs_tag.sample_rate"));
268270
}
269271

272+
@Test
273+
void openAiRequestSpanStaysInTheParentsAdoptedLlmObsTrace() throws Exception {
274+
// A parent that continued a trace from another service carries that service's LLMObs trace id,
275+
// which is not the local APM trace id. The LLM span must report the parent's id, or it lands
276+
// in a different LLMObs trace from its parent.
277+
String adoptedTraceId = "6ab41f6e00000000fea4f6b7295bd5e9";
278+
279+
AgentSpan parentSpan = AgentTracer.startSpan("test", "parent");
280+
try (ContextScope ignored1 = AgentTracer.activateSpan(parentSpan)) {
281+
try (ContextScope ignored2 =
282+
LLMObsContext.attach(
283+
parentSpan.spanContext(),
284+
adoptedTraceId,
285+
null,
286+
null,
287+
null,
288+
"1",
289+
LLMObsContext.SAMPLING_DECISION_SAMPLED,
290+
null,
291+
null)) {
292+
openAiClient.chat().completions().create(buildMinimalChatParams());
293+
}
294+
} finally {
295+
parentSpan.finish();
296+
}
297+
298+
writer.waitForTraces(1);
299+
DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
300+
assertNotNull(openAiSpan, "openai.request span should have been created");
301+
assertNotEquals(adoptedTraceId, openAiSpan.getTraceId().toHexString());
302+
assertEquals(adoptedTraceId, openAiSpan.getTag("_ml_obs_tag.trace_id"));
303+
}
304+
305+
@Test
306+
void openAiRequestSpanInheritsMlAppFromActiveContext() throws Exception {
307+
String expectedMlApp = "upstream-app";
308+
309+
AgentSpan parentSpan = AgentTracer.startSpan("test", "parent");
310+
try (ContextScope ignored1 = AgentTracer.activateSpan(parentSpan)) {
311+
try (ContextScope ignored2 =
312+
LLMObsContext.attach(
313+
parentSpan.spanContext(),
314+
null,
315+
expectedMlApp,
316+
null,
317+
null,
318+
"1",
319+
LLMObsContext.SAMPLING_DECISION_SAMPLED,
320+
null,
321+
null)) {
322+
openAiClient.chat().completions().create(buildMinimalChatParams());
323+
}
324+
} finally {
325+
parentSpan.finish();
326+
}
327+
328+
writer.waitForTraces(1);
329+
DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
330+
assertNotNull(openAiSpan, "openai.request span should have been created");
331+
assertEquals(expectedMlApp, openAiSpan.getTag("_ml_obs_tag.ml_app"));
332+
}
333+
334+
@Test
335+
void openAiRequestSpanUsesServiceDefaultsWithoutLlmObsContext() throws Exception {
336+
openAiClient.chat().completions().create(buildMinimalChatParams());
337+
338+
// No parent: no trace_id tag, so the mapper seeds it from the APM trace id, and ml_app is the
339+
// configured default.
340+
writer.waitForTraces(1);
341+
DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
342+
assertNotNull(openAiSpan, "openai.request span should have been created");
343+
assertNull(openAiSpan.getTag("_ml_obs_tag.trace_id"));
344+
assertEquals(Config.get().getLlmObsMlApp(), openAiSpan.getTag("_ml_obs_tag.ml_app"));
345+
}
346+
347+
@Test
348+
void openAiRequestSpanRecordsNoSamplingDecisionWhenItsParentHasNone() throws Exception {
349+
// The parent continued a trace whose caller sent no verdict, so it has none. The LLM span must
350+
// not roll its own: below a rate of 1 that could drop the span while its parent is kept.
351+
AgentSpan parentSpan = AgentTracer.startSpan("test", "parent");
352+
try (ContextScope ignored1 = AgentTracer.activateSpan(parentSpan)) {
353+
try (ContextScope ignored2 =
354+
LLMObsContext.attach(
355+
parentSpan.spanContext(), null, null, null, null, null, null, null, null)) {
356+
openAiClient.chat().completions().create(buildMinimalChatParams());
357+
}
358+
} finally {
359+
parentSpan.finish();
360+
}
361+
362+
writer.waitForTraces(1);
363+
DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
364+
assertNotNull(openAiSpan, "openai.request span should have been created");
365+
assertEquals(
366+
String.valueOf(parentSpan.getSpanId()), openAiSpan.getTag("_ml_obs_tag.parent_id"));
367+
assertNull(openAiSpan.getTag("_ml_obs_tag.sampling_decision"));
368+
assertNull(openAiSpan.getTag("_ml_obs_tag.sample_rate"));
369+
}
370+
270371
@Test
271372
void openAiRequestSpanInheritsNothingFromStaleCrossTraceContext() throws Exception {
272373
// Simulates a stale LLMObsContext leaked across an async boundary: the context is attached,

0 commit comments

Comments
 (0)