Propagate LLM Observability context across service boundaries

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParser.java

@@ -15,6 +15,9 @@ public final class DatadogAttributeParser {
 
   private static final Base64.Decoder BASE_64 = Base64.getDecoder();
 
+  private static final String TAGS_KEY = "x-datadog-tags";
+  private static final String LLMOBS_TAG_PREFIX = "_dd.p.llmobs_";
+
   /** Parses trace context properties from the given JSON and passes them to the classifier. */
   public static void forEachProperty(AgentPropagation.KeyClassifier classifier, String json) {
     if (null == json) {
@@ -24,6 +27,7 @@ public static void forEachProperty(AgentPropagation.KeyClassifier classifier, St
       if (acceptJsonProperty(classifier, json, "x-datadog-trace-id")) {
         acceptJsonProperty(classifier, json, "x-datadog-parent-id");
         acceptJsonProperty(classifier, json, "x-datadog-sampling-priority");
+        acceptLlmObsPropagationTags(classifier, json);
       }
       if (Config.get().isDataStreamsEnabled()) {
         acceptJsonProperty(classifier, json, "dd-pathway-ctx-base64");
@@ -56,22 +60,81 @@ public static void forEachProperty(AgentPropagation.KeyClassifier classifier, By
     }
   }
 
-  // Simple parser that assumes values are JSON strings that don't contain escaped quotes
+  /**
+   * Forwards the {@code _dd.p.llmobs_*} propagation tags out of {@code x-datadog-tags}, dropping
+   * the rest.
+   */
+  private static void acceptLlmObsPropagationTags(
+      AgentPropagation.KeyClassifier classifier, String json) {
+    String tags = jsonPropertyValue(json, TAGS_KEY);
+    // Nothing to forward is the common case, and it costs a scan rather than a parse.
+    if (null == tags || !tags.contains(LLMOBS_TAG_PREFIX)) {
+      return;
+    }
+    StringBuilder filtered = null;
+    int start = 0;
+    while (start < tags.length()) {
+      int end = tags.indexOf(',', start);
+      if (end < 0) {
+        end = tags.length();
+      }
+      if (tags.startsWith(LLMOBS_TAG_PREFIX, start)) {
+        if (null == filtered) {
+          filtered = new StringBuilder(tags.length());
+        } else {
+          filtered.append(',');
+        }
+        filtered.append(tags, start, end);
+      }
+      start = end + 1;
+    }
+    if (null != filtered) {
+      classifier.accept(TAGS_KEY, filtered.toString());
+    }
+  }
+
   private static boolean acceptJsonProperty(
       AgentPropagation.KeyClassifier classifier, String json, String key) {
-    int keyStart = json.indexOf(key);
+    String value = jsonPropertyValue(json, key);
+    return null != value && classifier.accept(key, value);
+  }
+
+  // Simple parser that assumes values are JSON strings that don't contain escaped quotes
+  private static String jsonPropertyValue(String json, String key) {
+    int keyStart = indexOfPropertyName(json, key);
     if (keyStart > 0) {
       int separator = json.indexOf(':', keyStart + key.length());
       if (separator > 0) {
         int valueStart = json.indexOf('"', separator + 1);
         if (valueStart > 0) {
           int valueEnd = json.indexOf('"', valueStart + 1);
           if (valueEnd > 0) {
-            return classifier.accept(key, json.substring(valueStart + 1, valueEnd));
+            return json.substring(valueStart + 1, valueEnd);
           }
         }
       }
     }
-    return false;
+    return null;
+  }
+
+  /**
+   * Finds where {@code key} is used as a property name, that is quoted on both sides, skipping any
+   * copy of the same text that sits inside a value. The {@code x-datadog-tags} header now carries
+   * application-supplied values, so a value can contain whatever the application named its ML app
+   * or session, including the name of another property.
+   */
+  private static int indexOfPropertyName(String json, String key) {
+    int from = 0;
+    while (true) {
+      int at = json.indexOf(key, from);
+      if (at < 1) {
+        return -1;
+      }
+      int after = at + key.length();
+      if ('"' == json.charAt(at - 1) && after < json.length() && '"' == json.charAt(after)) {
+        return at;
+      }
+      from = after;
+    }
   }
 }

dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/messaging/DatadogAttributeParserTest.java

@@ -0,0 +1,108 @@
+package datadog.trace.bootstrap.instrumentation.messaging;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+import org.tabletest.junit.TableTest;
+
+/**
+ * Covers the behaviour this branch adds to the {@code _datadog} message attribute parser shared by
+ * the AWS messaging instrumentations (SQS, SNS, EventBridge, Step Functions): the {@code
+ * _dd.p.llmobs_*} propagation tags are forwarded to the extractor out of {@code x-datadog-tags},
+ * and every other {@code _dd.p.*} tag stays dropped as it is on master.
+ */
+class DatadogAttributeParserTest {
+
+  private static final String TRACE_CONTEXT =
+      "\"x-datadog-trace-id\":\"1234567890\","
+          + "\"x-datadog-parent-id\":\"9876543210\","
+          + "\"x-datadog-sampling-priority\":\"1\"";
+
+  private static Map<String, String> parse(String json) {
+    Map<String, String> collected = new LinkedHashMap<>();
+    DatadogAttributeParser.forEachProperty(
+        (key, value) -> {
+          collected.put(key, value);
+          return true;
+        },
+        json);
+    return collected;
+  }
+
+  private static Map<String, String> parseWithTags(String tags) {
+    return parse("{" + TRACE_CONTEXT + ",\"x-datadog-tags\":\"" + tags + "\"}");
+  }
+
+  /**
+   * The LLMObs tags are kept wherever they sit in the header, and the tags other products key off
+   * are not — forwarding {@code _dd.p.ts} would change ASM consumer sampling and {@code _dd.p.tid}
+   * would change the {@code dd.trace_id} format in logs, neither of which belongs in this change.
+   */
+  @TableTest({
+    "scenario          | tags                                                                   | expected                                   ",
+    "llmobs only       | '_dd.p.llmobs_sid=sess-1'                                              | '_dd.p.llmobs_sid=sess-1'                  ",
+    "llmobs first      | '_dd.p.llmobs_sid=sess-1,_dd.p.dm=-1'                                  | '_dd.p.llmobs_sid=sess-1'                  ",
+    "llmobs last       | '_dd.p.dm=-1,_dd.p.llmobs_sid=sess-1'                                  | '_dd.p.llmobs_sid=sess-1'                  ",
+    "llmobs in between | '_dd.p.dm=-1,_dd.p.llmobs_sid=sess-1,_dd.p.tid=6aa01c5400000000'       | '_dd.p.llmobs_sid=sess-1'                  ",
+    "several llmobs    | '_dd.p.llmobs_ml_app=app,_dd.p.tid=6aa01c5400000000,_dd.p.llmobs_sr=1' | '_dd.p.llmobs_ml_app=app,_dd.p.llmobs_sr=1'",
+    "trace source      | '_dd.p.ts=02,_dd.p.llmobs_sid=sess-1'                                  | '_dd.p.llmobs_sid=sess-1'                  ",
+    "empty value       | '_dd.p.llmobs_sid=,_dd.p.dm=-1'                                        | '_dd.p.llmobs_sid='                        "
+  })
+  void forwardsOnlyTheLlmObsPropagationTags(String tags, String expected) {
+    assertEquals(expected, parseWithTags(tags).get("x-datadog-tags"));
+  }
+
+  /**
+   * A header with nothing to forward is not forwarded at all, rather than forwarded empty — an
+   * empty {@code x-datadog-tags} is not the same thing to the extractor as an absent one.
+   */
+  @TableTest({
+    "scenario                  | tags                                    ",
+    "no llmobs tags            | '_dd.p.dm=-1,_dd.p.tid=6aa01c5400000000'",
+    "empty header              | ''                                      ",
+    "prefix only in a value    | '_dd.p.dm=_dd.p.llmobs_sid'             ",
+    "prefix is not a whole key | '_dd.p.llmobsx=1'                       "
+  })
+  void doesNotForwardAHeaderWithNoLlmObsTags(String tags) {
+    Map<String, String> collected = parseWithTags(tags);
+    assertFalse(collected.containsKey("x-datadog-tags"), () -> "forwarded " + collected);
+  }
+
+  @Test
+  void forwardsNothingWhenTheHeaderIsAbsent() {
+    Map<String, String> collected = parse("{" + TRACE_CONTEXT + "}");
+    assertEquals("1234567890", collected.get("x-datadog-trace-id"));
+    assertNull(collected.get("x-datadog-tags"));
+  }
+
+  /**
+   * A property is located by its quoted name, not by a bare substring search, so a value that
+   * happens to contain another property's name does not shadow the real one. This matters now that
+   * {@code x-datadog-tags} carries application-supplied values — an ML app can be named anything.
+   */
+  @Test
+  void readsAPropertyWhoseNameAlsoAppearsInsideAnEarlierValue() {
+    Map<String, String> collected =
+        parse(
+            "{\"x-datadog-trace-id\":\"1234567890\","
+                + "\"x-datadog-tags\":\"_dd.p.llmobs_ml_app=x-datadog-sampling-priority:9\","
+                + "\"x-datadog-parent-id\":\"9876543210\","
+                + "\"x-datadog-sampling-priority\":\"1\"}");
+    assertEquals("1", collected.get("x-datadog-sampling-priority"));
+    assertEquals(
+        "_dd.p.llmobs_ml_app=x-datadog-sampling-priority:9", collected.get("x-datadog-tags"));
+  }
+
+  /** The tags are only read once a trace id has been found, which is what gates the whole block. */
+  @Test
+  void forwardsNothingWhenThereIsNoTraceId() {
+    assertEquals(
+        0,
+        parse("{\"x-datadog-tags\":\"_dd.p.llmobs_sid=sess-1\"}").size(),
+        "tags should not be read without a trace id");
+  }
+}

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsContextPropagationSource.java

@@ -0,0 +1,45 @@
+package datadog.trace.llmobs;
+
+import datadog.trace.api.llmobs.LLMObsContext;
+import datadog.trace.api.llmobs.LLMObsPropagationSource;
+import datadog.trace.api.llmobs.LLMObsPropagationValues;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.llmobs.domain.LLMObsTraceId;
+import javax.annotation.Nullable;
+
+/**
+ * Resolves the LLM Observability propagation tag values for an outbound request from the ambient
+ * {@link LLMObsContext}.
+ *
+ * <p>Runs at injection rather than when a span starts, so the innermost active LLMObs span always
+ * wins and leaving an LLMObs scope stops contributing its tags — with no save/restore bookkeeping,
+ * because nothing is stored: the values are returned to the caller and serialized immediately.
+ */
+public final class LLMObsContextPropagationSource implements LLMObsPropagationSource {
+
+  @Override
+  @Nullable
+  public LLMObsPropagationValues valuesFor(AgentSpanContext spanContext) {
+    if (spanContext == null) {
+      return null;
+    }
+    // Gate on trace-id consistency, the same way DDLLMObsSpan gates parent_id/session_id
+    // inheritance. An LLMObs context leaked across an async boundary must not tag an outbound
+    // request that belongs to an unrelated trace. Returning null here means the codec forwards the
+    // inbound values instead, which is what a pass-through service should do.
+    AgentSpanContext llmObsContext = LLMObsContext.current();
+    if (llmObsContext == null || !llmObsContext.getTraceId().equals(spanContext.getTraceId())) {
+      return null;
+    }
+    return new LLMObsPropagationValues(
+        // Stored as hex in-process and in the span payload, carried as decimal on the wire.
+        LLMObsTraceId.toWire(LLMObsContext.currentTraceId()),
+        LLMObsContext.currentMlApp(),
+        LLMObsContext.currentSessionId(),
+        LLMObsContext.currentParentAgentSpanId(),
+        LLMObsContext.currentParentAgentName(),
+        String.valueOf(llmObsContext.getSpanId()),
+        LLMObsContext.currentSampleRate(),
+        LLMObsContext.currentSamplingDecision());
+  }
+}