Attach SDS findings to ai_guard spans (#10721)

Attach SDS findings to ai_guard spans

Co-authored-by: oceane.bordeau <oceane.bordeau@datadoghq.com>

Author: obordeau

dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java

@@ -72,6 +72,7 @@ public BadConfigurationException(final String message) {
   static final String META_STRUCT_TAG = "ai_guard";
   static final String META_STRUCT_MESSAGES = "messages";
   static final String META_STRUCT_CATEGORIES = "attack_categories";
+  static final String META_STRUCT_SDS = "sds";
 
   public static void install() {
     final Config config = Config.get();
@@ -252,13 +253,18 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
         final String reason = (String) result.get("reason");
         @SuppressWarnings("unchecked")
         final List<String> tags = (List<String>) result.get("tags");
+        @SuppressWarnings("unchecked")
+        final List<?> sdsFindings = (List<?>) result.get("sds_findings");
         span.setTag(ACTION_TAG, action);
         if (reason != null) {
           span.setTag(REASON_TAG, reason);
         }
         if (tags != null && !tags.isEmpty()) {
           metaStruct.put(META_STRUCT_CATEGORIES, tags);
         }
+        if (sdsFindings != null && !sdsFindings.isEmpty()) {
+          metaStruct.put(META_STRUCT_SDS, sdsFindings);
+        }
         final boolean shouldBlock =
             isBlockingEnabled(options, result.get("is_blocking_enabled")) && action != Action.ALLOW;
         WafMetricCollector.get().aiGuardRequest(action, shouldBlock);

dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy

@@ -344,6 +344,57 @@ class AIGuardInternalTests extends DDSpecification {
     messages << [[], null]
   }
 
+  void 'test evaluate with sds findings'() {
+    given:
+    final sdsFindings = [
+      [
+        rule_display_name: 'Credit Card Number',
+        rule_tag: 'credit_card',
+        category: 'pii',
+        matched_text: '4111111111111111',
+        location: [start_index: 10, end_index_exclusive: 26, path: 'messages[0].content[0].text']
+      ],
+      [
+        rule_display_name: 'Social Security Number',
+        rule_tag: 'ssn',
+        category: 'pii',
+        matched_text: '123-45-6789',
+        location: [start_index: 30, end_index_exclusive: 41, path: 'messages[1].tool_calls[0].function.arguments']
+      ]
+    ]
+    final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine', sds_findings: sdsFindings]]])
+    Map<String, Object> receivedMeta
+
+    when:
+    aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT)
+
+    then:
+    1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> {
+      receivedMeta = it[1] as Map<String, Object>
+      return span
+    }
+    receivedMeta.sds == sdsFindings
+  }
+
+  void 'test evaluate with empty sds findings'() {
+    given:
+    final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine', sds_findings: sdsFindings]]])
+    Map<String, Object> receivedMeta
+
+    when:
+    aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT)
+
+    then:
+    1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> {
+      receivedMeta = it[1] as Map<String, Object>
+      return span
+    }
+    !receivedMeta.containsKey('sds')
+
+    where:
+    sdsFindings << [null, []]
+  }
+
   void 'test missing tool name'() {
     given:
     final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Just do it']]])