Prepare AI Guard tag probability coverage

smola · smola · commit fadbd2b16451 · 2026-03-30T15:08:40.000+02:00
Add a system test for tag probabilities in AI Guard responses and meta-structs.

APPSEC-61900
diff --git a/docs/understand/scenarios/ai_guard.md b/docs/understand/scenarios/ai_guard.md
@@ -96,7 +96,7 @@ Each language implements a `POST /ai_guard/evaluate` endpoint that:
 1. Reads messages from the request JSON body
 2. Reads the `X-AI-Guard-Block` header to determine blocking behavior
 3. Calls the AI Guard SDK `evaluate` method
-4. Returns the evaluation result (action, reason, tags)
+4. Returns the evaluation result (action, reason, tags, tag probabilities)
 
 See [weblogs](../weblogs/README.md) for details on weblog implementations.
 
diff --git a/docs/understand/weblogs/end-to-end_weblog.md b/docs/understand/weblogs/end-to-end_weblog.md
@@ -1236,7 +1236,10 @@ Successful evaluation:
 {
   "action": "ALLOW",
   "reason": "All looks good",
-  "tags": []
+  "tags": [],
+  "tag_probs": {
+    "jailbreak": 0.0
+  }
 }
 ```
 
diff --git a/manifests/cpp_httpd.yml b/manifests/cpp_httpd.yml
@@ -18,6 +18,7 @@ manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDK_Disabled: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDS_Findings_In_SDK_Response: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SensitiveDataScanning: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature
   tests/apm_tracing_e2e/: missing_feature (missing /e2e_otel_span endpoint on weblog)
   tests/appsec/: irrelevant (ASM is not implemented in C++)
   tests/appsec/api_security/test_endpoints.py: irrelevant (language not implementing this feature)
diff --git a/manifests/cpp_nginx.yml b/manifests/cpp_nginx.yml
@@ -16,6 +16,7 @@ manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDK_Disabled: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDS_Findings_In_SDK_Response: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SensitiveDataScanning: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span: missing_feature (missing /e2e_otel_span endpoint on weblog)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span::test_distributed_otel_trace: irrelevant (Golang specific test with OTel Go contrib package)
   tests/apm_tracing_e2e/test_process_tags.py::Test_Process_Tags: missing_feature
diff --git a/manifests/dotnet.yml b/manifests/dotnet.yml
@@ -8,6 +8,7 @@ manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDK_Disabled: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDS_Findings_In_SDK_Response: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SensitiveDataScanning: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span: missing_feature (missing /e2e_otel_span endpoint on weblog)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span::test_distributed_otel_trace: irrelevant (Golang specific test with OTel Go contrib package)
   tests/apm_tracing_e2e/test_process_tags.py::Test_Process_Tags: v3.34.0
diff --git a/manifests/golang.yml b/manifests/golang.yml
@@ -8,6 +8,7 @@ manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDK_Disabled: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDS_Findings_In_SDK_Response: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SensitiveDataScanning: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span:
     - weblog_declaration:
         "*": missing_feature (missing /e2e_otel_span endpoint on weblog)
diff --git a/manifests/java.yml b/manifests/java.yml
@@ -29,6 +29,7 @@ manifest:
     - weblog_declaration:
         "*": irrelevant (just one weblog is enough to test the SDK)
         "spring-boot": v1.61.0-SNAPSHOT
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature (APPSEC-61896)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span:
     - weblog_declaration:
         "*": v0.2.0  # real version not known
diff --git a/manifests/nodejs.yml b/manifests/nodejs.yml
@@ -114,6 +114,7 @@ manifest:
     - weblog_declaration:
         "*": irrelevant (just one weblog is enough to test the SDK)
         express4: *ref_5_89_0
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature (APPSEC-61899)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span: missing_feature (missing /e2e_otel_span endpoint on weblog)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span::test_distributed_otel_trace: irrelevant (Golang specific test with OTel Go contrib package)
   tests/apm_tracing_e2e/test_process_tags.py::Test_Process_Tags: missing_feature
diff --git a/manifests/php.yml b/manifests/php.yml
@@ -8,6 +8,7 @@ manifest:
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDK_Disabled: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SDS_Findings_In_SDK_Response: missing_feature
   tests/ai_guard/test_ai_guard_sdk.py::Test_SensitiveDataScanning: missing_feature
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span: missing_feature (missing /e2e_otel_span endpoint on weblog)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span::test_distributed_otel_trace: irrelevant (Golang specific test with OTel Go contrib package)
   tests/apm_tracing_e2e/test_process_tags.py::Test_Process_Tags: missing_feature
diff --git a/manifests/python.yml b/manifests/python.yml
@@ -35,6 +35,7 @@ manifest:
     - weblog_declaration:
         "*": irrelevant (just one weblog is enough to test the SDK)
         "flask-poc": v4.6.0-rc1
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature (APPSEC-61898)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span: missing_feature (missing /e2e_otel_span endpoint on weblog)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span::test_distributed_otel_trace: irrelevant (Golang specific test with OTel Go contrib package)
   tests/apm_tracing_e2e/test_process_tags.py::Test_Process_Tags: v4.1.0
diff --git a/manifests/ruby.yml b/manifests/ruby.yml
@@ -62,6 +62,7 @@ manifest:
         sinatra32: irrelevant
         sinatra41: irrelevant
         uds-sinatra: irrelevant
+  tests/ai_guard/test_ai_guard_sdk.py::Test_Tag_Probabilities: missing_feature (APPSEC-61897)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span: missing_feature (missing /e2e_otel_span endpoint on weblog)
   tests/apm_tracing_e2e/test_otel.py::Test_Otel_Span::test_distributed_otel_trace: irrelevant (Golang specific test with OTel Go contrib package)
   tests/apm_tracing_e2e/test_process_tags.py::Test_Process_Tags::test_remote_config_process_tags_svc: missing_feature
diff --git a/tests/ai_guard/test_ai_guard_sdk.py b/tests/ai_guard/test_ai_guard_sdk.py
@@ -1,4 +1,5 @@
 import json
+import math
 
 from utils import context, interfaces, scenarios, weblog, features
 from utils.dd_constants import SamplingMechanism, SamplingPriority
@@ -62,6 +63,22 @@ def _assert_key(values: dict, key: str, value: object | None = None):
     return result
 
 
+def _assert_tag_probabilities(values: dict) -> dict:
+    result = _assert_key(values, "tag_probs")
+    assert isinstance(result, dict), f"'tag_probs' should be a dictionary in '{values}'"
+    assert len(result) > 0, f"'tag_probs' should not be empty in '{values}'"
+    return result
+
+
+def _assert_probabilities_match(actual: dict, expected: dict):
+    assert actual.keys() == expected.keys(), f"Mismatched probability keys: {actual.keys()} != {expected.keys()}"
+    for key, expected_value in expected.items():
+        actual_value = actual[key]
+        assert math.isclose(actual_value, expected_value, rel_tol=1e-9, abs_tol=1e-12), (
+            f"Probability mismatch for '{key}': {actual_value} != {expected_value}"
+        )
+
+
 @features.ai_guard
 @scenarios.ai_guard
 class Test_Evaluation:
@@ -250,6 +267,51 @@ def test_evaluation(self):
         )
 
 
+@features.ai_guard
+@scenarios.ai_guard
+class Test_Tag_Probabilities:
+    def _assert_span(self, response: dict):
+        def validate(span: DataDogLibrarySpan):
+            if span["resource"] != "ai_guard":
+                return False
+
+            response_tags = _assert_key(response, "tags")
+            response_tag_probabilities = _assert_tag_probabilities(response)
+            for tag in response_tags:
+                assert tag in response_tag_probabilities, (
+                    f"Missing probability for '{tag}' in {response_tag_probabilities}"
+                )
+                assert response_tag_probabilities[tag] > 0, (
+                    f"Expected a positive probability for '{tag}' in {response_tag_probabilities}"
+                )
+
+            meta_struct = span["meta_struct"]
+            ai_guard = _assert_key(meta_struct, "ai_guard")
+            attack_categories = _assert_key(ai_guard, "attack_categories")
+            assert attack_categories == response_tags, (
+                f"Attack categories do not match the SDK response: {attack_categories} != {response_tags}"
+            )
+
+            span_tag_probabilities = _assert_tag_probabilities(ai_guard)
+            _assert_probabilities_match(span_tag_probabilities, response_tag_probabilities)
+            return True
+
+        return validate
+
+    def setup_tag_probabilities(self):
+        self.messages = MESSAGES["DENY"]
+        self.r = weblog.post("/ai_guard/evaluate", json=self.messages)
+
+    def test_tag_probabilities(self):
+        """Test AI Guard returns and stores tag probabilities.
+        Verifies the SDK response exposes tag probabilities and the ai_guard meta struct keeps the
+        same probability map received from the AI Guard REST API.
+        """
+        assert self.r.status_code == 200
+        body = json.loads(self.r.text)
+        interfaces.library.validate_one_span(self.r, validator=self._assert_span(response=body), full_trace=True)
+
+
 @features.ai_guard
 @scenarios.default
 class Test_SDK_Disabled:
diff --git a/utils/build/docker/python/flask/app.py b/utils/build/docker/python/flask/app.py
@@ -2193,14 +2193,16 @@ def ai_guard_evaluate():
 
     except Exception as e:
         if isinstance(e, AIGuardAbortError):
-            return jsonify(
-                {
-                    "action": getattr(e, "action", ""),
-                    "reason": getattr(e, "reason", ""),
-                    "tags": getattr(e, "tags", []),
-                    "sds": getattr(e, "sds", []),
-                }
-            ), 403
+            error_response = {
+                "action": getattr(e, "action", ""),
+                "reason": getattr(e, "reason", ""),
+                "tags": getattr(e, "tags", []),
+                "sds": getattr(e, "sds", []),
+            }
+            tag_probs = getattr(e, "tag_probs", None)
+            if tag_probs is not None:
+                error_response["tag_probs"] = tag_probs
+            return jsonify(error_response), 403
         else:
             return jsonify({"error": str(e), "type": e.__class__.__name__}), 500
 
diff --git a/utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb
@@ -29,10 +29,12 @@ def evaluate
       tags: result.tags,
       is_blocking_enabled: result.blocking_enabled?
     }
+    response_data[:tag_probs] = result.tag_probs if result.respond_to?(:tag_probs)
     response_data[:sds] = result.sds if result.respond_to?(:sds)
     render json: response_data
   rescue Datadog::AIGuard::AIGuardAbortError => e
     error_data = { action: e.action, reason: e.reason, tags: e.tags }
+    error_data[:tag_probs] = e.tag_probs if e.respond_to?(:tag_probs)
     error_data[:sds] = e.sds if e.respond_to?(:sds)
     render json: error_data, status: 403
   rescue => e
diff --git a/utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb
@@ -29,10 +29,12 @@ def evaluate
       tags: result.tags,
       is_blocking_enabled: result.blocking_enabled?
     }
+    response_data[:tag_probs] = result.tag_probs if result.respond_to?(:tag_probs)
     response_data[:sds] = result.sds if result.respond_to?(:sds)
     render json: response_data
   rescue Datadog::AIGuard::AIGuardAbortError => e
     error_data = { action: e.action, reason: e.reason, tags: e.tags }
+    error_data[:tag_probs] = e.tag_probs if e.respond_to?(:tag_probs)
     error_data[:sds] = e.sds if e.respond_to?(:sds)
     render json: error_data, status: 403
   rescue => e
diff --git a/utils/build/docker/ruby/rails72/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails72/app/controllers/ai_guard_controller.rb
@@ -29,10 +29,12 @@ def evaluate
       tags: result.tags,
       is_blocking_enabled: result.blocking_enabled?
     }
+    response_data[:tag_probs] = result.tag_probs if result.respond_to?(:tag_probs)
     response_data[:sds] = result.sds if result.respond_to?(:sds)
     render json: response_data
   rescue Datadog::AIGuard::AIGuardAbortError => e
     error_data = { action: e.action, reason: e.reason, tags: e.tags }
+    error_data[:tag_probs] = e.tag_probs if e.respond_to?(:tag_probs)
     error_data[:sds] = e.sds if e.respond_to?(:sds)
     render json: error_data, status: 403
   rescue => e
diff --git a/utils/build/docker/ruby/rails80/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails80/app/controllers/ai_guard_controller.rb
@@ -30,10 +30,12 @@ def evaluate
       tags: result.tags,
       is_blocking_enabled: result.blocking_enabled?
     }
+    response_data[:tag_probs] = result.tag_probs if result.respond_to?(:tag_probs)
     response_data[:sds] = result.sds if result.respond_to?(:sds)
     render json: response_data
   rescue Datadog::AIGuard::AIGuardAbortError => e
     error_data = { action: e.action, reason: e.reason, tags: e.tags }
+    error_data[:tag_probs] = e.tag_probs if e.respond_to?(:tag_probs)
     error_data[:sds] = e.sds if e.respond_to?(:sds)
     render json: error_data, status: 403
   rescue => e

Original file line number	Diff line number	Diff line change
`@@ -1236,7 +1236,10 @@ Successful evaluation:`
`1236`	`1236`	`{`
`1237`	`1237`	`"action": "ALLOW",`
`1238`	`1238`	`"reason": "All looks good",`
`1239`		`- "tags": []`
	`1239`	`+ "tags": [],`
	`1240`	`+ "tag_probs": {`
	`1241`	`+ "jailbreak": 0.0`
	`1242`	`+ }`
`1240`	`1243`	`}`
`1241`	`1244`	```
`1242`	`1245`