scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-09 08:18:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

fix(proxy): keep pass-through guardrail settings, gate frame policy picks, skip moderation-only hooks on socket turns

Browse source
This commit is contained in:
mateo-berri 2026-10-29 07:52:42 -07:00
parent 736f7c22d9
commit 7d5037cd79
6 changed files with 107 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/proxy/auth/auth_checks.py
View file

@ -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(

32
litellm/proxy/pass_through_endpoints/passthrough_guardrails.py
View file

@ -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(

29
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -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]:

2
litellm/proxy/utils.py
View file

@ -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
):

13
tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py
View file

@ -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
}

71
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -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