From ceb5b5df8c899e0754d815297872d23ca2a7258f Mon Sep 17 00:00:00 2001 From: Abhijoy Sarkar Date: Tue, 31 Mar 2026 14:02:07 +0530 Subject: [PATCH] Address Greptile review: redact path, null decision, error context - P1: Filter _extract_texts_from_messages to user-role messages only, preventing system/assistant content from being injected into texts - P1: Strengthen test_redact_updates_structured_messages assertion from weak `in` check to strict equality, catching the injection bug - P2: Use `result.get("decision") or "allow"` to handle explicit null decision values (not just absent keys) - P2: Wrap bare exception re-raise in GuardrailRaisedException so the caller knows which guardrail failed (block_on_error=True path) - P2: Add static Promptguard entry in guardrail_provider_map so the preset works before populateGuardrailProviderMap is called - Add test for explicit null decision treated as allow --- .../promptguard/promptguard.py | 19 ++++++++++-- .../guardrail_hooks/test_promptguard.py | 30 +++++++++++++++---- .../guardrails/guardrail_info_helpers.tsx | 1 + 3 files changed, 42 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index 009e29ef089..8c5001249ad 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -155,7 +155,13 @@ class PromptGuardGuardrail(CustomGuardrail): except Exception as exc: verbose_proxy_logger.error("PromptGuard API error: %s", str(exc)) if self.block_on_error: - raise + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=( + f"PromptGuard API unreachable " + f"(block_on_error=True): {exc}" + ), + ) from exc return inputs verbose_proxy_logger.debug( @@ -164,7 +170,7 @@ class PromptGuardGuardrail(CustomGuardrail): result.get("threat_type"), ) - decision = result.get("decision", "allow") + decision = result.get("decision") or "allow" if decision == "block": threat_type = result.get("threat_type", "unknown") @@ -196,9 +202,16 @@ class PromptGuardGuardrail(CustomGuardrail): @staticmethod def _extract_texts_from_messages(messages: list) -> List[str]: - """Extract text content strings from a list of chat messages.""" + """Extract text content from user-role messages only. + + Only user messages are extracted to avoid injecting system or + assistant content into the ``texts`` list, which should mirror + the original user-provided input. + """ texts: List[str] = [] for message in messages: + if message.get("role") != "user": + continue content = message.get("content") if isinstance(content, str): texts.append(content) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py index e5e2face752..efd14379ddd 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_promptguard.py @@ -373,7 +373,7 @@ class TestPromptGuardRedactAction: input_type="request", ) assert result["structured_messages"] == redacted - assert "My SSN is *********" in result["texts"] + assert result["texts"] == ["My SSN is *********"] @pytest.mark.asyncio async def test_redact_structured_only_does_not_create_texts( @@ -629,7 +629,7 @@ class TestPromptGuardErrorHandling: async def test_http_error_propagates_block_on_error( self, promptguard_guardrail, mock_request_data ): - """Default block_on_error=True re-raises HTTP errors.""" + """Default block_on_error=True wraps HTTP errors in GuardrailRaisedException.""" mock_request = httpx.Request("POST", "https://api.test.promptguard.co") mock_resp = httpx.Response(status_code=500, request=mock_request) with patch.object( @@ -641,29 +641,33 @@ class TestPromptGuardErrorHandling: response=mock_resp, ), ): - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(GuardrailRaisedException) as exc_info: await promptguard_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data, input_type="request", ) + assert "block_on_error=True" in str(exc_info.value) + assert exc_info.value.__cause__ is not None @pytest.mark.asyncio async def test_connection_error_propagates_block_on_error( self, promptguard_guardrail, mock_request_data ): - """Default block_on_error=True re-raises connection errors.""" + """Default block_on_error=True wraps connection errors in GuardrailRaisedException.""" with patch.object( promptguard_guardrail.async_handler, "post", side_effect=httpx.ConnectError("Connection refused"), ): - with pytest.raises(httpx.ConnectError): + with pytest.raises(GuardrailRaisedException) as exc_info: await promptguard_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data, input_type="request", ) + assert "block_on_error=True" in str(exc_info.value) + assert exc_info.value.__cause__ is not None @pytest.mark.asyncio async def test_fail_open_returns_inputs_on_http_error(self, mock_request_data): @@ -747,6 +751,22 @@ class TestPromptGuardErrorHandling: ) assert result["texts"] == ["test"] + @pytest.mark.asyncio + async def test_null_decision_treated_as_allow( + self, promptguard_guardrail, mock_request_data + ): + """Explicit null decision should be treated as allow.""" + resp = _make_response({"decision": None, "event_id": "evt-null"}) + with patch.object( + promptguard_guardrail.async_handler, "post", return_value=resp + ): + result = await promptguard_guardrail.apply_guardrail( + inputs={"texts": ["test"]}, + request_data=mock_request_data, + input_type="request", + ) + assert result["texts"] == ["test"] + # --------------------------------------------------------------------------- # Config model diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx index 088995dea8d..8ab8710d01a 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx @@ -48,6 +48,7 @@ export const guardrail_provider_map: Record = { LitellmContentFilter: "litellm_content_filter", ToolPermission: "tool_permission", BlockCodeExecution: "block_code_execution", + Promptguard: "promptguard", }; // Function to populate provider map from API response - updates the original map