feat(pillar): add automatic LiteLLM context headers (#17076)

- Automatically pass LiteLLM virtual key context as X-LiteLLM-* headers - Includes key_alias, user_id, team_id, org_id, and user_email - No configuration required - always enabled for application/user tracking - Excludes sensitive data (metadata, API tokens) for security - Add comprehensive tests (30 tests, all passing) - Update documentation with header details
2026-07-25 06:21:17 +00:00 · 2025-11-25 19:35:39 -08:00
parent 7c09187daf
commit e6e1e8fca4
3 changed files with 273 additions and 9 deletions
@@ -60,6 +60,8 @@ litellm_settings:
  set_verbose: true                          # Enable detailed logging
 ```

+**Note:** Virtual key context is **automatically passed** as headers - no additional configuration needed!
+
 ### 3. Start the Proxy

 ```bash
@@ -210,7 +212,7 @@ export PILLAR_API_KEY="your_api_key_here"
 export PILLAR_API_BASE="https://api.pillar.security"
 export PILLAR_ON_FLAGGED_ACTION="monitor"
 export PILLAR_FALLBACK_ON_ERROR="allow"
-export PILLAR_TIMEOUT="30.0"
+export PILLAR_TIMEOUT="5.0"
 ```

 ### Session Tracking
@@ -90,6 +90,10 @@ class PillarGuardrail(CustomGuardrail):
            fallback_on_error: Action when API errors occur ('allow' or 'block')
            timeout: Timeout for API calls in seconds
            **kwargs: Additional arguments passed to parent class
+
+        Note:
+            LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always
+            automatically passed as X-LiteLLM-* headers to enable application/user tracking.
        """
        self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
        self.api_key = api_key or os.environ.get("PILLAR_API_KEY")
@@ -222,7 +226,7 @@ class PillarGuardrail(CustomGuardrail):
            return data

        verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook")
-        result = await self.run_pillar_guardrail(data)
+        result = await self.run_pillar_guardrail(data, user_api_key_dict)

        # Add guardrail name to response headers
        add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@@ -265,7 +269,7 @@ class PillarGuardrail(CustomGuardrail):
            return data

        verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook")
-        result = await self.run_pillar_guardrail(data)
+        result = await self.run_pillar_guardrail(data, user_api_key_dict)

        # Add guardrail name to response headers
        add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@@ -315,7 +319,7 @@ class PillarGuardrail(CustomGuardrail):
        post_call_data["messages"] = data.get("messages", []) + response_messages

        # Reuse the existing guardrail logic - zero duplication!
-        await self.run_pillar_guardrail(post_call_data)
+        await self.run_pillar_guardrail(post_call_data, user_api_key_dict)

        # Add guardrail name to response headers
        add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@@ -326,12 +330,13 @@ class PillarGuardrail(CustomGuardrail):
    # CORE LOGIC METHOD
    # =========================================================================

-    async def run_pillar_guardrail(self, data: dict) -> dict:
+    async def run_pillar_guardrail(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> dict:
        """
        Core method to run the Pillar guardrail scan.

        Args:
            data: Request data containing messages and metadata
+            user_api_key_dict: User API key authentication info containing key context

        Returns:
            Original data if safe or in monitor mode
@@ -345,7 +350,7 @@ class PillarGuardrail(CustomGuardrail):
            return data

        try:
-            headers = self._prepare_headers()
+            headers = self._prepare_headers(user_api_key_dict)
            payload = self._prepare_payload(data)

            response = await self._call_pillar_api(
@@ -403,8 +408,16 @@ class PillarGuardrail(CustomGuardrail):
                },
            )

-    def _prepare_headers(self) -> Dict[str, str]:
-        """Prepare headers for the Pillar API request."""
+    def _prepare_headers(self, user_api_key_dict: UserAPIKeyAuth) -> Dict[str, str]:
+        """
+        Prepare headers for the Pillar API request.
+
+        Args:
+            user_api_key_dict: User API key authentication info containing key context
+
+        Returns:
+            Dictionary of headers to send to Pillar API
+        """
        if not self.api_key:
            msg = (
                "Couldn't get Pillar API key, either set the `PILLAR_API_KEY` in the environment or "
@@ -415,7 +428,7 @@ class PillarGuardrail(CustomGuardrail):
        headers: Dict[str, str] = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
-        }
+        }    

        # Add Pillar-specific headers based on configuration
        self._set_bool_header(headers, "plr_scanners", self.include_scanners)
@@ -423,6 +436,20 @@ class PillarGuardrail(CustomGuardrail):
        self._set_bool_header(headers, "plr_async", self.async_mode)
        self._set_bool_header(headers, "plr_persist", self.persist_session)

+        # Always add LiteLLM virtual key context headers (metadata excluded for security)
+        context_mapping = {
+            "X-LiteLLM-Key-Name": user_api_key_dict.key_name,
+            "X-LiteLLM-Key-Alias": user_api_key_dict.key_alias,
+            "X-LiteLLM-User-Id": user_api_key_dict.user_id,
+            "X-LiteLLM-User-Email": user_api_key_dict.user_email,
+            "X-LiteLLM-Team-Id": user_api_key_dict.team_id,
+            "X-LiteLLM-Team-Name": user_api_key_dict.team_alias,
+            "X-LiteLLM-Org-Id": user_api_key_dict.org_id,
+        }
+        for header_name, value in context_mapping.items():
+            if value:
+                headers[header_name] = str(value)
+
        return headers

    def _set_bool_header(self, headers: Dict[str, str], header_name: str, value: Optional[bool]) -> None:
@@ -517,6 +544,14 @@ class PillarGuardrail(CustomGuardrail):
        """
        Prepare the payload for the Pillar API request following the /api/v1/protect contract.

+        This method supports multi-modal content (images, files, audio, video, etc.) as messages
+        are passed through without modification. The messages array can contain any OpenAI-compatible
+        message structure including:
+        - Text content (string)
+        - Multi-modal content blocks (image_url, image_file, audio, video, document, file)
+        - Attachments
+        - Tool calls
+
        Args:
            data: Request data

@@ -234,6 +234,22 @@ def pillar_async_response():
    )


+@pytest.fixture
+def user_api_key_dict_with_context():
+    """Fixture providing UserAPIKeyAuth with complete context."""
+    return UserAPIKeyAuth(
+        token="hashed-test-token",
+        key_name="production-api-key",
+        key_alias="prod-key",
+        user_id="user-123",
+        user_email="test@example.com",
+        team_id="team-456",
+        team_alias="engineering-team",
+        org_id="org-789",
+        metadata={"environment": "production", "region": "us-east-1"},
+    )
+
+
@pytest.fixture
 def mock_llm_response_with_tools():
    """Fixture providing a mock LLM response with tool calls."""
@@ -502,6 +518,217 @@ async def test_pre_call_hook_custom_header_overrides(
    assert captured_headers.get("plr_evidence") == "false"


+# =========================================================================
+# LITELLM KEY CONTEXT HEADER TESTS
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_litellm_context_headers_automatically_added(
+    sample_request_data,
+    user_api_key_dict_with_context,
+    dual_cache,
+    pillar_clean_response,
+):
+    """Test that LiteLLM context headers are automatically added (always enabled)."""
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-context-enabled",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+    )
+
+    captured_headers: Dict[str, str] = {}
+
+    async def _mock_post(*args, **kwargs):
+        captured_headers.update(kwargs.get("headers", {}))
+        return pillar_clean_response
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        new=_mock_post,
+    ):
+        await guardrail.async_pre_call_hook(
+            data=sample_request_data,
+            cache=dual_cache,
+            user_api_key_dict=user_api_key_dict_with_context,
+            call_type="completion",
+        )
+
+    # Verify LiteLLM context headers are present
+    assert "X-LiteLLM-Key-Name" in captured_headers
+    assert captured_headers["X-LiteLLM-Key-Name"] == "production-api-key"
+    assert "X-LiteLLM-Key-Alias" in captured_headers
+    assert captured_headers["X-LiteLLM-Key-Alias"] == "prod-key"
+    assert "X-LiteLLM-User-Id" in captured_headers
+    assert captured_headers["X-LiteLLM-User-Id"] == "user-123"
+    assert "X-LiteLLM-User-Email" in captured_headers
+    assert captured_headers["X-LiteLLM-User-Email"] == "test@example.com"
+    assert "X-LiteLLM-Team-Id" in captured_headers
+    assert captured_headers["X-LiteLLM-Team-Id"] == "team-456"
+    assert "X-LiteLLM-Team-Name" in captured_headers
+    assert captured_headers["X-LiteLLM-Team-Name"] == "engineering-team"
+    assert "X-LiteLLM-Org-Id" in captured_headers
+    assert captured_headers["X-LiteLLM-Org-Id"] == "org-789"
+    
+    # Metadata is NOT sent (may contain sensitive information)
+    assert "X-LiteLLM-Metadata" not in captured_headers
+
+
+@pytest.mark.asyncio
+async def test_litellm_context_with_partial_fields(
+    sample_request_data,
+    dual_cache,
+    pillar_clean_response,
+):
+    """Test that partial LiteLLM context (only some fields present) is handled correctly."""
+    # Create UserAPIKeyAuth with only some fields populated
+    partial_context = UserAPIKeyAuth(
+        user_id="user-only",
+        team_id="team-only",
+    )
+
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-partial-context",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        pass_litellm_key_header=True,
+    )
+
+    captured_headers: Dict[str, str] = {}
+
+    async def _mock_post(*args, **kwargs):
+        captured_headers.update(kwargs.get("headers", {}))
+        return pillar_clean_response
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        new=_mock_post,
+    ):
+        await guardrail.async_pre_call_hook(
+            data=sample_request_data,
+            cache=dual_cache,
+            user_api_key_dict=partial_context,
+            call_type="completion",
+        )
+
+    # Verify only populated fields are present
+    assert "X-LiteLLM-User-Id" in captured_headers
+    assert captured_headers["X-LiteLLM-User-Id"] == "user-only"
+    assert "X-LiteLLM-Team-Id" in captured_headers
+    assert captured_headers["X-LiteLLM-Team-Id"] == "team-only"
+
+    # Verify empty fields are not present
+    assert "X-LiteLLM-Key-Name" not in captured_headers
+    assert "X-LiteLLM-User-Email" not in captured_headers
+
+
+# =========================================================================
+# MULTI-MODAL CONTENT TESTS
+# =========================================================================
+
+
+@pytest.mark.asyncio
+async def test_multimodal_image_url_support(
+    user_api_key_dict,
+    dual_cache,
+    pillar_clean_response,
+):
+    """Test that messages with image URLs are properly handled."""
+    multimodal_data = {
+        "model": "gpt-4-vision-preview",
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What's in this image?"},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": "https://example.com/image.jpg",
+                            "detail": "high",
+                        },
+                    },
+                ],
+            }
+        ],
+    }
+
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-multimodal",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+    )
+
+    captured_payload: Dict[str, Any] = {}
+
+    async def _mock_post(*args, **kwargs):
+        captured_payload.update(kwargs.get("json", {}))
+        return pillar_clean_response
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        new=_mock_post,
+    ):
+        result = await guardrail.async_pre_call_hook(
+            data=multimodal_data,
+            cache=dual_cache,
+            user_api_key_dict=user_api_key_dict,
+            call_type="completion",
+        )
+
+    # Verify multimodal message structure is preserved
+    assert result == multimodal_data
+    assert "messages" in captured_payload
+    assert len(captured_payload["messages"]) == 1
+    assert isinstance(captured_payload["messages"][0]["content"], list)
+    assert captured_payload["messages"][0]["content"][1]["type"] == "image_url"
+
+
+@pytest.mark.asyncio
+async def test_multimodal_with_attachments(
+    user_api_key_dict,
+    dual_cache,
+    pillar_clean_response,
+):
+    """Test that messages with file attachments are properly handled."""
+    multimodal_data = {
+        "model": "gpt-4",
+        "messages": [
+            {
+                "role": "user",
+                "content": "Analyze this document",
+                "attachments": [
+                    {
+                        "file_id": "file-abc123",
+                        "tools": [{"type": "code_interpreter"}],
+                    }
+                ],
+            }
+        ],
+    }
+
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-attachments",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+    )
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=pillar_clean_response,
+    ):
+        result = await guardrail.async_pre_call_hook(
+            data=multimodal_data,
+            cache=dual_cache,
+            user_api_key_dict=user_api_key_dict,
+            call_type="completion",
+        )
+
+    # Verify attachment structure is preserved
+    assert result == multimodal_data
+    assert result["messages"][0]["attachments"] is not None
+
+
 # ============================================================================
 # EDGE CASE TESTS
 # ============================================================================