From 1d39294b6312568fbc70a36483a93d8431c0fe0f Mon Sep 17 00:00:00 2001 From: jessicagamio Date: Wed, 13 May 2026 20:40:02 -0700 Subject: [PATCH 1/2] fix(llmobs): emit session_id as top-level field in LLMObs span event dd-trace-java was placing session_id only inside the `tags[]` array of each LLMObs span event on the wire, never as a top-level span field. The LLM Trace Explorer's Sessions filter keys off the top-level `session_id` field per the public HTTP intake API schema (https://docs.datadoghq.com/llm_observability/instrumentation/api?tab=model#span), so Java spans were invisible to that filter. dd-trace-py and dd-trace-js have conformed to this schema since the LLMObs v2 API launched (2024-06-21, MLOB-955). In LLMObsSpanMapper.map(), read the _ml_obs_tag.session_id tag from each span before opening the msgpack map; if present, extend the map size from 11 to 12 entries and emit the value as a top-level field. Mirrors the existing parent_id lift pattern, with one deliberate difference: the tag is not removed so the `session_id:` entry remains in the tags[] array, matching dd-trace-py and dd-trace-js behavior. Originating issue: MLOS-646. --- .../writer/ddintake/LLMObsSpanMapper.java | 20 ++- .../ddintake/LLMObsSpanMapperTest.groovy | 120 ++++++++++++++++++ 2 files changed, 138 insertions(+), 2 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java index cb8ebd3d8a1..287444c92bc 100644 --- a/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java +++ b/dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java @@ -55,6 +55,7 @@ public class LLMObsSpanMapper implements RemoteMapper { private static final byte[] DD = "_dd".getBytes(StandardCharsets.UTF_8); private static final byte[] APM_TRACE_ID = "apm_trace_id".getBytes(StandardCharsets.UTF_8); private static final byte[] PARENT_ID = "parent_id".getBytes(StandardCharsets.UTF_8); + private static final byte[] SESSION_ID = "session_id".getBytes(StandardCharsets.UTF_8); private static final byte[] NAME = "name".getBytes(StandardCharsets.UTF_8); private static final byte[] DURATION = "duration".getBytes(StandardCharsets.UTF_8); private static final byte[] START_NS = "start_ns".getBytes(StandardCharsets.UTF_8); @@ -88,6 +89,8 @@ public class LLMObsSpanMapper implements RemoteMapper { private static final byte[] LLM_TOOL_RESULT_RESULT = "result".getBytes(StandardCharsets.UTF_8); private static final String PARENT_ID_TAG_INTERNAL_FULL = LLMOBS_TAG_PREFIX + "parent_id"; + private static final String SESSION_ID_TAG_INTERNAL_FULL = + LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID; private final MetaWriter metaWriter = new MetaWriter(); private final int size; @@ -126,7 +129,13 @@ public void map(List> trace, Writable writable) { } for (CoreSpan span : llmobsSpans) { - writable.startMap(11); + // Read session_id off the span before opening the map so we can size it correctly. + // We deliberately do NOT remove the tag (unlike parent_id) — the session_id: + // entry must remain in the tags[] array to match dd-trace-py and dd-trace-js behavior. + String sessionId = span.getTag(SESSION_ID_TAG_INTERNAL_FULL); + boolean hasSessionId = sessionId != null && !sessionId.isEmpty(); + + writable.startMap(hasSessionId ? 12 : 11); // 1 writable.writeUTF8(SPAN_ID); writable.writeString(String.valueOf(span.getSpanId()), null); @@ -166,7 +175,14 @@ public void map(List> trace, Writable writable) { writable.writeUTF8(APM_TRACE_ID); writable.writeString(span.getTraceId().toHexString(), null); - /* 9 (metrics), 10 (tags), 11 meta */ + // 9 — optional top-level session_id field. Required by the LLMObs HTTP intake schema + // and by the LLM Trace Explorer's Sessions filter, which keys off this field. + if (hasSessionId) { + writable.writeUTF8(SESSION_ID); + writable.writeString(sessionId, null); + } + + /* metrics, tags, meta */ span.processTagsAndBaggage(metaWriter.withWritable(writable, getErrorsMap(span))); } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.groovy index 6ab958c3adc..e30f4fd2e68 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.groovy @@ -297,6 +297,126 @@ class LLMObsSpanMapperTest extends DDCoreSpecification { spanNames.contains("chat-completion-3") } + def "test LLMObsSpanMapper writes top-level session_id when set"() { + setup: + def mapper = new LLMObsSpanMapper() + def tracer = tracerBuilder().writer(new ListWriter()).build() + + def sessionId = "abc-123-session" + + def llmSpan = tracer.buildSpan("datadog", "openai.request") + .withResourceName("createCompletion") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.model_name", "gpt-4") + .withTag("_ml_obs_tag.model_provider", "openai") + .withTag("_ml_obs_tag.session_id", sessionId) + .start() + llmSpan.setSpanType(InternalSpanTypes.LLMOBS) + llmSpan.finish() + + def trace = [llmSpan] + CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer() + MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(16 * 1024, sink)) + + when: + packer.format(trace, mapper) + packer.flush() + + then: + sink.captured != null + def payload = mapper.newPayload() + payload.withBody(1, sink.captured) + + def channel = new ByteArrayOutputStream() + payload.writeTo(new WritableByteChannel() { + @Override + int write(ByteBuffer src) throws IOException { + def bytes = new byte[src.remaining()] + src.get(bytes) + channel.write(bytes) + return bytes.length + } + + @Override + boolean isOpen() { + return true + } + + @Override + void close() throws IOException { } + }) + + def result = objectMapper.readValue(channel.toByteArray(), Map) + def spanData = result["spans"][0] + + then: + // Top-level session_id field is present with the right value — this is what + // the LLM Trace Explorer's Sessions filter queries. + spanData.containsKey("session_id") + spanData["session_id"] == sessionId + + // The session_id: entry is ALSO present in the tags[] array, matching + // dd-trace-py and dd-trace-js wire-format behavior. + spanData["tags"].contains("session_id:${sessionId}".toString()) + } + + def "test LLMObsSpanMapper omits top-level session_id when not set"() { + setup: + def mapper = new LLMObsSpanMapper() + def tracer = tracerBuilder().writer(new ListWriter()).build() + + def llmSpan = tracer.buildSpan("datadog", "openai.request") + .withResourceName("createCompletion") + .withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND) + .withTag("_ml_obs_tag.model_name", "gpt-4") + .withTag("_ml_obs_tag.model_provider", "openai") + .start() + llmSpan.setSpanType(InternalSpanTypes.LLMOBS) + llmSpan.finish() + + def trace = [llmSpan] + CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer() + MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(16 * 1024, sink)) + + when: + packer.format(trace, mapper) + packer.flush() + + then: + sink.captured != null + def payload = mapper.newPayload() + payload.withBody(1, sink.captured) + + def channel = new ByteArrayOutputStream() + payload.writeTo(new WritableByteChannel() { + @Override + int write(ByteBuffer src) throws IOException { + def bytes = new byte[src.remaining()] + src.get(bytes) + channel.write(bytes) + return bytes.length + } + + @Override + boolean isOpen() { + return true + } + + @Override + void close() throws IOException { } + }) + + def result = objectMapper.readValue(channel.toByteArray(), Map) + def spanData = result["spans"][0] + + then: + // No top-level session_id field when the tag was never set. + !spanData.containsKey("session_id") + + // And no session_id entry leaks into tags[] either. + spanData["tags"].every { !it.startsWith("session_id:") } + } + static class CapturingByteBufferConsumer implements ByteBufferConsumer { ByteBuffer captured From c4ebbe227bde53c498b7810d2378af5bdf8b9c67 Mon Sep 17 00:00:00 2001 From: jessicagamio Date: Wed, 13 May 2026 20:40:22 -0700 Subject: [PATCH 2/2] fix(llmobs): propagate session_id from parent context to child LLMObs spans dd-trace-java's DDLLMObsSpan constructor only honored an explicitly passed sessionId argument. When a caller used the SDK's documented "set session_id on the root span only" pattern and started a child LLMObs span with sessionId=null, the child carried no session_id. dd-trace-py (ddtrace/llmobs/_llmobs.py:1911) and dd-trace-js (packages/dd-trace/src/llmobs/tagger.js:105-106) both auto-propagate session_id from parent context to descendant LLMObs spans. Add a SESSION_ID ContextKey to LLMObsContext alongside the existing span-context key, plus an attach(ctx, sessionId) overload and a currentSessionId() accessor. In DDLLMObsSpan's constructor, fall back to LLMObsContext.currentSessionId() when no explicit sessionId is passed. The constructor now also propagates the effective sessionId through the new context key so subsequent descendant spans can inherit it. Originating issue: MLOS-646. --- .../trace/llmobs/domain/DDLLMObsSpan.java | 13 ++++++- .../llmobs/domain/DDLLMObsSpanTest.groovy | 34 +++++++++++++++++++ .../trace/api/llmobs/LLMObsContext.java | 23 ++++++++++++- 3 files changed, 68 insertions(+), 2 deletions(-) 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 6532829cfa6..40db109da03 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 @@ -67,6 +67,16 @@ public DDLLMObsSpan( spanName = kind; } + // If no explicit session_id was passed, inherit it from the active LLMObs parent. + // This matches dd-trace-py and dd-trace-js, and the public SDK docs which state that + // session_id only needs to be set on the root span — descendants inherit it. + if (sessionId == null || sessionId.isEmpty()) { + String inherited = LLMObsContext.currentSessionId(); + if (inherited != null && !inherited.isEmpty()) { + sessionId = inherited; + } + } + AgentTracer.SpanBuilder spanBuilder = AgentTracer.get() .buildSpan(LLM_OBS_INSTRUMENTATION_NAME, spanName) @@ -109,7 +119,8 @@ public DDLLMObsSpan( } } span.setTag(LLMOBS_TAG_PREFIX + PARENT_ID_TAG_INTERNAL, parentSpanID); - scope = LLMObsContext.attach(span.context()); + // Propagate the effective sessionId to descendant LLMObs spans via the context. + scope = LLMObsContext.attach(span.context(), this.hasSessionId ? sessionId : null); } @Override diff --git a/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy b/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy index 0deb4e99a19..189b62a9a3f 100644 --- a/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy +++ b/dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy @@ -411,6 +411,40 @@ class DDLLMObsSpanTest extends DDSpecification{ null | "has_session_id:0" } + def "child LLMObs span inherits session_id from parent context when none is passed"() { + setup: + def expectedSessionId = "session-abc-123" + def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", expectedSessionId) + + when: + // Child created with null sessionId — should inherit from the parent's LLMObsContext. + def child = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "child-llm", null) + + then: + def innerChild = (AgentSpan) child.span + expectedSessionId == innerChild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID) + + cleanup: + child.finish() + parent.finish() + } + + def "child LLMObs span has no session_id when neither parent nor child passes one"() { + setup: + def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", null) + + when: + def child = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "child-llm", null) + + then: + def innerChild = (AgentSpan) child.span + null == innerChild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID) + + cleanup: + child.finish() + parent.finish() + } + def "global dd_tags are included in LLMObs span tags"() { setup: injectSysConfig("trace.global.tags", "team:backend,owner:ml-platform") 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 09d90417d92..83df56553b8 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 @@ -13,12 +13,33 @@ private LLMObsContext() { } private static final ContextKey CONTEXT_KEY = ContextKey.named("llmobs_span"); + private static final ContextKey SESSION_ID_KEY = ContextKey.named("llmobs_session_id"); public static ContextScope attach(AgentSpanContext ctx) { - return Context.current().with(CONTEXT_KEY, ctx).attach(); + return attach(ctx, null); + } + + /** + * Attach an LLMObs span context, optionally propagating a session_id to descendant LLMObs spans. + * When sessionId is non-null and non-empty, child LLMObs spans started under this context that do + * not specify their own sessionId will inherit it via {@link #currentSessionId()}. + */ + public static ContextScope attach(AgentSpanContext ctx, String sessionId) { + Context updated = Context.current().with(CONTEXT_KEY, ctx); + if (sessionId != null && !sessionId.isEmpty()) { + updated = updated.with(SESSION_ID_KEY, sessionId); + } + return updated.attach(); } public static AgentSpanContext current() { return Context.current().get(CONTEXT_KEY); } + + /** + * Return the session_id propagated from an enclosing LLMObs span, or null if no parent set one. + */ + public static String currentSessionId() { + return Context.current().get(SESSION_ID_KEY); + } }