diff --git a/docs/understand/scenarios/ai_guard.md b/docs/understand/scenarios/ai_guard.md index 38255999a06..32919c9a182 100644 --- 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 index ae114f8c750..c5fb2a076cc 100644 --- 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 index 8d7c6ba1fd5..a13295ea852 100644 --- 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 index 536dde96fc2..94e62bf0d21 100644 --- 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 index e49fb25082f..eed21016dff 100644 --- 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 index 6fb9dc15273..14ad2a523a0 100644 --- 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 index f502d1b3e06..6b7d16f7342 100644 --- 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 index 13a553b27e5..b518e40d28b 100644 --- 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 index 7d80cea05af..4899dcbf71f 100644 --- 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 index e0e054a183b..8e652765eeb 100644 --- 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 index db08f4c4cab..331c2d117e5 100644 --- 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 index 7813044d2c9..c736d4d4237 100644 --- 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 index 06b53c3cad8..1b34e92e2cb 100644 --- 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 index 24de2a409c9..9548b00f671 100644 --- 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 index 24de2a409c9..9548b00f671 100644 --- 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 index 24de2a409c9..9548b00f671 100644 --- 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 index 1ef7fef61d1..d73df32735c 100644 --- 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