From dd173a0b1b7099255c51812edb327dffddb121fc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:43:35 -0700 Subject: [PATCH] fix(guardrails): count the PANW latest-user scan over the hoisted system prompt --- .../panw_prisma_airs/panw_prisma_airs.py | 11 +--- .../guardrail_hooks/test_panw_prisma_airs.py | 55 ++++++------------- 2 files changed, 20 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 3bc0dfabefc..9002e2aea07 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): Args: texts: Flattened text entries from the framework. - messages: Original request messages (request_data["messages"]), - NOT structured_messages (which may have injected system content). + messages: The structured messages the framework flattened into ``texts``, + hoisted top-level system prompt included, so positions line up. Returns a set of scannable indices, or None on count mismatch or no user/developer message (safety fallback to existing role-filter behavior). @@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): structured_messages: Final = inputs.get("structured_messages") if structured_messages: # For Anthropic /v1/messages: default to latest-user-only scanning. - # Uses request_data["messages"] (original format), NOT structured_messages - # (which has injected system content from adapter translation). if self._use_latest_user_only(request_data, logging_obj): - original_messages: Final = request_data.get("messages") - if original_messages: - scannable_indices = self._get_latest_user_text_indices(texts, original_messages) + scannable_indices = self._get_latest_user_text_indices(texts, structured_messages) # Fall through to existing role filtering if: # - not Anthropic, OR flag explicitly False, OR - # - no original messages, OR # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: scannable_indices = self._get_scannable_text_indices(texts, structured_messages) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 3d7c6e06d94..f25727ebd9a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -4620,46 +4620,27 @@ class TestPanwAirsLatestRoleMessageOnly: @pytest.mark.asyncio async def test_anthropic_system_plus_multiturn_no_fallback(self): - """Anthropic with top-level system + multi-turn messages[] - — latest-user works, no scan-all fallback. + """Anthropic with a top-level system prompt and multi-turn messages[] + scans only the latest user turn, with no scan-all fallback. - Key scenario: Anthropic top-level `system` field causes - structured_messages to have an injected system entry, but - request_data["messages"] does NOT include it. + The Anthropic handler hoists the top-level `system` field into both + `texts` and `structured_messages`, so the latest-user walk has to + count the same entries the framework flattened. """ - handler = PanwPrismaAirsHandler( - guardrail_name="test_panw_airs", - api_key="test_api_key", - profile_name="test_profile", - default_on=True, + from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, ) - # Original Anthropic messages (no system in messages array) - original_messages = [ - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - # texts extracted from original_messages (3 text entries) - texts = ["First user turn", "First assistant turn", "Latest user turn"] - - # structured_messages has an INJECTED system message from translation - structured_messages = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - inputs: GenericGuardrailAPIInputs = { - "texts": texts, - "structured_messages": structured_messages, - } + handler = make_handler() request_data = { "litellm_call_id": "test-call-id", "model": "anthropic/claude-sonnet-4-20250514", - "messages": original_messages, + "system": "You are a helpful assistant.", + "messages": [ + {"role": "user", "content": "First user turn"}, + {"role": "assistant", "content": "First assistant turn"}, + {"role": "user", "content": "Latest user turn"}, + ], "proxy_server_request": { "url": "http://localhost:4000/v1/messages", }, @@ -4670,13 +4651,11 @@ class TestPanwAirsLatestRoleMessageOnly: ) as mock_api: mock_api.return_value = {"action": "allow", "category": "benign"} - await handler.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", + await AnthropicMessagesHandler().process_input_messages( + data=request_data, + guardrail_to_apply=handler, ) - # Should scan ONLY the latest user message, not fall back to scan-all assert mock_api.call_count == 1 assert mock_api.call_args.kwargs["content"] == "Latest user turn"