Merge pull request #40041 from BerriAI/litellm_presidio_ui_user_story_e2e

test(e2e/ui): automate the RC checklist's Presidio guardrail walk

.circleci/config.yml

@@ -2648,6 +2648,19 @@ jobs:
           name: Start mock LLM server
           command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
           background: true
+      - run:
+          name: Start mock Presidio server
+          command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
+          background: true
+      - run:
+          name: Wait for mock Presidio server
+          command: |
+            for i in $(seq 1 30); do
+              if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
+              sleep 1
+            done
+            echo "Mock Presidio server never answered /health on port 8091" >&2
+            exit 1
       - run:
           name: Start LiteLLM proxy
           environment:
@@ -2778,6 +2791,19 @@ jobs:
           name: Start mock LLM server
           command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
           background: true
+      - run:
+          name: Start mock Presidio server
+          command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
+          background: true
+      - run:
+          name: Wait for mock Presidio server
+          command: |
+            for i in $(seq 1 30); do
+              if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
+              sleep 1
+            done
+            echo "Mock Presidio server never answered /health on port 8091" >&2
+            exit 1
       - run:
           name: Start LiteLLM proxy under a server root path
           environment:

tests/e2e/ui/constants.ts

@@ -15,6 +15,8 @@ export const UI_BASE_URL = (
 // writable path (the image runner already exports TMPDIR) to relocate them.
 export const ARTIFACT_DIR = process.env.E2E_UI_ARTIFACT_DIR || ".";
 
+export const MOCK_PRESIDIO_URL = (process.env.E2E_MOCK_PRESIDIO_URL || "http://127.0.0.1:8091").replace(/\/+$/, "");
+
 const storagePath = (name: string): string => path.join(ARTIFACT_DIR, name);
 
 // Storage state paths for each role

tests/e2e/ui/fixtures/mock_presidio_server/server.py

@@ -0,0 +1,107 @@
+"""
+Mock Presidio analyzer + anonymizer for UI e2e tests.
+Serves POST /analyze and POST /anonymize over a fixed set of regex recognizers.
+"""
+
+import os
+import re
+
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+
+
+def _luhn_ok(candidate):
+    digits = [int(c) for c in candidate if c.isdigit()]
+    doubled = [d * 2 - 9 if d * 2 > 9 else d * 2 for d in digits[-2::-2]]
+    return len(digits) >= 13 and (sum(digits[::-1][::2]) + sum(doubled)) % 10 == 0
+
+
+RECOGNIZERS = (
+    ("EMAIL_ADDRESS", re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), 1.0, None),
+    ("CREDIT_CARD", re.compile(r"\b(?:\d[ -]?){13,19}\b"), 1.0, _luhn_ok),
+    ("US_SSN", re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), 0.85, None),
+    ("PHONE_NUMBER", re.compile(r"\b(?:\+?\d{1,2}[ .-]?)?(?:\(\d{3}\)|\d{3})[ .-]?\d{3}[ .-]?\d{4}\b"), 0.75, None),
+    ("IP_ADDRESS", re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), 0.6, None),
+    ("URL", re.compile(r"\bhttps?://[^\s]+"), 0.5, None),
+)
+
+app = FastAPI(title="Mock Presidio Server")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+def _detect(text, wanted):
+    found = [
+        {"entity_type": name, "start": m.start(), "end": m.end(), "score": score}
+        for name, pattern, score, validate in RECOGNIZERS
+        if wanted is None or name in wanted
+        for m in pattern.finditer(text)
+        if validate is None or validate(m.group())
+    ]
+    ranked = sorted(found, key=lambda r: (-r["score"], r["start"]))
+    kept = []
+    for candidate in ranked:
+        overlaps = any(candidate["start"] < k["end"] and k["start"] < candidate["end"] for k in kept)
+        if not overlaps:
+            kept.append(candidate)
+    return sorted(kept, key=lambda r: r["start"])
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+@app.post("/analyze")
+async def analyze(request: Request):
+    body = await request.json()
+    text = body.get("text", "") or ""
+    entities = body.get("entities")
+    wanted = set(entities) if entities else None
+    threshold = body.get("score_threshold") or 0.0
+    return [
+        {**result, "analysis_explanation": None, "recognition_metadata": {"recognizer_name": "MockRecognizer"}}
+        for result in _detect(text, wanted)
+        if result["score"] >= threshold
+    ]
+
+
+@app.post("/anonymize")
+async def anonymize(request: Request):
+    body = await request.json()
+    text = body.get("text", "") or ""
+    results = sorted(body.get("analyzer_results") or [], key=lambda r: r["start"])
+
+    pieces = []
+    items = []
+    cursor = 0
+    for result in results:
+        start, end = result["start"], result["end"]
+        if start < cursor:
+            continue
+        placeholder = f"<{result['entity_type']}>"
+        pieces.append(text[cursor:start])
+        masked_start = sum(len(p) for p in pieces)
+        pieces.append(placeholder)
+        items.append(
+            {
+                "operator": "replace",
+                "entity_type": result["entity_type"],
+                "start": masked_start,
+                "end": masked_start + len(placeholder),
+                "text": placeholder,
+            }
+        )
+        cursor = end
+    pieces.append(text[cursor:])
+
+    return {"text": "".join(pieces), "items": list(reversed(items))}
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("MOCK_PRESIDIO_PORT", "8091")))

tests/e2e/ui/run_e2e.sh

@@ -14,7 +14,7 @@ set -euo pipefail
 #
 # Ports default to 4000 / 5432 / 8090 and can be moved when another checkout
 # already holds them:
-#   PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 ./run_e2e.sh
+#   PROXY_PORT=4100 POSTGRES_PORT=5532 MOCK_LLM_PORT=8190 MOCK_PRESIDIO_PORT=8191 ./run_e2e.sh
 #
 # In CI (CI=true), expects:
 #   - PostgreSQL already running on 127.0.0.1:5432
@@ -29,6 +29,7 @@ DASHBOARD_DIR="$REPO_ROOT/ui/litellm-dashboard"
 IS_CI="${CI:-false}"
 CONTAINER_NAME="litellm-e2e-postgres-$$"
 MOCK_PID=""
+MOCK_PRESIDIO_PID=""
 PROXY_PID=""
 PROXY_LOG=""
 
@@ -40,7 +41,8 @@ PROXY_LOG=""
 PROXY_PORT="${PROXY_PORT:-4000}"
 POSTGRES_PORT="${POSTGRES_PORT:-5432}"
 MOCK_LLM_PORT="${MOCK_LLM_PORT:-8090}"
-export MOCK_LLM_PORT
+MOCK_PRESIDIO_PORT="${MOCK_PRESIDIO_PORT:-8091}"
+export MOCK_LLM_PORT MOCK_PRESIDIO_PORT
 
 # --- Ensure common tool paths are available (local dev only) ---
 if [ "$IS_CI" = "false" ]; then
@@ -82,6 +84,7 @@ fi
 cleanup() {
   echo "Cleaning up..."
   [ -n "$MOCK_PID" ] && kill "$MOCK_PID" 2>/dev/null || true
+  [ -n "$MOCK_PRESIDIO_PID" ] && kill "$MOCK_PRESIDIO_PID" 2>/dev/null || true
   [ -n "$PROXY_PID" ] && kill "$PROXY_PID" 2>/dev/null || true
   [ -n "$PROXY_LOG" ] && rm -f "$PROXY_LOG" || true
   if [ "$IS_CI" = "false" ]; then
@@ -110,9 +113,9 @@ if [ "$IS_CI" = "false" ]; then
   # to someone else's :5432 (a psql session, a running app, a Prisma engine
   # talking to a remote database) aborts the run with "port 5432 is in use"
   # while nothing is actually bound locally.
-  for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT"; do
+  for port in "$PROXY_PORT" "$POSTGRES_PORT" "$MOCK_LLM_PORT" "$MOCK_PRESIDIO_PORT"; do
     if lsof -nP -iTCP:"$port" -sTCP:LISTEN >/dev/null 2>&1; then
-      echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT)"
+      echo "Error: port $port is in use (override with PROXY_PORT / POSTGRES_PORT / MOCK_LLM_PORT / MOCK_PRESIDIO_PORT)"
       exit 1
     fi
   done
@@ -143,6 +146,7 @@ fi
 # --- Credentials ---
 export LITELLM_MASTER_KEY="sk-1234"
 export MOCK_LLM_URL="http://127.0.0.1:${MOCK_LLM_PORT}/v1"
+export E2E_MOCK_PRESIDIO_URL="http://127.0.0.1:${MOCK_PRESIDIO_PORT}"
 export DISABLE_SCHEMA_UPDATE="true"
 # The suite resolves its target from E2E_UI_BASE_URL (constants.ts), which
 # otherwise defaults to :4000 -- so without this a relocated stack would be
@@ -204,6 +208,20 @@ for i in $(seq 1 15); do
   sleep 1
 done
 
+echo "=== Starting mock Presidio server ==="
+uv run --no-sync python "$SCRIPT_DIR/fixtures/mock_presidio_server/server.py" &
+MOCK_PRESIDIO_PID=$!
+
+PRESIDIO_READY=0
+for i in $(seq 1 15); do
+  if curl -sf http://127.0.0.1:${MOCK_PRESIDIO_PORT}/health >/dev/null 2>&1; then PRESIDIO_READY=1; break; fi
+  sleep 1
+done
+if [ "$PRESIDIO_READY" -ne 1 ]; then
+  echo "Mock Presidio server never answered /health on port ${MOCK_PRESIDIO_PORT}" >&2
+  exit 1
+fi
+
 # --- LiteLLM proxy ---
 echo "=== Starting LiteLLM proxy ==="
 cd "$REPO_ROOT"

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

@@ -0,0 +1,157 @@
+import { test, expect, type Page as PlaywrightPage } from "@playwright/test";
+import { ADMIN_STORAGE_PATH, MOCK_PRESIDIO_URL } from "../../constants";
+import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation";
+import { Page } from "../../fixtures/pages";
+import { CHAT_MODEL_A, masterKey, rootPath, waitForSpendLogByPrompt } from "../../helpers/traffic";
+import { openPlayground, selectModel, sendButton, onlyVisible } from "../../helpers/playground";
+
+const RAW_EMAIL = "jane.doe@example.com";
+const RAW_PHONE = "555-867-5309";
+
+const visibleTestId = (page: PlaywrightPage, id: string) => page.getByTestId(id).filter({ visible: true });
+
+const requestLogsRows = (page: PlaywrightPage) =>
+  page.locator("table").filter({ visible: true }).first().locator("tbody tr");
+
+async function createPresidioGuardrail(page: PlaywrightPage, guardrailName: string): Promise<void> {
+  await navigateToPage(page, Page.Guardrails);
+  await dismissFeedbackPopup(page);
+
+  await page.getByRole("button", { name: /Add New Guardrail/i }).click();
+  await page.getByRole("menuitem", { name: "Add Provider Guardrail" }).click();
+
+  const dialog = page.getByRole("dialog", { name: "Create guardrail" });
+  await expect(dialog).toBeVisible({ timeout: 10_000 });
+
+  await dialog.getByLabel("Guardrail Name").fill(guardrailName);
+
+  const providerSelect = dialog.getByRole("combobox", { name: "Guardrail Provider" });
+  await providerSelect.click();
+  await providerSelect.fill("Presidio");
+  await page.getByRole("option", { name: "Presidio PII" }).click();
+
+  await dialog.getByLabel("Mode", { exact: true }).click();
+  await page.keyboard.type("pre_call");
+  await expect(page.getByRole("option", { name: "pre_call" })).toBeAttached({ timeout: 5_000 });
+  await page.keyboard.press("Enter");
+  await expect(dialog.getByText("pre_call", { exact: true })).toBeVisible({ timeout: 5_000 });
+  await dialog.getByText("Create guardrail", { exact: true }).click();
+
+  await dialog.getByLabel("presidio_analyzer_api_base").fill(MOCK_PRESIDIO_URL);
+  await dialog.getByLabel("presidio_anonymizer_api_base").fill(MOCK_PRESIDIO_URL);
+
+  await dialog.getByRole("button", { name: "Next" }).click();
+  await expect(dialog.getByText("Configure PII Protection")).toBeVisible({ timeout: 10_000 });
+  await dialog.getByRole("button", { name: "Select All & Mask" }).click();
+
+  await dialog.getByRole("button", { name: "Create Guardrail" }).click();
+  await expect(page.getByText("Guardrail created successfully").first()).toBeVisible({ timeout: 15_000 });
+  await expect(page.getByRole("row").filter({ hasText: guardrailName })).toHaveCount(1, { timeout: 15_000 });
+}
+
+async function deleteGuardrail(page: PlaywrightPage, guardrailName: string): Promise<void> {
+  await navigateToPage(page, Page.Guardrails);
+  await dismissFeedbackPopup(page);
+  const row = page.getByRole("row").filter({ hasText: guardrailName });
+  await expect(row).toHaveCount(1, { timeout: 15_000 });
+  await row.getByRole("button", { name: "Open guardrail actions" }).click();
+  await page.getByRole("menuitem", { name: "Delete" }).click();
+  const deleteModal = page.getByRole("dialog", { name: "Delete Guardrail" });
+  await expect(deleteModal).toBeVisible({ timeout: 5_000 });
+  await deleteModal.getByRole("button", { name: "Delete", exact: true }).click();
+  await expect(page.getByText(`Guardrail "${guardrailName}" deleted successfully`)).toBeVisible({ timeout: 10_000 });
+}
+
+test.describe("Presidio PII guardrail, end to end from the dashboard", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH });
+
+  test("masks PII sent from the Playground and shows the run in Logs", async ({ page, request }) => {
+    const guardrailName = `e2e-presidio-story-${Date.now()}`;
+    const marker = `case-ref-${Math.random().toString(36).slice(2, 10)}`;
+    const prompt = `${marker}. Email me at ${RAW_EMAIL} or call ${RAW_PHONE}.`;
+
+    await createPresidioGuardrail(page, guardrailName);
+
+    await openPlayground(page);
+    await selectModel(page, CHAT_MODEL_A);
+
+    const guardrailSelect = onlyVisible(page.getByPlaceholder("Select guardrails"));
+    await expect(guardrailSelect).toBeVisible({ timeout: 20_000 });
+    await guardrailSelect.click();
+    await guardrailSelect.fill(guardrailName);
+    await onlyVisible(page.getByRole("option", { name: guardrailName, exact: true })).click({ timeout: 20_000 });
+    await page.keyboard.press("Escape");
+
+    const input = onlyVisible(page.getByPlaceholder("Type your message", { exact: false }));
+    await expect(input).toBeVisible({ timeout: 15_000 });
+
+    await expect
+      .poll(
+        async () => {
+          await input.fill(prompt);
+          await sendButton(page).click();
+          const res = await request.get(`${rootPath()}/spend/logs`, {
+            headers: { Authorization: `Bearer ${masterKey()}` },
+          });
+          if (!res.ok()) return false;
+          const rows: { metadata?: { applied_guardrails?: string[] } }[] = await res.json();
+          return (Array.isArray(rows) ? rows : []).some((row) =>
+            (row.metadata?.applied_guardrails ?? []).includes(guardrailName),
+          );
+        },
+        {
+          message: `the playground never produced a request that ran ${guardrailName}`,
+          timeout: 90_000,
+          intervals: [5_000],
+        },
+      )
+      .toBe(true);
+
+    const requestId = await waitForSpendLogByPrompt(request, marker);
+
+    const stored = await request.get(`${rootPath()}/spend/logs?request_id=${requestId}`, {
+      headers: { Authorization: `Bearer ${masterKey()}` },
+    });
+    expect(stored.ok(), `spend log read failed: ${stored.status()}`).toBe(true);
+    const storedBody = JSON.stringify(await stored.json());
+    expect(storedBody, "the raw email reached the spend log, so the prompt was stored unmasked").not.toContain(
+      RAW_EMAIL,
+    );
+    expect(storedBody, "the raw phone number reached the spend log, so the prompt was stored unmasked").not.toContain(
+      RAW_PHONE,
+    );
+    expect(storedBody, "the stored prompt carries no EMAIL_ADDRESS placeholder, so nothing was masked").toContain(
+      "<EMAIL_ADDRESS>",
+    );
+    expect(storedBody, "the stored prompt carries no PHONE_NUMBER placeholder, so nothing was masked").toContain(
+      "<PHONE_NUMBER>",
+    );
+
+    await navigateToPage(page, Page.Logs);
+    await dismissFeedbackPopup(page);
+    const search = visibleTestId(page, "datatable-search");
+    await expect(search).toBeVisible({ timeout: 20_000 });
+    await search.fill(requestId);
+    const row = requestLogsRows(page).filter({ hasText: requestId });
+    await expect(row, `no logs row for request ${requestId}`).toHaveCount(1, { timeout: 30_000 });
+    await row.click();
+
+    const drawer = page.getByRole("dialog").first();
+    await expect(onlyVisible(drawer.getByText("Guardrails & Policy Compliance"))).toBeVisible({ timeout: 20_000 });
+    await expect(onlyVisible(drawer.getByText(`Pre-call guardrail: ${guardrailName}`))).toBeVisible({ timeout: 20_000 });
+    const maskedPrompt = drawer.getByText(`${marker}. Email me at <EMAIL_ADDRESS> or call <PHONE_NUMBER>.`);
+    await expect(onlyVisible(maskedPrompt)).toBeVisible({ timeout: 20_000 });
+
+    await onlyVisible(drawer.getByText("2 matched")).click();
+    await expect(onlyVisible(drawer.getByText("Detected Entities (2)"))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("EMAIL_ADDRESS", { exact: true }))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("PHONE_NUMBER", { exact: true }))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("Score: 1.00", { exact: true }))).toBeVisible({ timeout: 10_000 });
+    await expect(onlyVisible(drawer.getByText("Score: 0.75", { exact: true }))).toBeVisible({ timeout: 10_000 });
+
+    await expect(drawer.getByText(RAW_EMAIL)).toHaveCount(0);
+    await expect(drawer.getByText(RAW_PHONE)).toHaveCount(0);
+
+    await deleteGuardrail(page, guardrailName);
+  });
+});