From fb60f80c804b48e3702ca431ee4dad06799e88d7 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:02:49 +0000 Subject: [PATCH 1/2] fix(prompt_security): keep polling file sanitization through non-terminal statuses Prompt Security reports a queued sanitization job as status "created" before it moves to "in progress" and "done". The poller treated anything other than those two known strings as an error and returned HTTP 500 on the first poll, so every image or file request through the guardrail failed while the vendor job was still queued. Only "done" is terminal now. Every other status is logged and polled again until max_poll_attempts or the outer file_sanitization_timeout, after which the existing fail-open or fail-closed (408) policy applies. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_security/prompt_security.py | 17 +++-- .../test_prompt_security_guardrails.py | 63 +++++++++++++++++++ 2 files changed, 71 insertions(+), 9 deletions(-) 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 0954fe1698a..2a02560f6cc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -512,15 +512,14 @@ class PromptSecurityGuardrail(CustomGuardrail): "metadata": result.get("metadata", {}), "violations": result.get("metadata", {}).get("violations", []), } - elif status == "in progress": - verbose_proxy_logger.debug( - "Prompt Security Guardrail: File sanitization in progress (attempt %d/%d)", - attempt + 1, - self.max_poll_attempts, - ) - continue - else: - 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, + job_id, + attempt + 1, + self.max_poll_attempts, + ) raise HTTPException(status_code=408, detail="File sanitization timeout") 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 e650f796f29..2e44b4b91b8 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -497,6 +497,69 @@ async def test_file_sanitization_modify_can_rewrite_when_blocking_disabled(monke assert base64.b64decode(result["file"]["data"]) == b"name,email\nAlice,[REDACTED]\n" +@pytest.mark.asyncio +async def test_file_sanitization_keeps_polling_through_queued_statuses(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) + guardrail.poll_interval = 0 + upload_response = Response( + json={"jobId": "queued-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_request = Request(method="GET", url="https://test.prompt.security/api/sanitizeFile") + poll_responses = [ + Response(json={"status": "created"}, status_code=200, request=poll_request), + Response(json={"status": "in progress"}, status_code=200, request=poll_request), + Response( + json={"status": "done", "content": "clean", "metadata": {"action": "allow", "violations": []}}, + status_code=200, + request=poll_request, + ), + ] + + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=upload_response)): + with patch.object(guardrail.async_handler, "get", AsyncMock(side_effect=poll_responses)) as poll_mock: + result = await guardrail.sanitize_file_content(b"image-content", "image.png") + + assert poll_mock.await_count == 3 + assert result["action"] == "allow" + assert result["content"] == "clean" + + +@pytest.mark.asyncio +async def test_file_sanitization_never_finishing_job_times_out(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_fail_open=False + ) + guardrail.poll_interval = 0 + guardrail.max_poll_attempts = 3 + upload_response = Response( + json={"jobId": "stuck-job"}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/sanitizeFile"), + ) + poll_response = Response( + json={"status": "created"}, + 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 == 3 + assert exc_info.value.status_code == 408 + assert exc_info.value.detail == "File sanitization timeout" + + @pytest.mark.asyncio @pytest.mark.parametrize( "timeout", From 3a3075b8a22df70dabc54b24a6ea6c7ac2826207 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:01:39 +0000 Subject: [PATCH 2/2] fix(prompt_security): poll only on queued statuses, keep 500 for terminal or missing status Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_security/prompt_security.py | 4 +++ .../test_prompt_security_guardrails.py | 29 +++++++++++++++++++ 2 files changed, 33 insertions(+) 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",