From 6ff0f4a03c06159b25564d79fd7c668f4da04a66 Mon Sep 17 00:00:00 2001 From: scthornton Date: Mon, 31 Aug 2026 12:50:46 -0400 Subject: [PATCH] fix(panw_prisma_airs): apply experimental_use_latest_role_message_only to every request shape The flag was silently ignored for any request not classified as a native Anthropic /v1/messages call, so OpenAI-compatible /v1/chat/completions traffic (the default shape for Vertex/Gemini, Azure OpenAI, and Bedrock backends) rescanned the full user/system history on every turn of a multi-turn conversation. Per-turn scan size grows linearly with turn count, so cumulative scanned tokens grow quadratically with session length (measured 7.46x over 8 turns against the live AIRS Scan API). An explicit true/false now governs scan scope for every request shape, matching the bedrock guardrail's semantics for the same flag. The unset default is unchanged: latest-only for Anthropic /v1/messages, full history for everything else. The existing safety fallbacks in _get_latest_user_text_indices (count mismatch, no user message) engage identically on the new path. Claude-Session: https://claude.ai/code/session_019e3TmBntsNQTBbXB9iu9uZ --- .../panw_prisma_airs/panw_prisma_airs.py | 24 +- .../guardrail_hooks/panw_prisma_airs.py | 8 +- .../guardrail_hooks/test_panw_prisma_airs.py | 323 +++++++++++++++++- 3 files changed, 336 insertions(+), 19 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 5f07f529e7a..92296e6cde9 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 @@ -194,7 +194,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): # internal '<=' comparison and surfaces as a misleading api_error. self.timeout = float(timeout) if timeout is not None else 10.0 - # Tri-state: None = not set (default-on for Anthropic), True = explicit on, False = explicit off + # Tri-state: None = not set (latest-only for Anthropic /v1/messages, full history + # otherwise), True = latest-only for every request shape, False = full history everywhere self.experimental_use_latest_role_message_only: bool | None = kwargs.get( "experimental_use_latest_role_message_only" ) @@ -1568,16 +1569,14 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) -> bool: """Resolve whether to scan only the latest user message. - - Non-Anthropic requests: always False (existing behavior) - - Anthropic requests: - - Flag explicitly True/False: respect it - - Flag None (not set): default to True + - Flag explicitly True/False: respect it for every request shape, + matching the bedrock guardrail's semantics for the same flag + - Flag None (not set): keep the historical defaults - True for + Anthropic /v1/messages requests, False for everything else """ - if not self._is_anthropic_request(request_data, logging_obj): - return False - if self.experimental_use_latest_role_message_only is None: - return True # Default-on for Anthropic - return self.experimental_use_latest_role_message_only + if self.experimental_use_latest_role_message_only is not None: + return self.experimental_use_latest_role_message_only + return self._is_anthropic_request(request_data, logging_obj) @staticmethod def _get_latest_user_text_indices( @@ -1775,7 +1774,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): if input_type == "request": structured_messages: Final = inputs.get("structured_messages") if structured_messages: - # For Anthropic /v1/messages: default to latest-user-only scanning. + # Latest-user-only scanning: flag-controlled for every request shape, + # default-on only for Anthropic /v1/messages (see _use_latest_user_only). # 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): @@ -1783,7 +1783,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): if original_messages: scannable_indices = self._get_latest_user_text_indices(texts, original_messages) # Fall through to existing role filtering if: - # - not Anthropic, OR flag explicitly False, OR + # - latest-only not enabled for this request, OR # - no original messages, OR # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py b/litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py index 606210d3b8b..f12dd927661 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py @@ -54,9 +54,11 @@ class PanwPrismaAirsGuardrailConfigModel(GuardrailConfigModel): experimental_use_latest_role_message_only: bool | None = Field( default=None, - description="Anthropic /v1/messages only. When unset: scans only latest user/developer " - "message on request side. Set false to scan all user/system/developer messages. " - "Non-Anthropic unaffected.", + description="Scan only the latest user/developer message on the request side instead of " + "the full conversation history. Set true to enable for every request shape (avoids " + "rescanning prior turns on every call of a multi-turn conversation); set false to always " + "scan all user/system/developer messages. When unset: defaults to latest-only for " + "Anthropic /v1/messages requests and full-history for all other request shapes.", ) @staticmethod 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 8f29ba66814..87ddb46638d 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 @@ -4538,10 +4538,19 @@ class TestPanwAirsLatestRoleMessageOnly: assert mock_api.call_args.kwargs["content"] == "Latest user message" @pytest.mark.asyncio - async def test_non_anthropic_any_flag_unchanged(self): - """Non-Anthropic + any flag state: existing role-filter behavior.""" - # Even with flag explicitly True, non-Anthropic should not change - handler = make_handler(experimental_use_latest_role_message_only=True) + @pytest.mark.parametrize("flag_value", [None, False]) + async def test_non_anthropic_flag_unset_or_false_full_scan(self, flag_value): + """Non-Anthropic + flag unset or False: existing role-filter behavior. + + The unset default for non-Anthropic requests must stay full-history: + only an explicit True opts a request shape into latest-only scanning. + """ + overrides = ( + {} + if flag_value is None + else {"experimental_use_latest_role_message_only": flag_value} + ) + handler = make_handler(**overrides) inputs: GenericGuardrailAPIInputs = { "texts": ["user prompt", "assistant reply", "system instruction"], @@ -4924,6 +4933,312 @@ class TestPanwAirsLatestRoleMessageOnly: ) +class TestPanwAirsLatestRoleMessageOnlyOpenAIShape: + """Latest-user-only scanning for OpenAI-shaped /v1/chat/completions requests. + + Before this behavior was added, experimental_use_latest_role_message_only was + silently ignored for any request the handler did not classify as a native + Anthropic /v1/messages call, so every turn of a multi-turn conversation + rescanned the full user/system history (per-turn scan size grows linearly + with turn count, cumulative scanned tokens grow quadratically). An explicit + True now scopes scanning to the latest user/developer message for every + request shape; the unset default is unchanged (full history for + non-Anthropic requests). + """ + + @pytest.fixture + def openai_multiturn_request_data(self): + """Multi-turn OpenAI chat-completions request data (system + history).""" + return { + "litellm_call_id": "test-call-id", + "model": "gemini-2.5-flash", + "messages": [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": "First user message"}, + {"role": "assistant", "content": "First assistant reply"}, + {"role": "user", "content": "Second user message"}, + {"role": "assistant", "content": "Second assistant reply"}, + {"role": "user", "content": "Latest user message"}, + ], + "proxy_server_request": { + "url": "http://localhost:4000/v1/chat/completions", + }, + } + + @pytest.fixture + def openai_multiturn_inputs(self): + """Inputs matching openai_multiturn_request_data (texts in message order).""" + return GenericGuardrailAPIInputs( + texts=[ + "You are a helpful assistant", + "First user message", + "First assistant reply", + "Second user message", + "Second assistant reply", + "Latest user message", + ], + structured_messages=[ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": "First user message"}, + {"role": "assistant", "content": "First assistant reply"}, + {"role": "user", "content": "Second user message"}, + {"role": "assistant", "content": "Second assistant reply"}, + {"role": "user", "content": "Latest user message"}, + ], + ) + + @pytest.mark.asyncio + async def test_flag_true_openai_scans_latest_user_only( + self, openai_multiturn_request_data, openai_multiturn_inputs + ): + """OpenAI shape + flag True: only the latest user message is scanned.""" + handler = make_handler(experimental_use_latest_role_message_only=True) + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = {"action": "allow", "category": "benign"} + + result = await handler.apply_guardrail( + inputs=openai_multiturn_inputs, + request_data=openai_multiturn_request_data, + input_type="request", + ) + + assert mock_api.call_count == 1 + assert mock_api.call_args.kwargs["content"] == "Latest user message" + # All texts preserved in output + assert result["texts"] == list(openai_multiturn_inputs["texts"]) + + @pytest.mark.asyncio + @pytest.mark.parametrize("flag_value", [None, False]) + async def test_flag_unset_or_false_openai_full_history_scan( + self, flag_value, openai_multiturn_request_data, openai_multiturn_inputs + ): + """OpenAI shape + flag unset or False: full-history role-filter scan. + + Guards against silently flipping the default: users who have not set the + flag must keep today's behavior (system + every user turn scanned). + """ + overrides = ( + {} + if flag_value is None + else {"experimental_use_latest_role_message_only": flag_value} + ) + handler = make_handler(**overrides) + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = {"action": "allow", "category": "benign"} + + await handler.apply_guardrail( + inputs=openai_multiturn_inputs, + request_data=openai_multiturn_request_data, + input_type="request", + ) + + # system + 3 user messages scanned, 2 assistant replies skipped + assert mock_api.call_count == 4 + scanned = [call.kwargs["content"] for call in mock_api.call_args_list] + assert "You are a helpful assistant" in scanned + assert "First user message" in scanned + assert "Second user message" in scanned + assert "Latest user message" in scanned + assert "First assistant reply" not in scanned + assert "Second assistant reply" not in scanned + + @pytest.mark.asyncio + async def test_flag_true_count_mismatch_falls_back_to_role_filter(self): + """OpenAI shape + flag True + texts/messages count mismatch: safety fallback. + + When request_data["messages"] does not walk to the same text count as + the flattened texts (e.g. a message was scoped out upstream), + _get_latest_user_text_indices returns None and the existing role-filter + scan over structured_messages must engage. + """ + handler = make_handler(experimental_use_latest_role_message_only=True) + + inputs: GenericGuardrailAPIInputs = { + "texts": ["system text", "user one", "user two"], + "structured_messages": [ + {"role": "system", "content": "system text"}, + {"role": "user", "content": "user one"}, + {"role": "user", "content": "user two"}, + ], + } + # messages has one extra entry vs texts → count mismatch → fallback + request_data = { + "litellm_call_id": "test-call-id", + "model": "gemini-2.5-flash", + "messages": [ + {"role": "system", "content": "system text"}, + {"role": "user", "content": "user one"}, + {"role": "user", "content": "user two"}, + {"role": "user", "content": "not represented in texts"}, + ], + } + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = {"action": "allow", "category": "benign"} + + await handler.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + # Fallback role-filter scan: all system/user texts scanned + assert mock_api.call_count == 3 + scanned = [call.kwargs["content"] for call in mock_api.call_args_list] + assert scanned == ["system text", "user one", "user two"] + + @pytest.mark.asyncio + async def test_flag_true_no_user_message_falls_back(self): + """OpenAI shape + flag True + no user/developer message: safety fallback.""" + handler = make_handler(experimental_use_latest_role_message_only=True) + + inputs: GenericGuardrailAPIInputs = { + "texts": ["system text", "assistant text"], + "structured_messages": [ + {"role": "system", "content": "system text"}, + {"role": "assistant", "content": "assistant text"}, + ], + } + request_data = { + "litellm_call_id": "test-call-id", + "model": "gemini-2.5-flash", + "messages": [ + {"role": "system", "content": "system text"}, + {"role": "assistant", "content": "assistant text"}, + ], + } + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = {"action": "allow", "category": "benign"} + + await handler.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + # Fallback role-filter scan: system scanned, assistant skipped + assert mock_api.call_count == 1 + assert mock_api.call_args.kwargs["content"] == "system text" + + @pytest.mark.asyncio + async def test_flag_true_tool_calls_still_scanned( + self, openai_multiturn_request_data, openai_multiturn_inputs + ): + """OpenAI shape + flag True: tool-call scanning is unaffected by scan scope.""" + handler = make_handler(experimental_use_latest_role_message_only=True) + + inputs = dict(openai_multiturn_inputs) + inputs["tool_calls"] = [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ] + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = {"action": "allow", "category": "benign"} + + await handler.apply_guardrail( + inputs=inputs, + request_data=openai_multiturn_request_data, + input_type="request", + ) + + # 1 text scan (latest user message) + 1 tool-call scan + assert mock_api.call_count == 2 + scanned = [call.kwargs["content"] for call in mock_api.call_args_list] + assert "Latest user message" in scanned + assert 'get_weather\n{"city": "Paris"}' in scanned + + @pytest.mark.asyncio + async def test_flag_true_masking_applies_to_latest_message_only( + self, openai_multiturn_request_data, openai_multiturn_inputs + ): + """OpenAI shape + flag True + DLP masking: masked text replaces only the + scanned (latest) message; unscanned history texts pass through untouched.""" + handler = make_handler(experimental_use_latest_role_message_only=True) + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = { + "action": "allow", + "category": "dlp", + "prompt_masked_data": {"data": "Latest user message [MASKED]"}, + } + + result = await handler.apply_guardrail( + inputs=openai_multiturn_inputs, + request_data=openai_multiturn_request_data, + input_type="request", + ) + + assert mock_api.call_count == 1 + expected = list(openai_multiturn_inputs["texts"]) + expected[-1] = "Latest user message [MASKED]" + assert result["texts"] == expected + + @pytest.mark.asyncio + async def test_flag_true_block_on_latest_message_still_blocks( + self, openai_multiturn_request_data, openai_multiturn_inputs + ): + """OpenAI shape + flag True: a block verdict on the latest message raises.""" + handler = make_handler(experimental_use_latest_role_message_only=True) + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = {"action": "block", "category": "injection"} + + with pytest.raises(HTTPException) as exc_info: + await handler.apply_guardrail( + inputs=openai_multiturn_inputs, + request_data=openai_multiturn_request_data, + input_type="request", + ) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_flag_true_response_side_unaffected( + self, openai_multiturn_request_data + ): + """Response-side scanning ignores the request-side scope flag.""" + handler = make_handler(experimental_use_latest_role_message_only=True) + + inputs: GenericGuardrailAPIInputs = { + "texts": ["response text one", "response text two"], + } + + with patch.object( + handler, "_call_panw_api", new_callable=AsyncMock + ) as mock_api: + mock_api.return_value = {"action": "allow", "category": "benign"} + + await handler.apply_guardrail( + inputs=inputs, + request_data=openai_multiturn_request_data, + input_type="response", + ) + + # Both response texts scanned; scope filtering is request-side only + assert mock_api.call_count == 2 + + class TestPanwAirsMcpToolCallWithoutCallId: """Tests for MCP tool invocations flowing through apply_guardrail without litellm_call_id — the bug fix for _convert_mcp_to_llm_format synthetic data."""