From ff2709f5c33c9d843593bed2be63926c4dbb6811 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 5 Nov 2026 13:00:21 -0800 Subject: [PATCH] fix(zscaler_ai_guard)!: scan only the last user turn on requests (#50139) * fix(zscaler_ai_guard): scan only the last user turn on requests * fix(zscaler_ai_guard): keep scanning A2A requests and refuse scan_only_tool_results * fix(zscaler_ai_guard): join the last user turn's text parts with spaces like before * fix(zscaler_ai_guard): keep scanning tool results when scan_only_tool_results is set * fix(zscaler_ai_guard): scan the latest user turn when assistant or tool messages follow it * fix(zscaler_ai_guard): skip the scan when the latest user turn is only whitespace A blank latest user turn was sent to Zscaler on its own, which answers Content is empty, so the proxy returned 500 where main returned 200 --------- Co-authored-by: Arnav Maiti --- .../zscaler_ai_guard/zscaler_ai_guard.py | 45 +++- .../zscaler_ai_guard/test_zscaler_ai_guard.py | 212 ++++++++++++++++++ 2 files changed, 252 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index e8f0dc68d83..326f6400ba3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -4,7 +4,8 @@ # # +-------------------------------------------------------------+ import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from itertools import chain, dropwhile, takewhile from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException @@ -14,6 +15,10 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -218,15 +223,14 @@ class ZscalerAIGuard(CustomGuardrail): zscaler_ai_guard_result = None direction: Final = "OUT" if input_type == "response" else "IN" verbose_proxy_logger.debug("direction: %s", direction) - # Concatenate all texts and send to Zscaler AI Guard - if texts: - concatenated_text: Final = " ".join(texts) + content: Final = self._content_to_scan(inputs, input_type) + if content: zscaler_ai_guard_result = await self.make_zscaler_ai_guard_api_call( zscaler_ai_guard_url=self.zscaler_ai_guard_url, api_key=self.api_key, policy_id=policy_id, direction=direction, - content=concatenated_text, + content=content, user=user_for_payload, **kwargs, ) @@ -244,6 +248,37 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.") return inputs + def _content_to_scan( + self, + inputs: GenericGuardrailAPIInputs, + input_type: Literal["request", "response"], + ) -> str | None: + structured_messages: Final = inputs.get("structured_messages") or [] + scans_last_user_turn: Final = ( + input_type == "request" + and not effective_scan_only_tool_results_for_guardrail(self) + and any("content" in message for message in structured_messages) + ) + texts: Final = ( + self._last_user_turn_texts(structured_messages) if scans_last_user_turn else inputs.get("texts", []) + ) + return " ".join(texts) if texts else None + + @staticmethod + def _last_user_turn_texts(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + from_last_user_message: Final = dropwhile(lambda message: message.get("role") != "user", reversed(messages)) + last_user_turn: Final = tuple(takewhile(lambda message: message.get("role") == "user", from_last_user_message)) + texts: Final = tuple(chain.from_iterable(map(ZscalerAIGuard._content_texts, reversed(last_user_turn)))) + return texts if any(text.strip() for text in texts) else () + + @staticmethod + def _content_texts(message: Mapping[str, object]) -> tuple[str, ...]: + content: Final = message.get("content") + if isinstance(content, str): + return (content,) + parts: Final = (as_str_mapping(part) for part in as_object_list(content) or ()) + return tuple(text for part in parts if part is not None and isinstance(text := part.get("text"), str)) + def extract_blocking_info(self, response): """ Extracts transaction ID and blocking detector details from a response. diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py index bccc531793f..712750f702e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py @@ -23,6 +23,8 @@ from fastapi import HTTPException import litellm from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ZscalerAIGuard +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.types.utils import GenericGuardrailAPIInputs from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -141,6 +143,216 @@ async def test_apply_guardrail_text_concatenation(mock_api_call): assert call_args.kwargs["content"] == "Hello world" +_CONVERSATION: Final = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is two plus two?"}, + {"role": "assistant", "content": "Four."}, +] + + +async def _execute_policy_bodies( + inputs: GenericGuardrailAPIInputs, + input_type: Literal["request", "response"], + scan_only_tool_results: bool = False, +) -> tuple[object, ...]: + received: Final = asyncio.Queue[object]() + + async def execute_policy(request: web.Request) -> web.Response: + await received.put(await request.json()) + return web.json_response({"statusCode": 200, "action": "ALLOW"}) + + app: Final = web.Application() + app.router.add_post("/v1/detection/execute-policy", execute_policy) + handler: Final = InMemoryGuardrailHandler() + async with TestServer(app) as server: + guardrail_id: Final = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "zscaler-last-user-turn", + "litellm_params": { + "guardrail": "zscaler_ai_guard", + "mode": "pre_call", + "api_key": "test_key", + "api_base": str(server.make_url("/v1/detection/execute-policy")), + "policy_id": 1, + "scan_only_tool_results": scan_only_tool_results, + }, + }, + )["guardrail_id"] + try: + guardrail: Final = handler.guardrail_id_to_custom_guardrail[guardrail_id] + assert guardrail is not None + returned: Final = await guardrail.apply_guardrail(inputs, {}, input_type) + finally: + handler.delete_in_memory_guardrail(guardrail_id) + + assert returned is inputs + return tuple(received.get_nowait() for _ in range(received.qsize())) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inputs", "expected_content"), + [ + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [*_CONVERSATION, {"role": "user", "content": "And three plus three?"}], + }, + "And three plus three?", + ), + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [ + *_CONVERSATION, + { + "role": "user", + "content": [{"type": "text", "text": "And three"}, {"type": "text", "text": "plus three?"}], + }, + ], + }, + "And three plus three?", + ), + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [ + *_CONVERSATION, + {"role": "user", "content": "And three plus three?"}, + {"role": "user", "content": "Answer in words."}, + ], + }, + "And three plus three? Answer in words.", + ), + ( + { + "texts": ["Summarize", "this report"], + "structured_messages": [ + { + "role": "user", + "parts": [{"kind": "text", "text": "Summarize"}, {"kind": "text", "text": "this report"}], + "messageId": "msg-1", + } + ], + }, + "Summarize this report", + ), + ({"texts": ["first input", "second input"]}, "first input second input"), + ( + { + "texts": ["every", "message"], + "structured_messages": [*_CONVERSATION, {"role": "assistant", "content": "Prefilled answer:"}], + }, + "What is two plus two?", + ), + ( + { + "texts": ["every", "message"], + "structured_messages": [ + *_CONVERSATION, + {"role": "tool", "tool_call_id": "call_1", "content": "tool output"}, + ], + }, + "What is two plus two?", + ), + ( + { + "texts": ["every", "message"], + "structured_messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "assistant", "content": "Prefilled answer:"}, + ], + }, + None, + ), + ( + { + "texts": ["You are a helpful assistant.", " "], + "structured_messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": " "}, + ], + }, + None, + ), + ( + { + "texts": ["every", "message", "\n"], + "structured_messages": [ + *_CONVERSATION, + { + "role": "user", + "content": [ + {"type": "text", "text": "\n"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + ], + }, + ], + }, + None, + ), + ], + ids=[ + "string-content", + "content-parts", + "consecutive-user-messages", + "a2a-message-parts", + "texts-only-route", + "ends-on-assistant", + "ends-on-tool", + "no-user-message", + "whitespace-only-latest-turn", + "whitespace-text-beside-image", + ], +) +async def test_request_scan_sends_only_the_last_user_turn( + inputs: GenericGuardrailAPIInputs, expected_content: str | None +) -> None: + bodies: Final = await _execute_policy_bodies(inputs, "request") + + assert bodies == ( + () if expected_content is None else ({"direction": "IN", "content": expected_content, "policyId": 1},) + ) + + +@pytest.mark.asyncio +async def test_response_scan_sends_every_text() -> None: + inputs: Final = GenericGuardrailAPIInputs( + texts=["First choice.", "Second choice."], + structured_messages=[*_CONVERSATION, {"role": "user", "content": "And three plus three?"}], + ) + + bodies: Final = await _execute_policy_bodies(inputs, "response") + + assert bodies == ({"direction": "OUT", "content": "First choice. Second choice.", "policyId": 1},) + + +@pytest.mark.asyncio +async def test_scan_only_tool_results_sends_the_tool_results() -> None: + inputs: Final = GenericGuardrailAPIInputs( + texts=["Sunny, 22C"], + structured_messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Sunny, 22C"}, + ], + ) + + bodies: Final = await _execute_policy_bodies(inputs, "request", scan_only_tool_results=True) + + assert bodies == ({"direction": "IN", "content": "Sunny, 22C", "policyId": 1},) + + @pytest.mark.asyncio @patch( "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",