From 3d5591e0b0b759ffe5ecfb762bbfa1af898ee016 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:31:52 -0700 Subject: [PATCH] fix(responses): check frame guardrail picks against the team permission, honor picks beside a nested response, prefer a guardrail entry with settings --- litellm/integrations/custom_guardrail.py | 37 ++-- litellm/proxy/auth/auth_checks.py | 75 ++++---- .../passthrough_guardrails.py | 4 +- .../proxy/response_api_endpoints/endpoints.py | 126 +++++++++----- .../integrations/test_custom_guardrail.py | 7 +- ...st_auth_checks_object_access_and_lookup.py | 6 +- .../test_passthrough_guardrails.py | 15 ++ .../response_api_endpoints/test_endpoints.py | 162 ++++++++++++++++-- .../utils/helpers/test_guardrail_merge.py | 8 + 9 files changed, 330 insertions(+), 110 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 15a29ddfb7e..1b597da8a24 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -213,23 +213,40 @@ _GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: - """Merges guardrail lists in order, dropping an entry whose guardrails an earlier one already names, so the - first entry for a guardrail, plain or ``{name: settings}``, is the one that stays""" + """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over + any plain name in either order, else its first plain name stays""" entries: Final = tuple(chain.from_iterable(entry_lists)) - names: Final = tuple(_guardrail_entry_names(entry) for entry in entries) - return [entry for index, entry in enumerate(entries) if not _named_earlier(names, index)] + settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries) + named_with_settings: Final = frozenset[str]().union(*settings_names) + return [ + entry + for index, entry in enumerate(entries) + if not _superseded_guardrail_entry(entries, settings_names, named_with_settings, index) + ] -def _guardrail_entry_names(entry: object) -> frozenset[str]: +def guardrail_entry_names(entry: object) -> tuple[str, ...]: if isinstance(entry, str): - return frozenset((entry,)) + return (entry,) if isinstance(entry, dict): - return frozenset(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) - return frozenset() + return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return () -def _named_earlier(names: tuple[frozenset[str], ...], index: int) -> bool: - return bool(names[index]) and names[index] <= frozenset[str]().union(*names[:index]) +def _settings_entry_names(entry: object) -> frozenset[str]: + return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset() + + +def _superseded_guardrail_entry( + entries: tuple[object, ...], + settings_names: tuple[frozenset[str], ...], + named_with_settings: frozenset[str], + index: int, +) -> bool: + entry: Final = entries[index] + if isinstance(entry, str): + return entry in named_with_settings or entry in entries[:index] + return bool(settings_names[index]) and settings_names[index] <= frozenset[str]().union(*settings_names[:index]) def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 448643ef59b..ebd4fb3e691 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -930,13 +930,45 @@ _GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = ( ) -def _guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: +def requests_guardrail_modification(request_body: Mapping[str, object]) -> bool: + """Whether a request body sets the plural ``guardrails`` list or a per-request toggle that influences whether + default-on guardrails run (``disable_global_guardrails``, ``disable_global_guardrail`` singular, and + ``opted_out_global_guardrails``), at the top level or in either metadata key""" + return ( + _requests_guardrail_modification_in(request_body.get("metadata")) + or _requests_guardrail_modification_in(request_body.get("litellm_metadata")) + or _requests_guardrail_modification_in(request_body) + ) + + +def _requests_guardrail_modification_in(container: object) -> bool: + coerced: Final = _guardrail_container_as_dict(container) + if coerced is None: + return False + return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) + + +def _guardrail_container_as_dict(container: object) -> Mapping[str, object] | None: + """Accept dict or JSON-string (from multipart/form-data or extra_body). + + Without this, an attacker can smuggle guardrail keys past the check by + sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — + ``isinstance(dict)`` on the string returns False, the check returns + no-modification, and ``add_litellm_data_to_request`` parses the string + to a dict downstream. + """ + if isinstance(container, Mapping): + return container + if isinstance(container, str): + parsed: Final = _safe_json_loads_obj(container) + return parsed if isinstance(parsed, dict) else None + return None + + +def guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: """ Reject user-supplied metadata flags that would modify guardrail behavior - unless the team has explicit permission. Checked keys include the plural - ``guardrails`` list plus the per-request toggles that influence whether - default-on guardrails run (``disable_global_guardrails``, - ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``). + unless the team has explicit permission. User-supplied values for the bypass toggles are also silently ignored by ``_get_admin_metadata`` at read time; this check adds defense in depth by @@ -945,36 +977,7 @@ def _guardrail_modification_check(request_body: Mapping[str, object], team_objec """ from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails - def _coerce_to_dict(container: object) -> dict | None: - """Accept dict or JSON-string (from multipart/form-data or extra_body). - - Without this, an attacker can smuggle guardrail keys past the check by - sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — - ``isinstance(dict)`` on the string returns False, the check returns - no-modification, and ``add_litellm_data_to_request`` parses the string - to a dict downstream. - """ - if isinstance(container, dict): - return container - if isinstance(container, str): - parsed: Final = _safe_json_loads_obj(container) - return parsed if isinstance(parsed, dict) else None - return None - - def _user_requested_modification(container: object) -> bool: - coerced: Final = _coerce_to_dict(container) - if coerced is None: - return False - return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) - - # Check both metadata keys — callers can populate either depending on the - # endpoint. Cover the top-level too so root-level injection is rejected. - modifies: Final = ( - _user_requested_modification(request_body.get("metadata")) - or _user_requested_modification(request_body.get("litellm_metadata")) - or _user_requested_modification(request_body) - ) - if not modifies: + if not requests_guardrail_modification(request_body): return if not can_modify_guardrails(team_object): @@ -1465,7 +1468,7 @@ async def common_checks( _enforce_user_param_check(general_settings, request, request_body, route) _global_proxy_budget_check(global_proxy_spend, skip_all_budget_checks, route) - _guardrail_modification_check(_typed_request_body(request_body), team_object) + guardrail_modification_check(_typed_request_body(request_body), team_object) # 10 [OPTIONAL] Organization RBAC checks organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body) diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index 1c5d270542f..f35d4346d87 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -7,9 +7,11 @@ 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.proxy._types import ( PassThroughGuardrailsConfig, PassThroughGuardrailSettings, @@ -261,7 +263,7 @@ class PassthroughGuardrailHandler: # Merge inherited guardrails into guardrails_to_run inherited_guardrails: Final = temp_data["metadata"].get("guardrails", []) - for guardrail_name in inherited_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) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 3ebf99d176a..8c94b604423 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -47,6 +47,11 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( build_blocked_response, ) from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import ( + get_team_object, + guardrail_modification_check, + requests_guardrail_modification, +) from litellm.proxy.auth.auth_utils import carries_client_side_overrides from litellm.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, @@ -2062,7 +2067,6 @@ _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") _METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config") _PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") -_NESTED_BODY_ONLY_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, "guardrails", "policies") _UNFORWARDED_FRAME_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS, "litellm_metadata") _OBJECTS: Final = TypeAdapter(list[object]) _TURN_COPIED_CONTAINERS: Final = (list, dict) @@ -2087,15 +2091,33 @@ class _GuardrailWalk(Protocol): ) -> object: ... +async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + 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, + ) + + def _responses_ws_frame_guard( admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: _GuardrailWalk, record_refused_turn: Callable[[tuple[Mapping[str, object], ...], Exception], None], + *, + team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]] = _key_team_object, ) -> ResponsesWsFrameGuard: """Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every ``response.create`` frame gets the pre_call and during_call guardrail walks the HTTP ``/v1/responses`` path runs on - its body. during_call runs before the frame is forwarded, since a forwarded frame cannot be recalled.""" + its body. during_call runs before the frame is forwarded, since a forwarded frame cannot be recalled, and like HTTP, + which sends the model call beside during_call, the frame carries the pre_call output and not during_call rewrites.""" metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) admitted_bucket: Final = admitted.get(metadata_key) bucket: Final = { @@ -2110,16 +2132,17 @@ def _responses_ws_frame_guard( # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) + async def rejection(exception: Exception, data: dict[str, object]) -> ResponsesWsFrameRejection: # mutable-ok: billed + await _bill_rejected_frame(rejected_frame_auth, exception, data) + proxy_exception: Final = proxy_exception_from_route_error(exception) + return ResponsesWsFrameRejection(proxy_exception.message, proxy_exception.provider_specific_fields) + async def guard(frame: str) -> GuardedResponsesWsFrame: frame_obj: Final = _WS_FRAME.validate_json(frame) nested: Final = frame_obj.get("response") - stray: Final = _field_beside_nested_response(frame_obj) if isinstance(nested, dict) else None - if stray is not None: - raise ValueError( - f"response.create sends {stray} at the top level next to response; send it inside response" - ) body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj - frame_data: Final = guardrail_request_data_with_streaming( + selection_sources: Final = (frame_obj, body) if isinstance(nested, dict) else (body,) + turn_data: Final = guardrail_request_data_with_streaming( { **{ key: _without_untrusted_metadata(key, value) @@ -2128,23 +2151,35 @@ def _responses_ws_frame_guard( }, "model": model, "litellm_call_id": str(uuid.uuid4()), - metadata_key: await _metadata_with_frame_guardrails( - bucket, body, model, metadata_key, user_api_key_dict - ), + metadata_key: {**bucket}, }, is_streaming=True, ) + try: + _reject_guarded_field_beside_nested_response(frame_obj, nested) + await _check_frame_may_modify_guardrails(selection_sources, user_api_key_dict, team_object) + frame_data: Final = { + **turn_data, + metadata_key: await _metadata_with_frame_guardrails( + bucket, selection_sources, model, metadata_key, user_api_key_dict + ), + } + except Exception as e: + raise await rejection(e, turn_data) from e try: guarded_data: Final = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True ) + # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them + guarded_fields: Final = { + key: copy.deepcopy(guarded_data[key]) for key in _GUARDED_FRAME_FIELDS if key in guarded_data + } + 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" ) except Exception as e: - await _bill_rejected_frame(rejected_frame_auth, e, frame_data) - proxy_exception: Final = proxy_exception_from_route_error(e) - raise ResponsesWsFrameRejection(proxy_exception.message, proxy_exception.provider_specific_fields) from e + raise await rejection(e, frame_data) from e if guarded_data.get("model") != model: rerouted: Final = ValueError( f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" @@ -2152,39 +2187,47 @@ def _responses_ws_frame_guard( ) await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data) raise rerouted - # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them - guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} - kept: Final = { - key: _without_guardrail_selection(value) if key == "metadata" else value - for key, value in body.items() - if key not in _UNFORWARDED_FRAME_FIELDS - } rewritten: Final = ( - { - **{key: value for key, value in frame_obj.items() if key not in _PROXY_ONLY_FRAME_FIELDS}, - "response": {**kept, **guarded_fields}, - } + {**_forwarded_frame_fields(frame_obj), "response": {**_forwarded_frame_fields(body), **guarded_fields}} if isinstance(nested, dict) - else {**kept, **guarded_fields} + else {**_forwarded_frame_fields(body), **guarded_fields} ) guardrail_information: Final = _recorded_guardrail_information(guarded_data.get(metadata_key)) return GuardedResponsesWsFrame( frame=json.dumps(rewritten), guardrail_information=guardrail_information, - pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + pii_tokens=pii_tokens, bill_refusal=partial(record_refused_turn, guardrail_information), ) return guard -def _field_beside_nested_response(frame: Mapping[str, object]) -> str | None: - metadata_picks: Final = ( - f"{key}.guardrails" - for key, value in frame.items() - if key in _CLIENT_METADATA_KEYS and isinstance(value, dict) and "guardrails" in value - ) - return next(chain((key for key in _NESTED_BODY_ONLY_FIELDS if key in frame), metadata_picks), None) +def _reject_guarded_field_beside_nested_response(frame: Mapping[str, object], nested: object) -> None: + """Only the nested ``response`` body runs through the guardrails, so a guarded field beside it would go upstream + unguarded""" + stray: Final = next((key for key in _GUARDED_FRAME_FIELDS if key in frame), None) if isinstance(nested, dict) else None + if stray is not None: + raise ValueError(f"response.create sends {stray} at the top level next to response; send it inside response") + + +async def _check_frame_may_modify_guardrails( + sources: tuple[Mapping[str, object], ...], + 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)) + + +def _forwarded_frame_fields(source: Mapping[str, object]) -> dict[str, object]: + return { + key: _without_guardrail_selection(value) if key == "metadata" else value + for key, value in source.items() + if key not in _UNFORWARDED_FRAME_FIELDS + } class _FrameGuardrailSelection(TypedDict, total=False): @@ -2198,18 +2241,21 @@ _FRAME_GUARDRAIL_SELECTION: Final = TypeAdapter(_FrameGuardrailSelection) async def _metadata_with_frame_guardrails( bucket: Mapping[str, object], - body: Mapping[str, object], + sources: tuple[Mapping[str, object], ...], model: object, metadata_key: str, user_api_key_dict: UserAPIKeyAuth, ) -> Mapping[str, object]: """A frame selects guardrails the way an HTTP ``/v1/responses`` body does, adding to the key, team and policy - guardrails the socket was admitted with, so a frame can never shrink that list""" - selections: Final = _frame_guardrail_selections(body) - policies: Final = selections[-1].get("policies") + guardrails the socket was admitted with, so a frame can never shrink that list. A nested frame selects at its top + level and in its ``response`` body, the body last so its ``guardrail_config`` wins""" + selections: Final = tuple(chain.from_iterable(_frame_guardrail_selections(source) for source in sources)) + picked_policies: Final = tuple(selection["policies"] for selection in selections if "policies" in selection) metadata: Final = ( - await _metadata_with_frame_policies(bucket, policies, model, metadata_key, user_api_key_dict) - if policies is not None + await _metadata_with_frame_policies( + bucket, tuple(chain.from_iterable(picked_policies)), model, metadata_key, user_api_key_dict + ) + if picked_policies else bucket ) selected_guardrails: Final = tuple(chain.from_iterable(selection.get("guardrails", ()) for selection in selections)) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index de5f6995dcc..1a04f649642 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3788,11 +3788,14 @@ def test_default_on_guardrail() -> None: ) -def test_merged_guardrail_entries_keeps_the_first_entry_per_guardrail_and_any_entry_that_names_a_new_one(): +def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): from litellm.integrations.custom_guardrail import merged_guardrail_entries assert merged_guardrail_entries(["a"], [{"a": {}, "b": {}}, {"a": {"x": 1}}], ["b", "c", "c"]) == [ - "a", {"a": {}, "b": {}}, "c", ] + assert merged_guardrail_entries(["pii", {"pii": {"extra_body": {"threshold": 1}}}], ["model-guard"]) == [ + {"pii": {"extra_body": {"threshold": 1}}}, + "model-guard", + ] diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index effdd9725c2..e116c07dc11 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -5636,7 +5636,7 @@ async def test_team_member_budget_check_reads_from_spend_counter(): class TestGuardrailModificationCheck: - """Defense-in-depth: `_guardrail_modification_check` must 403 when the + """Defense-in-depth: `guardrail_modification_check` must 403 when the caller's metadata attempts to modify any guardrail-related key and the team lacks the `modify_guardrails` permission. Checks both the historically-covered `guardrails` list and the bypass toggles that @@ -5644,11 +5644,11 @@ class TestGuardrailModificationCheck: """ def _call(self, request_body): - from litellm.proxy.auth.auth_checks import _guardrail_modification_check + from litellm.proxy.auth.auth_checks import guardrail_modification_check team_object = MagicMock() team_object.metadata = {} # no permission - return _guardrail_modification_check(request_body=request_body, team_object=team_object) + return guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): # no-op — should return silently 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 40403de846f..7ebb82386a8 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py @@ -73,6 +73,21 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: assert set(result) == {"pii-detection", "content-moderation"} +class TestPassthroughGuardrailHandlerCollectGuardrails: + def test_inherits_a_key_guardrail_listed_with_settings_by_its_name(self): + from litellm.proxy._types import UserAPIKeyAuth + + key = UserAPIKeyAuth( + 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, + } + + class TestPassthroughGuardrailHandlerNormalizeConfig: """Tests for PassthroughGuardrailHandler.normalize_config method.""" diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 5431c633623..a4e14a11748 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1996,7 +1996,7 @@ class _FrameGuardrail(CustomGuardrail): async def async_moderation_hook( self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral ) -> None: - self._scan(data) + data["input"] = self._scan(data)["input"] def _scan(self, data: dict[str, object]) -> dict[str, object]: self.inputs.append(data["input"]) @@ -2245,7 +2245,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( - {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() ) metadata = { "user_api_key_metadata": {"disable_global_guardrails": True}, @@ -2275,7 +2275,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, - MagicMock(), + _key(), proxy_logging_obj, MagicMock(), ) @@ -2422,12 +2422,17 @@ class TestResponsesWSFirstFrameModelAuth: ), ], ) - async def test_frame_guard_refuses_a_frame_whose_guardrail_selection_is_malformed(self, selection, message): + async def test_frame_guard_refuses_and_bills_a_frame_whose_guardrail_selection_is_malformed( + self, selection, message + ): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard key_required = _FrameGuardrail("key-required") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() - with _guardrails_on(key_required) as proxy_logging_obj: + with _guardrails_on(key_required, spend_logger) as proxy_logging_obj: guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, _key(), @@ -2438,6 +2443,7 @@ class TestResponsesWSFirstFrameModelAuth: await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) assert key_required.inputs == [] + spend_logger.async_post_call_failure_hook.assert_awaited_once() @pytest.mark.asyncio async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self): @@ -2613,30 +2619,150 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - @pytest.mark.parametrize( - ("top_level", "field"), - [ - ({"input": "forbidden"}, "input"), - ({"guardrails": ["frame-pick"]}, "guardrails"), - ({"policies": ["frame-policy"]}, "policies"), - ({"metadata": {"guardrails": ["frame-pick"]}}, "metadata.guardrails"), - ({"litellm_metadata": {"guardrails": ["frame-pick"]}}, "litellm_metadata.guardrails"), - ], - ) - async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self, top_level, field): + @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) + async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() ) message = f"response.create sends {field} at the top level next to response; send it inside response" - with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - await guard(json.dumps({"type": "response.create", **top_level, "response": {}})) + with ( + patch.object(litellm, "callbacks", [spend_logger]), + pytest.raises(ValueError, match=f"^{re.escape(message)}$"), + ): + await guard(json.dumps({"type": "response.create", field: "unguarded", "response": {"input": "hi"}})) proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"]["input"] == "hi" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "top_level", + [ + {"guardrails": ["frame-pick"]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "metadata", "litellm-metadata"], + ) + async def test_frame_guard_runs_the_guardrails_a_nested_frame_selects_beside_its_response(self, top_level): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick", rewrites=True) + unselected = _FrameGuardrail("unselected") + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps({"type": "response.create", **top_level, "response": {"input": "hi"}})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in top_level else {} + assert json.loads(guarded.frame) == { + "type": "response.create", + **kept_metadata, + "response": {"input": "hi [frame-pick]"}, + } + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}, + {"type": "response.create", "input": "hi", "metadata": {"guardrails": ["frame-pick"]}}, + {"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"}}, + ], + ids=["top-level", "metadata", "global-toggle", "nested", "beside-nested"], + ) + async def test_frame_guard_refuses_guardrail_changes_from_a_team_without_permission_like_http(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 + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + team_lookups: list[UserAPIKeyAuth] = [] + key = _key() + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + team_lookups.append(user_api_key_dict) + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick, spend_logger) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + key, + proxy_logging_obj, + MagicMock(), + team_object=locked_team, + ) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(json.dumps(frame)) + + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "plain"} + assert key_required.inputs == ["plain"] + assert frame_pick.inputs == [] + assert team_lookups == [key] + spend_logger.async_post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_frame_guard_runs_a_frame_s_picks_for_a_team_allowed_to_modify_guardrails(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + + async def open_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": True}}) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, + _key(), + proxy_logging_obj, + MagicMock(), + team_object=open_team, + ) + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]})) + + assert frame_pick.inputs == ["hi"] + + @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 + + pre_pick = _FrameGuardrail("pre-pick", rewrites=True) + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, rewrites=True) + + with _guardrails_on(pre_pick, during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["pre-pick", "during-pick"]}) + ) + + assert during_pick.inputs == ["hi [pre-pick]"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi [pre-pick]"} @pytest.mark.asyncio async def test_frame_guard_forwards_a_nested_turn_without_top_level_fields_that_pick_no_guardrail(self): diff --git a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py index 4ae87dda343..b65bc535a9d 100644 --- a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py +++ b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py @@ -366,3 +366,11 @@ def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_ov result = _merge_guardrails_with_existing({"metadata": {"guardrails": [configured]}}, ["pii-redact", "toxic-filter"]) assert result["metadata"]["guardrails"] == [configured, "toxic-filter"] + + +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_listed_after_the_key_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": ["pii-redact", configured]}}, ["toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"]