Fix sending multiple LLM Obs traces per request.

ygree · ygree · commit 774603321184 · 2026-01-16T17:07:29.000-08:00
Previously, llm-obs-events-processor only recognized the first trace per request and ignored the rest.

Fix sizeInBytes calculation for metrics and logging.
Previously, it was the size of the entire buffer rather than of the actual number of bytes for the request
diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/PayloadDispatcherImpl.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/PayloadDispatcherImpl.java
@@ -112,19 +112,21 @@ public void accept(int messageCount, ByteBuffer buffer) {
       if (response.success()) {
         if (log.isDebugEnabled()) {
           log.debug(
-              "Successfully sent {} traces {} bytes to the API {}",
+              "Successfully sent {} traces of size {} bytes to the API {}",
               messageCount,
-              buffer.position(),
+              sizeInBytes,
               mapper.endpoint());
         }
         healthMetrics.onSend(messageCount, sizeInBytes, response);
       } else {
         if (log.isDebugEnabled()) {
           log.debug(
-              "Failed to send {} traces of size {} bytes to the API {}",
+              "Failed to send {} traces of size {} bytes to the API {} status {} response {}",
               messageCount,
               sizeInBytes,
-              mapper.endpoint());
+              mapper.endpoint(),
+              response.status(),
+              response.response());
         }
         healthMetrics.onFailedSend(messageCount, sizeInBytes, response);
       }
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
@@ -2,7 +2,9 @@
 
 import static datadog.communication.http.OkHttpUtils.gzippedMsgpackRequestBodyOf;
 
+import datadog.communication.serialization.GrowableBuffer;
 import datadog.communication.serialization.Writable;
+import datadog.communication.serialization.msgpack.MsgPackWriter;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.intake.TrackType;
 import datadog.trace.api.llmobs.LLMObs;
@@ -77,15 +79,30 @@ public class LLMObsSpanMapper implements RemoteMapper {
 
   private static final String PARENT_ID_TAG_INTERNAL_FULL = LLMOBS_TAG_PREFIX + "parent_id";
 
-  private final LLMObsSpanMapper.MetaWriter metaWriter = new MetaWriter();
+  private final MetaWriter metaWriter = new MetaWriter();
   private final int size;
 
+  private final ByteBuffer header;
+  private int spansWritten;
+
   public LLMObsSpanMapper() {
     this(5 << 20);
   }
 
   private LLMObsSpanMapper(int size) {
     this.size = size;
+
+    GrowableBuffer header = new GrowableBuffer(64);
+    MsgPackWriter headerWriter = new MsgPackWriter(header);
+
+    headerWriter.startMap(3);
+    headerWriter.writeUTF8(EVENT_TYPE);
+    headerWriter.writeString("span", null);
+    headerWriter.writeUTF8(STAGE);
+    headerWriter.writeString("raw", null);
+    headerWriter.writeUTF8(SPANS);
+
+    this.header = header.slice();
   }
 
   @Override
@@ -98,16 +115,6 @@ public void map(List<? extends CoreSpan<?>> trace, Writable writable) {
       return;
     }
 
-    writable.startMap(3);
-
-    writable.writeUTF8(EVENT_TYPE);
-    writable.writeString("span", null);
-
-    writable.writeUTF8(STAGE);
-    writable.writeString("raw", null);
-
-    writable.writeUTF8(SPANS);
-    writable.startArray(llmobsSpans.size());
     for (CoreSpan<?> span : llmobsSpans) {
       writable.startMap(11);
       // 1
@@ -148,6 +155,10 @@ public void map(List<? extends CoreSpan<?>> trace, Writable writable) {
       /* 9 (metrics), 10 (tags), 11 meta */
       span.processTagsAndBaggage(metaWriter.withWritable(writable, getErrorsMap(span)));
     }
+
+    // Increase only after all spans have been written. This way, if it rolls back because of a
+    // buffer overflow, the counter won't be skewed.
+    spansWritten += llmobsSpans.size();
   }
 
   private static boolean isLLMObsSpan(CoreSpan<?> span) {
@@ -157,7 +168,7 @@ private static boolean isLLMObsSpan(CoreSpan<?> span) {
 
   @Override
   public Payload newPayload() {
-    return new PayloadV1();
+    return new PayloadV1(header, spansWritten);
   }
 
   @Override
@@ -166,7 +177,10 @@ public int messageBufferSize() {
   }
 
   @Override
-  public void reset() {}
+  public void reset() {
+    // Reset the number of spans per message with each flush.
+    spansWritten = 0;
+  }
 
   @Override
   public String endpoint() {
@@ -206,7 +220,7 @@ private static final class MetaWriter implements MetadataConsumer {
                     LLMOBS_TAG_PREFIX + LLMObsTags.MODEL_VERSION,
                     LLMOBS_TAG_PREFIX + LLMObsTags.METADATA)));
 
-    LLMObsSpanMapper.MetaWriter withWritable(Writable writable, Map<String, String> errorInfo) {
+    MetaWriter withWritable(Writable writable, Map<String, String> errorInfo) {
       this.writable = writable;
       this.errorInfo = errorInfo;
       return this;
@@ -348,14 +362,20 @@ public void accept(Metadata metadata) {
   }
 
   private static class PayloadV1 extends Payload {
+    private final ByteBuffer header;
+    private final int spansWritten;
+
+    public PayloadV1(ByteBuffer header, int spansWritten) {
+      this.spansWritten = spansWritten;
+      this.header = header;
+    }
 
     @Override
     public int sizeInBytes() {
       if (traceCount() == 0) {
         return msgpackMapHeaderSize(0);
       }
-
-      return body.array().length;
+      return header.remaining() + msgpackArrayHeaderSize(spansWritten) + body.remaining();
     }
 
     @Override
@@ -368,6 +388,8 @@ public void writeTo(WritableByteChannel channel) throws IOException {
         }
       } else {
         while (body.hasRemaining()) {
+          channel.write(header.slice());
+          channel.write(msgpackArrayHeader(spansWritten));
           channel.write(body);
         }
       }
@@ -379,9 +401,13 @@ public RequestBody toRequest() {
       if (traceCount() == 0) {
         buffers = Collections.singletonList(msgpackMapHeader(0));
       } else {
-        buffers = Collections.singletonList(body);
+        buffers =
+            Arrays.asList(
+                header.slice(),
+                // Third Value: is an array of spans serialized into the body
+                msgpackArrayHeader(spansWritten),
+                body);
       }
-
       return gzippedMsgpackRequestBodyOf(buffers);
     }
   }
