fix(llmobs): emit session_id as top-level field + propagate via context (#11371)

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:<value>` entry
remains in the tags[] array, matching dd-trace-py and dd-trace-js
behavior.

Originating issue: MLOS-646.

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.

test(llmobs): consolidate session_id mapper tests, add transitive inheritance test

fix(llmobs/openai-java): inherit session_id on auto-instrumented openai.request span

OpenAiDecorator.afterStart() already consulted LLMObsContext.current()
to set the LLMObs parent_id tag but never read session_id from the
context. Auto-instrumented openai.request spans therefore carried no
_ml_obs_tag.session_id even when a manual LLMObs workflow parent had
one. dd-trace-py and dd-trace-js both auto-propagate session_id to
auto-instrumented LLM spans via parent context.

Add a SESSION_ID constant to CommonTags. In OpenAiDecorator.afterStart(),
after the existing parent_id block, read LLMObsContext.currentSessionId()
and stamp it on the span as CommonTags.SESSION_ID when present. With
this change, auto-instrumented openai.request spans now appear under
their session in the LLM Trace Explorer's Sessions view, matching
Python and Node behavior.

Depends on the LLMObsContext.currentSessionId() API added in the
preceding fix(llmobs): propagate session_id from parent context commit.

Originating issue: MLOS-646.

test(llmobs/openai-java): migrate session_id propagation test to JUnit 5 + Java

refactor(llmobs): drop redundant hasSessionId ternary on LLMObsContext.attach

Merge remote-tracking branch 'origin/master' into jessica.gamio/llmobs-session-id-mlos-646

style(llmobs/openai-java): apply spotless formatting to SessionIdPropagationForkedTest

fix(llmobs): address Codex review — safe session_id cast, gate inheritance on trace match

- LLMObsSpanMapper: read session_id tag as Object and instanceof-check before casting,
  so a non-string value can't ClassCastException and drop the entire LLMObs payload.
- DDLLMObsSpan: move session_id inheritance from the constructor preamble to inside
  the existing trace-id consistency check that already gates parent_id. Prevents a
  stale LLMObsContext from leaking session attribution across traces.
- DDLLMObsSpanTest: existing 2-level + 3-level inheritance tests now activate the
  parent's AgentScope so they truly share a trace. New test asserts session_id is
  NOT inherited when the LLMObsContext belongs to a different trace.

test(llmobs): rename cross-trace test to mention async-boundary leak as the canonical scenario

Merge branch 'master' into jessica.gamio/llmobs-session-id-mlos-646

Co-authored-by: yury.gribkov <yury.gribkov@datadoghq.com>

Author: jessicagamio

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java

@@ -89,11 +89,9 @@ public DDLLMObsSpan(
     span.setTag(SPAN_KIND, kind);
     spanKind = kind;
     span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.ML_APP, mlApp);
-    this.hasSessionId = sessionId != null && !sessionId.isEmpty();
-    if (this.hasSessionId) {
-      span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
-    }
-
+    // Resolve effective parent_id and session_id from the LLMObs context, both gated on
+    // trace-id consistency. A stale context from a different trace (e.g. async boundary
+    // leakage) must not contribute either tag.
     AgentSpanContext parent = LLMObsContext.current();
     String parentSpanID = LLMObsContext.ROOT_SPAN_ID;
     if (null != parent) {
@@ -106,10 +104,25 @@ public DDLLMObsSpan(
             span.getSpanId());
       } else {
         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
+        // span; descendants inherit transitively via context propagation.
+        if (sessionId == null || sessionId.isEmpty()) {
+          String inherited = LLMObsContext.currentSessionId();
+          if (inherited != null && !inherited.isEmpty()) {
+            sessionId = inherited;
+          }
+        }
       }
     }
+
+    this.hasSessionId = sessionId != null && !sessionId.isEmpty();
+    if (this.hasSessionId) {
+      span.setTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID, sessionId);
+    }
     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(), sessionId);
   }
 
   @Override

dd-java-agent/agent-llmobs/src/test/groovy/datadog/trace/llmobs/domain/DDLLMObsSpanTest.groovy

@@ -411,6 +411,96 @@ 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)
+    // Activate the parent's AgentScope so the child span is created in the same trace.
+    // Without this, the child gets a fresh trace_id and the trace-consistency gate in
+    // DDLLMObsSpan would (correctly) skip session_id inheritance.
+    def parentScope = AgentTracer.activateSpan((AgentSpan) parent.span)
+
+    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()
+    parentScope.close()
+    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 "grandchild LLMObs span transitively inherits session_id through intermediate span"() {
+    setup:
+    def expectedSessionId = "session-grandparent-xyz"
+    def grandparent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "grandparent-workflow", expectedSessionId)
+    // Activate each ancestor's AgentScope so descendants stay in the same trace —
+    // session_id inheritance is gated on trace-id consistency in DDLLMObsSpan.
+    def grandparentScope = AgentTracer.activateSpan((AgentSpan) grandparent.span)
+    def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", null)
+    def parentScope = AgentTracer.activateSpan((AgentSpan) parent.span)
+
+    when:
+    // Grandchild created with null sessionId — should inherit transitively
+    // through parent's re-attached LLMObsContext (which itself inherited from grandparent).
+    def grandchild = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "grandchild-llm", null)
+
+    then:
+    def innerGrandchild = (AgentSpan) grandchild.span
+    expectedSessionId == innerGrandchild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID)
+
+    cleanup:
+    grandchild.finish()
+    parentScope.close()
+    parent.finish()
+    grandparentScope.close()
+    grandparent.finish()
+  }
+
+  def "child does NOT inherit session_id when stale LLMObsContext is from a different trace (e.g. async boundary leak)"() {
+    setup:
+    // Simulates a stale LLMObsContext (e.g. leaked across an async boundary). The parent's
+    // LLMObsContext is attached, but its AgentScope is deliberately NOT activated — so the
+    // next span we create starts a fresh trace and the trace-consistency gate must skip
+    // session_id inheritance.
+    def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "stale-workflow", "stale-session-id")
+
+    when:
+    def child = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "child-llm", null)
+
+    then:
+    def innerParent = (AgentSpan) parent.span
+    def innerChild = (AgentSpan) child.span
+    // Sanity: traces differ — confirms the scenario is set up correctly.
+    innerParent.getTraceId() != innerChild.getTraceId()
+    // Session_id from the stale-trace context must NOT leak into the new 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")

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

@@ -30,6 +30,7 @@ interface CommonTags {
   String ENV = TAG_PREFIX + "env";
   String SERVICE = TAG_PREFIX + "service";
   String PARENT_ID = TAG_PREFIX + "parent_id";
+  String SESSION_ID = TAG_PREFIX + LLMObsTags.SESSION_ID;
 
   String TOOL_DEFINITIONS = TAG_PREFIX + "tool_definitions";
 

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

@@ -115,6 +115,16 @@ public AgentSpan afterStart(AgentSpan span) {
         parentSpanId = String.valueOf(parent.getSpanId());
       }
       span.setTag(CommonTags.PARENT_ID, parentSpanId);
+
+      // Inherit session_id from the active LLMObs parent (e.g. a manual workflow span).
+      // Matches dd-trace-py / dd-trace-js, where auto-instrumented LLM spans inherit
+      // session_id from the workflow root via context propagation. Without this, the
+      // auto-instrumented openai.request span would not appear under its session in
+      // the LLM Trace Explorer's Sessions view.
+      String sessionId = LLMObsContext.currentSessionId();
+      if (sessionId != null && !sessionId.isEmpty()) {
+        span.setTag(CommonTags.SESSION_ID, sessionId);
+      }
     }
     return super.afterStart(span);
   }

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

@@ -0,0 +1,129 @@
+package datadog.trace.instrumentation.openai_java;
+
+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 com.openai.client.OpenAIClient;
+import com.openai.client.okhttp.OpenAIOkHttpClient;
+import com.openai.credential.BearerTokenCredential;
+import com.openai.models.ChatModel;
+import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import com.sun.net.httpserver.HttpServer;
+import datadog.context.ContextScope;
+import datadog.trace.agent.test.AbstractInstrumentationTest;
+import datadog.trace.api.llmobs.LLMObsContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.core.DDSpan;
+import datadog.trace.junit.utils.config.WithConfig;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.List;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Verifies that auto-instrumented openai.request spans inherit session_id from an active LLMObs
+ * parent context. Forked + @WithConfig used together so the LLMObs system property is in place
+ * before the agent installs and there's no leakage from prior test state.
+ *
+ * <p>The mock OpenAI backend returns a minimal 200 response — the test asserts on the span tag set
+ * by OpenAiDecorator.afterStart(), which runs before the HTTP response is parsed, so the response
+ * body shape doesn't matter for what's being tested.
+ */
+@WithConfig(key = "llmobs.enabled", value = "true")
+class SessionIdPropagationForkedTest extends AbstractInstrumentationTest {
+
+  private static HttpServer mockServer;
+  private static OpenAIClient openAiClient;
+
+  @BeforeAll
+  static void setupMockOpenAi() throws IOException {
+    mockServer = HttpServer.create(new InetSocketAddress("localhost", 0), 0);
+    mockServer.createContext(
+        "/v1/",
+        exchange -> {
+          exchange.sendResponseHeaders(200, -1);
+          exchange.close();
+        });
+    mockServer.start();
+
+    openAiClient =
+        OpenAIOkHttpClient.builder()
+            .baseUrl(
+                "http://"
+                    + mockServer.getAddress().getHostString()
+                    + ":"
+                    + mockServer.getAddress().getPort()
+                    + "/v1")
+            .credential(BearerTokenCredential.create(""))
+            .build();
+  }
+
+  @AfterAll
+  static void tearDownMockOpenAi() {
+    if (mockServer != null) {
+      mockServer.stop(0);
+      mockServer = null;
+    }
+    openAiClient = null;
+  }
+
+  @Test
+  void openAiRequestSpanInheritsSessionIdFromActiveContext() throws Exception {
+    String expectedSessionId = "session-propagation-test-abc";
+
+    AgentSpan parentSpan = AgentTracer.startSpan("test", "parent");
+    try (AgentScope ignored1 = AgentTracer.activateSpan(parentSpan)) {
+      try (ContextScope ignored2 = LLMObsContext.attach(parentSpan.context(), expectedSessionId)) {
+        try {
+          openAiClient.chat().completions().create(buildMinimalChatParams());
+        } catch (Exception ignored) {
+          // Mock server returns no body — the SDK may throw on parse. The span we care about
+          // is already created by the instrumentation advice before this point.
+        }
+      }
+    } finally {
+      parentSpan.finish();
+    }
+
+    writer.waitForTraces(1);
+    DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
+    assertNotNull(openAiSpan, "openai.request span should have been created");
+    assertEquals(expectedSessionId, openAiSpan.getTag("_ml_obs_tag.session_id"));
+  }
+
+  @Test
+  void openAiRequestSpanHasNoSessionIdWhenNoLlmObsContext() throws Exception {
+    try {
+      openAiClient.chat().completions().create(buildMinimalChatParams());
+    } catch (Exception ignored) {
+      // Mock server returns no body — the SDK may throw on parse. The span we care about
+      // is already created by the instrumentation advice before this point.
+    }
+
+    writer.waitForTraces(1);
+    DDSpan openAiSpan = findSpanByOperationName(writer, "openai.request");
+    assertNotNull(openAiSpan, "openai.request span should have been created");
+    assertNull(openAiSpan.getTag("_ml_obs_tag.session_id"));
+  }
+
+  private static ChatCompletionCreateParams buildMinimalChatParams() {
+    return ChatCompletionCreateParams.builder()
+        .model(ChatModel.GPT_4O_MINI)
+        .addSystemMessage("")
+        .addUserMessage("")
+        .build();
+  }
+
+  private static DDSpan findSpanByOperationName(List<List<DDSpan>> traces, String operationName) {
+    return traces.stream()
+        .flatMap(List::stream)
+        .filter(s -> operationName.equals(s.getOperationName().toString()))
+        .findFirst()
+        .orElse(null);
+  }
+}

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,17 @@ public void map(List<? extends CoreSpan<?>> 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:<value>
+      // entry must remain in the tags[] array to match dd-trace-py and dd-trace-js behavior.
+      // span.getTag returns Object — guard against generic tag APIs setting a non-string
+      // session_id value, which would otherwise throw ClassCastException here and drop
+      // the entire LLMObs payload for the trace.
+      Object rawSessionId = span.getTag(SESSION_ID_TAG_INTERNAL_FULL);
+      String sessionId = rawSessionId instanceof String ? (String) rawSessionId : null;
+      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 +179,14 @@ public void map(List<? extends CoreSpan<?>> 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);
+      }
+
+      /* 10 (metrics), 11 (tags), 12 meta — shift down 1 if session_id absent */
       span.processTagsAndBaggage(metaWriter.withWritable(writable, getErrorsMap(span)));
     }