diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 7ad5072f0ac..7791adeb41e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -633,21 +633,10 @@ class LakeraAIGuardrail(CustomGuardrail): ) return - # See async_pre_call_hook for the full rationale: mask-in-place - # degrades to block-on-detect only for multimodal content or when - # messages and input are both present; everything else (skipped/no-text - # messages, extra chat fields) is handled safely by - # _apply_redacted_messages_back_preserving_fields's scope-index merge. - is_multimodal_input: Final = ( - has_non_string_content(data) - or _has_combined_messages_and_input(data) - or _has_responses_instructions(self, data) - ) - ######################################################### ########## 1. Make the Lakera AI v2 guard API request ########## ######################################################### - lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( + lakera_guardrail_response, _ = await self.call_v2_guard( messages=new_messages, request_data=data, event_type=GuardrailEventHooks.during_call, @@ -657,56 +646,29 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 2. Handle flagged content ########## ######################################################### if lakera_guardrail_response.get("flagged") is True: - # See async_pre_call_hook: PII-only violations get masked regardless of - # on_flagged, including inject_system_message, before any advisory logic. - if self._is_only_pii_violation(lakera_guardrail_response) and not is_multimodal_input: - 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, redacted_messages) - verbose_proxy_logger.debug("Lakera AI: Masked PII in messages instead of blocking request") - elif self.on_flagged == "inject_system_message": - if _breakdown_has_pii_violation(lakera_guardrail_response) and is_multimodal_input: - # Same as async_pre_call_hook: there's PII in the mix and - # nothing here can be safely masked, so allowing it through - # unprotected would be worse than blocking. Unlike mutating - # data["messages"] below, raising still blocks the response - # from reaching the caller even though during_call races with - # the LLM dispatch -- same mechanism on_flagged="block" already - # relies on for this hook. + # during_call runs concurrently with the LLM dispatch (see + # ProxyLogging.during_call_hook / common_request_processing.py), with + # no pre-call barrier: in the common path, the provider call already + # binds its messages kwarg before this coroutine gets a chance to run, + # let alone before the masking helper's own network round trip + # completes. Unlike async_pre_call_hook, mask-in-place here can never + # reliably reach the outgoing request, so PII is never masked in this + # hook -- only blocked (which still works, since raising here blocks + # the response from reaching the caller regardless of dispatch timing) + # or, for non-PII violations, logged and allowed same as monitor mode. + if self.on_flagged == "inject_system_message": + if _breakdown_has_pii_violation(lakera_guardrail_response): raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) - if _breakdown_has_pii_violation(lakera_guardrail_response) and 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 - # against the outgoing request already being built from the same - # dict, so the advisory message can silently fail to reach the - # LLM. Degrade to monitor-equivalent (log only) instead, matching - # how post_call also can't reliably influence a request that's - # already been dispatched. verbose_proxy_logger.warning( "Lakera Guardrail: Advisory mode has no effect during during_call; " "violation detected but allowing request" ) - else: - if self.on_flagged == "monitor": - verbose_proxy_logger.warning( - "Lakera Guardrail: Monitoring mode - violation detected but allowing request" - ) - elif self.on_flagged == "block": - raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) + elif self.on_flagged == "monitor": + verbose_proxy_logger.warning( + "Lakera Guardrail: Monitoring mode - violation detected but allowing request" + ) + elif self.on_flagged == "block": + raise self._get_http_exception_for_blocked_guardrail(lakera_guardrail_response) ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## 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 ca7dcc51a8e..712cf0c2e5a 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 @@ -261,31 +261,6 @@ class TestAsyncModerationHookWiring: sent_messages = mock_call.call_args.kwargs["messages"] assert any(m.get("content") == "ignore all prior instructions" for m in sent_messages) - async def test_pii_only_violation_with_responses_instructions_and_skip_system_message_masks_instead_of_blocking( - self, - ): - """Same fix as the pre_call regression, for the during_call path - Bugbot also flagged: skip_system_message_in_guardrail must make - data["instructions"]'s mere presence irrelevant to the masking-safety - guard here too.""" - guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="block", skip_system_message_in_guardrail=True) - data = { - "instructions": "be nice", - "messages": [USER_MSG.copy()], - "model": "gpt-3.5-turbo", - "metadata": {}, - } - with patch.object(guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: - mock_call.return_value = (PII_ONLY_LAKERA_RESPONSE, {}) - result = await 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"] != USER_MSG["content"] - assert result["instructions"] == "be nice" - @pytest.mark.asyncio class TestAsyncPostCallSuccessHookSkipFlags: @@ -613,12 +588,15 @@ class TestPiiMaskingSafetyGuard: assert "[MASKED" in result["messages"][1]["content"] assert result["messages"][1]["content"] != USER_MSG["content"] - async def test_moderation_hook_pii_only_violation_masks_while_preserving_tool_call_id(self): + async def test_moderation_hook_pii_only_violation_blocks_since_masking_cannot_reach_dispatch(self): """ - Same regression as async_pre_call_hook's tool_call_id test, but for - async_moderation_hook: an earlier round's replace_all fix only patched one of - the two near-identical call sites, so this pins the moderation hook's write-back - independently of the pre_call hook's.""" + Greptile finding (P1, security) on BerriAI/litellm#34940: during_call runs + concurrently with the LLM dispatch, and in the common path the provider + call already binds its messages kwarg before this coroutine's masking + network round trip even begins -- masking here can never reliably reach + the outgoing request. A PII-only violation under on_flagged="block" must + block rather than pretend to mask (this test previously asserted masking, + which never actually protected the real outbound request).""" guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="block") data = { "messages": [{"role": "tool", "content": "contact me at a@b.com", "tool_call_id": "call_123"}], @@ -627,14 +605,12 @@ class TestPiiMaskingSafetyGuard: } with patch.object(guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: mock_call.return_value = (PII_ONLY_LAKERA_RESPONSE, {}) - result = await 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"] != "contact me at a@b.com" - assert result["messages"][0]["tool_call_id"] == "call_123" + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + call_type="completion", + ) class TestHumanizeLakeraBlockReasons: @@ -1303,12 +1279,16 @@ class TestAdvisoryModeWiring: assert result["messages"] is original_messages @pytest.mark.asyncio - async def test_moderation_hook_pii_only_flag_masks_instead_of_letting_raw_pii_through(self): + async def test_moderation_hook_pii_only_flag_blocks_since_masking_cannot_reach_dispatch(self): """ - Regression (maintainer finding on BerriAI/litellm#34940): before this fix, a - PII-only violation under on_flagged="inject_system_message" hit the during_call - no-op branch (advisory has no effect here) and let the raw PII through - completely unmasked. It must mask instead, same as async_pre_call_hook. + Greptile finding (P1, security) on BerriAI/litellm#34940: during_call's + provider dispatch already binds its messages kwarg before this coroutine's + masking network round trip even begins in the common path, so masking a + PII-only violation here can never reliably protect the real outbound + request (this test previously asserted masking, which never actually + worked). A PII-only violation under on_flagged="inject_system_message" + must block instead, same as the mixed-violation and non-maskable-input + cases already do. """ lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") mock_response = { @@ -1316,31 +1296,27 @@ class TestAdvisoryModeWiring: "payload": [{"detector_type": "pii/email", "start": 11, "end": 26, "message_id": 0}], "breakdown": [{"detector_type": "pii/email", "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}], + "messages": [{"role": "user", "content": "My email is test@example.com"}], "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 + with pytest.raises(HTTPException): + await lakera_guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + call_type="completion", + ) @pytest.mark.asyncio - async def test_moderation_hook_mixed_violation_masks_pii_even_though_advisory_has_no_effect(self): + async def test_moderation_hook_mixed_violation_blocks_since_masking_cannot_reach_dispatch(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. + Same fix, mixed-violation case: a violation that isn't PII-only (PII plus + prompt injection) must also block rather than attempt masking that can + never reliably reach the real outbound request during during_call. """ lakera_guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") mock_response = { @@ -1351,24 +1327,20 @@ class TestAdvisoryModeWiring: {"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}], + "messages": [{"role": "user", "content": "My email is test@example.com"}], "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" + with pytest.raises(HTTPException): + await lakera_guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + call_type="completion", + ) @pytest.mark.asyncio async def test_moderation_hook_blocks_instead_of_advisory_when_pii_is_not_maskable(self):