test: add guardrail e2e coverage (presidio masking, bedrock post and during call, moderation on messages) (#38553)

* test: add guardrail e2e coverage (presidio masking, bedrock post/during, moderation on messages)

* test: require the phone placeholder positively in the presidio masking predicate

* test: count only the 400 verdict body as a bedrock post_call block

tests/e2e/guardrails/guardrails_client.py

@@ -68,11 +68,27 @@ class BlockCodeExecutionParamsBody(GuardrailParamsBase):
     guardrail: Literal["block_code_execution"] = "block_code_execution"
 
 
+class PresidioParamsBody(GuardrailParamsBase):
+    """Presidio PII guardrail params. `presidio_filter_scope="input"` keeps the
+    registration to a single callback on the configured mode; the default
+    ("both") also registers a second post_call output-masking callback, which a
+    pre_call- or logging_only-scoped test must not drag in. `output_parse_pii`
+    stays unset/False: True would unmask the response back to the caller."""
+
+    guardrail: Literal["presidio"] = "presidio"
+    presidio_analyzer_api_base: str
+    presidio_anonymizer_api_base: str
+    presidio_filter_scope: Literal["input", "output", "both"] | None = None
+    presidio_language: str | None = None
+    output_parse_pii: bool | None = None
+
+
 GuardrailParamsBody = (
     ContentFilterParamsBody
     | BedrockGuardrailParamsBody
     | OpenAIModerationParamsBody
     | BlockCodeExecutionParamsBody
+    | PresidioParamsBody
 )
 
 
@@ -253,6 +269,30 @@ class GuardrailsClient:
             ),
         )
 
+    def chat_stream_raw(
+        self,
+        key: str,
+        model: str,
+        text: str,
+        *,
+        guardrails: list[str] | None = None,
+        max_tokens: int = 64,
+    ) -> StreamingResponse:
+        """Drive /chat/completions with stream=true, returning the raw HTTP
+        outcome (status, headers, SSE events) via the shared ProxyClient stream
+        sender - a streamed guardrail block is judged on status and stream
+        shape, not a typed body."""
+        return self.proxy.chat_stream(
+            key,
+            ChatBody(
+                model=model,
+                messages=[ChatMessage(role="user", content=text)],
+                max_tokens=max_tokens,
+                stream=True,
+                guardrails=guardrails,
+            ),
+        )
+
     def messages(
         self,
         key: str,
@@ -318,7 +358,7 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient:
     return GuardrailsClient(proxy=proxy)
 
 
-def poll_until_blocked(call: Callable[[], Result[ChatResponse]]) -> Result[ChatResponse]:
+def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]:
     """Retry a call that a guardrail should reject until it is, returning the last result.
 
     Registering a guardrail is a control-plane write; the data-plane worker that
@@ -337,3 +377,25 @@ def poll_until_blocked(call: Callable[[], Result[ChatResponse]]) -> Result[ChatR
         time.sleep(POLL_INTERVAL)
         last = call()
     return last
+
+
+#: Statuses a stream poll keeps retrying through instead of returning as "the
+#: block": network failures (-1), key propagation (401), rate limits (429) -
+#: transient rig noise, not a guardrail verdict.
+_TRANSIENT_STREAM_STATUSES = frozenset({-1, 401, 429})
+
+
+def poll_until_blocked_stream(call: Callable[[], StreamingResponse]) -> StreamingResponse:
+    """poll_until_blocked for raw/streamed sends, which return a StreamingResponse
+    instead of a Result: retry while the call still succeeds (the data-plane worker
+    has not picked the new guardrail up yet) or fails with a transient status,
+    returning the first guardrail-shaped non-2xx outcome or the last result at
+    the deadline."""
+    deadline = time.monotonic() + POLL_TIMEOUT
+    last = call()
+    while time.monotonic() < deadline:
+        if not last.ok and last.status_code not in _TRANSIENT_STREAM_STATUSES:
+            return last
+        time.sleep(POLL_INTERVAL)
+        last = call()
+    return last

tests/e2e/guardrails/test_bedrock_guardrail_e2e.py

@@ -1,9 +1,12 @@
-"""Live e2e: Bedrock ApplyGuardrail pre_call blocks denied input on chat.
+"""Live e2e: Bedrock ApplyGuardrail blocks on chat, pre_call and post_call.
 
-Registers a default-on bedrock guardrail via POST /guardrails with identifier/
+pre_call registers a bedrock guardrail via POST /guardrails with identifier/
 version from env, then sends a prompt the guardrail's configured policy denies.
 HTTP 400 (or other non-2xx block) with a guardrail-shaped body is the contract;
-a 200 means the guardrail never ran.
+a 200 means the guardrail never ran. post_call scans the MODEL OUTPUT only, so
+its test makes the model echo the word the guardrail's word policy denies
+(BEDROCK_GUARDRAIL_BLOCKED_WORD, default FORBIDDENWORD) and the block must
+arrive without leaking the model's text.
 
 No AWS keys are passed: the gateway signs ApplyGuardrail with its own
 pod-identity role, since the static AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY
@@ -18,7 +21,11 @@ import pytest
 
 from e2e_config import unique_marker
 from e2e_http import UnknownApiError
-from guardrails_client import GuardrailsClient, poll_until_blocked
+from guardrails_client import (
+    BedrockGuardrailParamsBody,
+    GuardrailsClient,
+    poll_until_blocked,
+)
 from lifecycle import ResourceManager
 
 pytestmark = pytest.mark.e2e
@@ -42,23 +49,17 @@ class TestBedrockGuardrail:
         version = os.environ["BEDROCK_GUARDRAIL_VERSION"]
 
         name = f"e2e-bedrock-guard-{unique_marker()}"
-        guardrail_id = client.create_bedrock_guardrail(
-            name, identifier=identifier, version=version
-        )
+        guardrail_id = client.create_bedrock_guardrail(name, identifier=identifier, version=version)
         resources.defer(lambda: client.delete_guardrail(guardrail_id))
 
         # Selected per request rather than registered default_on, so an upstream
         # ApplyGuardrail failure surfaces here instead of 403ing every other suite
         # running against this proxy.
-        result = poll_until_blocked(
-            lambda: client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name])
-        )
+        result = poll_until_blocked(lambda: client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name]))
 
         match result:
             case UnknownApiError(status_code=status, body=body):
-                assert status in {400, 403}, (
-                    f"expected a guardrail block status, got {status}: {body[:400]}"
-                )
+                assert status in {400, 403}, f"expected a guardrail block status, got {status}: {body[:400]}"
                 body_lower = body.lower()
                 assert any(
                     token in body_lower
@@ -72,6 +73,49 @@ class TestBedrockGuardrail:
                     )
                 ), f"block body should name the guardrail reason; got: {body[:400]}"
             case _:
-                pytest.fail(
-                    f"bedrock default-on guardrail did not block harmful prompt; got {result}"
+                pytest.fail(f"bedrock default-on guardrail did not block harmful prompt; got {result}")
+
+    @pytest.mark.covers(
+        "guardrail.bedrock.post_call.blocks",
+        exercised_on=["chat_completions"],
+    )
+    def test_bedrock_post_call_blocks_denied_model_output(
+        self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
+        version = os.environ["BEDROCK_GUARDRAIL_VERSION"]
+        blocked_word = os.environ.get("BEDROCK_GUARDRAIL_BLOCKED_WORD", "FORBIDDENWORD")
+
+        name = f"e2e-bedrock-post-{unique_marker()}"
+        guardrail_id = client.register(
+            name,
+            BedrockGuardrailParamsBody(
+                mode="post_call",
+                default_on=False,
+                guardrailIdentifier=identifier,
+                guardrailVersion=version,
+            ),
+        )
+        resources.defer(lambda: client.delete_guardrail(guardrail_id))
+
+        # post_call scans OUTPUT only, so the denied word has to come out of the
+        # model: ask it to echo the word verbatim. The word in the prompt itself
+        # is not scanned in this mode.
+        prompt = f"Reply with exactly this one word and nothing else: {blocked_word}"
+        result = poll_until_blocked(lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128))
+
+        match result:
+            case UnknownApiError(status_code=status, body=body):
+                # A policy block is a 400 naming the verdict; a failed
+                # ApplyGuardrail call surfaces as 403 "guardrail request
+                # failed", which must not count as a block.
+                assert status == 400, f"expected the guardrail block status 400, got {status}: {body[:400]}"
+                body_lower = body.lower()
+                assert any(token in body_lower for token in ("violated", "blocked", "intervened")), (
+                    f"block body should name the guardrail verdict; got: {body[:400]}"
                 )
+                assert blocked_word not in body, (
+                    f"the blocked model output must not leak into the error body; got: {body[:400]}"
+                )
+            case _:
+                pytest.fail(f"bedrock post_call guardrail did not block denied model output; got {result}")

tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py

@@ -7,7 +7,9 @@ before the upstream model runs; a prompt that trips the policy must be rejected
 with HTTP 400 naming the moderation policy, and the same guardrail must let a
 benign prompt through. The chat backend is a gemini deployment created for the
 test (and torn down); moderation runs independently of it, so the block is
-attributable to the guardrail, not the model.
+attributable to the guardrail, not the model. The same pre_call contract is
+also exercised through /v1/messages (Anthropic format): a flagged prompt is
+rejected with a 400 naming moderation and a benign one passes.
 """
 
 from __future__ import annotations
@@ -69,3 +71,46 @@ class TestOpenAIModerationGuardrail:
             "the same moderation guardrail must let a benign prompt through, but the "
             f"call returned no choices: {allowed}"
         )
+
+    @pytest.mark.covers(
+        "guardrail.openai_moderations.pre_call.blocks",
+        exercised_on=["messages"],
+    )
+    def test_moderation_blocks_flagged_input_on_messages(
+        self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        model = client.create_backend_model(resources, prefix="e2e-moderation-msg-backend")
+
+        name = f"e2e-openai-moderation-msg-{unique_marker()}"
+        guardrail_id = client.register(
+            name,
+            OpenAIModerationParamsBody(
+                mode="pre_call", default_on=False, api_key="os.environ/OPENAI_API_KEY"
+            ),
+        )
+        resources.defer(lambda: client.delete_guardrail(guardrail_id))
+
+        blocked = poll_until_blocked(
+            lambda: client.messages(scoped_key, model, FLAGGED_PROMPT, guardrails=[name])
+        )
+        match blocked:
+            case UnknownApiError(status_code=400, body=body):
+                assert "moderation" in body.lower(), (
+                    f"the block body must name the moderation policy, got: {body[:400]}"
+                )
+            case UnknownApiError(status_code=status, body=body):
+                pytest.fail(
+                    f"expected a 400 moderation block on /v1/messages, got {status}: {body[:400]}"
+                )
+            case _:
+                pytest.fail(
+                    f"openai moderation did not block a flagged /v1/messages prompt; got {blocked}"
+                )
+
+        allowed = unwrap(
+            client.messages(scoped_key, model, BENIGN_PROMPT, guardrails=[name], max_tokens=64)
+        )
+        assert allowed.content or allowed.choices, (
+            "the same moderation guardrail must let a benign /v1/messages prompt through, but "
+            f"the response carried neither content nor choices: {allowed}"
+        )

tests/e2e/guardrails/test_presidio_masking_e2e.py

@@ -0,0 +1,184 @@
+"""Live e2e: the Presidio PII guardrail masks, per its configured hook point.
+
+pre_call: the guardrail calls the Presidio analyzer/anonymizer on the request
+messages BEFORE the model runs, so the model only ever sees placeholders like
+<EMAIL_ADDRESS>. A prompt asking the model to repeat a fake email + phone back
+must come back with the placeholders echoed and the raw PII absent, on
+/chat/completions and on /v1/messages (Anthropic format).
+
+The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE /
+PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip.
+Each guardrail registers with presidio_filter_scope="input" so only the
+configured hook's callback exists (the default "both" adds a second post_call
+output masker), and is deleted on teardown.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from collections.abc import Callable
+
+import pytest
+from pydantic import BaseModel
+
+from e2e_config import unique_marker
+from e2e_http import Result, Success
+from guardrails_client import GuardrailsClient, PresidioParamsBody
+from lifecycle import ResourceManager
+from models import AnthropicMessagesResponse, ChatResponse
+
+pytestmark = pytest.mark.e2e
+
+MODEL = "gemini-2.5-flash"
+
+# A guardrail created via POST /guardrails reaches the worker that served the
+# create immediately, but every other worker only picks it up on its next
+# periodic DB sync (~30s), so the first requests can be served unguarded.
+GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0
+GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0
+
+# Presidio's anonymizer replaces a detected entity with its unnumbered type
+# placeholder, e.g. <EMAIL_ADDRESS>. The pre_call assertions match on the bare
+# token because the model is echoing the masked prompt and may not preserve the
+# angle brackets; the logged payload keeps the placeholder verbatim.
+MASKED_EMAIL_TOKEN = "EMAIL_ADDRESS"
+MASKED_PHONE_TOKEN = "PHONE_NUMBER"
+
+# Fictional NANP 555 number; a standard format Presidio's phone recognizer detects.
+FAKE_PHONE = "+1 415-555-0134"
+
+
+def _presidio_bases() -> tuple[str, str]:
+    analyzer = os.environ.get("PRESIDIO_ANALYZER_API_BASE", "").strip()
+    anonymizer = os.environ.get("PRESIDIO_ANONYMIZER_API_BASE", "").strip()
+    if not analyzer or not anonymizer:
+        pytest.fail(
+            "Presidio e2e requires PRESIDIO_ANALYZER_API_BASE and PRESIDIO_ANONYMIZER_API_BASE "
+            "(the running Presidio analyzer/anonymizer services); missing env is a hard failure, not a skip"
+        )
+    return analyzer, anonymizer
+
+
+def _register_presidio(
+    client: GuardrailsClient,
+    resources: ResourceManager,
+    *,
+    name: str,
+) -> None:
+    analyzer, anonymizer = _presidio_bases()
+    guardrail_id = client.register(
+        name,
+        PresidioParamsBody(
+            mode="pre_call",
+            default_on=False,
+            presidio_analyzer_api_base=analyzer,
+            presidio_anonymizer_api_base=anonymizer,
+            presidio_filter_scope="input",
+        ),
+    )
+    resources.defer(lambda: client.delete_guardrail(guardrail_id))
+
+
+def _fake_email() -> str:
+    return f"jane.doe.{unique_marker()}@example.com"
+
+
+def _pii_prompt(marker: str, email: str) -> str:
+    return (
+        f"{marker} Repeat this sentence back to me exactly, word for word: "
+        f"My email address is {email} and my phone number is {FAKE_PHONE}."
+    )
+
+
+def _first_content(response: ChatResponse) -> str:
+    if not response.choices:
+        return ""
+    message = response.choices[0].message
+    return (message.content if message else None) or ""
+
+
+def _messages_text(response: AnthropicMessagesResponse) -> str:
+    """The text of a /v1/messages answer, whichever shape the proxy produced
+    (Anthropic-native content blocks or OpenAI-normalized choices)."""
+    parts: list[str] = []
+    for block in response.content or []:
+        if block.text:
+            parts.append(block.text)
+    for choice in response.choices or []:
+        if choice.message and choice.message.content:
+            parts.append(choice.message.content)
+    return "\n".join(parts)
+
+
+def _assert_eventually_masked[R: BaseModel](
+    fetch: Callable[[], Result[R]], extract: Callable[[R], str], *, email: str
+) -> None:
+    """Retry the call until the response comes back masked, to the propagation
+    deadline. An unmasked early response is in-flight guardrail propagation, not
+    a failure, and neither is a transient non-Success (a replica that has not
+    reloaded the guardrail answers 404, the live model can rate-limit) - only a
+    response that still carries the raw PII at the deadline is."""
+    deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
+    last: str = "<no successful response yet>"
+    while True:
+        result = fetch()
+        match result:
+            case Success(data=data):
+                content = extract(data)
+                last = content
+                masked = MASKED_EMAIL_TOKEN in content and MASKED_PHONE_TOKEN in content and email not in content
+                if masked:
+                    assert FAKE_PHONE not in content, (
+                        f"the raw phone number must be masked before the model sees it, but the "
+                        f"response echoed it: {content[:300]!r}"
+                    )
+                    return
+            case _:
+                last = f"<non-Success result: {result}>"
+        if time.monotonic() >= deadline:
+            pytest.fail(
+                f"presidio pre_call guardrail never masked the PII within "
+                f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
+            )
+        time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
+
+
+class TestPresidioPreCallMasking:
+    @pytest.mark.covers(
+        "guardrail.presidio.pre_call.masks",
+        exercised_on=["chat_completions"],
+    )
+    def test_pre_call_masks_pii_on_chat_completions(
+        self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        name = f"e2e-presidio-pre-chat-{unique_marker()}"
+        _register_presidio(client, resources, name=name)
+
+        email = _fake_email()
+        prompt = _pii_prompt(unique_marker(), email)
+
+        _assert_eventually_masked(
+            lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
+            _first_content,
+            email=email,
+        )
+
+    @pytest.mark.covers(
+        "guardrail.presidio.pre_call.masks",
+        exercised_on=["messages"],
+    )
+    def test_pre_call_masks_pii_on_messages(
+        self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        name = f"e2e-presidio-pre-msg-{unique_marker()}"
+        _register_presidio(client, resources, name=name)
+
+        email = _fake_email()
+        prompt = _pii_prompt(unique_marker(), email)
+
+        _assert_eventually_masked(
+            lambda: client.messages(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
+            _messages_text,
+            email=email,
+        )

tests/e2e/guardrails/test_streaming_guardrail_e2e.py

@@ -0,0 +1,87 @@
+"""Live e2e: a Bedrock guardrail in during_call mode blocks a streamed chat.
+
+during_call runs the Bedrock ApplyGuardrail INPUT scan in an asyncio.gather
+alongside the LLM call (common_request_processing.py); when the scan flags the
+prompt, the raised block cancels the LLM task before the stream ever starts, so
+the client sees a non-2xx JSON error - not an SSE stream, not an in-stream
+error frame - and zero content chunks are delivered.
+
+The prompt deliberately contains the exact word the guardrail's word policy
+denies (BEDROCK_GUARDRAIL_BLOCKED_WORD, default FORBIDDENWORD), so the INPUT
+scan intervenes deterministically. Identifier/version come from
+BEDROCK_GUARDRAIL_IDENTIFIER / BEDROCK_GUARDRAIL_VERSION like the rest of the
+bedrock suite; no AWS keys are passed (the gateway signs with pod identity).
+The guardrail registers default_on=False and is selected per request, so an
+upstream ApplyGuardrail failure surfaces here instead of 403ing other suites.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from e2e_config import unique_marker
+from guardrails_client import (
+    BedrockGuardrailParamsBody,
+    GuardrailsClient,
+    poll_until_blocked_stream,
+)
+from lifecycle import ResourceManager
+
+pytestmark = pytest.mark.e2e
+
+MODEL = "gemini-2.5-flash"
+
+
+class TestBedrockDuringCallStreaming:
+    @pytest.mark.covers(
+        "guardrail.bedrock.during.blocks",
+        exercised_on=["chat_completions"],
+    )
+    def test_during_call_blocks_stream_before_first_chunk(
+        self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
+        version = os.environ["BEDROCK_GUARDRAIL_VERSION"]
+        blocked_word = os.environ.get("BEDROCK_GUARDRAIL_BLOCKED_WORD", "FORBIDDENWORD")
+
+        name = f"e2e-bedrock-during-{unique_marker()}"
+        guardrail_id = client.register(
+            name,
+            BedrockGuardrailParamsBody(
+                mode="during_call",
+                default_on=False,
+                guardrailIdentifier=identifier,
+                guardrailVersion=version,
+            ),
+        )
+        resources.defer(lambda: client.delete_guardrail(guardrail_id))
+
+        # The denied word sits in the INPUT: during_call scans the request
+        # messages while the model call runs, and the flag must win the race
+        # by cancelling the stream outright.
+        prompt = f"Please use the word {blocked_word} in a sentence."
+        result = poll_until_blocked_stream(
+            lambda: client.chat_stream_raw(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=64)
+        )
+
+        assert not result.ok, (
+            f"the during_call guardrail never blocked the streamed request; got a "
+            f"{result.status_code} with {result.chunks} chunks"
+        )
+        assert result.status_code == 400, (
+            f"a during_call block surfaces as HTTP 400 before the stream starts, got "
+            f"{result.status_code}: {result.body[:400]}"
+        )
+        assert result.chunks == 0 and not result.stream_events, (
+            f"no content chunk may be delivered on a during_call block, but "
+            f"{result.chunks} chunks arrived: {result.stream_events[:3]}"
+        )
+        assert "text/event-stream" not in (result.content_type or ""), (
+            f"the block must be a JSON error response, not an SSE stream; got content-type {result.content_type!r}"
+        )
+        body_lower = result.body.lower()
+        assert any(token in body_lower for token in ("guardrail", "violated", "blocked", "bedrock", "intervened")), (
+            f"block body should name the guardrail reason; got: {result.body[:400]}"
+        )