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

jessicagamio · jessicagamio · commit 4bdfc39123c3 · 2026-05-19T15:18:41.000-07:00
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.
diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/CommonTags.java b/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";
 
diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java b/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);
   }
diff --git a/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/SessionIdPropagationTest.groovy b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/SessionIdPropagationTest.groovy
@@ -0,0 +1,63 @@
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+
+import datadog.context.ContextScope
+import datadog.trace.api.llmobs.LLMObsContext
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+
+/**
+ * Verifies that auto-instrumented openai.request spans inherit session_id from an active
+ * LLMObs parent context, matching the behavior of dd-trace-py and dd-trace-js.
+ *
+ * Background: customer apps typically set session_id on a manual LLMObs workflow root
+ * via LLMObs.startWorkflowSpan(..., sessionId) and then call OpenAI without manually
+ * wrapping in LLMObs.startLLMSpan. The auto-instrumented LLM span needs to inherit the
+ * session_id from the workflow parent so the trace appears under its session in the
+ * LLM Trace Explorer's Sessions view.
+ *
+ * Note: this test attaches the LLMObsContext directly rather than going through
+ * LLMObs.startWorkflowSpan() — the latter resolves to a no-op factory in this
+ * test module since agent-llmobs isn't loaded here. We're specifically validating
+ * OpenAiDecorator's session_id inheritance behavior, which only depends on
+ * LLMObsContext.currentSessionId() being set.
+ */
+class SessionIdPropagationTest extends OpenAiTest {
+
+  def "openai.request span inherits session_id from active LLMObs context"() {
+    setup:
+    def expectedSessionId = "session-propagation-test-abc"
+
+    when:
+    runUnderTrace("parent") {
+      def parentCtx = AgentTracer.activeSpan().context()
+      ContextScope scope = LLMObsContext.attach(parentCtx, expectedSessionId)
+      try {
+        openAiClient.chat().completions().create(chatCompletionCreateParams(false))
+      } finally {
+        scope.close()
+      }
+    }
+    TEST_WRITER.waitForTraces(1)
+    def openAiSpan = TEST_WRITER.flatten().find {
+      it.operationName.toString() == "openai.request"
+    }
+
+    then:
+    openAiSpan != null
+    expectedSessionId == openAiSpan.getTag("_ml_obs_tag.session_id")
+  }
+
+  def "openai.request span has no session_id when no LLMObs parent context is active"() {
+    when:
+    runUnderTrace("non-llmobs-parent") {
+      openAiClient.chat().completions().create(chatCompletionCreateParams(false))
+    }
+    TEST_WRITER.waitForTraces(1)
+    def openAiSpan = TEST_WRITER.flatten().find {
+      it.operationName.toString() == "openai.request"
+    }
+
+    then:
+    openAiSpan != null
+    null == openAiSpan.getTag("_ml_obs_tag.session_id")
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,16 @@ public AgentSpan afterStart(AgentSpan span) {`
`115`	`115`	`parentSpanId = String.valueOf(parent.getSpanId());`
`116`	`116`	`}`
`117`	`117`	`span.setTag(CommonTags.PARENT_ID, parentSpanId);`
	`118`	`+`
	`119`	`+ // Inherit session_id from the active LLMObs parent (e.g. a manual workflow span).`
	`120`	`+ // Matches dd-trace-py / dd-trace-js, where auto-instrumented LLM spans inherit`
	`121`	`+ // session_id from the workflow root via context propagation. Without this, the`
	`122`	`+ // auto-instrumented openai.request span would not appear under its session in`
	`123`	`+ // the LLM Trace Explorer's Sessions view.`
	`124`	`+ String sessionId = LLMObsContext.currentSessionId();`
	`125`	`+ if (sessionId != null && !sessionId.isEmpty()) {`
	`126`	`+ span.setTag(CommonTags.SESSION_ID, sessionId);`
	`127`	`+ }`
`118`	`128`	`}`
`119`	`129`	`return super.afterStart(span);`
`120`	`130`	`}`