diff --git a/format.sh b/format.sh index 6f634787937..d8ed6b1f2aa 100755 --- a/format.sh +++ b/format.sh @@ -210,6 +210,14 @@ if ! gen_redaction_output="$(python utils/scripts/gen_redaction_cassettes.py $ge exit 1 fi +echo "Checking AI Guard extra_tags fixtures..." +if ! gen_extra_tags_output="$(python utils/scripts/gen_extra_tags_cassettes.py $gen_redaction_args 2>&1)"; then + echo "$gen_extra_tags_output" + echo "AI Guard extra_tags fixtures are out of date or invalid. Regenerate them with" + echo "python utils/scripts/gen_extra_tags_cassettes.py 💥 💔 💥" + exit 1 +fi + echo "Running shellcheck checks..." if ! ./utils/scripts/shellcheck.sh; then echo "shellcheck checks failed. Please fix the errors above. 💥 💔 💥" diff --git a/manifests/cpp_httpd.yml b/manifests/cpp_httpd.yml index a39091927ae..be228ae9558 100644 --- a/manifests/cpp_httpd.yml +++ b/manifests/cpp_httpd.yml @@ -21,6 +21,8 @@ manifest: 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 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse: missing_feature diff --git a/manifests/cpp_nginx.yml b/manifests/cpp_nginx.yml index fd2f877a7bc..6e39341455b 100644 --- a/manifests/cpp_nginx.yml +++ b/manifests/cpp_nginx.yml @@ -19,6 +19,8 @@ manifest: 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 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse: missing_feature diff --git a/manifests/dotnet.yml b/manifests/dotnet.yml index 7b197898c10..1e4c311591f 100644 --- a/manifests/dotnet.yml +++ b/manifests/dotnet.yml @@ -11,6 +11,8 @@ manifest: 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 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse: missing_feature diff --git a/manifests/golang.yml b/manifests/golang.yml index 3c8d348a356..51bc6ac27a5 100644 --- a/manifests/golang.yml +++ b/manifests/golang.yml @@ -11,6 +11,8 @@ manifest: 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 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse: missing_feature diff --git a/manifests/java.yml b/manifests/java.yml index a93333e7f8f..1fa8227e037 100644 --- a/manifests/java.yml +++ b/manifests/java.yml @@ -47,6 +47,8 @@ manifest: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) "spring-boot": v1.54.0 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) diff --git a/manifests/nodejs.yml b/manifests/nodejs.yml index d7acdfbdbd7..9d8050fe990 100644 --- a/manifests/nodejs.yml +++ b/manifests/nodejs.yml @@ -152,6 +152,8 @@ manifest: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) express5: *ref_5_69_0 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) diff --git a/manifests/php.yml b/manifests/php.yml index 7f61142f192..2f03c1e7c48 100644 --- a/manifests/php.yml +++ b/manifests/php.yml @@ -11,6 +11,8 @@ manifest: 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 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_NoRedaction: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_RedactedMessagesInSDKResponse: missing_feature diff --git a/manifests/python.yml b/manifests/python.yml index 8cce8aec05b..80a995718fc 100644 --- a/manifests/python.yml +++ b/manifests/python.yml @@ -60,6 +60,14 @@ manifest: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) "flask-poc": v3.18.1 + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: + - weblog_declaration: + "*": irrelevant (just one weblog is enough to test the SDK) + "flask-poc": v4.17.0-dev + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: + - weblog_declaration: + "*": irrelevant (just one weblog is enough to test the SDK) + "flask-poc": v4.17.0-dev tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) diff --git a/manifests/ruby.yml b/manifests/ruby.yml index 53c45d32ea4..b091d5e3630 100644 --- a/manifests/ruby.yml +++ b/manifests/ruby.yml @@ -88,6 +88,8 @@ manifest: sinatra32: irrelevant sinatra41: irrelevant uds-sinatra: irrelevant + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTags: missing_feature (APPSEC-70724) + tests/ai_guard/test_ai_guard_sdk.py::Test_ExtraTagsFailSafe: missing_feature (APPSEC-70724) tests/ai_guard/test_ai_guard_sdk.py::Test_Full_Response_And_Tags: - weblog_declaration: "*": v2.25.0-dev diff --git a/tests/ai_guard/extra_tags_scenarios.json b/tests/ai_guard/extra_tags_scenarios.json new file mode 100644 index 00000000000..d3b51ec3654 --- /dev/null +++ b/tests/ai_guard/extra_tags_scenarios.json @@ -0,0 +1,360 @@ +{ + "EXTRA_TAGS": { + "evaluation_id": "13c92662-52c8-5959-8fec-fb7d6bb5b5d0", + "messages": [ + { + "role": "user", + "content": "Summarize the latest newsroom style guide changes." + } + ], + "action": "ALLOW", + "tags": [], + "extra_tags": [ + { + "key": "ai_guard.policy_name", + "value": "newsroom-default" + }, + { + "key": "ai_guard.policy_revision", + "value": "7" + } + ], + "expected_tags": { + "ai_guard.policy_name": "newsroom-default", + "ai_guard.policy_revision": "7" + }, + "expected_metastruct": {}, + "skipped_tags": [], + "cassette": "aiguard_evaluate_post_f0ef2bd8.json" + }, + "EXTRA_METASTRUCT": { + "evaluation_id": "73606e72-26b6-52f8-8652-7f321db98dda", + "messages": [ + { + "role": "user", + "content": "Draft a headline for the city council budget story." + } + ], + "action": "ALLOW", + "tags": [], + "extra_metastruct": [ + { + "key": "policy_match", + "value": { + "rule": "credit_card", + "categories": [ + "pii", + "financial" + ] + } + }, + { + "key": "matched_rules", + "value": [ + { + "id": "r1", + "action": "DENY" + }, + { + "id": "r2", + "action": "ALLOW" + } + ] + }, + { + "key": "policy_version", + "value": "v3" + } + ], + "expected_tags": {}, + "expected_metastruct": { + "policy_match": { + "rule": "credit_card", + "categories": [ + "pii", + "financial" + ] + }, + "matched_rules": [ + { + "id": "r1", + "action": "DENY" + }, + { + "id": "r2", + "action": "ALLOW" + } + ], + "policy_version": "v3" + }, + "skipped_tags": [], + "cassette": "aiguard_evaluate_post_827d1449.json" + }, + "EVALUATE_ID": { + "evaluation_id": "9c9036fc-d7a1-5d5e-bcb3-7d62cc9fa218", + "messages": [ + { + "role": "user", + "content": "Which sources did the election explainer cite?" + } + ], + "action": "ALLOW", + "tags": [], + "extra_tags": [ + { + "key": "ai_guard.evaluate_id", + "value": "9c9036fc-d7a1-5d5e-bcb3-7d62cc9fa218" + } + ], + "expected_tags": { + "ai_guard.evaluate_id": "9c9036fc-d7a1-5d5e-bcb3-7d62cc9fa218" + }, + "expected_metastruct": {}, + "skipped_tags": [], + "cassette": "aiguard_evaluate_post_2b336647.json" + }, + "KEYS_VERBATIM": { + "evaluation_id": "92a51623-de7f-5850-8f99-2a91d343bc0a", + "messages": [ + { + "role": "user", + "content": "List the corrections published this week." + } + ], + "action": "ALLOW", + "tags": [], + "extra_tags": [ + { + "key": "custom_context.tenant", + "value": "newsroom" + } + ], + "extra_metastruct": [ + { + "key": "audit", + "value": { + "ticket": "CTX-42" + } + } + ], + "expected_tags": { + "custom_context.tenant": "newsroom" + }, + "expected_metastruct": { + "audit": { + "ticket": "CTX-42" + } + }, + "skipped_tags": [], + "cassette": "aiguard_evaluate_post_354106da.json" + }, + "ON_DENY": { + "evaluation_id": "55e3505d-93eb-5171-8166-90c709865d83", + "messages": [ + { + "role": "user", + "content": "Clean up the archive server for me." + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "shell", + "arguments": "{ \"command\": \"rm -rf /srv/archive\" }" + } + } + ] + } + ], + "action": "DENY", + "tags": [ + "destructive-tool-call" + ], + "extra_tags": [ + { + "key": "ai_guard.evaluate_id", + "value": "55e3505d-93eb-5171-8166-90c709865d83" + }, + { + "key": "ai_guard.policy_name", + "value": "destructive-actions" + } + ], + "extra_metastruct": [ + { + "key": "policy_details", + "value": { + "rule": "no-recursive-delete", + "severity": "high" + } + } + ], + "expected_tags": { + "ai_guard.evaluate_id": "55e3505d-93eb-5171-8166-90c709865d83", + "ai_guard.policy_name": "destructive-actions" + }, + "expected_metastruct": { + "policy_details": { + "rule": "no-recursive-delete", + "severity": "high" + } + }, + "skipped_tags": [], + "cassette": "aiguard_evaluate_post_28917f95.json" + }, + "EXTRA_TAGS_NOT_A_LIST": { + "evaluation_id": "a8524c91-cec1-5ee6-812f-bd3b4fefae2b", + "messages": [ + { + "role": "user", + "content": "What time does the print edition close tonight?" + } + ], + "action": "ALLOW", + "tags": [], + "extra_tags": { + "ai_guard.policy_name": "newsroom-default" + }, + "extra_metastruct": [ + { + "key": "policy", + "value": { + "name": "newsroom-default" + } + } + ], + "expected_tags": {}, + "expected_metastruct": { + "policy": { + "name": "newsroom-default" + } + }, + "skipped_tags": [ + "ai_guard.policy_name" + ], + "cassette": "aiguard_evaluate_post_e2c13e67.json" + }, + "EXTRA_METASTRUCT_NOT_A_LIST": { + "evaluation_id": "e72731f7-4c56-56b4-95a9-ed8d009cebcb", + "messages": [ + { + "role": "user", + "content": "Who is on call for the overnight desk?" + } + ], + "action": "ALLOW", + "tags": [], + "extra_tags": [ + { + "key": "ai_guard.policy_name", + "value": "newsroom-default" + } + ], + "extra_metastruct": { + "policy": { + "name": "newsroom-default" + } + }, + "expected_tags": { + "ai_guard.policy_name": "newsroom-default" + }, + "expected_metastruct": {}, + "skipped_tags": [], + "cassette": "aiguard_evaluate_post_bf6868f9.json" + }, + "MALFORMED_ENTRIES": { + "evaluation_id": "f55f1430-86f7-5441-96b5-9938f4a10261", + "messages": [ + { + "role": "user", + "content": "When does the weekend magazine go to print?" + } + ], + "action": "ALLOW", + "tags": [], + "extra_tags": [ + "ai_guard.policy_name=newsroom-default", + [ + "ai_guard.policy_name", + "newsroom-default" + ], + { + "value": "missing key" + }, + { + "key": 1, + "value": "integer key" + }, + { + "key": "", + "value": "empty key" + }, + { + "key": "ai_guard.policy_revision", + "value": 7 + }, + { + "key": "ai_guard.policy_details", + "value": { + "rule": "embargo-check" + } + }, + { + "key": "ai_guard.policy_enforced", + "value": true + }, + { + "key": "ai_guard.policy_name", + "value": "newsroom-default" + } + ], + "extra_metastruct": [ + "policy=newsroom-default", + { + "value": { + "name": "missing key" + } + }, + { + "key": 1, + "value": { + "name": "integer key" + } + }, + { + "key": "", + "value": { + "name": "empty key" + } + }, + { + "key": "policy_note", + "value": null + }, + { + "key": "policy", + "value": { + "name": "newsroom-default" + } + } + ], + "expected_tags": { + "ai_guard.policy_name": "newsroom-default" + }, + "expected_metastruct": { + "policy": { + "name": "newsroom-default" + } + }, + "skipped_tags": [ + "ai_guard.policy_revision", + "ai_guard.policy_details", + "ai_guard.policy_enforced", + "1", + "" + ], + "cassette": "aiguard_evaluate_post_19cfc9ce.json" + } +} diff --git a/tests/ai_guard/test_ai_guard_sdk.py b/tests/ai_guard/test_ai_guard_sdk.py index 3d06752c00d..717d80eab2b 100644 --- a/tests/ai_guard/test_ai_guard_sdk.py +++ b/tests/ai_guard/test_ai_guard_sdk.py @@ -1299,6 +1299,192 @@ def test_ai_guard_event(self): interfaces.library.validate_one_trace(self.r, validator=self._assert_trace) +# Extra tags scenarios and the matching VCR cassettes are generated together by +# utils/scripts/gen_extra_tags_cassettes.py: each entry keeps the messages we send, the exact +# extra_tags and extra_metastruct values the backend returns, and what the tracer must apply from +# them. Regenerate both after editing. +EXTRA_TAGS_SCENARIOS: dict = json.loads((Path(__file__).parent / "extra_tags_scenarios.json").read_text()) + +# What the tracer itself writes to the ai_guard meta struct. Anything else there came from extra_metastruct. +TRACER_META_STRUCT_KEYS = frozenset(("messages", "attack_categories", "sds", "tag_probs")) + + +def _post_extra_tags_scenario(scenario_key: str, *, block: str = "false") -> tuple[dict, HttpResponse]: + scenario = EXTRA_TAGS_SCENARIOS[scenario_key] + response = weblog.post("/ai_guard/evaluate", headers={BLOCKING_HEADER: block}, json=scenario["messages"]) + return scenario, response + + +def _assert_extra_tags_applied(scenario: dict): + """Validator asserting the backend-provided entries landed exactly where the contract puts them. + + Each extra_tags entry is a string tag on the ai_guard span and each extra_metastruct entry is a + value in the ai_guard meta struct, both under the key exactly as returned: the backend owns the + ai_guard prefix, so the tracer neither adds it nor filters keys by namespace. Nothing crosses + over: a tag never reaches the meta struct, and a meta struct value is never flattened into a tag. + """ + expected_tags: dict[str, str] = scenario["expected_tags"] + expected_metastruct: dict = scenario["expected_metastruct"] + + def validate(span: DataDogLibrarySpan): + if span["resource"] != "ai_guard": + return False + + meta = span.get("meta") or {} + for key, value in expected_tags.items(): + assert meta.get(key) == value, ( + f"Extra tag '{key}' should be '{value}' on the ai_guard span, got '{meta.get(key)}'" + ) + assert _span_tag(span, f"ai_guard.{key}") is None, f"The tracer must not prefix extra tag '{key}'" + + ai_guard = _assert_key(span["meta_struct"], "ai_guard") + for key, value in expected_metastruct.items(): + assert key in ai_guard, f"Extra meta struct key '{key}' not found in the ai_guard meta struct {ai_guard}" + assert ai_guard[key] == value, f"Extra meta struct key '{key}': {ai_guard[key]} != {value}" + assert f"ai_guard.{key}" not in ai_guard, f"The tracer must not prefix extra meta struct key '{key}'" + assert _span_tag(span, key) is None, f"Extra meta struct key '{key}' should not be a span tag" + + unexpected = set(ai_guard) - TRACER_META_STRUCT_KEYS - set(expected_metastruct) + assert not unexpected, f"Unexpected keys in the ai_guard meta struct: {unexpected}" + + for key in scenario["skipped_tags"]: + assert _span_tag(span, key) is None, ( + f"Malformed extra tag '{key}' should be skipped, got '{_span_tag(span, key)}'" + ) + return True + + return validate + + +@rfc("https://datadoghq.atlassian.net/browse/APPSEC-70724") +@features.ai_guard +@scenarios.ai_guard +class Test_ExtraTags: + """The AI Guard backend enriches the ai_guard span through extra_tags and extra_metastruct. + + Lets an AI Guard policy attach its own context to traces without a tracer change per tag. Both + attributes of the evaluate response are lists of {key, value} objects. + """ + + def setup_extra_tags(self): + self.scenario, self.r = _post_extra_tags_scenario("EXTRA_TAGS") + + def test_extra_tags(self): + """Each extra_tags entry is set as a string tag on the ai_guard span.""" + assert self.r.status_code == 200 + interfaces.library.validate_one_span( + self.r, validator=_assert_extra_tags_applied(self.scenario), full_trace=True + ) + + def setup_extra_metastruct(self): + self.scenario, self.r = _post_extra_tags_scenario("EXTRA_METASTRUCT") + + def test_extra_metastruct(self): + """Each extra_metastruct entry, object, array or scalar, is added to the ai_guard meta struct.""" + assert self.r.status_code == 200 + interfaces.library.validate_one_span( + self.r, validator=_assert_extra_tags_applied(self.scenario), full_trace=True + ) + + def setup_evaluate_id(self): + self.scenario, self.r = _post_extra_tags_scenario("EVALUATE_ID") + + def test_evaluate_id(self): + """The evaluation ID lands only as the ai_guard.evaluate_id tag, correlating the span with the backend.""" + assert self.r.status_code == 200 + evaluation_id = self.scenario["evaluation_id"] + assert self.scenario["expected_tags"] == {"ai_guard.evaluate_id": evaluation_id} + assert self.scenario["expected_metastruct"] == {} + extra_tags_validator = _assert_extra_tags_applied(self.scenario) + + def validate(span: DataDogLibrarySpan): + if not extra_tags_validator(span): + return False + _assert_key(span["meta"], "ai_guard.evaluate_id", evaluation_id) + ai_guard = _assert_key(span["meta_struct"], "ai_guard") + for key in ("evaluate_id", "ai_guard.evaluate_id"): + assert key not in ai_guard, f"The evaluation ID must not be in the ai_guard meta struct: {ai_guard}" + return True + + interfaces.library.validate_one_span(self.r, validator=validate, full_trace=True) + + def setup_keys_not_prefixed_or_filtered(self): + self.scenario, self.r = _post_extra_tags_scenario("KEYS_VERBATIM") + + def test_keys_not_prefixed_or_filtered(self): + """Keys land verbatim, even outside the ai_guard namespace: namespacing is the backend's job.""" + assert self.r.status_code == 200 + interfaces.library.validate_one_span( + self.r, validator=_assert_extra_tags_applied(self.scenario), full_trace=True + ) + + def setup_extra_tags_on_block(self): + self.scenario, self.r = _post_extra_tags_scenario("ON_DENY", block="true") + + def test_extra_tags_on_block(self): + """A blocked evaluation still carries both kinds of entries, alongside ai_guard.blocked.""" + assert self.r.status_code == 403 + extra_tags_validator = _assert_extra_tags_applied(self.scenario) + + def validate(span: DataDogLibrarySpan): + if not extra_tags_validator(span): + return False + _assert_key(span["meta"], "ai_guard.action", "DENY") + assert is_same_boolean(actual=span["meta"].get("ai_guard.blocked"), expected="true"), ( + f"'ai_guard.blocked' with value 'true' not found in '{span['meta']}'" + ) + return True + + interfaces.library.validate_one_span(self.r, validator=validate, full_trace=True) + + +@rfc("https://datadoghq.atlassian.net/browse/APPSEC-70724") +@features.ai_guard +@scenarios.ai_guard +class Test_ExtraTagsFailSafe: + """A malformed extra_tags or extra_metastruct is skipped and never changes the evaluation. + + Both fields only enrich the span, so a shape the tracer does not understand must leave the + action, the SDK response and the span as they would be without it, while every well-formed + part of the response is still applied. + """ + + def _assert_skipped(self, response: HttpResponse, scenario: dict) -> None: + assert response.status_code == 200 + _assert_key(json.loads(response.text), "action", scenario["action"]) + extra_tags_validator = _assert_extra_tags_applied(scenario) + + def validate(span: DataDogLibrarySpan): + if not extra_tags_validator(span): + return False + assert span.get("error", 0) == 0, f"A malformed payload must not flag the span as an error: {span}" + _assert_key(span["meta"], "ai_guard.action", scenario["action"]) + return True + + interfaces.library.validate_one_span(response, validator=validate, full_trace=True) + + def setup_extra_tags_not_a_list(self): + self.scenario, self.r = _post_extra_tags_scenario("EXTRA_TAGS_NOT_A_LIST") + + def test_extra_tags_not_a_list(self): + """An extra_tags object instead of a list is skipped, and extra_metastruct is still applied.""" + self._assert_skipped(self.r, self.scenario) + + def setup_extra_metastruct_not_a_list(self): + self.scenario, self.r = _post_extra_tags_scenario("EXTRA_METASTRUCT_NOT_A_LIST") + + def test_extra_metastruct_not_a_list(self): + """An extra_metastruct object instead of a list is skipped, and extra_tags is still applied.""" + self._assert_skipped(self.r, self.scenario) + + def setup_malformed_entries(self): + self.scenario, self.r = _post_extra_tags_scenario("MALFORMED_ENTRIES") + + def test_malformed_entries(self): + """Malformed entries and non-string extra_tags values are skipped one by one, the rest still lands.""" + self._assert_skipped(self.r, self.scenario) + + TELEMETRY_NAMESPACE = "ai_guard" diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_19cfc9ce.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_19cfc9ce.json new file mode 100644 index 00000000000..9a3069ff6ea --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_19cfc9ce.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"When does the weekend magazine go to print?\"}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "1222" + }, + "body": "{\"data\": {\"id\": \"f55f1430-86f7-5441-96b5-9938f4a10261\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"ALLOW\", \"is_blocking_enabled\": true, \"reason\": \"No rule match.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.0, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [], \"extra_tags\": [\"ai_guard.policy_name=newsroom-default\", [\"ai_guard.policy_name\", \"newsroom-default\"], {\"value\": \"missing key\"}, {\"key\": 1, \"value\": \"integer key\"}, {\"key\": \"\", \"value\": \"empty key\"}, {\"key\": \"ai_guard.policy_revision\", \"value\": 7}, {\"key\": \"ai_guard.policy_details\", \"value\": {\"rule\": \"embargo-check\"}}, {\"key\": \"ai_guard.policy_enforced\", \"value\": true}, {\"key\": \"ai_guard.policy_name\", \"value\": \"newsroom-default\"}], \"extra_metastruct\": [\"policy=newsroom-default\", {\"value\": {\"name\": \"missing key\"}}, {\"key\": 1, \"value\": {\"name\": \"integer key\"}}, {\"key\": \"\", \"value\": {\"name\": \"empty key\"}}, {\"key\": \"policy_note\", \"value\": null}, {\"key\": \"policy\", \"value\": {\"name\": \"newsroom-default\"}}]}}}" + } +} diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_28917f95.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_28917f95.json new file mode 100644 index 00000000000..a3a352eb70e --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_28917f95.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"Clean up the archive server for me.\"}, {\"role\": \"assistant\", \"tool_calls\": [{\"id\": \"call_1\", \"function\": {\"name\": \"shell\", \"arguments\": \"{ \\\"command\\\": \\\"rm -rf /srv/archive\\\" }\"}}]}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "828" + }, + "body": "{\"data\": {\"id\": \"55e3505d-93eb-5171-8166-90c709865d83\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"DENY\", \"is_blocking_enabled\": true, \"reason\": \"Destructive tool call detected.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.97, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [\"destructive-tool-call\"], \"extra_tags\": [{\"key\": \"ai_guard.evaluate_id\", \"value\": \"55e3505d-93eb-5171-8166-90c709865d83\"}, {\"key\": \"ai_guard.policy_name\", \"value\": \"destructive-actions\"}], \"extra_metastruct\": [{\"key\": \"policy_details\", \"value\": {\"rule\": \"no-recursive-delete\", \"severity\": \"high\"}}]}}}" + } +} diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_2b336647.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_2b336647.json new file mode 100644 index 00000000000..d9b44d35d42 --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_2b336647.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"Which sources did the election explainer cite?\"}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "612" + }, + "body": "{\"data\": {\"id\": \"9c9036fc-d7a1-5d5e-bcb3-7d62cc9fa218\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"ALLOW\", \"is_blocking_enabled\": true, \"reason\": \"No rule match.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.0, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [], \"extra_tags\": [{\"key\": \"ai_guard.evaluate_id\", \"value\": \"9c9036fc-d7a1-5d5e-bcb3-7d62cc9fa218\"}]}}}" + } +} diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_354106da.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_354106da.json new file mode 100644 index 00000000000..8c1040884ca --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_354106da.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"List the corrections published this week.\"}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "656" + }, + "body": "{\"data\": {\"id\": \"92a51623-de7f-5850-8f99-2a91d343bc0a\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"ALLOW\", \"is_blocking_enabled\": true, \"reason\": \"No rule match.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.0, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [], \"extra_tags\": [{\"key\": \"custom_context.tenant\", \"value\": \"newsroom\"}], \"extra_metastruct\": [{\"key\": \"audit\", \"value\": {\"ticket\": \"CTX-42\"}}]}}}" + } +} diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_827d1449.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_827d1449.json new file mode 100644 index 00000000000..3ee1eee7dc9 --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_827d1449.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"Draft a headline for the city council budget story.\"}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "775" + }, + "body": "{\"data\": {\"id\": \"73606e72-26b6-52f8-8652-7f321db98dda\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"ALLOW\", \"is_blocking_enabled\": true, \"reason\": \"No rule match.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.0, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [], \"extra_metastruct\": [{\"key\": \"policy_match\", \"value\": {\"rule\": \"credit_card\", \"categories\": [\"pii\", \"financial\"]}}, {\"key\": \"matched_rules\", \"value\": [{\"id\": \"r1\", \"action\": \"DENY\"}, {\"id\": \"r2\", \"action\": \"ALLOW\"}]}, {\"key\": \"policy_version\", \"value\": \"v3\"}]}}}" + } +} diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_bf6868f9.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_bf6868f9.json new file mode 100644 index 00000000000..6c4c92ca100 --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_bf6868f9.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"Who is on call for the overnight desk?\"}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "654" + }, + "body": "{\"data\": {\"id\": \"e72731f7-4c56-56b4-95a9-ed8d009cebcb\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"ALLOW\", \"is_blocking_enabled\": true, \"reason\": \"No rule match.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.0, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [], \"extra_tags\": [{\"key\": \"ai_guard.policy_name\", \"value\": \"newsroom-default\"}], \"extra_metastruct\": {\"policy\": {\"name\": \"newsroom-default\"}}}}}" + } +} diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_e2c13e67.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_e2c13e67.json new file mode 100644 index 00000000000..1e34054357e --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_e2c13e67.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"What time does the print edition close tonight?\"}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "654" + }, + "body": "{\"data\": {\"id\": \"a8524c91-cec1-5ee6-812f-bd3b4fefae2b\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"ALLOW\", \"is_blocking_enabled\": true, \"reason\": \"No rule match.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.0, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [], \"extra_tags\": {\"ai_guard.policy_name\": \"newsroom-default\"}, \"extra_metastruct\": [{\"key\": \"policy\", \"value\": {\"name\": \"newsroom-default\"}}]}}}" + } +} diff --git a/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_f0ef2bd8.json b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_f0ef2bd8.json new file mode 100644 index 00000000000..73e7dce2329 --- /dev/null +++ b/utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_f0ef2bd8.json @@ -0,0 +1,22 @@ +{ + "request": { + "method": "POST", + "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate", + "headers": { + "Content-Type": "application/json", + "DD-AI-GUARD-SOURCE": "SDK" + }, + "body": "{\"data\": {\"attributes\": {\"messages\": [{\"role\": \"user\", \"content\": \"Summarize the latest newsroom style guide changes.\"}], \"meta\": {\"service\": \"weblog\", \"env\": \"system-tests\"}}}}" + }, + "response": { + "status": { + "code": 200, + "message": "OK" + }, + "headers": { + "content-type": "application/vnd.api+json", + "content-length": "643" + }, + "body": "{\"data\": {\"id\": \"13c92662-52c8-5959-8fec-fb7d6bb5b5d0\", \"type\": \"evaluations\", \"attributes\": {\"action\": \"ALLOW\", \"is_blocking_enabled\": true, \"reason\": \"No rule match.\", \"sds_findings\": [], \"tag_probs\": {\"authority-override\": 0.0, \"data-exfiltration\": 0.0, \"denial-of-service-tool-call\": 0.0, \"destructive-tool-call\": 0.0, \"indirect-prompt-injection\": 0.0, \"instruction-override\": 0.0, \"jailbreak\": 0.0, \"obfuscation\": 0.0, \"role-play\": 0.0, \"security-exploit\": 0.0, \"system-prompt-extraction\": 0.0}, \"tags\": [], \"extra_tags\": [{\"key\": \"ai_guard.policy_name\", \"value\": \"newsroom-default\"}, {\"key\": \"ai_guard.policy_revision\", \"value\": \"7\"}]}}}" + } +} diff --git a/utils/scripts/gen_extra_tags_cassettes.py b/utils/scripts/gen_extra_tags_cassettes.py new file mode 100755 index 00000000000..3a05d8e2df2 --- /dev/null +++ b/utils/scripts/gen_extra_tags_cassettes.py @@ -0,0 +1,308 @@ +#!/usr/bin/env python3 +"""Generate hand-crafted AI Guard VCR cassettes for backend-provided extra_tags and extra_metastruct. + +The evaluate response can carry two optional attributes so an AI Guard policy can enrich the +ai_guard span without a tracer release per tag (APPSEC-70724). Both are lists of {key, value} +objects rather than maps, because dotted keys such as ai_guard.evaluate_id are not valid JSON:API +member names: + + - extra_tags holds string values, each set as an ai_guard span tag under the key exactly as + returned; + - extra_metastruct holds any value (object, array or scalar), each added to the ai_guard meta + struct under the key exactly as returned. + +The backend owns the ai_guard prefix and the namespace, so tracers never add the prefix or filter +keys. The backend sends the evaluation ID (the same value as data.id) only in extra_tags, as +ai_guard.evaluate_id; extra_metastruct is empty, and so omitted, in today's responses. A malformed +field or entry is skipped and never changes the evaluation. + +The backend does not return these attributes in the recorded cassettes, so this script writes them +by hand, reusing the request hashing and rendering of gen_redaction_cassettes.py, and writes +tests/ai_guard/extra_tags_scenarios.json, the sidecar the tests read the scenarios from. The two +are generated together because a cassette is addressed by the hash of its request body: a scenario +whose messages drift from its cassette would never be replayed again. + +Run from the repo root: python3 utils/scripts/gen_extra_tags_cassettes.py +Pass --check to report drift without writing anything (used by format.sh --check). +""" + +from __future__ import annotations + +import json +import sys +import uuid +from pathlib import Path + +from gen_redaction_cassettes import ( + CASSETTES_DIR, + FOREIGN_CASSETTES, + NO_MATCH_REASON, + _cassette_name, + _tag_probs_for, + _wrap_request_body, + render_cassette, +) + +SIDECAR = Path("tests/ai_guard/extra_tags_scenarios.json") +REDACTION_SIDECAR = Path("tests/ai_guard/redaction_scenarios.json") + +DENY_REASON = "Destructive tool call detected." + +# Fixed namespace so every scenario gets a stable evaluation ID (data.id) across regenerations. +EVALUATION_ID_NAMESPACE = uuid.UUID("6f0e8c1a-3b9d-4c55-9a51-0c2f4e7d8b13") + + +def _evaluation_id(name: str) -> str: + return str(uuid.uuid5(EVALUATION_ID_NAMESPACE, name)) + + +def _kv(**entries: object) -> list[dict]: + return [{"key": key, "value": value} for key, value in entries.items()] + + +def _evaluate_id_tag(name: str) -> list[dict]: + """The extra_tags entry carrying the evaluation ID, the only place the backend sends it.""" + return [{"key": "ai_guard.evaluate_id", "value": _evaluation_id(name)}] + + +# Every scenario produces one cassette and one sidecar entry: +# messages what the weblog POSTs, and therefore what the cassette is keyed on +# action the evaluation outcome the backend returns +# tags the attack categories the backend returns (a DENY needs at least one) +# extra_tags the extra_tags value the backend returns, verbatim (omitted when absent) +# extra_metastruct the extra_metastruct value the backend returns, verbatim (omitted when absent) +# expected_tags, expected_metastruct, skipped_tags +# what the tracer must apply, only spelled out for the fail-safe scenarios, whose +# payloads are malformed on purpose; for the others it is derived from the lists +SCENARIOS: dict[str, dict] = { + "EXTRA_TAGS": { + "messages": [{"role": "user", "content": "Summarize the latest newsroom style guide changes."}], + "action": "ALLOW", + "tags": [], + "extra_tags": _kv(**{"ai_guard.policy_name": "newsroom-default", "ai_guard.policy_revision": "7"}), + }, + # The illustrative extra_metastruct from the backend contract: an object, an array of objects and + # a scalar. No null anywhere and no boolean directly inside an array, shapes the backend avoids + # and that some tracers' meta struct encoders drop. + "EXTRA_METASTRUCT": { + "messages": [{"role": "user", "content": "Draft a headline for the city council budget story."}], + "action": "ALLOW", + "tags": [], + "extra_metastruct": _kv( + policy_match={"rule": "credit_card", "categories": ["pii", "financial"]}, + matched_rules=[{"id": "r1", "action": "DENY"}, {"id": "r2", "action": "ALLOW"}], + policy_version="v3", + ), + }, + # Mirrors today's backend response: the evaluation ID in extra_tags only, no extra_metastruct. + "EVALUATE_ID": { + "messages": [{"role": "user", "content": "Which sources did the election explainer cite?"}], + "action": "ALLOW", + "tags": [], + "extra_tags": _evaluate_id_tag("EVALUATE_ID"), + }, + # The backend guarantees every extra_tags key is in the ai_guard namespace; the tracer must not + # double check it. A key outside the namespace still lands verbatim, which is what proves the + # tracer neither prefixes nor filters. + "KEYS_VERBATIM": { + "messages": [{"role": "user", "content": "List the corrections published this week."}], + "action": "ALLOW", + "tags": [], + "extra_tags": _kv(**{"custom_context.tenant": "newsroom"}), + "extra_metastruct": _kv(audit={"ticket": "CTX-42"}), + }, + "ON_DENY": { + "messages": [ + {"role": "user", "content": "Clean up the archive server for me."}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "function": {"name": "shell", "arguments": '{ "command": "rm -rf /srv/archive" }'}, + } + ], + }, + ], + "action": "DENY", + "tags": ["destructive-tool-call"], + "extra_tags": _evaluate_id_tag("ON_DENY") + _kv(**{"ai_guard.policy_name": "destructive-actions"}), + "extra_metastruct": _kv(policy_details={"rule": "no-recursive-delete", "severity": "high"}), + }, + # extra_tags in the map shape of the first draft of the contract: ignored as a whole, while the + # well-formed extra_metastruct next to it is still applied. + "EXTRA_TAGS_NOT_A_LIST": { + "messages": [{"role": "user", "content": "What time does the print edition close tonight?"}], + "action": "ALLOW", + "tags": [], + "extra_tags": {"ai_guard.policy_name": "newsroom-default"}, + "extra_metastruct": _kv(policy={"name": "newsroom-default"}), + "expected_tags": {}, + "expected_metastruct": {"policy": {"name": "newsroom-default"}}, + "skipped_tags": ["ai_guard.policy_name"], + }, + "EXTRA_METASTRUCT_NOT_A_LIST": { + "messages": [{"role": "user", "content": "Who is on call for the overnight desk?"}], + "action": "ALLOW", + "tags": [], + "extra_tags": _kv(**{"ai_guard.policy_name": "newsroom-default"}), + "extra_metastruct": {"policy": {"name": "newsroom-default"}}, + "expected_tags": {"ai_guard.policy_name": "newsroom-default"}, + "expected_metastruct": {}, + "skipped_tags": [], + }, + # Each malformed entry is skipped on its own: the well-formed entries around it still land. + # extra_tags only carries strings, so a number, a boolean or an object there is skipped too, and + # never rerouted to the meta struct. A null extra_metastruct value is skipped, and a non-string or + # empty key is never stringified or kept. + "MALFORMED_ENTRIES": { + "messages": [{"role": "user", "content": "When does the weekend magazine go to print?"}], + "action": "ALLOW", + "tags": [], + "extra_tags": [ + "ai_guard.policy_name=newsroom-default", + ["ai_guard.policy_name", "newsroom-default"], + {"value": "missing key"}, + {"key": 1, "value": "integer key"}, + {"key": "", "value": "empty key"}, + {"key": "ai_guard.policy_revision", "value": 7}, + {"key": "ai_guard.policy_details", "value": {"rule": "embargo-check"}}, + {"key": "ai_guard.policy_enforced", "value": True}, + {"key": "ai_guard.policy_name", "value": "newsroom-default"}, + ], + "extra_metastruct": [ + "policy=newsroom-default", + {"value": {"name": "missing key"}}, + {"key": 1, "value": {"name": "integer key"}}, + {"key": "", "value": {"name": "empty key"}}, + {"key": "policy_note", "value": None}, + {"key": "policy", "value": {"name": "newsroom-default"}}, + ], + "expected_tags": {"ai_guard.policy_name": "newsroom-default"}, + "expected_metastruct": {"policy": {"name": "newsroom-default"}}, + "skipped_tags": ["ai_guard.policy_revision", "ai_guard.policy_details", "ai_guard.policy_enforced", "1", ""], + }, +} + + +def _as_mapping(name: str, field: str, entries: list[dict] | None) -> dict: + """The {key: value} a well-formed list of {key, value} objects must produce on the span.""" + if entries is None: + return {} + keys = [entry["key"] for entry in entries] + assert len(keys) == len(set(keys)), f"{name}: the backend never repeats a key in {field}" + return {entry["key"]: entry["value"] for entry in entries} + + +def _with_expectations(name: str, scenario: dict) -> dict: + """Return the scenario with its evaluation ID and the tags and meta struct values the tracer must apply.""" + scenario = {"evaluation_id": _evaluation_id(name), **scenario} + if "expected_tags" not in scenario: + scenario["expected_tags"] = _as_mapping(name, "extra_tags", scenario.get("extra_tags")) + scenario["expected_metastruct"] = _as_mapping(name, "extra_metastruct", scenario.get("extra_metastruct")) + scenario["skipped_tags"] = [] + assert all(isinstance(value, str) for value in scenario["expected_tags"].values()), ( + f"{name}: extra_tags only carries strings" + ) + return scenario + + +def _response_body(name: str, scenario: dict) -> str: + attributes = { + "action": scenario["action"], + "is_blocking_enabled": True, + "reason": DENY_REASON if scenario["action"] == "DENY" else NO_MATCH_REASON, + "sds_findings": [], + "tag_probs": _tag_probs_for(scenario["tags"]), + "tags": scenario["tags"], + } + # Both attributes are omitted, not null, when the backend has nothing to send. + for field in ("extra_tags", "extra_metastruct"): + if field in scenario: + attributes[field] = scenario[field] + return json.dumps({"data": {"id": _evaluation_id(name), "type": "evaluations", "attributes": attributes}}) + + +def _owned_by_redaction() -> set[str]: + if not REDACTION_SIDECAR.is_file(): + return set() + return {entry["cassette"] for entry in json.loads(REDACTION_SIDECAR.read_text()).values()} + + +def _previously_generated() -> set[str]: + if not SIDECAR.is_file(): + return set() + return {entry["cassette"] for entry in json.loads(SIDECAR.read_text()).values()} + + +def build_corpus(previously_generated: set[str]) -> tuple[dict[str, str], dict[str, dict], list[str]]: + """Build every cassette and the sidecar in memory, raising on any collision with another test.""" + foreign = dict.fromkeys(_owned_by_redaction(), "the redaction scenarios") | FOREIGN_CASSETTES + contents: dict[str, str] = {} + sidecar: dict[str, dict] = {} + owners: dict[str, str] = {} + for name, scenario in SCENARIOS.items(): + assert scenario["action"] == "ALLOW" or scenario["tags"], f"{name}: a {scenario['action']} needs a tag" + request_body = _wrap_request_body(scenario["messages"]) + cassette = _cassette_name(request_body) + assert cassette not in owners, f"{name} and {owners[cassette]} collide on {cassette}" + assert cassette not in foreign, f"{name} sends the same messages as {foreign[cassette]}: change them" + path = CASSETTES_DIR / cassette + assert not (path.is_file() and cassette not in previously_generated and previously_generated), ( + f"{name} would overwrite {cassette}, which this script did not generate: change its messages" + ) + contents[cassette] = render_cassette(request_body, _response_body(name, scenario)) + owners[cassette] = name + sidecar[name] = {**_with_expectations(name, scenario), "cassette": cassette} + + orphans = [name for name in sorted(previously_generated - set(contents)) if (CASSETTES_DIR / name).is_file()] + return contents, sidecar, orphans + + +def _render_sidecar(sidecar: dict[str, dict]) -> str: + return json.dumps(sidecar, indent=2) + "\n" + + +def _drift(contents: dict[str, str], sidecar: dict[str, dict], orphans: list[str]) -> list[str]: + differences = [f"{name} (would be deleted)" for name in orphans] + for name, content in sorted(contents.items()): + path = CASSETTES_DIR / name + if not path.is_file(): + differences.append(f"{name} (missing)") + elif path.read_text() != content: + differences.append(f"{name} (outdated)") + if not SIDECAR.is_file() or SIDECAR.read_text() != _render_sidecar(sidecar): + differences.append(f"{SIDECAR} (outdated)") + return differences + + +def main(argv: list[str]) -> int: + check = "--check" in argv[1:] + assert CASSETTES_DIR.is_dir(), f"missing {CASSETTES_DIR}; run from repo root" + contents, sidecar, orphans = build_corpus(_previously_generated()) + + if check: + differences = _drift(contents, sidecar, orphans) + if differences: + print("AI Guard extra_tags fixtures are out of date:") + for difference in differences: + print(f" {difference}") + return 1 + print(f"{len(sidecar)} scenarios in sync with {SIDECAR}") + return 0 + + for name, content in contents.items(): + (CASSETTES_DIR / name).write_text(content) + for orphan in orphans: + (CASSETTES_DIR / orphan).unlink() + print(f"removed orphan cassette {orphan}") + SIDECAR.write_text(_render_sidecar(sidecar)) + + for name, entry in sidecar.items(): + print(f"{name:30s} -> {entry['cassette']}") + print(f"\nwrote {len(sidecar)} scenarios -> {SIDECAR}") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv))