diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 2a02560f6cc..6f75c74405c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0 +_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"}) class PromptSecurityGuardrailMissingSecrets(Exception): @@ -513,6 +514,9 @@ class PromptSecurityGuardrail(CustomGuardrail): "violations": result.get("metadata", {}).get("violations", []), } + if status not in _SANITIZE_FILE_QUEUED_STATUSES: + raise HTTPException(status_code=500, detail=f"Unexpected sanitization status: {status}") + verbose_proxy_logger.debug( "Prompt Security Guardrail: File sanitization status=%s for jobId=%s (attempt %d/%d)", status, diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 2e44b4b91b8..ee12bc4d223 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -560,6 +560,35 @@ async def test_file_sanitization_never_finishing_job_times_out(monkeypatch: pyte assert exc_info.value.detail == "File sanitization timeout" +@pytest.mark.asyncio +@pytest.mark.parametrize("poll_body", [{"status": "failed"}, {}]) +async def test_file_sanitization_terminal_failure_does_not_fail_open(monkeypatch: pytest.MonkeyPatch, poll_body): + 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) + guardrail.poll_interval = 0 + upload_response = Response( + json={"jobId": "failed-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json=poll_body, + status_code=200, + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(return_value=poll_response)) as poll_mock: + with pytest.raises(HTTPException) as exc_info: + await guardrail.sanitize_file_content(b"file-content", "document.pdf") + + assert poll_mock.await_count == 1 + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == f"Unexpected sanitization status: {poll_body.get('status')}" + + @pytest.mark.asyncio @pytest.mark.parametrize( "timeout",