feat(pillar): add masking support and MCP call support (#17959)

- Add 'mask' action to SUPPORTED_ON_FLAGGED_ACTIONS - Automatically sanitizes sensitive content using masked_session_messages - Allows requests to proceed with masked content instead of blocking - Add MCP call support - Add pre_mcp_call and during_mcp_call to supported_event_hooks - Verify mcp_call is supported in call_type Literal types - Control exception details based on config - Conditionally include scanners/evidence in exceptions based on include_scanners and include_evidence settings - Reduces payload size when detailed exception info isn't needed - Add comprehensive test coverage - Tests for masking functionality - Tests for conditional exception details - Tests for MCP call support - Update documentation - Add Mask section explaining masking functionality - Clarify exception details control All changes maintain backward compatibility.
2026-09-14 23:21:35 +00:00 · 2025-12-16 05:19:03 +02:00 · 2025-12-16 05:19:03 +02:00 · 636efb7795
commit 636efb7795
parent ae7c994526
3 changed files with 418 additions and 10 deletions
--- a/docs/my-website/docs/proxy/guardrails/pillar_security.md
+++ b/docs/my-website/docs/proxy/guardrails/pillar_security.md
@ -72,13 +72,15 @@ litellm --config config.yaml --port 4000

 ### Overview

-Pillar Security supports three execution modes for comprehensive protection:
+Pillar Security supports five execution modes for comprehensive protection:

 | Mode | When It Runs | What It Protects | Use Case
 |------|-------------|------------------|----------
 | **`pre_call`** | Before LLM call | User input only | Block malicious prompts, prevent prompt injection
 | **`during_call`** | Parallel with LLM call | User input only | Input monitoring with lower latency
 | **`post_call`** | After LLM response | Full conversation context | Output filtering, PII detection in responses
+| **`pre_mcp_call`** | Before MCP tool call | MCP tool inputs | Validate and sanitize MCP tool call arguments
+| **`during_mcp_call`** | During MCP tool call | MCP tool inputs | Real-time monitoring of MCP tool calls

 ### Why Dual Mode is Recommended

@ -198,6 +200,85 @@ litellm_settings:
  set_verbose: true                          # Enable detailed logging
 ```

+</TabItem>
+<TabItem value="masking" label="Masking Mode - Auto-Sanitize PII">
+
+**Best for:**
+- 🔒 **PII Protection**: Automatically sanitize sensitive data before sending to LLM
+- ✅ **Continue Workflows**: Allow requests to proceed with masked content
+- 🛡️ **Zero Trust**: Never expose sensitive data to LLM models
+- 📊 **Compliance**: Meet data privacy requirements without blocking legitimate requests
+
+```yaml
+model_list:
+  - model_name: gpt-4.1-mini
+    litellm_params:
+      model: openai/gpt-4.1-mini
+      api_key: os.environ/OPENAI_API_KEY
+
+guardrails:
+  - guardrail_name: "pillar-masking"
+    litellm_params:
+      guardrail: pillar
+      mode: "pre_call"                       # Scan input before LLM call
+      api_key: os.environ/PILLAR_API_KEY     # Your Pillar API key
+      api_base: os.environ/PILLAR_API_BASE   # Pillar API endpoint
+      on_flagged_action: "mask"             # Mask sensitive content instead of blocking
+      persist_session: true                  # Keep records for investigation
+      include_scanners: true                 # Understand which scanners triggered
+      include_evidence: true                 # Capture evidence for analysis
+      default_on: true                       # Enable for all requests
+
+general_settings:
+  master_key: "YOUR_LITELLM_PROXY_MASTER_KEY"
+
+litellm_settings:
+  set_verbose: true
+```
+
+**How it works:**
+1. User sends request with sensitive data: `"My email is john@example.com"`
+2. Pillar detects PII and returns masked version: `"My email is [MASKED_EMAIL]"`
+3. LiteLLM replaces original messages with masked messages
+4. Request proceeds to LLM with sanitized content
+5. User receives response without exposing sensitive data
+
+</TabItem>
+<TabItem value="mcp" label="MCP Call Protection">
+
+**Best for:**
+- 🤖 **Agent Workflows**: Protect MCP (Model Context Protocol) tool calls
+- 🔒 **Tool Input Validation**: Scan arguments passed to MCP tools
+- 🛡️ **Comprehensive Coverage**: Extend security to all LLM endpoints
+
+```yaml
+model_list:
+  - model_name: gpt-4.1-mini
+    litellm_params:
+      model: openai/gpt-4.1-mini
+      api_key: os.environ/OPENAI_API_KEY
+
+guardrails:
+  - guardrail_name: "pillar-mcp-guard"
+    litellm_params:
+      guardrail: pillar
+      mode: "pre_mcp_call"                   # Scan MCP tool call inputs
+      api_key: os.environ/PILLAR_API_KEY     # Your Pillar API key
+      api_base: os.environ/PILLAR_API_BASE   # Pillar API endpoint
+      on_flagged_action: "block"             # Block malicious MCP calls
+      default_on: true                       # Enable for all MCP calls
+
+general_settings:
+  master_key: "YOUR_LITELLM_PROXY_MASTER_KEY"
+
+litellm_settings:
+  set_verbose: true
+```
+
+**MCP Modes:**
+- `pre_mcp_call`: Scan MCP tool call inputs before execution
+- `during_mcp_call`: Monitor MCP tool calls in real-time
+
 </TabItem>
 </Tabs>

@ -251,6 +332,15 @@ Logs the violation but allows the request to proceed:
 on_flagged_action: "monitor"
 ```

+#### Mask
+Automatically sanitizes sensitive content (PII, secrets, etc.) in your messages before sending them to the LLM:
+
+```yaml
+on_flagged_action: "mask"
+```
+
+When masking is enabled, sensitive information is automatically replaced with masked versions, allowing requests to proceed safely without exposing sensitive data to the LLM.
+
 **Response Headers:**

 You can opt in to receiving detection details in response headers by configuring `include_scanners: true` and/or `include_evidence: true`. When enabled, these headers are included for **every request**—not just flagged ones—enabling comprehensive metrics, false positive analysis, and threat investigation.
@ -383,7 +473,8 @@ export PILLAR_TIMEOUT="5.0"
 **Quick takeaways**
 - Every request still runs *all* Pillar scanners; these options only change what comes back.
 - Choose richer responses when you need audit trails, lighter responses when latency or cost matters.
- Blocking is controlled by LiteLLM’s `on_flagged_action` configuration—Pillar headers do not change block/monitor behaviour.
+- Actions (block/monitor/mask) are controlled by LiteLLM's `on_flagged_action` configuration—Pillar headers are automatically set based on your config.
+- When blocking (`on_flagged_action: "block"`), the `include_scanners` and `include_evidence` settings control what details are included in the exception response.

 Pillar Security executes the full scanner suite on each call. The settings below tune the Protect response headers LiteLLM sends, letting you balance fidelity, retention, and latency.

@ -415,9 +506,10 @@ include_evidence: true    # → plr_evidence (default true in LiteLLM)
  ```
  Use when you only care about whether Pillar detected a threat.

-  > **📝 Note:** `flagged: true` means Pillar’s scanners recommend blocking. Pillar only reports this verdict—LiteLLM enforces your policy via the `on_flagged_action` configuration (no Pillar header controls it):
-  > - `on_flagged_action: "block"` → LiteLLM raises a 400 guardrail error
+  > **📝 Note:** `flagged: true` means Pillar's scanners recommend blocking. Pillar only reports this verdict—LiteLLM enforces your policy via the `on_flagged_action` configuration:
+  > - `on_flagged_action: "block"` → LiteLLM raises a 400 guardrail error (exception includes scanners/evidence based on `include_scanners`/`include_evidence` settings)
  > - `on_flagged_action: "monitor"` → LiteLLM logs the threat but still returns the LLM response
+  > - `on_flagged_action: "mask"` → LiteLLM replaces messages with masked versions and allows the request to proceed

 - **Scanner breakdown** (`include_scanners=true`)
  ```json
--- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py
@ -164,7 +164,7 @@ class PillarGuardrail(CustomGuardrail):
    using the Pillar Security API.
    """

-    SUPPORTED_ON_FLAGGED_ACTIONS = ["block", "monitor"]
+    SUPPORTED_ON_FLAGGED_ACTIONS = ["block", "monitor", "mask"]
    DEFAULT_ON_FLAGGED_ACTION = "monitor"
    SUPPORTED_FALLBACK_ACTIONS = ["allow", "block"]
    DEFAULT_FALLBACK_ACTION = "allow"
@ -280,6 +280,8 @@ class PillarGuardrail(CustomGuardrail):
            GuardrailEventHooks.pre_call,
            GuardrailEventHooks.during_call,
            GuardrailEventHooks.post_call,
+            GuardrailEventHooks.pre_mcp_call,
+            GuardrailEventHooks.during_mcp_call,
        ]

        super().__init__(
@ -773,6 +775,15 @@ class PillarGuardrail(CustomGuardrail):
            verbose_proxy_logger.warning("Pillar Guardrail: Threat detected")
            if self.on_flagged_action == "block":
                self._raise_pillar_detection_exception(pillar_response)
+            elif self.on_flagged_action == "mask":
+                verbose_proxy_logger.info("Pillar Guardrail: Masking mode - masking flagged content")
+                masked_messages = pillar_response.get("masked_session_messages", [])
+                if masked_messages:
+                    original_data["messages"] = masked_messages
+                else:
+                    verbose_proxy_logger.warning(
+                        "Pillar Guardrail: Masking requested but no masked_session_messages in response"
+                    )
            elif self.on_flagged_action == "monitor":
                verbose_proxy_logger.info("Pillar Guardrail: Monitoring mode - allowing flagged content to proceed")

@ -788,14 +799,20 @@ class PillarGuardrail(CustomGuardrail):
        Raises:
            HTTPException: Always raises with security detection details
        """
+        pillar_response_dict = {
+            "session_id": pillar_response.get("session_id"),
+        }
+        
+        # Conditionally include scanners and evidence based on config
+        if self.include_scanners:
+            pillar_response_dict["scanners"] = pillar_response.get("scanners", {})
+        if self.include_evidence:
+            pillar_response_dict["evidence"] = pillar_response.get("evidence", [])
+
        error_detail = {
            "error": "Blocked by Pillar Security Guardrail",
            "detection_message": "Security threats detected",
-            "pillar_response": {
-                "session_id": pillar_response.get("session_id"),
-                "scanners": pillar_response.get("scanners", {}),
-                "evidence": pillar_response.get("evidence", []),
-            },
+            "pillar_response": pillar_response_dict,
        }

        verbose_proxy_logger.warning("Pillar Guardrail: Request blocked - Security threats detected")
--- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py
+++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py
@ -1142,6 +1142,305 @@ def test_get_config_model():
    assert hasattr(config_model, "ui_friendly_name")


+# ============================================================================
+# MASKING TESTS
+# ============================================================================
+
+
+@pytest.fixture
+def pillar_masked_response():
+    """Fixture providing a Pillar API response with masked messages."""
+    return Response(
+        json={
+            "session_id": "test-session-123",
+            "flagged": True,
+            "masked_session_messages": [
+                {"role": "user", "content": "My email is [MASKED_EMAIL]"}
+            ],
+            "evidence": [
+                {
+                    "category": "pii",
+                    "type": "email",
+                    "evidence": "test@example.com",
+                }
+            ],
+            "scanners": {
+                "jailbreak": False,
+                "prompt_injection": False,
+                "pii": True,
+                "toxic_language": False,
+            },
+        },
+        status_code=200,
+        request=Request(
+            method="POST", url="https://api.pillar.security/api/v1/protect"
+        ),
+    )
+
+
+@pytest.fixture
+def pillar_mask_guardrail(env_setup):
+    """Fixture providing a PillarGuardrail instance in mask mode."""
+    return PillarGuardrail(
+        guardrail_name="pillar-mask",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        on_flagged_action="mask",
+    )
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_masking_mode(
+    pillar_mask_guardrail,
+    sample_request_data,
+    user_api_key_dict,
+    dual_cache,
+    pillar_masked_response,
+):
+    """Test pre-call hook masks content when action is 'mask'."""
+    original_messages = sample_request_data["messages"].copy()
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=pillar_masked_response,
+    ):
+        result = await pillar_mask_guardrail.async_pre_call_hook(
+            data=sample_request_data,
+            cache=dual_cache,
+            user_api_key_dict=user_api_key_dict,
+            call_type="completion",
+        )
+
+    # Messages should be replaced with masked messages
+    assert result["messages"] == pillar_masked_response.json()["masked_session_messages"]
+    assert result["messages"] != original_messages
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_masking_no_masked_messages(
+    pillar_mask_guardrail,
+    sample_request_data,
+    user_api_key_dict,
+    dual_cache,
+):
+    """Test masking mode when API doesn't return masked_session_messages."""
+    response_no_mask = Response(
+        json={
+            "session_id": "test-session-123",
+            "flagged": True,
+            # No masked_session_messages
+        },
+        status_code=200,
+        request=Request(
+            method="POST", url="https://api.pillar.security/api/v1/protect"
+        ),
+    )
+
+    original_messages = sample_request_data["messages"].copy()
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=response_no_mask,
+    ):
+        result = await pillar_mask_guardrail.async_pre_call_hook(
+            data=sample_request_data,
+            cache=dual_cache,
+            user_api_key_dict=user_api_key_dict,
+            call_type="completion",
+        )
+
+    # Messages should remain unchanged if no masked messages provided
+    assert result["messages"] == original_messages
+
+
+# ============================================================================
+# CONDITIONAL EXCEPTION DETAILS TESTS
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_exception_without_scanners(
+    sample_request_data,
+    user_api_key_dict,
+    dual_cache,
+    pillar_flagged_response,
+):
+    """Test exception excludes scanners when include_scanners is False."""
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-no-scanners",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        on_flagged_action="block",
+        include_scanners=False,
+        include_evidence=True,
+    )
+
+    with pytest.raises(HTTPException) as excinfo:
+        with patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            return_value=pillar_flagged_response,
+        ):
+            await guardrail.async_pre_call_hook(
+                data=sample_request_data,
+                cache=dual_cache,
+                user_api_key_dict=user_api_key_dict,
+                call_type="completion",
+            )
+
+    error_detail = excinfo.value.detail
+    assert "pillar_response" in error_detail
+    assert "scanners" not in error_detail["pillar_response"]
+    assert "evidence" in error_detail["pillar_response"]
+
+
+@pytest.mark.asyncio
+async def test_exception_without_evidence(
+    sample_request_data,
+    user_api_key_dict,
+    dual_cache,
+    pillar_flagged_response,
+):
+    """Test exception excludes evidence when include_evidence is False."""
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-no-evidence",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        on_flagged_action="block",
+        include_scanners=True,
+        include_evidence=False,
+    )
+
+    with pytest.raises(HTTPException) as excinfo:
+        with patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            return_value=pillar_flagged_response,
+        ):
+            await guardrail.async_pre_call_hook(
+                data=sample_request_data,
+                cache=dual_cache,
+                user_api_key_dict=user_api_key_dict,
+                call_type="completion",
+            )
+
+    error_detail = excinfo.value.detail
+    assert "pillar_response" in error_detail
+    assert "scanners" in error_detail["pillar_response"]
+    assert "evidence" not in error_detail["pillar_response"]
+
+
+@pytest.mark.asyncio
+async def test_exception_without_scanners_or_evidence(
+    sample_request_data,
+    user_api_key_dict,
+    dual_cache,
+    pillar_flagged_response,
+):
+    """Test exception excludes both scanners and evidence when both are False."""
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-minimal",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        on_flagged_action="block",
+        include_scanners=False,
+        include_evidence=False,
+    )
+
+    with pytest.raises(HTTPException) as excinfo:
+        with patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            return_value=pillar_flagged_response,
+        ):
+            await guardrail.async_pre_call_hook(
+                data=sample_request_data,
+                cache=dual_cache,
+                user_api_key_dict=user_api_key_dict,
+                call_type="completion",
+            )
+
+    error_detail = excinfo.value.detail
+    assert "pillar_response" in error_detail
+    pillar_response = error_detail["pillar_response"]
+    assert "scanners" not in pillar_response
+    assert "evidence" not in pillar_response
+    assert "session_id" in pillar_response  # session_id should always be present
+
+
+# ============================================================================
+# MCP CALL SUPPORT TESTS
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_mcp_call(
+    pillar_guardrail_instance,
+    sample_request_data,
+    user_api_key_dict,
+    dual_cache,
+    pillar_clean_response,
+):
+    """Test pre-call hook works with MCP call type."""
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=pillar_clean_response,
+    ):
+        result = await pillar_guardrail_instance.async_pre_call_hook(
+            data=sample_request_data,
+            cache=dual_cache,
+            user_api_key_dict=user_api_key_dict,
+            call_type="mcp_call",
+        )
+
+    assert result == sample_request_data
+
+
+@pytest.mark.asyncio
+async def test_moderation_hook_mcp_call(
+    pillar_guardrail_instance,
+    sample_request_data,
+    user_api_key_dict,
+    pillar_clean_response,
+):
+    """Test moderation hook works with MCP call type."""
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=pillar_clean_response,
+    ):
+        result = await pillar_guardrail_instance.async_moderation_hook(
+            data=sample_request_data,
+            user_api_key_dict=user_api_key_dict,
+            call_type="mcp_call",
+        )
+
+    assert result == sample_request_data
+
+
+@pytest.mark.asyncio
+async def test_mcp_call_masking(
+    pillar_mask_guardrail,
+    sample_request_data,
+    user_api_key_dict,
+    dual_cache,
+    pillar_masked_response,
+):
+    """Test masking works with MCP call type."""
+    original_messages = sample_request_data["messages"].copy()
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=pillar_masked_response,
+    ):
+        result = await pillar_mask_guardrail.async_pre_call_hook(
+            data=sample_request_data,
+            cache=dual_cache,
+            user_api_key_dict=user_api_key_dict,
+            call_type="mcp_call",
+        )
+
+    # Messages should be replaced with masked messages
+    assert result["messages"] == pillar_masked_response.json()["masked_session_messages"]
+    assert result["messages"] != original_messages
+
+
 if __name__ == "__main__":
    # Run the tests
    pytest.main([__file__, "-v"])