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..b02e02817a0 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(), sessionId); } @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..505408b5924 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,61 @@ 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 "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) + def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", null) + + 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() + parent.finish() + grandparent.finish() + } + def "global dd_tags are included in LLMObs span tags"() { setup: injectSysConfig("trace.global.tags", "team:backend,owner:ml-platform") 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/java/datadog/trace/instrumentation/openai_java/SessionIdPropagationForkedTest.java b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/SessionIdPropagationForkedTest.java new file mode 100644 index 00000000000..96d6ea83639 --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/java/datadog/trace/instrumentation/openai_java/SessionIdPropagationForkedTest.java @@ -0,0 +1,130 @@ +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. + * + *

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> traces, String operationName) { + return traces.stream() + .flatMap(List::stream) + .filter(s -> operationName.equals(s.getOperationName().toString())) + .findFirst() + .orElse(null); + } +} 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..2863a9bf610 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); + } + + /* 10 (metrics), 11 (tags), 12 meta — shift down 1 if session_id absent */ 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..a62c5315aff 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 @@ -36,6 +36,7 @@ class LLMObsSpanMapperTest extends DDCoreSpecification { .withTag("_ml_obs_metric.input_tokens", 50) .withTag("_ml_obs_metric.output_tokens", 25) .withTag("_ml_obs_metric.total_tokens", 75) + .withTag("_ml_obs_tag.session_id", "abc-123-session") .start() llmSpan.setSpanType(InternalSpanTypes.LLMOBS) @@ -132,6 +133,10 @@ class LLMObsSpanMapperTest extends DDCoreSpecification { spanData["_dd"]["trace_id"] == spanData["trace_id"] spanData["_dd"]["apm_trace_id"] == spanData["trace_id"] + // Top-level session_id field — what the LLM Trace Explorer's Sessions filter queries. + spanData.containsKey("session_id") + spanData["session_id"] == "abc-123-session" + spanData.containsKey("meta") spanData["meta"]["span.kind"] == "llm" spanData["meta"].containsKey("error") @@ -176,6 +181,7 @@ class LLMObsSpanMapperTest extends DDCoreSpecification { spanData.containsKey("tags") spanData["tags"].contains("language:jvm") + spanData["tags"].contains("session_id:abc-123-session") } def "test LLMObsSpanMapper writes no spans when none are LLMObs spans"() { @@ -297,6 +303,63 @@ class LLMObsSpanMapperTest extends DDCoreSpecification { spanNames.contains("chat-completion-3") } + 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 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); + } }