From 7d5037cd79c8b314d3e8ec568ab7bddd563adfe9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:52:42 -0700 Subject: [PATCH] fix(proxy): keep pass-through guardrail settings, gate frame policy picks, skip moderation-only hooks on socket turns --- litellm/proxy/auth/auth_checks.py | 8 ++- .../passthrough_guardrails.py | 32 ++------- .../proxy/response_api_endpoints/endpoints.py | 29 ++++---- litellm/proxy/utils.py | 2 + .../test_passthrough_guardrails.py | 13 ++-- .../response_api_endpoints/test_endpoints.py | 71 ++++++++++++++++++- 6 files changed, 107 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3044caca42e..beb8f614f4f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -981,10 +981,12 @@ def guardrail_modification_check(request_body: Mapping[str, object], team_object failing loudly at the auth layer so operators see an explicit 403 instead of a confusing silent-ignore. """ - from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails + if requests_guardrail_modification(request_body): + require_guardrail_modification_permission(team_object) - if not requests_guardrail_modification(request_body): - return + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: + from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails if not can_modify_guardrails(team_object): raise HTTPException( diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index f35d4346d87..bd8390d704e 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -7,11 +7,10 @@ Handles guardrail execution for passthrough endpoints with: - Automatic inheritance from org/team/key levels when enabled """ -from itertools import chain from typing import Any, Final from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import guardrail_entry_names +from litellm.integrations.custom_guardrail import merged_guardrail_entries from litellm.proxy._types import ( PassThroughGuardrailsConfig, PassThroughGuardrailSettings, @@ -207,7 +206,7 @@ class PassthroughGuardrailHandler: def collect_guardrails( user_api_key_dict: UserAPIKeyAuth, passthrough_guardrails_config: PassThroughGuardrailsConfigInput | None, - ) -> dict[str, bool] | None: + ) -> list[object] | None: """ Collect guardrails for a passthrough endpoint. @@ -227,7 +226,7 @@ class PassthroughGuardrailHandler: passthrough_guardrails_config: List or Dict of guardrail names/settings Returns: - Dict of guardrail names to run (format: {guardrail_name: True}), or None + Guardrail entries to run, each a name or a ``{name: settings}`` entry, or None """ from litellm.proxy.litellm_pre_call_utils import ( add_guardrails_from_key_or_team_metadata, @@ -244,15 +243,6 @@ class PassthroughGuardrailHandler: verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None - # Passthrough is enabled - collect guardrails - guardrails_to_run: Final[dict[str, bool]] = {} - - # Add passthrough-specific guardrails - for guardrail_name in normalized_config: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added passthrough-specific guardrail: %s", guardrail_name) - - # Add org/team/key level guardrails using shared helper temp_data: Final[dict[str, Any]] = {"metadata": {}} add_guardrails_from_key_or_team_metadata( key_metadata=user_api_key_dict.metadata, @@ -260,20 +250,12 @@ class PassthroughGuardrailHandler: data=temp_data, metadata_variable_name="metadata", ) + inherited_guardrails: Final[list[object]] = temp_data["metadata"].get("guardrails", []) + guardrails_to_run: Final = merged_guardrail_entries(normalized_config, inherited_guardrails) - # Merge inherited guardrails into guardrails_to_run - inherited_guardrails: Final = temp_data["metadata"].get("guardrails", []) - for guardrail_name in chain.from_iterable(guardrail_entry_names(entry) for entry in inherited_guardrails): - if guardrail_name not in guardrails_to_run: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added inherited guardrail (key/team level): %s", guardrail_name) + verbose_proxy_logger.debug("Collected guardrails for passthrough endpoint: %s", guardrails_to_run) - verbose_proxy_logger.debug( - "Collected guardrails for passthrough endpoint: %s", - list(guardrails_to_run.keys()), - ) - - return guardrails_to_run if guardrails_to_run else None + return guardrails_to_run or None @staticmethod def get_field_targeted_text( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5bf55c71e88..685f1086eb2 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -50,8 +50,8 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( get_team_object, - guardrail_modification_check, requests_guardrail_modification, + require_guardrail_modification_permission, ) from litellm.proxy.auth.auth_utils import carries_client_side_overrides from litellm.proxy.auth.user_api_key_auth import ( @@ -2183,6 +2183,7 @@ class _GuardrailWalk(Protocol): data: dict[str, object], # mutable-ok: ProxyLogging.during_call_hook hands the guardrails a dict they edit user_api_key_dict: UserAPIKeyAuth | None, call_type: Literal["aresponses"], + guardrails_only: bool, ) -> object: ... @@ -2192,13 +2193,16 @@ async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTab team_id: Final = user_api_key_dict.team_id if team_id is None or team_id == UI_TEAM_ID: return None - return await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_dict.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + return LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata) def _responses_ws_frame_guard( @@ -2274,7 +2278,7 @@ def _responses_ws_frame_guard( } pii_tokens: Final = _recorded_pii_tokens(guarded_data.get("metadata")) await proxy_logging_obj.during_call_hook( - data=guarded_data, user_api_key_dict=user_api_key_dict, call_type="aresponses" + data=guarded_data, user_api_key_dict=user_api_key_dict, call_type="aresponses", guardrails_only=True ) except Exception as e: raise await rejection(e, frame_data) from e @@ -2316,10 +2320,9 @@ async def _check_frame_may_modify_guardrails( user_api_key_dict: UserAPIKeyAuth, team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]], ) -> None: - """A frame that picks guardrails needs the team permission an HTTP body that picks them needs""" - modifying_source: Final = next((source for source in sources if requests_guardrail_modification(source)), None) - if modifying_source is not None: - guardrail_modification_check(modifying_source, await team_object(user_api_key_dict)) + """A frame that picks guardrails or policies needs the team's permission to modify guardrails""" + if any("policies" in source or requests_guardrail_modification(source) for source in sources): + require_guardrail_modification_permission(await team_object(user_api_key_dict)) def _forwarded_frame_fields(source: Mapping[str, object]) -> dict[str, object]: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8f51d14f14a..47159c81e71 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3497,6 +3497,7 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ): caps: Final = ProxyLogging.callback_capabilities() if not caps.has_guardrail and not caps.has_moderation_override: @@ -3508,6 +3509,7 @@ class ProxyLogging: if ( isinstance(callback, CustomLogger) and not isinstance(callback, CustomGuardrail) + and (not guardrails_only or callback.enforces_request_content) and _overrides_moderation_hook(callback) and user_api_key_dict is not None ): diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py index ddf22b920b5..ef19482b30a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py @@ -74,7 +74,8 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: class TestPassthroughGuardrailHandlerCollectGuardrails: - def test_inherits_a_key_guardrail_listed_with_settings_by_its_name(self, monkeypatch): + def test_inherits_a_key_guardrail_with_its_settings(self, monkeypatch: pytest.MonkeyPatch): + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) @@ -82,10 +83,12 @@ class TestPassthroughGuardrailHandlerCollectGuardrails: api_key="hashed", metadata={"guardrails": [{"pii": {"extra_body": {"threshold": 1}}}, "toxic"]} ) - assert PassthroughGuardrailHandler.collect_guardrails(key, ["passthrough-guard"]) == { - "passthrough-guard": True, - "pii": True, - "toxic": True, + guardrails = PassthroughGuardrailHandler.collect_guardrails(key, ["passthrough-guard", "toxic"]) + + assert guardrails == ["passthrough-guard", "toxic", {"pii": {"extra_body": {"threshold": 1}}}] + request_data = {"metadata": {"guardrails": guardrails}} + assert CustomGuardrail(guardrail_name="pii").get_guardrail_dynamic_request_body_params(request_data) == { + "threshold": 1 } diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 0fc88fc7f3b..9dea7fa0af9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2628,6 +2628,34 @@ class TestResponsesWSFirstFrameModelAuth: assert json.loads(unselected_turn.frame) == {"type": "response.create", "input": "plain"} assert during_pick.inputs == ["hi"] + @pytest.mark.asyncio + async def test_frame_guard_runs_only_the_moderation_hooks_that_enforce_request_content(self): + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + class _ModerationHook(CustomLogger): + def __init__(self, enforces_request_content: bool) -> None: + super().__init__() + self.enforces_request_content = enforces_request_content + self.inputs: list[object] = [] + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + self.inputs.append(data.get("input")) + + vendor_moderation = _ModerationHook(enforces_request_content=False) + content_enforcer = _ModerationHook(enforces_request_content=True) + + with _guardrails_on(vendor_moderation, content_enforcer) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert vendor_moderation.inputs == [] + assert content_enforcer.inputs == ["hi"] + @pytest.mark.asyncio @pytest.mark.parametrize( ("selection", "message"), @@ -2918,10 +2946,12 @@ class TestResponsesWSFirstFrameModelAuth: {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": True}}, {"type": "response.create", "response": {"input": "hi", "guardrails": ["frame-pick"]}}, {"type": "response.create", "litellm_metadata": {"guardrails": ["frame-pick"]}, "response": {"input": "hi"}}, + {"type": "response.create", "input": "hi", "policies": ["frame-policy"]}, + {"type": "response.create", "response": {"input": "hi", "policies": ["frame-policy"]}}, ], - ids=["top-level", "metadata", "global-toggle", "nested", "beside-nested"], + ids=["top-level", "metadata", "global-toggle", "nested", "beside-nested", "policies", "nested-policies"], ) - async def test_frame_guard_refuses_guardrail_changes_from_a_team_without_permission_like_http(self, frame): + async def test_frame_guard_refuses_guardrail_and_policy_picks_from_a_team_without_permission(self, frame): from litellm.proxy._types import LiteLLM_TeamTable from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard @@ -2977,6 +3007,43 @@ class TestResponsesWSFirstFrameModelAuth: assert frame_pick.inputs == ["hi"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("modify_guardrails", "picked_inputs"), + [(True, ["hi"]), (False, [])], + ids=["open", "locked"], + ) + async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_row_is_gone( + self, modify_guardrails: bool, picked_inputs: list[str] + ): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + key = UserAPIKeyAuth( + api_key="hashed", + team_id="deleted-team", + team_metadata={"guardrails": {"modify_guardrails": modify_guardrails}}, + ) + frame = json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + + with ( + _guardrails_on(frame_pick) as proxy_logging_obj, + patch( + "litellm.proxy.response_api_endpoints.endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})), + ), + ): + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, key, proxy_logging_obj, MagicMock() + ) + if modify_guardrails: + await guard(frame) + else: + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(frame) + + assert frame_pick.inputs == picked_inputs + @pytest.mark.asyncio async def test_frame_guard_forwards_the_pre_call_output_and_not_a_during_call_rewrite_like_http(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard