Skip to content

Commit cd29846

Browse files
jessicagamiojessicagamio
committed
test(llmobs): consolidate session_id mapper tests, add transitive inheritance test
1 parent c4ebbe2 commit cd29846

3 files changed

Lines changed: 28 additions & 64 deletions

File tree

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,27 @@ class DDLLMObsSpanTest extends DDSpecification{
445445
parent.finish()
446446
}
447447

448+
def "grandchild LLMObs span transitively inherits session_id through intermediate span"() {
449+
setup:
450+
def expectedSessionId = "session-grandparent-xyz"
451+
def grandparent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "grandparent-workflow", expectedSessionId)
452+
def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", null)
453+
454+
when:
455+
// Grandchild created with null sessionId — should inherit transitively
456+
// through parent's re-attached LLMObsContext (which itself inherited from grandparent).
457+
def grandchild = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "grandchild-llm", null)
458+
459+
then:
460+
def innerGrandchild = (AgentSpan) grandchild.span
461+
expectedSessionId == innerGrandchild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID)
462+
463+
cleanup:
464+
grandchild.finish()
465+
parent.finish()
466+
grandparent.finish()
467+
}
468+
448469
def "global dd_tags are included in LLMObs span tags"() {
449470
setup:
450471
injectSysConfig("trace.global.tags", "team:backend,owner:ml-platform")

‎dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,7 @@ public void map(List<? extends CoreSpan<?>> trace, Writable writable) {
182182
writable.writeString(sessionId, null);
183183
}
184184

185-
/* metrics, tags, meta */
185+
/* 10 (metrics), 11 (tags), 12 meta — shift down 1 if session_id absent */
186186
span.processTagsAndBaggage(metaWriter.withWritable(writable, getErrorsMap(span)));
187187
}
188188

‎dd-trace-core/src/test/groovy/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapperTest.groovy‎

Lines changed: 6 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ class LLMObsSpanMapperTest extends DDCoreSpecification {
3636
.withTag("_ml_obs_metric.input_tokens", 50)
3737
.withTag("_ml_obs_metric.output_tokens", 25)
3838
.withTag("_ml_obs_metric.total_tokens", 75)
39+
.withTag("_ml_obs_tag.session_id", "abc-123-session")
3940
.start()
4041

4142
llmSpan.setSpanType(InternalSpanTypes.LLMOBS)
@@ -132,6 +133,10 @@ class LLMObsSpanMapperTest extends DDCoreSpecification {
132133
spanData["_dd"]["trace_id"] == spanData["trace_id"]
133134
spanData["_dd"]["apm_trace_id"] == spanData["trace_id"]
134135

136+
// Top-level session_id field — what the LLM Trace Explorer's Sessions filter queries.
137+
spanData.containsKey("session_id")
138+
spanData["session_id"] == "abc-123-session"
139+
135140
spanData.containsKey("meta")
136141
spanData["meta"]["span.kind"] == "llm"
137142
spanData["meta"].containsKey("error")
@@ -176,6 +181,7 @@ class LLMObsSpanMapperTest extends DDCoreSpecification {
176181

177182
spanData.containsKey("tags")
178183
spanData["tags"].contains("language:jvm")
184+
spanData["tags"].contains("session_id:abc-123-session")
179185
}
180186

181187
def "test LLMObsSpanMapper writes no spans when none are LLMObs spans"() {
@@ -297,69 +303,6 @@ class LLMObsSpanMapperTest extends DDCoreSpecification {
297303
spanNames.contains("chat-completion-3")
298304
}
299305

300-
def "test LLMObsSpanMapper writes top-level session_id when set"() {
301-
setup:
302-
def mapper = new LLMObsSpanMapper()
303-
def tracer = tracerBuilder().writer(new ListWriter()).build()
304-
305-
def sessionId = "abc-123-session"
306-
307-
def llmSpan = tracer.buildSpan("datadog", "openai.request")
308-
.withResourceName("createCompletion")
309-
.withTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND)
310-
.withTag("_ml_obs_tag.model_name", "gpt-4")
311-
.withTag("_ml_obs_tag.model_provider", "openai")
312-
.withTag("_ml_obs_tag.session_id", sessionId)
313-
.start()
314-
llmSpan.setSpanType(InternalSpanTypes.LLMOBS)
315-
llmSpan.finish()
316-
317-
def trace = [llmSpan]
318-
CapturingByteBufferConsumer sink = new CapturingByteBufferConsumer()
319-
MsgPackWriter packer = new MsgPackWriter(new FlushingBuffer(16 * 1024, sink))
320-
321-
when:
322-
packer.format(trace, mapper)
323-
packer.flush()
324-
325-
then:
326-
sink.captured != null
327-
def payload = mapper.newPayload()
328-
payload.withBody(1, sink.captured)
329-
330-
def channel = new ByteArrayOutputStream()
331-
payload.writeTo(new WritableByteChannel() {
332-
@Override
333-
int write(ByteBuffer src) throws IOException {
334-
def bytes = new byte[src.remaining()]
335-
src.get(bytes)
336-
channel.write(bytes)
337-
return bytes.length
338-
}
339-
340-
@Override
341-
boolean isOpen() {
342-
return true
343-
}
344-
345-
@Override
346-
void close() throws IOException { }
347-
})
348-
349-
def result = objectMapper.readValue(channel.toByteArray(), Map)
350-
def spanData = result["spans"][0]
351-
352-
then:
353-
// Top-level session_id field is present with the right value — this is what
354-
// the LLM Trace Explorer's Sessions filter queries.
355-
spanData.containsKey("session_id")
356-
spanData["session_id"] == sessionId
357-
358-
// The session_id:<value> entry is ALSO present in the tags[] array, matching
359-
// dd-trace-py and dd-trace-js wire-format behavior.
360-
spanData["tags"].contains("session_id:${sessionId}".toString())
361-
}
362-
363306
def "test LLMObsSpanMapper omits top-level session_id when not set"() {
364307
setup:
365308
def mapper = new LLMObsSpanMapper()

0 commit comments

Comments
 (0)