Merge pull request #40038 from BerriAI/litellm_/guardrail-automation-testing-3ecd3d

test(guardrails): pin the presidio spend-log record and the UI's masked-entity persistence

tests/e2e/CONTRIBUTING.md

@@ -52,6 +52,15 @@ The suites run against a live proxy, so bring one up first by running the litell
 
 Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE`
 
+The Presidio spend-log audit test also requires prompt storage on the proxy, because its assertions inspect the detected entities and their scores. Add the following to the proxy config before starting it:
+
+```yaml
+general_settings:
+  store_prompts_in_spend_logs: true
+```
+
+Alternatively, start the proxy with `STORE_PROMPTS_IN_SPEND_LOGS=true litellm --config <your-e2e-config>.yml --port 4000`. Setting the variable only on the pytest process does not configure the proxy. With prompt storage disabled, the proxy correctly redacts `guardrail_response`, so that configuration cannot exercise this test's entity-detail assertions. Enable this only on a test stack using synthetic prompts
+
 A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships
 
 ### The pull request check

tests/e2e/coverage_registry/guardrail.yaml

@@ -3,6 +3,7 @@
 - {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"}
 - {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"}
 - {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"}
+- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"}
 - {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
 - {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
 - {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"}

tests/e2e/guardrails/test_presidio_masking_e2e.py

@@ -19,6 +19,12 @@ PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip.
 Each guardrail registers with an explicit presidio_filter_scope so only the
 configured hook's callback exists (the default "both" registers input masking
 AND a post_call output masker), and is deleted on teardown.
+
+The spend-log audit test requires general_settings.store_prompts_in_spend_logs:
+true in the proxy config, or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy
+process environment before startup. Setting it only on pytest has no effect.
+Without this opt-in, redacting guardrail_response is expected proxy behavior;
+this suite deliberately requires the detected-entity details to remain visible.
 """
 
 from __future__ import annotations
@@ -26,16 +32,22 @@ from __future__ import annotations
 import os
 import time
 from collections.abc import Callable
-from typing import Literal
+from typing import Final, Literal
 
 import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 from e2e_config import unique_marker
-from e2e_http import Result, Success
+from e2e_http import Result, StreamingResponse, Success
 from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody
 from lifecycle import ResourceManager
-from models import AnthropicMessagesResponse, ChatResponse
+from models import (
+    AnthropicMessagesResponse,
+    ChatResponse,
+    GuardrailEntityMatch,
+    GuardrailRunRecord,
+    SpendLogRow,
+)
 
 pytestmark = pytest.mark.e2e
 
@@ -276,3 +288,114 @@ class TestPresidioPostCallMasking:
                     f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
                 )
             time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
+
+
+_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
+
+_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch])
+
+
+def _applied_guardrails(outcome: StreamingResponse) -> str:
+    return outcome.headers.get("x-litellm-applied-guardrails", "")
+
+
+def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]:
+    metadata = row.metadata
+    if metadata is None:
+        return ()
+    return tuple(metadata.guardrail_information or ())
+
+
+def _poll_until_guardrail_applied(
+    client: GuardrailsClient, key: str, guardrail_name: str, prompt: str
+) -> StreamingResponse:
+    deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
+    last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
+    while time.monotonic() < deadline:
+        if last.ok and guardrail_name in _applied_guardrails(last):
+            return last
+        time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
+        last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
+    return last
+
+
+class TestPresidioSpendLogRecord:
+    @pytest.mark.covers(
+        "guardrail.presidio.pre_call.logs_masked_entities",
+        exercised_on=["chat_completions"],
+    )
+    def test_masking_run_is_recorded_on_the_spend_log(
+        self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
+    ) -> None:
+        name = f"e2e-presidio-log-{unique_marker()}"
+        _register_presidio(client, resources, name=name, entities=_LOGGED_ENTITIES)
+
+        email = _fake_email()
+        prompt = _pii_prompt(unique_marker(), email)
+
+        outcome = _poll_until_guardrail_applied(client, scoped_key, name, prompt)
+        assert outcome.ok, f"the guarded call must be served, got {outcome.status_code}: {outcome.body[:400]}"
+        assert name in _applied_guardrails(outcome), (
+            "the response must carry x-litellm-applied-guardrails naming the guardrail; without it "
+            f"the 200 only proves the guardrail never attached. Got {_applied_guardrails(outcome)!r}"
+        )
+
+        request_id = ChatResponse.model_validate_json(outcome.body).id
+        assert request_id, f"the served response must carry an id to look the spend log up by: {outcome.body[:400]}"
+
+        rows = client.proxy.poll_logs_for_request_id(
+            request_id,
+            predicate=lambda logged: bool(_guardrail_records(logged[0])),
+        )
+        assert rows, f"no spend log row ever appeared for request {request_id}"
+
+        records = _guardrail_records(rows[0])
+        assert records, (
+            f"the spend log for {request_id} carries no guardrail_information, so the dashboard's "
+            "guardrail panel would render nothing for a request the guardrail demonstrably ran on"
+        )
+
+        record = next(
+            (entry for entry in records if entry.guardrail_name == name and entry.guardrail_mode == "pre_call"),
+            None,
+        )
+        assert record is not None, (
+            "guardrail_information carries no pre_call record for this guardrail, only "
+            f"{[(entry.guardrail_name, entry.guardrail_mode) for entry in records]}"
+        )
+        assert record.guardrail_status == "success", (
+            f"the recorded status must be success for a run that masked and served, got {record.guardrail_status!r}"
+        )
+        assert record.guardrail_provider == "presidio", (
+            f"the record must attribute the run to presidio so the dashboard picks the right "
+            f"renderer, got {record.guardrail_provider!r}"
+        )
+
+        counts = record.masked_entity_count or {}
+        assert _LOGGED_ENTITIES.keys() <= counts.keys(), (
+            f"every entity the guardrail was configured to mask must appear in masked_entity_count, got {counts}"
+        )
+        assert all(counts[entity] >= 1 for entity in _LOGGED_ENTITIES), (
+            f"each masked entity must be counted at least once, got {counts}"
+        )
+
+        assert not isinstance(record.guardrail_response, str), (
+            "This audit test requires general_settings.store_prompts_in_spend_logs: true in the proxy config "
+            "or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy process environment before startup "
+            "(not just the pytest environment). Default prompt redaction is valid proxy behavior, "
+            f"but prevents entity-detail assertions; got guardrail_response={record.guardrail_response!r}"
+        )
+        entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response)
+        assert entities, (
+            "guardrail_response must carry the detected entities; the dashboard's Detected Entities "
+            "list and its per-entity scores are rendered from exactly this array"
+        )
+        assert {entity.entity_type for entity in entities} >= _LOGGED_ENTITIES.keys(), (
+            f"the detected entities must cover what was masked, got "
+            f"{sorted(entity.entity_type for entity in entities)}"
+        )
+        for entity in entities:
+            assert 0.0 < entity.score <= 1.0, f"{entity.entity_type} carries an out-of-range score: {entity.score}"
+            assert 0 <= entity.start < entity.end, (
+                f"{entity.entity_type} carries a degenerate span: {entity.start}-{entity.end}"
+            )

tests/e2e/models.py

@@ -584,6 +584,27 @@ class OcrResponse(BaseModel):
 # ---------- spend logs ----------
 
 
+class GuardrailEntityMatch(BaseModel):
+    entity_type: str
+    score: float
+    start: int
+    end: int
+
+
+class GuardrailRunRecord(BaseModel):
+    guardrail_name: str | None = None
+    guardrail_mode: str | None = None
+    guardrail_status: str | None = None
+    guardrail_provider: str | None = None
+    masked_entity_count: dict[str, int] | None = None
+    guardrail_response: object | None = None
+
+
+class SpendLogMetadata(BaseModel):
+    applied_guardrails: list[str] | None = None
+    guardrail_information: list[GuardrailRunRecord] | None = None
+
+
 class SpendLogRow(BaseModel):
     request_id: str | None = None
     api_key: str | None = None
@@ -600,6 +621,7 @@ class SpendLogRow(BaseModel):
     completion_tokens: int | None = None
     total_tokens: int | None = None
     request_tags: list[str] | None = None
+    metadata: SpendLogMetadata | None = None
 
 
 class SpendLogs(RootModel[list[SpendLogRow]]):

tests/e2e/ui/run_e2e.sh

@@ -151,6 +151,7 @@ export DISABLE_SCHEMA_UPDATE="true"
 export E2E_UI_BASE_URL="${E2E_UI_BASE_URL:-http://127.0.0.1:${PROXY_PORT}}"
 # Ensure the proxy serves UI at /ui (not behind a subpath)
 export SERVER_ROOT_PATH=""
+export PROXY_BASE_URL="$E2E_UI_BASE_URL"
 # Boot with an external logout URL so proxyLogoutUrl.spec.ts can assert the
 # redirect. This same value is exported to the Playwright process below (the
 # spec's skip guard reads it). Safe for the rest of the suite — nothing else

tests/e2e/ui/tests/guardrails/guardrails.spec.ts

@@ -7,6 +7,7 @@ import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traff
 interface StoredGuardrail {
   guardrail_id: string;
   guardrail_name: string | null;
+  litellm_params?: { pii_entities_config?: Record<string, string> | null } | null;
 }
 
 async function listGuardrails(page: PlaywrightPage): Promise<StoredGuardrail[]> {
@@ -249,6 +250,17 @@ test.describe("Guardrails", () => {
     const row = page.getByRole("row").filter({ hasText: guardrailName });
     await expect(row).toHaveCount(1, { timeout: 15_000 });
 
+    const stored = await findGuardrail(page, guardrailName);
+    const piiConfig = stored?.litellm_params?.pii_entities_config ?? {};
+    expect(
+      Object.keys(piiConfig).length,
+      "Select All & Mask persisted no PII entities, so the guardrail would mask nothing",
+    ).toBeGreaterThan(0);
+    expect(
+      Object.entries(piiConfig).filter(([, action]) => action !== "MASK"),
+      "Select All & Mask must persist every entity with the MASK action",
+    ).toEqual([]);
+
     await navigateToPage(page, Page.Teams);
     await dismissFeedbackPopup(page);
     await clickTeamId(page, E2E_TEAM_NO_ADMIN_ID);