From 0870af74990e04dae2cd71d62280fccfec0b2acb Mon Sep 17 00:00:00 2001 From: jessicagamio Date: Wed, 13 May 2026 20:40:57 -0700 Subject: [PATCH] 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. --- .../openai_java/CommonTags.java | 1 + .../openai_java/OpenAiDecorator.java | 10 +++ .../groovy/SessionIdPropagationTest.groovy | 63 +++++++++++++++++++ 3 files changed, 74 insertions(+) create mode 100644 dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/SessionIdPropagationTest.groovy 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 index a992c85400c..af206ac341d 100644 --- 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 index 83a2971953f..f1adcfcaeb7 100644 --- 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 new file mode 100644 index 00000000000..9e5848debea --- /dev/null +++ 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") + } +}