Add Test_AnomalyDetectionTags for AI Guard span enrichment

docs/understand/weblogs/end-to-end_weblog.md

@@ -1197,6 +1197,8 @@ This endpoint triggers AI Guard SDK evaluations using the `evaluate()` method fr
 - **Content-Type:** `application/json`
 - **Body:** List of AI Guard `Message` objects to be evaluated
 - **Header:** `X-AI-Guard-Block` (optional, default: `false`) - Controls whether blocking is enabled
+- **Header:** `X-User-Id` (optional) - If present and non-empty, sets `usr.id` on the local root span
+- **Header:** `X-Session-Id` (optional) - If present and non-empty, sets `session.id` on the local root span
 
 **Request Body Example:**
 ```json

manifests/cpp_httpd.yml

@@ -12,6 +12,7 @@
 # use `>1.0.0` to indicate that requirement.
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature

manifests/cpp_nginx.yml

@@ -10,6 +10,7 @@
 # use `>1.0.0` to indicate that requirement.
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature

manifests/dotnet.yml

@@ -2,6 +2,7 @@
 ---
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature

manifests/golang.yml

@@ -2,6 +2,7 @@
 ---
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62216)
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature

manifests/java.yml

@@ -2,6 +2,10 @@
 ---
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62215)
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags:
+    - weblog_declaration:
+        "*": irrelevant (just one weblog is enough to test the SDK)
+        "spring-boot": missing_feature (APPSEC-62449)
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected:
     - weblog_declaration:
         "*": irrelevant (just one weblog is enough to test the SDK)

manifests/nodejs.yml

@@ -91,6 +91,10 @@ refs:
   - &ref_6_0_0 '>=6.0.0-pre'
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62217)
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags:
+    - weblog_declaration:
+        "*": irrelevant (just one weblog is enough to test the SDK)
+        express4: missing_feature (APPSEC-62452)
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected:
     - weblog_declaration:
         "*": irrelevant (just one weblog is enough to test the SDK)

manifests/php.yml

@@ -2,6 +2,7 @@
 ---
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature

manifests/python.yml

@@ -8,6 +8,10 @@ manifest:
     - weblog_declaration:
         tornado: v4.4.0-rc2
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62216)
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags:
+    - weblog_declaration:
+        "*": irrelevant (just one weblog is enough to test the SDK)
+        "flask-poc": missing_feature (APPSEC-62450)
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected:
     - weblog_declaration:
         "*": irrelevant (just one weblog is enough to test the SDK)

manifests/ruby.yml

@@ -2,6 +2,16 @@
 ---
 manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62218)
+  tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags:
+    - weblog_declaration:
+        "*": missing_feature (APPSEC-62451)
+        rails42: irrelevant
+        rack: irrelevant
+        sinatra14: irrelevant
+        sinatra22: irrelevant
+        sinatra32: irrelevant
+        sinatra41: irrelevant
+        uds-sinatra: irrelevant
   tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected:
     - weblog_declaration:
         "*": missing_feature (APPSEC-62200)

tests/ai_guard/test_ai_guard_sdk.py

@@ -493,6 +493,82 @@ def test_sds_in_response(self):
             assert _assert_key(location, "path")
 
 
+@rfc("https://datadoghq.atlassian.net/wiki/x/KIApiQE")
+@features.ai_guard
+@scenarios.ai_guard
+class Test_AnomalyDetectionTags:
+    """Test that anomaly detection attributes are propagated from the root span into every AI Guard span."""
+
+    PUBLIC_IP = "5.6.7.9"
+    USER_ID = "u12345"
+    SESSION_ID = "s12345"
+
+    def _assert_span(self, root_span: DataDogLibrarySpan):
+        def validate(span: DataDogLibrarySpan):
+            if span["resource"] != "ai_guard":
+                return False
+
+            meta = span["meta"]
+
+            # Tags copied from the root span must be present on every AI Guard span
+            _assert_key(meta, "ai_guard.http.client_ip")
+            _assert_key(meta, "ai_guard.network.client.ip")
+            _assert_key(meta, "ai_guard.http.useragent")
+            _assert_key(meta, "ai_guard.usr.id", self.USER_ID)
+            _assert_key(meta, "ai_guard.session.id", self.SESSION_ID)
+
+            # Values must match what is on the root span
+            root_meta = root_span["meta"]
+            assert meta["ai_guard.http.client_ip"] == root_meta.get("http.client_ip"), (
+                f"ai_guard.http.client_ip mismatch: {meta['ai_guard.http.client_ip']} != {root_meta.get('http.client_ip')}"
+            )
+            assert meta["ai_guard.network.client.ip"] == root_meta.get("network.client.ip"), (
+                f"ai_guard.network.client.ip mismatch: {meta['ai_guard.network.client.ip']} != {root_meta.get('network.client.ip')}"
+            )
+            assert meta["ai_guard.http.useragent"] == root_meta.get("http.useragent"), (
+                f"ai_guard.http.useragent mismatch: {meta['ai_guard.http.useragent']} != {root_meta.get('http.useragent')}"
+            )
+            assert meta["ai_guard.usr.id"] == root_meta.get("usr.id"), (
+                f"ai_guard.usr.id mismatch: {meta['ai_guard.usr.id']} != {root_meta.get('usr.id')}"
+            )
+            assert meta["ai_guard.session.id"] == root_meta.get("session.id"), (
+                f"ai_guard.session.id mismatch: {meta['ai_guard.session.id']} != {root_meta.get('session.id')}"
+            )
+
+            return True
+
+        return validate
+
+    def setup_anomaly_detection_tags(self):
+        self.r = weblog.post(
+            "/ai_guard/evaluate",
+            headers={
+                "X-Forwarded-For": self.PUBLIC_IP,
+                "X-User-Id": self.USER_ID,
+                "X-Session-Id": self.SESSION_ID,
+            },
+            json=MESSAGES["ALLOW"],
+        )
+
+    def test_anomaly_detection_tags(self):
+        """Test that AI Guard spans carry anomaly detection attributes copied from the root span.
+
+        Verifies that http.client_ip, network.client.ip, http.useragent, usr.id and session.id
+        are all present on the AI Guard span with the ai_guard. prefix, and that their values
+        match the corresponding tags on the local root span.
+        """
+        assert self.r.status_code == 200
+
+        root_span = interfaces.library.get_root_span(self.r)
+        assert root_span, "No root span found"
+
+        interfaces.library.validate_one_span(
+            self.r,
+            validator=self._assert_span(root_span=root_span),
+            full_trace=True,
+        )
+
+
 @features.ai_guard
 @scenarios.ai_guard
 class Test_AIGuardEvent_Tag:

utils/build/docker/java/spring-boot/src/main/java/com/datadoghq/system_tests/springboot/ai_guard/AIGuardController.java

@@ -5,6 +5,9 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import datadog.trace.api.aiguard.AIGuard;
 import datadog.trace.api.aiguard.AIGuard.Evaluation;
+import datadog.trace.api.interceptor.MutableSpan;
+import io.opentracing.Span;
+import io.opentracing.util.GlobalTracer;
 import org.springframework.boot.autoconfigure.jackson.Jackson2ObjectMapperBuilderCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -39,7 +42,19 @@ public Jackson2ObjectMapperBuilderCustomizer mixInCustomizer() {
     @PostMapping("/ai_guard/evaluate")
     public ResponseEntity<?> evaluate(
             @RequestHeader(name = "X-AI-Guard-Block", defaultValue = "false") final boolean block,
+            @RequestHeader(name = "X-User-Id", required = false) final String userId,
+            @RequestHeader(name = "X-Session-Id", required = false) final String sessionId,
             @RequestBody final List<Message> data) {
+        final Span activeSpan = GlobalTracer.get().activeSpan();
+        if (activeSpan instanceof MutableSpan) {
+            final MutableSpan rootSpan = ((MutableSpan) activeSpan).getLocalRootSpan();
+            if (userId != null && !userId.isEmpty()) {
+                rootSpan.setTag("usr.id", userId);
+            }
+            if (sessionId != null && !sessionId.isEmpty()) {
+                rootSpan.setTag("session.id", sessionId);
+            }
+        }
         try {
             final List<AIGuard.Message> messages = data.stream().map(Message::toAIGuard).collect(Collectors.toList());
             final Evaluation result = AIGuard.evaluate(messages, new AIGuard.Options().block(block));

utils/build/docker/nodejs/express/app.js

@@ -759,6 +759,17 @@ app.post('/ai_guard/evaluate', async (req, res) => {
   const renameAttrs = ({ tagProbabilities: tag_probs, ...rest }) => ({ ...rest, tag_probs })
   const block = req.headers['x-ai-guard-block'] === 'true'
   const messages = req.body
+  const userId = req.headers['x-user-id']
+  const sessionId = req.headers['x-session-id']
+  const rootSpan = tracer.scope().active()?.context()._trace.started[0]
+  if (rootSpan) {
+    if (userId) {
+      rootSpan.setTag('usr.id', userId)
+    }
+    if (sessionId) {
+      rootSpan.setTag('session.id', sessionId)
+    }
+  }
   try {
     const evaluation = await tracer.aiguard.evaluate(messages, { block })
     res.status(200).json(renameAttrs(evaluation))

utils/build/docker/python/flask/app.py

@@ -2187,6 +2187,15 @@ def ai_guard_evaluate():
         should_block = flask_request.headers.get("X-AI-Guard-Block", "false").lower() == "true"
         messages = flask_request.get_json()
 
+        user_id = flask_request.headers.get("X-User-Id")
+        session_id = flask_request.headers.get("X-Session-Id")
+        root_span = tracer.current_root_span()
+        if root_span:
+            if user_id:
+                root_span.set_tag("usr.id", user_id)
+            if session_id:
+                root_span.set_tag("session.id", session_id)
+
         client = new_ai_guard_client(endpoint=os.environ.get("DD_AI_GUARD_ENDPOINT"))
         evaluation = client.evaluate(messages, Options(block=should_block))
         return jsonify(evaluation), 200

utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb

@@ -19,9 +19,9 @@ def evaluate
         Datadog::AIGuard.message(role: message_data[:role]) do |m|
           message_data[:content].each do |part|
             case part[:type]
-            when "text"
+            when 'text'
               m.text(part[:text])
-            when "image_url"
+            when 'image_url'
               m.image_url(part.dig(:image_url, :url))
             end
           end
@@ -31,7 +31,15 @@ def evaluate
       end
     end
 
-    allow_raise = request.headers['X-AI-Guard-Block']&.downcase == "true"
+    trace = Datadog::Tracing.active_trace
+    if trace
+      user_id = request.headers['X-User-Id']
+      session_id = request.headers['X-Session-Id']
+      trace.set_tag('usr.id', user_id) if user_id.present?
+      trace.set_tag('session.id', session_id) if session_id.present?
+    end
+
+    allow_raise = request.headers['X-AI-Guard-Block']&.downcase == 'true'
     result = Datadog::AIGuard.evaluate(*messages, allow_raise: allow_raise)
 
     response_data = {
@@ -48,7 +56,7 @@ def evaluate
     error_data[:tag_probabilities] = e.tag_probabilities if e.respond_to?(:tag_probabilities)
     error_data[:sds_findings] = e.sds_findings if e.respond_to?(:sds_findings)
     render json: error_data, status: 403
-  rescue => e
-    render json: {error: e.to_s, type: e.class.name}, status: 500
+  rescue StandardError => e
+    render json: { error: e.to_s, type: e.class.name }, status: 500
   end
 end

utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb

@@ -19,9 +19,9 @@ def evaluate
         Datadog::AIGuard.message(role: message_data[:role]) do |m|
           message_data[:content].each do |part|
             case part[:type]
-            when "text"
+            when 'text'
               m.text(part[:text])
-            when "image_url"
+            when 'image_url'
               m.image_url(part.dig(:image_url, :url))
             end
           end
@@ -31,7 +31,15 @@ def evaluate
       end
     end
 
-    allow_raise = request.headers['X-AI-Guard-Block']&.downcase == "true"
+    trace = Datadog::Tracing.active_trace
+    if trace
+      user_id = request.headers['X-User-Id']
+      session_id = request.headers['X-Session-Id']
+      trace.set_tag('usr.id', user_id) if user_id.present?
+      trace.set_tag('session.id', session_id) if session_id.present?
+    end
+
+    allow_raise = request.headers['X-AI-Guard-Block']&.downcase == 'true'
     result = Datadog::AIGuard.evaluate(*messages, allow_raise: allow_raise)
 
     response_data = {
@@ -48,7 +56,7 @@ def evaluate
     error_data[:tag_probabilities] = e.tag_probabilities if e.respond_to?(:tag_probabilities)
     error_data[:sds_findings] = e.sds_findings if e.respond_to?(:sds_findings)
     render json: error_data, status: 403
-  rescue => e
-    render json: {error: e.to_s, type: e.class.name}, status: 500
+  rescue StandardError => e
+    render json: { error: e.to_s, type: e.class.name }, status: 500
   end
 end

utils/build/docker/ruby/rails72/app/controllers/ai_guard_controller.rb

@@ -19,9 +19,9 @@ def evaluate
         Datadog::AIGuard.message(role: message_data[:role]) do |m|
           message_data[:content].each do |part|
             case part[:type]
-            when "text"
+            when 'text'
               m.text(part[:text])
-            when "image_url"
+            when 'image_url'
               m.image_url(part.dig(:image_url, :url))
             end
           end
@@ -31,7 +31,15 @@ def evaluate
       end
     end
 
-    allow_raise = request.headers['X-AI-Guard-Block']&.downcase == "true"
+    trace = Datadog::Tracing.active_trace
+    if trace
+      user_id = request.headers['X-User-Id']
+      session_id = request.headers['X-Session-Id']
+      trace.set_tag('usr.id', user_id) if user_id.present?
+      trace.set_tag('session.id', session_id) if session_id.present?
+    end
+
+    allow_raise = request.headers['X-AI-Guard-Block']&.downcase == 'true'
     result = Datadog::AIGuard.evaluate(*messages, allow_raise: allow_raise)
 
     response_data = {
@@ -48,7 +56,7 @@ def evaluate
     error_data[:tag_probabilities] = e.tag_probabilities if e.respond_to?(:tag_probabilities)
     error_data[:sds_findings] = e.sds_findings if e.respond_to?(:sds_findings)
     render json: error_data, status: 403
-  rescue => e
-    render json: {error: e.to_s, type: e.class.name}, status: 500
+  rescue StandardError => e
+    render json: { error: e.to_s, type: e.class.name }, status: 500
   end
 end