From aab8edde67ebaa6a24205cbd57016410df7f894e Mon Sep 17 00:00:00 2001 From: datzscaler Date: Fri, 13 Feb 2026 11:28:14 -0800 Subject: [PATCH] fix(guardrails): Zscaler AI Guard bug fixes and support during post-call (#20801) * fix(guardrails): fixed post-call issue with Zscaler guardrail and invalid headers. Added unittests * fix(guardrails): Addressed greptil comments. Make policyid hanlding more clear * fix(guardrails): Address greptile comment * Update litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../zscaler_ai_guard/__init__.py | 4 + .../zscaler_ai_guard/zscaler_ai_guard.py | 88 +++++++++++------ .../guardrails_tests/test_zscaler_ai_guard.py | 97 +++++++++++++++++++ 3 files changed, 161 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py index c987ace7ed..23f0970eba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py @@ -14,6 +14,10 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" _zscaler_ai_guard_callback = ZscalerAIGuard( api_base=litellm_params.api_base, api_key=litellm_params.api_key, + policy_id=litellm_params.policy_id, + send_user_api_key_alias=litellm_params.send_user_api_key_alias, + send_user_api_key_user_id=litellm_params.send_user_api_key_user_id, + send_user_api_key_team_id=litellm_params.send_user_api_key_team_id, guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, 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 b0b18166d2..fbf9e1f29f 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 @@ -32,9 +32,9 @@ class ZscalerAIGuard(CustomGuardrail): api_key: Optional[str] = None, api_base: Optional[str] = None, policy_id: Optional[int] = None, - send_user_api_key_alias: Optional[bool] = False, - send_user_api_key_user_id: Optional[bool] = False, - send_user_api_key_team_id: Optional[bool] = False, + send_user_api_key_alias: Optional[bool] = None, + send_user_api_key_user_id: Optional[bool] = None, + send_user_api_key_team_id: Optional[bool] = None, **kwargs, ): self.optional_params = kwargs @@ -42,15 +42,15 @@ class ZscalerAIGuard(CustomGuardrail): "ZSCALER_AI_GUARD_URL", "https://api.us1.zseclipse.net/v1/detection/execute-policy", ) - self.policy_id = policy_id or int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1)) + self.policy_id = policy_id if policy_id is not None else int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1)) self.api_key = api_key or os.getenv("ZSCALER_AI_GUARD_API_KEY") - self.send_user_api_key_alias = send_user_api_key_alias or os.getenv( + self.send_user_api_key_alias = send_user_api_key_alias if send_user_api_key_alias is not None else os.getenv( "SEND_USER_API_KEY_ALIAS", "False" ).lower() in ("true", "1") - self.send_user_api_key_user_id = send_user_api_key_user_id or os.getenv( + self.send_user_api_key_user_id = send_user_api_key_user_id if send_user_api_key_user_id is not None else os.getenv( "SEND_USER_API_KEY_USER_ID", "False" - ).lower() in ("true", "1,") - self.send_user_api_key_team_id = send_user_api_key_team_id or os.getenv( + ).lower() in ("true", "1") + self.send_user_api_key_team_id = send_user_api_key_team_id if send_user_api_key_team_id is not None else os.getenv( "SEND_USER_API_KEY_TEAM_ID", "False" ).lower() in ("true", "1") @@ -60,19 +60,48 @@ class ZscalerAIGuard(CustomGuardrail): send_user_api_key_team_id:{self.send_user_api_key_team_id}""" ) - super().__init__(default_on=True) + super().__init__(**kwargs) verbose_proxy_logger.debug("ZscalerAIGuard Initializing ...") - def _get_stripped_metadata_value( - self, request_data: Optional[dict], key: str - ) -> Optional[str]: + @staticmethod + def _resolve_metadata_value(request_data: Optional[dict], key: str) -> Optional[str]: + """ + Resolve metadata value from request_data, checking both metadata locations. + + During pre-call: metadata is at request_data["metadata"][key] + During post-call: metadata is at request_data["litellm_metadata"][key] + (set by transform_user_api_key_dict_to_metadata which prefixes keys with 'user_api_key_') + + Also handles key name mapping for UserAPIKeyAuth fields: + - key_alias -> user_api_key_key_alias (in litellm_metadata) + - user_id -> user_api_key_user_id + - team_id -> user_api_key_team_id + """ if request_data is None: - return "N/A" - value = request_data.get("metadata", {}).get(key, "N/A") - if value is not None: - return str(value).strip() - return "N/A" + return None + + # Check litellm_metadata first (set during post-call by guardrail framework) + litellm_metadata = request_data.get("litellm_metadata", {}) + if litellm_metadata: + value = litellm_metadata.get(key) + if value is not None: + return str(value).strip() + # Handle key_alias -> user_api_key_key_alias mapping + # transform_user_api_key_dict_to_metadata prefixes "key_alias" -> "user_api_key_key_alias" + if key == "user_api_key_alias": + value = litellm_metadata.get("user_api_key_key_alias") + if value is not None: + return str(value).strip() + + # Then check regular metadata (set during pre-call by proxy_server) + metadata = request_data.get("metadata", {}) + if metadata: + value = metadata.get(key) + if value is not None: + return str(value).strip() + + return None @log_guardrail_information async def apply_guardrail( @@ -128,17 +157,17 @@ class ZscalerAIGuard(CustomGuardrail): kwargs = {} if self.send_user_api_key_alias: - kwargs["user_api_key_alias"] = self._get_stripped_metadata_value( + kwargs["user_api_key_alias"] = self._resolve_metadata_value( request_data, "user_api_key_alias" - ) + ) or "N/A" if self.send_user_api_key_team_id: - kwargs["user_api_key_team_id"] = self._get_stripped_metadata_value( + kwargs["user_api_key_team_id"] = self._resolve_metadata_value( request_data, "user_api_key_team_id" - ) + ) or "N/A" if self.send_user_api_key_user_id: - kwargs["user_api_key_user_id"] = self._get_stripped_metadata_value( + kwargs["user_api_key_user_id"] = self._resolve_metadata_value( request_data, "user_api_key_user_id" - ) + ) or "N/A" verbose_proxy_logger.debug(f"inside apply_guardrail kwargs: {kwargs}") zscaler_ai_guard_result = None @@ -155,6 +184,7 @@ class ZscalerAIGuard(CustomGuardrail): content=concatenated_text, **kwargs, ) + verbose_proxy_logger.debug(f"response from zscaler ai guards: {zscaler_ai_guard_result}") if ( zscaler_ai_guard_result and zscaler_ai_guard_result.get("action") == "BLOCK" @@ -217,7 +247,7 @@ class ZscalerAIGuard(CustomGuardrail): extra_headers.update({"user-api-key-team-id": user_api_key_team_id}) if self.send_user_api_key_user_id: - user_api_key_user_id = kwargs.get("user-api-key-user-id", "N/A") + user_api_key_user_id = kwargs.get("user_api_key_user_id", "N/A") extra_headers.update({"user-api-key-user-id": user_api_key_user_id}) verbose_proxy_logger.debug(f"extra_headers: {extra_headers}") @@ -320,11 +350,14 @@ class ZscalerAIGuard(CustomGuardrail): extra_headers = self._prepare_headers(api_key, **kwargs) data = { - "policyId": policy_id, "direction": direction, "content": content, } - + # Only include policyId when explicitly configured (policy_id >= 1) + # When policy_id is None, 0, or -1 (default), use resolve-and-execute-policy which infers + # the policy from headers (e.g., user-api-key-alias) + if policy_id is not None and policy_id >= 1: + data["policyId"] = policy_id try: response = await self._send_request( zscaler_ai_guard_url, extra_headers, data @@ -332,8 +365,7 @@ class ZscalerAIGuard(CustomGuardrail): return self._handle_response(response, direction) except Exception as e: verbose_proxy_logger.error(f"{e}. Blocking request.") - user_facing_error = self._create_user_facing_error(f"{str(e)})") - # This exception will be caught by the proxy and returned to the user + user_facing_error = self._create_user_facing_error(f"{str(e)}") raise HTTPException(status_code=500, detail=user_facing_error) @staticmethod diff --git a/tests/guardrails_tests/test_zscaler_ai_guard.py b/tests/guardrails_tests/test_zscaler_ai_guard.py index 9d519c17f6..fe58c2be9e 100644 --- a/tests/guardrails_tests/test_zscaler_ai_guard.py +++ b/tests/guardrails_tests/test_zscaler_ai_guard.py @@ -1,3 +1,15 @@ +""" +Unit tests for Zscaler AI Guard guardrail + +Tests covering: +- API call handling (ALLOW, BLOCK, errors) +- Policy ID precedence (metadata > user_api_key > team > init) +- Boolean config parameter handling (send_user_api_key_*) +- Metadata resolution (pre-call vs post-call) +- Header preparation with kwargs +- resolve-and-execute-policy endpoint (policyId omission) +""" + import pytest from unittest.mock import AsyncMock, Mock, patch from fastapi import HTTPException @@ -244,3 +256,88 @@ async def test_policy_id_zero_from_request_metadata(mock_api_call): await guardrail.apply_guardrail(inputs, request_data, "request") mock_api_call.assert_called_once() assert mock_api_call.call_args.kwargs["policy_id"] == 0 + +@pytest.mark.asyncio +async def test_should_use_config_send_user_api_key_alias_when_true(): + """Test that send_user_api_key_alias=True from config is used (not overridden by env)""" + guardrail = ZscalerAIGuard( + api_key="test_key", + send_user_api_key_alias=True, + ) + assert guardrail.send_user_api_key_alias is True + + +@pytest.mark.asyncio +async def test_should_preserve_policy_id_zero_in_init(): + """Test that policy_id=0 is preserved (not treated as falsy and overridden by env)""" + guardrail = ZscalerAIGuard( + api_key="test_key", + policy_id=0, + ) + assert guardrail.policy_id == 0 + + +@pytest.mark.asyncio +async def test_should_resolve_from_litellm_metadata_during_post_call(): + """Test that user_api_key_alias is resolved from litellm_metadata during post-call""" + request_data = { + "litellm_metadata": { + "user_api_key_alias": "test-alias-post-call" + } + } + result = ZscalerAIGuard._resolve_metadata_value(request_data, "user_api_key_alias") + assert result == "test-alias-post-call" + + +@pytest.mark.asyncio +async def test_should_resolve_user_api_key_key_alias_mapping(): + """Test key_alias -> user_api_key_key_alias mapping in litellm_metadata""" + request_data = { + "litellm_metadata": { + "user_api_key_key_alias": "test-key-alias" + } + } + result = ZscalerAIGuard._resolve_metadata_value(request_data, "user_api_key_alias") + assert result == "test-key-alias" + + +@pytest.mark.asyncio +async def test_should_include_user_api_key_alias_header(): + """Test that user-api-key-alias header is included when send_user_api_key_alias is True""" + guardrail = ZscalerAIGuard( + api_key="test_key", + send_user_api_key_alias=True, + ) + headers = guardrail._prepare_headers("test_key", user_api_key_alias="test-alias") + assert headers.get("user-api-key-alias") == "test-alias" + + +@pytest.mark.asyncio +async def test_should_omit_policy_id_when_zero_or_negative(): + """Test that policyId is omitted from request body when policy_id <= 0 (for resolve-and-execute-policy)""" + guardrail = ZscalerAIGuard( + api_key="test_key", + policy_id=-1, + ) + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "statusCode": 200, + "action": "ALLOW", + } + + with patch.object(guardrail, "_send_request", new_callable=AsyncMock) as mock_send: + mock_send.return_value = mock_response + + await guardrail.make_zscaler_ai_guard_api_call( + zscaler_ai_guard_url="http://example.com", + api_key="test_key", + policy_id=-1, + direction="OUT", + content="test content", + ) + + call_args = mock_send.call_args + data = call_args[0][2] # Third positional arg is data + assert "policyId" not in data