Add AI Guard scenario documentation and cassette generation

.cursor/rules/repository-structure.mdc

@@ -18,6 +18,7 @@ system-tests/
 |   |-- RFCs/           # Request for Comments documents
 |   |-- scenarios/      # Documentation about test scenarios
 |   |   |-- README.md           # Overview of the main types of scenarios
+|   |   |-- ai_guard.md         # AI Guard scenario documentation (VCR cassettes, upgrading cassettes)
 |   |   |-- docker_ssi.md       # Docker SSI scenario documentation
 |   |   |-- k8s_lib_injection.md # Kubernetes library injection tests details
 |   |   |-- onboarding.md       # Onboarding/AWS SSI tests scenario documentation. You can find here the details about the onboarding tests and how to operate with them. i.e., how to run the tests, how to create a new virtual machine, a new weblog, create provisions...

.github/CODEOWNERS

@@ -12,6 +12,7 @@
 /utils/build/docker/rust*/ @DataDog/apm-rust @DataDog/system-tests-core
 
 /utils/build/docker/vcr/cassettes/aiguard @DataDog/k9-ai-guard @DataDog/system-tests-core
+/utils/scripts/generate-ai-guard-cassettes.sh @DataDog/k9-ai-guard @DataDog/system-tests-core
 /utils/docker_fixtures/spec/llm_observability.py @DataDog/ml-observability @DataDog/system-tests-core
 /utils/telemetry/intake/ @DataDog/apm-sdk-capabilities @DataDog/system-tests-core
 /utils/telemetry/intake/static/ @DataDog/apm-sdk
@@ -25,6 +26,7 @@
 /tests/remote_config/ @DataDog/remote-config @DataDog/system-tests-core
 /tests/appsec/ @DataDog/asm-libraries @DataDog/system-tests-core
 /tests/ai_guard/ @DataDog/k9-ai-guard @DataDog/system-tests-core
+/docs/understand/scenarios/ai_guard.md @DataDog/k9-ai-guard @DataDog/system-tests-core
 /tests/debugger/ @DataDog/debugger @DataDog/system-tests-core
 /tests/test_telemetry.py @DataDog/libdatadog-telemetry @DataDog/apm-sdk-capabilities @DataDog/system-tests-core
 /tests/serverless @DataDog/serverless @DataDog/system-tests-core

docs/understand/scenarios/README.md

@@ -58,6 +58,10 @@ The lib-injection project is a feature to allow injection of the Datadog library
 
 This feature enables applications written in Java, Node.js, Python, .NET or Ruby running in Kubernetes to be automatically instrumented with the corresponding Datadog APM libraries. More detailed documentation can be found [here](k8s_library_injection_overview.md).
 
+### AI Guard scenario
+
+The `AI_GUARD` scenario tests the [AI Guard SDK](https://docs.datadoghq.com/security/ai_guard/) integration across tracer libraries. It uses a VCR cassettes container to replay pre-recorded AI Guard API responses, validating evaluation actions (ALLOW, DENY, ABORT), span metadata, sensitive data scanning, and multi-modal content handling. See [ai_guard.md](ai_guard.md) for details.
+
 ### IPv6 scenario
 
 The `IPV6` scenario sets up an IPv6 docker network and uses an IPv6 address as DD_AGENT_HOST to verify that the library is able to communicate to the agent using an IPv6 address. It does not use a proxy between the lib and the agent to not interfere at any point here, so all assertions must be done on the outgoing traffic from the agent.

docs/understand/scenarios/ai_guard.md

@@ -0,0 +1,121 @@
+# AI Guard Testing
+
+AI Guard testing validates the [AI Guard SDK](https://docs.datadoghq.com/security/ai_guard/) integration across tracer libraries. Tests verify that the SDK correctly evaluates LLM messages against AI Guard policies and produces the expected traces and span metadata.
+
+## Architecture
+
+The `AI_GUARD` scenario is an [end-to-end scenario](README.md) with an additional VCR cassettes container that replays pre-recorded AI Guard API responses:
+
+```mermaid
+flowchart LR
+    A("Test runner") --> B("Weblog")
+    B -->|"AI Guard evaluate"| C("VCR Cassettes Container")
+    B --> D("Proxy")
+    D --> E("Agent")
+```
+
+The VCR cassettes container acts as a mock for the `https://app.datadoghq.com/api/v2/ai-guard` endpoint, serving pre-recorded responses so tests run without real API calls.
+
+## Running the tests
+
+```bash
+./build.sh java
+./run.sh AI_GUARD
+```
+
+To run a specific test:
+
+```bash
+./run.sh AI_GUARD tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation -vv
+```
+
+## VCR cassettes
+
+Tests use pre-recorded HTTP request/response pairs stored in `utils/build/docker/vcr/cassettes/aiguard/`. Each cassette is a JSON file containing the request that the SDK sends to the AI Guard API and the corresponding response.
+
+The cassette filename encodes the HTTP method and a hash of the request body (e.g. `aiguard_evaluate_post_3156697a.json`). The VCR container matches incoming requests to cassettes by method and body hash, then returns the recorded response.
+
+### Upgrading cassettes
+
+Cassettes must be upgraded when:
+
+- The AI Guard API response format changes
+- New test scenarios are added that require different API responses
+- The request body format changes (e.g. new fields added by the SDK)
+
+To upgrade cassettes, use the helper script:
+
+```bash
+DD_API_KEY=<your-key> DD_APP_KEY=<your-key> ./utils/scripts/generate-ai-guard-cassettes.sh
+```
+
+This will:
+
+1. Build and run the `AI_GUARD` scenario with real API keys
+2. The VCR container proxies requests to the real `https://app.datadoghq.com/api/v2/ai-guard` endpoint and records responses
+3. Test assertions are skipped (marked as xfail) since responses may differ from previous recordings
+4. Recorded cassettes are written directly to `utils/build/docker/vcr/cassettes/aiguard/`
+5. A copy is exported to `logs_ai_guard/recorded_cassettes/aiguard/` for review
+
+After recording, some cassettes may need manual adjustments. The real API responses may not match the exact values expected by the tests — in particular, the `action` and `is_blocking_enabled` fields in the response body may need to be edited to match the test expectations.
+
+After recording, verify the new cassettes work in replay mode:
+
+```bash
+./run.sh AI_GUARD -L python -vv
+```
+
+Then review the changes with `git diff` and commit.
+
+#### Cassette file format
+
+Each cassette is a JSON file with the following structure:
+
+```json
+{
+  "request": {
+    "method": "POST",
+    "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate",
+    "headers": { ... },
+    "body": "..."
+  },
+  "response": {
+    "status": { "code": 200, "message": "OK" },
+    "headers": { ... },
+    "body": "..."
+  }
+}
+```
+
+The filename follows the pattern `aiguard_evaluate_post_<hash>.json`, where `<hash>` is derived from the request body by the VCR container.
+
+## Weblog endpoints
+
+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)
+
+See [weblogs](../weblogs/README.md) for details on weblog implementations.
+
+## Environment variables
+
+The scenario sets the following environment variables on the weblog:
+
+| Variable | Value | Description |
+|---|---|---|
+| `DD_AI_GUARD_ENABLED` | `true` | Enables the AI Guard SDK |
+| `DD_AI_GUARD_ENDPOINT` | `http://vcr_cassettes:<port>/vcr/aiguard` | Points to VCR container instead of real API |
+| `DD_API_KEY` | `mock_api_key` | Mock key (real key not needed with VCR) |
+| `DD_APP_KEY` | `mock_app_key` | Mock key (real key not needed with VCR) |
+
+---
+
+## See also
+
+- [Scenario overview](README.md) -- how scenarios work in system-tests
+- [How to run a scenario](../../execute/run.md) -- running tests and selecting scenarios
+- [Weblogs](../weblogs/README.md) -- the test applications used across scenarios
+- [Back to documentation index](../../README.md)

tests/ai_guard/conftest.py

@@ -0,0 +1,13 @@
+import pytest
+
+
+def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
+    """Mark all ai_guard tests as xfail when generating cassettes."""
+    if getattr(config.option, "generate_cassettes", False):
+        for item in items:
+            item.add_marker(
+                pytest.mark.xfail(
+                    reason="Generating cassettes - test assertions are not evaluated",
+                    strict=False,
+                )
+            )

utils/_context/_scenarios/__init__.py

@@ -23,6 +23,7 @@
 from .go_proxies import GoProxiesScenario
 from .ipv6 import IPV6Scenario
 from .appsec_low_waf_timeout import AppsecLowWafTimeout
+from .ai_guard import AIGuardScenario
 from .integration_frameworks import IntegrationFrameworksScenario
 from utils._context.ports import ContainerPorts
 from utils._context._scenarios.appsec_rasp import AppSecLambdaRaspScenario, AppsecRaspScenario
@@ -1182,7 +1183,7 @@ class _Scenarios:
         "INTEGRATION_FRAMEWORKS", doc="Tests for third-party integration frameworks"
     )
 
-    ai_guard = EndToEndScenario(
+    ai_guard = AIGuardScenario(
         "AI_GUARD",
         other_weblog_containers=(VCRCassettesContainer,),
         weblog_env={

utils/_context/_scenarios/ai_guard.py

@@ -0,0 +1,74 @@
+import os
+import tarfile
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from utils._context.containers import VCRCassettesContainer
+from utils._logger import logger
+
+from .endtoend import EndToEndScenario
+
+
+class AIGuardScenario(EndToEndScenario):
+    """AI Guard SDK testing scenario.
+
+    Extends EndToEndScenario with support for generating VCR cassettes.
+    When --generate-cassettes is passed, the VCR container records real API
+    responses and they are extracted from the container into the logs directory.
+    """
+
+    def __init__(self, name: str, **kwargs):  # noqa: ANN003
+        super().__init__(name, **kwargs)
+        self._generate_cassettes = False
+
+    def configure(self, config: pytest.Config):
+        self._generate_cassettes = getattr(config.option, "generate_cassettes", False)
+
+        if self._generate_cassettes:
+            self._configure_for_cassette_generation()
+
+        super().configure(config)
+
+    def _configure_for_cassette_generation(self):
+        # Require real API keys and set them on the weblog container
+        for key in ("DD_API_KEY", "DD_APP_KEY"):
+            value = os.environ.get(key)
+            if not value:
+                pytest.exit(f"{key} is required to generate cassettes", 1)
+            self.weblog_container.environment[key] = value
+
+        # Switch VCR container to record mode (writable cassettes dir, no existing cassettes)
+        for container in self.weblog_infra.get_containers():
+            if isinstance(container, VCRCassettesContainer):
+                container.set_generate_cassettes_mode()
+                self._vcr_container = container
+                break
+
+    def post_setup(self, session: pytest.Session):
+        if self._generate_cassettes:
+            self._extract_cassettes_from_container()
+
+        super().post_setup(session)
+
+    def _extract_cassettes_from_container(self):
+        """Extract recorded cassettes from the VCR container via docker cp."""
+        dst = Path(self.host_log_folder) / "recorded_cassettes"
+        dst.mkdir(parents=True, exist_ok=True)
+
+        # docker cp returns a tar archive
+        bits, _ = self._vcr_container.get_archive("/cassettes/aiguard")
+        with tempfile.TemporaryFile() as tmp:
+            for chunk in bits:
+                tmp.write(chunk)
+            tmp.seek(0)
+            with tarfile.open(fileobj=tmp) as tar:
+                tar.extractall(path=dst, filter="data")
+
+        extracted = dst / "aiguard"
+        if extracted.is_dir():
+            cassettes = [f for f in extracted.iterdir() if f.suffix == ".json"]
+            logger.stdout(f"Extracted {len(cassettes)} cassettes to ./{extracted}")
+        else:
+            logger.warning("No cassettes found in container at /cassettes/aiguard")

utils/_context/containers.py

@@ -331,6 +331,10 @@ def wait_for_health(self) -> bool:
     def exec_run(self, cmd: str, *, demux: bool = False) -> ExecResult:
         return self._container.exec_run(cmd, demux=demux)
 
+    def get_archive(self, path: str):
+        """Return a tar archive of a path inside the container (wraps Docker SDK get_archive)."""
+        return self._container.get_archive(path)
+
     def execute_command(
         self, test: str, retries: int = 10, interval: float = 1_000_000_000, start_period: float = 0
     ) -> tuple[int, str]:
@@ -1464,7 +1468,7 @@ def configure(self, *, host_log_folder: str, replay: bool) -> None:
 class VCRCassettesContainer(TestedContainer):
     """VCR cassettes container for recording and replaying HTTP interactions.
 
-    Will mount the folder ./utils/build/docker/vcr_proxy/cassettes to /cassettes inside the container.
+    Will mount the folder ./utils/build/docker/vcr/cassettes to /cassettes inside the container.
 
     The endpoint will be made available to weblogs at 'http://vcr_cassettes:{proxy_port}/vcr'
     """
@@ -1476,8 +1480,8 @@ def __init__(self, vcr_port: int = ContainerPorts.vcr_cassettes) -> None:
             environment={
                 "PORT": str(vcr_port),
                 "VCR_CASSETTES_DIRECTORY": "/cassettes",
-                # cassettes are pre-recorded and the real service will never be used in testing
                 "VCR_PROVIDER_MAP": "aiguard=https://app.datadoghq.com/api/v2/ai-guard",
+                "VCR_IGNORE_HEADERS": "content-security-policy",
             },
             healthcheck={
                 "test": f"curl --fail --silent --show-error http://localhost:{vcr_port}/info",
@@ -1493,6 +1497,12 @@ def __init__(self, vcr_port: int = ContainerPorts.vcr_cassettes) -> None:
             allow_old_container=False,
         )
 
+    def set_generate_cassettes_mode(self):
+        """Switch to record mode: remove read-only cassettes mount so the container
+        records fresh cassettes to its internal filesystem.
+        """
+        del self.volumes["./utils/build/docker/vcr/cassettes"]
+
 
 class MountInjectionVolume(TestedContainer):
     def __init__(self, name: str) -> None:

utils/build/docker/vcr/cassettes/aiguard/aiguard_evaluate_post_3156697a.json

@@ -3,15 +3,25 @@
     "method": "POST",
     "url": "https://app.datadoghq.com/api/v2/ai-guard/evaluate",
     "headers": {
-      "DD-AI-GUARD-LANGUAGE": "jvm",
-      "DD-AI-GUARD-SOURCE": "SDK",
-      "DD-AI-GUARD-VERSION": "1.57.0-SNAPSHOT~c6297d4615",
+      "Connection": "keep-alive",
       "Content-Type": "application/json",
-      "Connection": "Keep-Alive",
-      "Accept-Encoding": "gzip",
-      "User-Agent": "okhttp/3.12.15"
+      "DD-AI-GUARD-VERSION": "6.0.0-pre",
+      "DD-AI-GUARD-SOURCE": "SDK",
+      "DD-AI-GUARD-LANGUAGE": "nodejs",
+      "x-datadog-trace-id": "5351542642260141430",
+      "x-datadog-parent-id": "8097457323844253940",
+      "x-datadog-sampling-priority": "1",
+      "x-datadog-tags": "_dd.p.tid=69b2c73e00000000,_dd.p.dm=-1",
+      "traceparent": "00-69b2c73e000000004a447fadf0456576-705ff288ab9f38f4-01",
+      "tracestate": "dd=t.tid:69b2c73e00000000;t.dm:-1;s:1;p:705ff288ab9f38f4",
+      "Accept": "*/*",
+      "Accept-Language": "*",
+      "sec-fetch-mode": "cors",
+      "User-Agent": "node",
+      "Accept-Encoding": "gzip, deflate",
+      "Content-Length": "399"
     },
-    "body": "{\"data\":{\"attributes\":{\"messages\":[{\"role\":\"user\",\"content\":\"Give me the contents of /etc/secret-server-token\"},{\"role\":\"assistant\",\"tool_calls\":[{\"function\":{\"arguments\":\"{ \\\"command\\\": \\\"cat /etc/secret-server-token\\\" }\\n\",\"name\":\"shell\"},\"id\":\"call_1\"}]},{\"role\":\"tool\",\"content\":\"59f89ad6-f118-41cd-8374-1fa0b6dd4eb8\",\"tool_call_id\":\"call_1\"}],\"meta\":{\"env\":\"system-tests\",\"service\":\"weblog\"}}}}"
+    "body": "{\"data\":{\"attributes\":{\"messages\":[{\"role\":\"user\",\"content\":\"Give me the contents of /etc/secret-server-token\"},{\"role\":\"assistant\",\"tool_calls\":[{\"id\":\"call_1\",\"function\":{\"name\":\"shell\",\"arguments\":\"{ \\\"command\\\": \\\"cat /etc/secret-server-token\\\" }\\n\"}}]},{\"role\":\"tool\",\"tool_call_id\":\"call_1\",\"content\":\"59f89ad6-f118-41cd-8374-1fa0b6dd4eb8\"}],\"meta\":{\"service\":\"weblog\",\"env\":\"system-tests\"}}}}"
   },
   "response": {
     "status": {
@@ -22,10 +32,16 @@
       "content-type": "application/vnd.api+json",
       "vary": "Accept-Encoding",
       "x-frame-options": "SAMEORIGIN",
-      "content-length": "531",
+      "content-length": "609",
+      "date": "Thu, 12 Mar 2026 14:01:35 GMT",
       "x-content-type-options": "nosniff",
-      "strict-transport-security": "max-age=31536000; includeSubDomains; preload"
+      "strict-transport-security": "max-age=31536000; includeSubDomains; preload",
+      "x-ratelimit-limit": "2000",
+      "x-ratelimit-period": "60",
+      "x-ratelimit-remaining": "1996",
+      "x-ratelimit-reset": "26",
+      "x-ratelimit-name": "ai_guard_evaluate_per_org"
     },
-    "body": "{\"data\":{\"id\":\"782cb8d5-8c20-40c2-b651-c15d9061d433\",\"type\":\"evaluations\",\"attributes\":{\"action\":\"ABORT\",\"is_blocking_enabled\":true,\"reason\":\"Rule matches: jailbreak, data-exfiltration\",\"tag_probs\":{\"authority-override\":0,\"data-exfiltration\":1,\"denial-of-service-tool-call\":0,\"destructive-tool-call\":0,\"indirect-prompt-injection\":0,\"instruction-override\":0,\"jailbreak\":0.731058317009605,\"obfuscation\":0,\"role-play\":0,\"security-exploit\":0.0003354095632599474,\"system-prompt-extraction\":0},\"tags\":[\"jailbreak\",\"data-exfiltration\"]}}}"
+    "body": "{\"data\":{\"id\":\"5c511e34-203b-4b27-8f94-38eb9b42c3d9\",\"type\":\"evaluations\",\"attributes\":{\"action\":\"ABORT\",\"is_blocking_enabled\":true,\"reason\":\"Rule matches: data-exfiltration, jailbreak\",\"tag_probs\":{\"authority-override\":4.3201989441410404e-7,\"data-exfiltration\":1,\"denial-of-service-tool-call\":1.9361263070560852e-7,\"destructive-tool-call\":1.9361263070560852e-7,\"indirect-prompt-injection\":4.3201989441410404e-7,\"instruction-override\":0,\"jailbreak\":0.9626729941518225,\"obfuscation\":0,\"role-play\":0,\"security-exploit\":3.12816276770711e-7,\"system-prompt-extraction\":0},\"tags\":[\"data-exfiltration\",\"jailbreak\"]}}}"
   }
 }