fix(guardrails): fail open on Prompt Security file timeouts

litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py

@@ -4,10 +4,12 @@ import os
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Final, Literal, Optional
 
+import httpx
 from fastapi import HTTPException
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
+from litellm.exceptions import Timeout as LiteLLMTimeout
 from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     log_guardrail_information,
@@ -24,6 +26,9 @@ if TYPE_CHECKING:
     from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
 
 
+_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
+
+
 class PromptSecurityGuardrailMissingSecrets(Exception):
     pass
 
@@ -63,6 +68,13 @@ class _SanitizeStatusResponse(TypedDict, total=False):
     metadata: ReadOnly[_SanitizeMetadata]
 
 
+class _SanitizeResult(TypedDict):
+    action: ReadOnly[str]
+    content: ReadOnly[str | None]
+    metadata: ReadOnly[_SanitizeMetadata]
+    violations: ReadOnly[Sequence[str]]
+
+
 class PromptSecurityGuardrail(CustomGuardrail):
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@@ -79,6 +91,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         user: str | None = None,
         system_prompt: str | None = None,
         check_tool_results: bool | None = None,
+        file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
         **kwargs,
     ):
         kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
@@ -108,6 +121,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         # Configuration for file sanitization
         self.max_poll_attempts = 30  # Maximum number of polling attempts
         self.poll_interval = 2  # Seconds between polling attempts
+        self.file_sanitization_timeout = file_sanitization_timeout
 
         super().__init__(**kwargs)
 
@@ -397,6 +411,31 @@ class PromptSecurityGuardrail(CustomGuardrail):
         Sanitize file content using Prompt Security API.
         Returns: dict with keys 'action', 'content', 'metadata'
         """
+        try:
+            return await asyncio.wait_for(
+                self._sanitize_file_content(file_data, filename, user_api_key_alias),
+                timeout=self.file_sanitization_timeout,
+            )
+        except (asyncio.TimeoutError, httpx.TimeoutException, LiteLLMTimeout) as exc:
+            verbose_proxy_logger.error(
+                "Prompt Security Guardrail: file sanitization for %s timed out with %s; failing open",
+                filename,
+                type(exc).__name__,
+            )
+            fail_open_result: Final[_SanitizeResult] = {
+                "action": "allow",
+                "content": None,
+                "metadata": {},
+                "violations": (),
+            }
+            return fail_open_result
+
+    async def _sanitize_file_content(
+        self,
+        file_data: bytes,
+        filename: str,
+        user_api_key_alias: str | None,
+    ) -> _SanitizeResult:
         headers: Final = {"APP-ID": self.api_key}
         if user_api_key_alias:
             headers["X-LiteLLM-Key-Alias"] = user_api_key_alias

tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py

@@ -1,16 +1,16 @@
-from fastapi.exceptions import HTTPException
-from unittest.mock import patch, AsyncMock
-from httpx import Response, Request
+import asyncio
 import base64
+from unittest.mock import AsyncMock, patch
 
 import pytest
-
-from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import (
-    PromptSecurityGuardrailMissingSecrets,
-    PromptSecurityGuardrail,
-)
+from fastapi.exceptions import HTTPException
+from httpx import ReadTimeout, Request, Response
 
 import litellm
+from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import (
+    PromptSecurityGuardrail,
+    PromptSecurityGuardrailMissingSecrets,
+)
 from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
 
 
@@ -374,6 +374,62 @@ async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch):
     assert result is not None
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "timeout",
+    (
+        litellm.Timeout(
+            message="Prompt Security upload timed out",
+            model="default-model-name",
+            llm_provider="litellm-httpx-handler",
+        ),
+        ReadTimeout(
+            "Prompt Security poll timed out",
+            request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"),
+        ),
+    ),
+    ids=("litellm", "httpx"),
+)
+async def test_file_sanitization_fails_open_on_request_timeout(monkeypatch: pytest.MonkeyPatch, timeout: Exception):
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+
+    guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True)
+
+    with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=timeout)):
+        result = await guardrail.sanitize_file_content(b"file-content", "document.pdf")
+
+    assert result == {
+        "action": "allow",
+        "content": None,
+        "metadata": {},
+        "violations": (),
+    }
+
+
+@pytest.mark.asyncio
+async def test_file_sanitization_fails_open_on_overall_timeout(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+
+    guardrail = PromptSecurityGuardrail(
+        guardrail_name="test-guard",
+        event_hook="pre_call",
+        default_on=True,
+        file_sanitization_timeout=0.01,
+    )
+
+    async def hanging_post(*args, **kwargs):
+        await asyncio.sleep(60)
+        raise AssertionError("sanitization request should have been cancelled")
+
+    with patch.object(guardrail.async_handler, "post", side_effect=hanging_post):
+        result = await guardrail.sanitize_file_content(b"file-content", "document.pdf")
+
+    assert result["action"] == "allow"
+    assert result["content"] is None
+
+
 @pytest.mark.asyncio
 async def test_file_sanitization_block(monkeypatch: pytest.MonkeyPatch):
     """Test that file sanitization blocks malicious files"""
@@ -544,7 +600,7 @@ async def test_role_filtering(monkeypatch: pytest.MonkeyPatch):
         return mock_response
 
     with patch.object(guardrail.async_handler, "post", side_effect=mock_post):
-        result = await guardrail.apply_guardrail(
+        await guardrail.apply_guardrail(
             inputs=inputs,
             request_data=request_data,
             input_type="request",