From de63702f0835ecc606a72a6d4b761d48053b3c2b Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 17:27:38 -0800 Subject: [PATCH] fix(guardrails): tolerate pass-through bodies request setup hasn't normalized --- .../generic_guardrail_api.py | 5 +- .../guardrails/guardrail_hooks/lakera_ai.py | 5 +- .../guardrail_hooks/lakera_ai_v2.py | 3 + .../guardrail_hooks/noma/noma_v2.py | 10 ++- litellm/proxy/litellm_pre_call_utils.py | 3 +- ...guardrails_on_unexpected_request_fields.py | 83 +++++++++++++++++++ .../guardrail_hooks/test_lakera_ai_v2.py | 29 +++++++ .../unit/proxy/test_litellm_pre_call_utils.py | 27 ++++++ 8 files changed, 160 insertions(+), 5 deletions(-) create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 7e052796545..f3309b98462 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -365,7 +365,10 @@ class GenericGuardrailAPI(CustomGuardrail): top_level_metadata: Final = request_data.get("metadata", {}) # Merge both sources, preferring litellm_metadata if both exist - metadata_dict: Final = {**top_level_metadata, **litellm_metadata} + metadata_dict: Final = { + **(top_level_metadata if isinstance(top_level_metadata, dict) else {}), + **(litellm_metadata if isinstance(litellm_metadata, dict) else {}), + } if not metadata_dict: return result_metadata diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..f7eb05e2ab6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -219,7 +219,10 @@ class lakeraAI_Moderation(CustomGuardrail): } ) elif "input" in data and isinstance(data["input"], list): - text = "\n".join(data["input"]) + text = "\n".join(item for item in data["input"] if isinstance(item, str)) + if not text: + verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text") + return _json_data = json.dumps( { "input": text, diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 3e5c54b4501..d127e10e685 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -781,6 +781,9 @@ class LakeraAIGuardrail(CustomGuardrail): # Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"] post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages + if not post_call_messages: + verbose_proxy_logger.debug("Lakera AI: not running guardrail. No inspectable text in data") + return response # Call Lakera guardrail lakera_guardrail_response, _ = await self.call_v2_guard( diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 43eb54d1464..71772ad6c46 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( @@ -63,6 +64,11 @@ class _Action(str, enum.Enum): GUARDRAIL_INTERVENED = "GUARDRAIL_INTERVENED" +def _user_api_key_alias(request_data: Mapping[str, object], metadata_key: str) -> object: + """The key alias recorded under ``metadata_key``, which a pass-through client can set to anything.""" + return (as_str_mapping(request_data.get(metadata_key)) or {}).get("user_api_key_alias") + + class NomaV2Guardrail(CustomGuardrail): reads_stream_marks: ClassVar[bool] = True """Noma keeps one record per streamed reply when every scan carries ``stream`` and only the last one carries the @@ -336,8 +342,8 @@ class NomaV2Guardrail(CustomGuardrail): # (ports v1 fallback from PR #16832). if application_id is None: application_id = self._get_non_empty_str( - request_data.get("litellm_metadata", {}).get("user_api_key_alias") - ) or self._get_non_empty_str(request_data.get("metadata", {}).get("user_api_key_alias")) + _user_api_key_alias(request_data, "litellm_metadata") + ) or self._get_non_empty_str(_user_api_key_alias(request_data, "metadata")) try: payload: Final = self._build_scan_payload( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..501c8352b10 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4377,11 +4377,12 @@ async def add_guardrails_from_policy_engine( all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias + _model: Final = data.get("model") context: Final = PolicyMatchContext.model_validate( { "team_alias": _team_alias if isinstance(_team_alias, str) else None, "key_alias": _key_alias if isinstance(_key_alias, str) else None, - "model": data.get("model"), + "model": _model if isinstance(_model, str) else None, "tags": all_tags, } ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py new file mode 100644 index 00000000000..f1e1dca04c3 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py @@ -0,0 +1,83 @@ +"""Guardrails on built-in pass-through routes read client JSON that request setup hasn't normalized.""" + +from typing import Final +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( + GenericGuardrailAPI, +) +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation +from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail + +NOT_AN_OBJECT: Final = (None, 0, False, "", [], "text", 7) + + +@pytest.mark.parametrize("value", NOT_AN_OBJECT) +def test_generic_guardrail_reads_metadata_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + + metadata: Final = guardrail._extract_user_api_key_metadata( + {"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value} + ) + + assert metadata["user_api_key_hash"] == "hash" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", NOT_AN_OBJECT) +async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None: + guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma") + verdict: Final = httpx.Response( + 200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid") + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)): + result: Final = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}, "litellm_metadata": value}, + input_type="request", + ) + + assert result["texts"] == ["hello"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]]) +async def test_legacy_lakera_has_no_text_to_scan_in_token_array_embeddings(token_arrays: list[object]) -> None: + guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid") + vendor_error: Final = httpx.HTTPStatusError( + "vendor called", + request=httpx.Request("POST", "http://vendor.invalid"), + response=httpx.Response(500, text="vendor called"), + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)): + await guardrail._check( + data={"input": token_arrays}, + user_api_key_dict=UserAPIKeyAuth(api_key="key"), + call_type="pass_through_endpoint", + ) + + +@pytest.mark.asyncio +async def test_legacy_lakera_still_scans_string_inputs() -> None: + guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid") + vendor_error: Final = httpx.HTTPStatusError( + "vendor called", + request=httpx.Request("POST", "http://vendor.invalid"), + response=httpx.Response(500, text="vendor called"), + ) + + with ( + patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)), + pytest.raises(Exception, match="vendor called"), + ): + await guardrail._check( + data={"input": ["hello"]}, + user_api_key_dict=UserAPIKeyAuth(api_key="key"), + call_type="pass_through_endpoint", + ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 23288a6136f..52c3c9eb92d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -2413,6 +2413,35 @@ async def test_lakera_post_call_allows_clean_content(): assert result is llm_response +@pytest.mark.asyncio +async def test_post_call_with_no_messages_and_no_reply_text_returns_the_response(): + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + response: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}}]} + rejected: Final = HTTPException(status_code=400, detail="either content, messages, or tools must be provided") + with patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)): + result: Final = await guardrail.async_post_call_success_hook( + data={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response=response, + ) + assert result == response + + +@pytest.mark.asyncio +async def test_post_call_with_messages_still_reaches_lakera(): + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + rejected: Final = HTTPException(status_code=400, detail="rejected") + with ( + patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)), + pytest.raises(HTTPException), + ): + await guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hello"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response={}, + ) + + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_lakera_post_call_masks_pii_and_allows(): diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..eeffba68969 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5327,6 +5327,33 @@ async def test_embedding_header_forwarding_without_model_group_config(): litellm.model_group_settings = original_model_group_settings +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [0, False, [], {}, 7, True]) +async def test_policy_engine_leaves_a_model_that_is_not_a_string_alone( + model: object, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails + + policy_registry = get_policy_registry() + monkeypatch.setattr( + policy_registry, "_policies", {"baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"]))} + ) + monkeypatch.setattr(policy_registry, "_initialized", True) + attachment_registry = get_attachment_registry() + monkeypatch.setattr(attachment_registry, "_attachments", [PolicyAttachment(policy="baseline", scope="*")]) + monkeypatch.setattr(attachment_registry, "_initialized", True) + data = {"model": model, "metadata": {}} + + await add_guardrails_from_policy_engine( + data=data, metadata_variable_name="metadata", user_api_key_dict=UserAPIKeyAuth(api_key="test-key") + ) + + assert data["model"] == model + assert data["metadata"]["guardrails"] == ["pii_blocker"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine(): """