From 4c7aac48a4e9696d7eb2c8976c44ba11c653315e Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 21:46:38 -0400 Subject: [PATCH] fix(guardrails): mask maskable PII in mixed violations before advising in Lakera v2 on_flagged="inject_system_message" only masked when a violation was PII-only; a mixed violation (PII plus a non-PII flag like prompt injection) fell straight through to the advisory branch with the raw PII still in place, in both async_pre_call_hook and async_moderation_hook. Mask whatever Lakera returned location data for before appending or logging the advisory, so a mixed violation never ships raw PII just because something else was also flagged. Also degrade to blocking, same as block mode already does, when nothing can be safely masked at all (multimodal content, or messages combined with a Responses API input field) instead of showing an advisory note next to raw, unredacted content. Widened call_v2_guard/_mask_pii_in_messages/the write-back helper's message parameters from list to Sequence to match what's actually passed through from _filter_skipped_messages, instead of duplicating list(...) casts at every call site. --- .../guardrail_hooks/lakera_ai_v2.py | 39 +++++- .../guardrail_hooks/test_lakera_ai_v2.py | 113 ++++++++++++++++++ 2 files changed, 147 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 1b3571cf242..0c0637aa8f0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -138,7 +138,7 @@ def _pre_masking_scope_indices( def _apply_redacted_messages_back_preserving_fields( guardrail: "LakeraAIGuardrail", data: dict[str, object], # mutable-ok: writes the redacted result back into the caller's request dict in place - redacted_messages: Sequence[Mapping[str, str]], + redacted_messages: Sequence[AllMessageValues], ) -> None: """Write masked content back to ``data["messages"]`` without losing fields the synthetic role/content-only ``redacted_messages`` never carried (e.g. a @@ -346,7 +346,7 @@ class LakeraAIGuardrail(CustomGuardrail): async def call_v2_guard( self, - messages: list[AllMessageValues], + messages: Sequence[AllMessageValues], request_data: dict, event_type: GuardrailEventHooks, ) -> tuple[LakeraAIResponse, dict]: @@ -408,10 +408,10 @@ class LakeraAIGuardrail(CustomGuardrail): def _mask_pii_in_messages( self, - messages: list[AllMessageValues], + messages: Sequence[AllMessageValues], lakera_response: LakeraAIResponse | None, masked_entity_count: dict, - ) -> list[AllMessageValues]: + ) -> Sequence[AllMessageValues]: """ Return a copy of messages with any detected PII replaced by “[MASKED ]” tokens. @@ -540,12 +540,30 @@ class LakeraAIGuardrail(CustomGuardrail): _apply_redacted_messages_back_preserving_fields(self, data, redacted_messages) verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request") elif self.on_flagged == "inject_system_message": + if is_multimodal_input: + # Nothing here can be safely masked, so an advisory note next + # to this raw, unredacted content would be no safer than a + # note next to nothing. Degrade to blocking instead, same as + # this on_flagged setting already does when the advisory + # itself has no field it can be delivered into. + raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) + # A mixed violation (PII plus something else, e.g. prompt injection): + # mask whatever Lakera returned location data for before advising + # about what remains, so the advisory is never shown next to raw + # PII that could have been redacted. + mixed_redacted_messages: Final = self._mask_pii_in_messages( + messages=new_messages, + lakera_response=lakera_guardrail_response, + masked_entity_count=masked_entity_count, + ) + _apply_redacted_messages_back_preserving_fields(self, data, mixed_redacted_messages) advisory_delivered: Final = self.inject_advisory_message( data, self._build_advisory_message(lakera_guardrail_response) ) if advisory_delivered: verbose_proxy_logger.warning( - "Lakera Guardrail: Advisory mode - violation detected, appended advisory system message" + "Lakera Guardrail: Advisory mode - violation detected, masked PII and appended advisory " + "system message" ) else: # Structured Responses-API input (a list, not a plain string) @@ -636,6 +654,17 @@ class LakeraAIGuardrail(CustomGuardrail): _apply_redacted_messages_back_preserving_fields(self, data, redacted_messages) verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request") elif self.on_flagged == "inject_system_message": + if not is_multimodal_input: + # A mixed violation (PII plus something else): mask whatever's + # maskable even though the advisory note below has no effect + # here, so raw PII doesn't pass through untouched just because + # this violation wasn't PII-only. + mixed_redacted_messages: Final = self._mask_pii_in_messages( + messages=new_messages, + lakera_response=lakera_guardrail_response, + masked_entity_count=masked_entity_count, + ) + _apply_redacted_messages_back_preserving_fields(self, data, mixed_redacted_messages) # during_call runs concurrently with the LLM dispatch (see # ProxyLogging.during_call_hook / common_request_processing.py), # with no pre-call barrier -- mutating data["messages"] here races diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 0cc27ca169f..977d26e1a40 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -930,6 +930,83 @@ class TestAdvisoryModeWiring: assert result["messages"][0]["content"] != original_content assert len(result["messages"]) == 1, "no advisory note should be appended once PII is masked" + @pytest.mark.asyncio + async def test_pre_call_mixed_violation_masks_pii_before_appending_advisory(self): + """ + Bugbot finding on BerriAI/litellm#34940: a mixed violation (PII plus a + non-PII flag like prompt injection) isn't PII-only, so it fell straight + through to the advisory branch with the raw PII still in place. It must + mask the maskable PII first, then still append the advisory note for the + remaining, non-PII concern. + """ + lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") + mock_response = { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 0}], + "breakdown": [ + {"detector_type": "pii/email", "detected": True}, + {"detector_type": "prompt_injection", "detected": True}, + ], + } + original_content = "My email is test@example.com" + + with patch.object(lakera_guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + mock_call.return_value = (mock_response, {}) + data = { + "messages": [{"role": "user", "content": original_content}], + "model": "gpt-5-mini", + "metadata": {}, + } + + result = await lakera_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert "[MASKED" in result["messages"][0]["content"] + assert result["messages"][0]["content"] != original_content + assert len(result["messages"]) == 2, "the remaining, non-PII concern still gets an advisory note" + assert result["messages"][1]["role"] == "system" + + @pytest.mark.asyncio + async def test_pre_call_blocks_instead_of_advisory_when_pii_is_not_maskable(self): + """ + Bugbot finding on BerriAI/litellm#34940: a PII-only or mixed violation on + input that can't be safely masked (combined messages+input, multimodal + content) fell through to the advisory branch with raw, unredacted content. + It must degrade to blocking instead, same as block mode already does for + this exact case, rather than showing an advisory note next to raw content. + """ + lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") + mock_response = { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 0}], + "breakdown": [{"detector_type": "pii/email", "detected": True}], + } + + with patch.object(lakera_guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + mock_call.return_value = (mock_response, {}) + data = { + "messages": [{"role": "user", "content": "My email is test@example.com"}], + "input": "responses-api content", + "model": "gpt-5-mini", + "metadata": {}, + } + with pytest.raises(HTTPException): + await lakera_guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert "messages" in data + assert data["messages"][0]["content"] == "My email is test@example.com", ( + "the raw content must be untouched, not partially rewritten before the block" + ) + @pytest.mark.asyncio async def test_moderation_hook_inspects_all_message_roles_not_just_user(self): """See test_pre_call_inspects_all_message_roles_not_just_user.""" @@ -1022,6 +1099,42 @@ class TestAdvisoryModeWiring: assert "[MASKED" in result["messages"][0]["content"] assert result["messages"][0]["content"] != original_content + @pytest.mark.asyncio + async def test_moderation_hook_mixed_violation_masks_pii_even_though_advisory_has_no_effect(self): + """ + Bugbot finding on BerriAI/litellm#34940: a mixed violation isn't PII-only, + so it fell through to the during_call no-op branch with the raw PII still + in place. It must mask the maskable PII even though the advisory note + itself still has no effect during during_call. + """ + lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") + mock_response = { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 0}], + "breakdown": [ + {"detector_type": "pii/email", "detected": True}, + {"detector_type": "prompt_injection", "detected": True}, + ], + } + original_content = "My email is test@example.com" + + with patch.object(lakera_guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + mock_call.return_value = (mock_response, {}) + data = { + "messages": [{"role": "user", "content": original_content}], + "model": "gpt-5-mini", + "metadata": {}, + } + result = await lakera_guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + call_type="completion", + ) + + assert "[MASKED" in result["messages"][0]["content"] + assert result["messages"][0]["content"] != original_content + assert len(result["messages"]) == 1, "no advisory note is appended during during_call" + class TestAdvisoryModePostCall: """