From 7e37c207482dde6b17c19783e9d8aeafadaecd11 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:12:51 -0700 Subject: [PATCH] fix(responses): run the guardrails a Responses WebSocket frame selects (#48725) * fix(responses): run the guardrails a Responses WebSocket frame selects * fix(responses): name the field a malformed frame guardrail selection got wrong A frame whose guardrails or guardrail_config has the wrong shape now gets an error event naming the field, instead of pydantic's raw validation dump * fix(responses): accept every guardrail selection shape a Responses WebSocket frame sends, and run during_call guardrails A frame can pick guardrails as a list of names or name-to-config dicts, inside metadata or litellm_metadata, with policies validated as a list of strings. during_call guardrails now run before the turn goes upstream, and selection fields sent next to a nested response are rejected instead of dropped * fix(proxy): keep a key's configured guardrails when policies add guardrails A policy's guardrails were merged with dict.fromkeys, so a configured guardrail entry ({name: {extra_body}}) in key metadata raised unhashable type: 'dict' and every request naming policies failed with a 500 on HTTP and on the Responses WebSocket. The merge now keeps existing entries and skips policy guardrails already selected by name. Also types the WebSocket test guardrail's hooks. * fix(responses): run streaming-scoped guardrails on socket turns, keep guardrail fields on blocks, merge guardrail lists by name * fix(responses): check frame guardrail picks against the team permission, honor picks beside a nested response, prefer a guardrail entry with settings * style(responses): format the frame guard and give the pass-through guardrail test an enterprise license * fix(proxy): keep pass-through guardrail settings, gate frame policy picks, skip moderation-only hooks on socket turns * fix(responses): keep string metadata guardrail values as client metadata on socket turns * fix(responses): ask no team permission for string metadata named like a guardrail pick on socket turns * fix(proxy): socket guardrail picks read the key's team metadata when the proxy has no database --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 43 +- litellm/proxy/auth/auth_checks.py | 81 +-- litellm/proxy/litellm_pre_call_utils.py | 67 +- .../passthrough_guardrails.py | 30 +- .../proxy/response_api_endpoints/endpoints.py | 294 ++++++++- litellm/proxy/utils.py | 6 +- litellm/responses/streaming_iterator.py | 36 +- .../integrations/test_custom_guardrail.py | 13 + ...st_auth_checks_object_access_and_lookup.py | 6 +- .../test_passthrough_guardrails.py | 19 + .../response_api_endpoints/test_endpoints.py | 615 +++++++++++++++++- .../unit/proxy/test_litellm_pre_call_utils.py | 82 +++ .../utils/helpers/test_guardrail_merge.py | 18 +- 13 files changed, 1134 insertions(+), 176 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a9fb5c29bad..736ebd58130 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,8 +3,9 @@ import copy import hashlib import os import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -208,6 +209,46 @@ def guardrail_request_data_with_streaming( } +_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) + + +def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: + """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)) + 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) -> tuple[str, ...]: + if isinstance(entry, str): + return (entry,) + if isinstance(entry, dict): + return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return () + + +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: if event_type in _REALTIME_STREAMING_HOOKS: return True diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 79a46ef355e..7d379a6f042 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -937,53 +937,58 @@ _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 failing loudly at the auth layer so operators see an explicit 403 instead of a confusing silent-ignore. """ + if requests_guardrail_modification(request_body): + require_guardrail_modification_permission(team_object) + + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: 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: - return - if not can_modify_guardrails(team_object): raise HTTPException( status_code=403, @@ -1485,7 +1490,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/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ae9a6d4b959..a9cffcc6566 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -48,7 +48,7 @@ from litellm.constants import ( VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) -from litellm.integrations.custom_guardrail import without_server_streaming_classification +from litellm.integrations.custom_guardrail import merged_guardrail_entries, without_server_streaming_classification from litellm.litellm_core_utils.core_helpers import is_codex_user_agent, normalize_drop_params from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -3948,30 +3948,14 @@ def add_guardrails_from_key_or_team_metadata( """ from litellm.proxy.utils import premium_user_check - # Initialize guardrails set (avoiding duplicates) - combined_guardrails: Final = set() - - # Add key-level guardrails first - if key_metadata and "guardrails" in key_metadata: - if isinstance(key_metadata["guardrails"], list) and len(key_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(key_metadata["guardrails"]) - - # Add team-level guardrails (set automatically handles duplicates) - if team_metadata and "guardrails" in team_metadata: - if isinstance(team_metadata["guardrails"], list) and len(team_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(team_metadata["guardrails"]) - - # Add project-level guardrails (set automatically handles duplicates) - if project_metadata and "guardrails" in project_metadata: - if isinstance(project_metadata["guardrails"], list) and len(project_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(project_metadata["guardrails"]) - - # Set combined guardrails in metadata as list - if combined_guardrails: - data[metadata_variable_name]["guardrails"] = list(combined_guardrails) + configured_guardrails: Final = tuple( + metadata["guardrails"] + for metadata in (key_metadata, team_metadata, project_metadata) + if metadata and isinstance(metadata.get("guardrails"), list) and len(metadata["guardrails"]) > 0 + ) + if configured_guardrails: + premium_user_check() + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries(*configured_guardrails) _add_guardrails_from_key_or_team_metadata: Final = add_guardrails_from_key_or_team_metadata @@ -4066,14 +4050,10 @@ def _add_guardrails_from_policies_in_metadata( if metadata_variable_name not in data: data[metadata_variable_name] = {} - existing_guardrails = data[metadata_variable_name].get("guardrails", []) - if not isinstance(existing_guardrails, list): - existing_guardrails = [] - - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = set(existing_guardrails) - combined.update(resolved_guardrails) - data[metadata_variable_name]["guardrails"] = list(combined) + existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries( + existing_guardrails if isinstance(existing_guardrails, list) else (), resolved_guardrails + ) # Store applied policies in metadata for tracking if "applied_policies" not in data[metadata_variable_name]: @@ -4308,10 +4288,11 @@ def _apply_resolved_guardrails_to_metadata( return existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) - existing_guardrails_list: Final = existing_guardrails if isinstance(existing_guardrails, list) else [] + existing_guardrails_list: Final = ( + _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () + ) - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = list(dict.fromkeys((*existing_guardrails_list, *resolved_guardrails))) + combined: Final = merged_guardrail_entries(existing_guardrails_list, resolved_guardrails) data[metadata_variable_name]["guardrails"] = combined verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) @@ -4328,7 +4309,7 @@ def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_nam async def add_guardrails_from_policy_engine( - data: dict, + data: dict[str, object], # mutable-ok: pops the request's policies and records the resolved guardrails in place metadata_variable_name: str, user_api_key_dict: UserAPIKeyAuth, ) -> None: @@ -4371,11 +4352,13 @@ 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 - context: Final = PolicyMatchContext( - 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"), - tags=all_tags, + 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"), + "tags": all_tags, + } ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index 1c5d270542f..bd8390d704e 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -10,6 +10,7 @@ Handles guardrail execution for passthrough endpoints with: from typing import Any, Final from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import merged_guardrail_entries from litellm.proxy._types import ( PassThroughGuardrailsConfig, PassThroughGuardrailSettings, @@ -205,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. @@ -225,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, @@ -242,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, @@ -258,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 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 ef59deb2e5e..089cb75a795 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,11 +1,13 @@ import asyncio import contextlib +import copy import json import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum from functools import partial +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -35,7 +37,7 @@ from litellm.anthropic_interface.exceptions.exception_mapping_utils import Anthr from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.exceptions import EndpointNotSupportedByModelError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import ModifyResponseException, guardrail_request_data_with_streaming from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -47,6 +49,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, + 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 ( UserAPIKeyAuth, @@ -57,6 +64,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + proxy_exception_from_route_error, resolve_counted_tool_fields, resolve_litellm_call_id, ) @@ -79,6 +87,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + add_guardrails_from_policy_engine, is_chat_completions_body, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -89,7 +98,12 @@ from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_ from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.retrieve_stream_query import RETRIEVE_INCLUDE_QUERY_KEYS, RETRIEVE_STREAM_QUERY_KEYS -from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task +from litellm.responses.streaming_iterator import ( + GuardedResponsesWsFrame, + ResponsesWsFrameGuard, + ResponsesWsFrameRejection, + spawn_logging_task, +) from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -2149,11 +2163,15 @@ _PII_TOKENS: Final = TypeAdapter(dict[str, str]) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" _GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") -# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it _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") +_UNFORWARDED_FRAME_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS, "litellm_metadata") +_OBJECTS: Final = TypeAdapter(list[object]) +_TURN_COPIED_CONTAINERS: Final = (list, dict) -class _PreCallGuardrailWalk(Protocol): +class _GuardrailWalk(Protocol): async def pre_call_hook( self, *, @@ -2163,15 +2181,49 @@ class _PreCallGuardrailWalk(Protocol): guardrails_only: bool, ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict + async def during_call_hook( + self, + *, + 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: ... + + +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 + team_from_key: Final = LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata) + if prisma_client is None: + return team_from_key + 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 team_from_key + def _responses_ws_frame_guard( admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - proxy_logging_obj: _PreCallGuardrailWalk, + 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 guardrail walk the HTTP ``/v1/responses`` path runs on its body.""" + ``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, 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 = { @@ -2186,32 +2238,57 @@ 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], # mutable-ok: billed + ) -> ResponsesWsFrameRejection: + 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") - if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): - raise ValueError( - "response.create carries input, instructions or tools both at the top level and under response" - ) body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj - frame_data: Final = { - **{ - key: _without_untrusted_metadata(key, value) - for key, value in body.items() - if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS + 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) + for key, value in body.items() + if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + metadata_key: {**bucket}, }, - "model": model, - "litellm_call_id": str(uuid.uuid4()), - # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied - 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", guardrails_only=True + ) except Exception as e: - await _bill_rejected_frame(rejected_frame_auth, e, frame_data) - raise + 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" @@ -2219,25 +2296,186 @@ 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: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} rewritten: Final = ( - {**frame_obj, "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 _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 or policies needs the team's permission to modify guardrails""" + if any( + "policies" in source or requests_guardrail_modification(_without_client_metadata(source)) for source in sources + ): + require_guardrail_modification_permission(await team_object(user_api_key_dict)) + + +def _without_client_metadata(source: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = source.get("metadata") + return {**source, "metadata": _metadata_picks("metadata", metadata)} if isinstance(metadata, dict) else source + + +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): + guardrails: ReadOnly[Sequence[str | Mapping[str, Mapping[str, object]]]] + guardrail_config: ReadOnly[Mapping[str, object]] + policies: ReadOnly[Sequence[str]] + + +_FRAME_GUARDRAIL_SELECTION: Final = TypeAdapter(_FrameGuardrailSelection) + + +async def _metadata_with_frame_guardrails( + bucket: 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. 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, 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)) + selected_config: Final = dict( + chain.from_iterable( + selection["guardrail_config"].items() for selection in selections if "guardrail_config" in selection + ) + ) + existing_guardrails: Final = metadata.get("guardrails") + existing_config: Final = metadata.get("guardrail_config") + return { + **metadata, + **( + { + "guardrails": [ + *(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()), + *selected_guardrails, + ] + } + if selected_guardrails + else {} + ), + **( + { + "guardrail_config": { + **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), + **selected_config, + } + } + if selected_config + else {} + ), + } + + +def _frame_guardrail_selections(body: Mapping[str, object]) -> tuple[_FrameGuardrailSelection, ...]: + """Like HTTP ``/v1/responses``, a frame selects guardrails at its top level and in its ``metadata`` and + ``litellm_metadata``, the top level last so its ``guardrail_config`` wins""" + return ( + *( + _frame_guardrail_selection((key,), _metadata_picks(key, value), _METADATA_GUARDRAIL_SELECTORS) + for key, value in body.items() + if key in _CLIENT_METADATA_KEYS and isinstance(value, dict) + ), + _frame_guardrail_selection((), body, _CLIENT_GUARDRAIL_SELECTORS), + ) + + +def _metadata_picks(key: str, metadata: object) -> Mapping[str, object]: + """OpenAI ``metadata`` values are strings, so a string there is the client's own metadata and never a pick""" + validated: Final = _WS_FRAME.validate_python(metadata) + if key != "metadata": + return validated + return { + name: item + for name, item in validated.items() + if name not in _METADATA_GUARDRAIL_SELECTORS or not isinstance(item, str) + } + + +def _frame_guardrail_selection( + location: tuple[str, ...], source: Mapping[str, object], selectors: tuple[str, ...] +) -> _FrameGuardrailSelection: + try: + return _FRAME_GUARDRAIL_SELECTION.validate_python({key: source[key] for key in selectors if key in source}) + except ValidationError as e: + raise ValueError(f"response.create {first_validation_error(e, prefix=location)}") from e + + +def _without_guardrail_selection(metadata: object) -> object: + if not isinstance(metadata, dict): + return metadata + return { + name: item + for name, item in _WS_FRAME.validate_python(metadata).items() + if name not in _METADATA_GUARDRAIL_SELECTORS or isinstance(item, str) + } + + +async def _metadata_with_frame_policies( + bucket: Mapping[str, object], + policies: Sequence[str], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + selection: Final = { + "model": model, + "policies": list(policies), + metadata_key: {key: _turn_copy(value) for key, value in bucket.items()}, + } + await add_guardrails_from_policy_engine( + data=selection, metadata_variable_name=metadata_key, user_api_key_dict=user_api_key_dict + ) + return _WS_FRAME.validate_python(selection[metadata_key]) + + +def _turn_copy(value: object) -> object: + """Shallow, since the bucket carries proxy-owned objects (the OTEL parent span) that cannot be copied, and the + policy engine writes its lists in place""" + return copy.copy(value) if type(value) in _TURN_COPIED_CONTAINERS else value + + def _without_untrusted_metadata(key: str, value: object) -> object: """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cc72dd03a40..0b0698faa9a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -147,6 +147,7 @@ from litellm.exceptions import ( from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, + merged_guardrail_entries, ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -3506,6 +3507,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: @@ -3517,6 +3519,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 ): @@ -9100,8 +9103,7 @@ def _merge_guardrails_with_existing(data: dict, model_level_guardrails: object) if not isinstance(model_level_guardrails, list): model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] - # Combine existing and model-level guardrails - metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) + metadata["guardrails"] = merged_guardrail_entries(existing_guardrails, model_level_guardrails) return modified_data diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d42f1af380e..a8e55a5e94d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2277,6 +2277,29 @@ class GuardedResponsesWsFrame: ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] +class ResponsesWsFrameRejection(ValueError): + def __init__(self, message: str, provider_specific_fields: Mapping[str, object] | None = None) -> None: + super().__init__(message) + self.provider_specific_fields: Final = provider_specific_fields + + +def _frame_rejection_event(rejection: Exception) -> str: + provider_fields: Final = ( + rejection.provider_specific_fields if isinstance(rejection, ResponsesWsFrameRejection) else None + ) + return json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": str(rejection), + **({"provider_specific_fields": dict(provider_fields)} if provider_fields else {}), + }, + }, + default=str, + ) + + async def _reject_guarded_turn( websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception ) -> None: @@ -2562,7 +2585,8 @@ class ResponsesWebSocketStreaming: return await self.frame_guard(message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error_event("invalid_request_error", str(e)) + with contextlib.suppress(Exception): + await self.websocket.send_text(_frame_rejection_event(e)) return None async def _forward_guarded_frame( @@ -2995,14 +3019,6 @@ class ResponsesWebSocketStreaming: except Exception: # noqa: BLE001, S110 # client may already be gone pass - async def _send_error_event(self, error_type: str, message: str) -> None: - try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass - async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: @@ -3416,7 +3432,7 @@ class ManagedResponsesWebSocketHandler: return await self.frame_guard(raw_message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error(str(e), error_type="invalid_request_error") + await self._send_frame(_frame_rejection_event(e)) return None @staticmethod diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 0483888533e..c32c1b71f75 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3800,3 +3800,16 @@ def test_default_on_guardrail() -> None: ) is True ) + + +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": {}, "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 6b59f6d40dc..dc7bd8cb312 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 @@ -5884,7 +5884,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 @@ -5892,11 +5892,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..ef19482b30a 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,25 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: assert set(result) == {"pii-detection", "content-moderation"} +class TestPassthroughGuardrailHandlerCollectGuardrails: + 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) + key = UserAPIKeyAuth( + api_key="hashed", metadata={"guardrails": [{"pii": {"extra_body": {"threshold": 1}}}, "toxic"]} + ) + + 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 + } + + 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 1e4dded801e..07f85cb9628 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -7,6 +7,7 @@ import contextlib import dataclasses import datetime import json +import re import unittest from collections.abc import Callable, Mapping from typing import Any, Final, Literal, cast @@ -22,12 +23,17 @@ from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.caching.caching import DualCache from litellm.constants import EMPTY_MAPPING -from litellm.proxy._types import ProxyException +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.proxy_server import app +from litellm.proxy.utils import ProxyLogging from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors +from litellm.types.utils import CallTypesLiteral from tests._master_key import MASTER_KEY @@ -2213,6 +2219,59 @@ class TestResponsesWSFirstFrameValidation: ws.close.assert_awaited_once_with(code=1011, reason="Internal server error") +class _FrameGuardrail(CustomGuardrail): + def __init__( + self, + name: str, + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + rewrites: bool = False, + blocks: bool = False, + ): + super().__init__(guardrail_name=name, event_hook=event_hook, default_on=False) + self.inputs: list[object] = [] + self.configs: list[object] = [] + self.rewrites = rewrites + self.blocks = blocks + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + return self._scan(data) + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + data["input"] = self._scan(data)["input"] + + def _scan(self, data: dict[str, object]) -> dict[str, object]: + self.inputs.append(data["input"]) + self.configs.append(cast(dict[str, object], data["litellm_metadata"]).get("guardrail_config")) + if self.blocks: + raise HTTPException( + status_code=400, + detail={"error": f"Content blocked by {self.guardrail_name}", "guardrail_name": self.guardrail_name}, + ) + return {**data, "input": f"{data['input']} [{self.guardrail_name}]"} if self.rewrites else data + + +@contextlib.contextmanager +def _guardrails_on(*guardrails: CustomGuardrail): + with patch.object(litellm, "callbacks", list(guardrails)): + ProxyLogging._callback_capabilities_cache.clear() + try: + yield ProxyLogging(user_api_key_cache=DualCache()) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + +def _key() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="hashed") + + class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio async def test_endpoint_enforces_auth_after_model_from_first_frame(self): @@ -2344,7 +2403,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}] return {**data, "input": data["input"] + " [guarded]"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) admitted = { "model": "gpt-4o-mini", @@ -2375,7 +2434,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock() spend_logger = MagicMock(spec=ProxyDBLogger) @@ -2413,7 +2472,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2432,10 +2491,10 @@ class TestResponsesWSFirstFrameModelAuth: async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + 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}, @@ -2454,17 +2513,18 @@ class TestResponsesWSFirstFrameModelAuth: [ {"type": "response.create", "input": "forbidden", "guardrails": []}, {"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}}, + {"type": "response.create", "input": "forbidden", "litellm_metadata": {"guardrails": []}}, {"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}}, ], ) async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + 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": {"guardrails": ["key-required"]}}, - MagicMock(), + _key(), proxy_logging_obj, MagicMock(), ) @@ -2476,13 +2536,258 @@ class TestResponsesWSFirstFrameModelAuth: assert "guardrails" not in (scanned.get("metadata") or {}) assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "selected", + [ + {"guardrails": ["frame-pick"]}, + {"guardrails": [{"frame-pick": {"extra_body": {"threshold": 0.8}}}]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + {"response": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "with-extra-body", "metadata", "litellm-metadata", "nested"], + ) + async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, selected: dict[str, object]): + 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") + admitted = {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}} + nested = "response" in selected + body = {"input": "hi", "store": False, **(selected["response"] if nested else selected)} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard(admitted, _key(), proxy_logging_obj, MagicMock()) + selected_turn = await guard(json.dumps(frame)) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "again"})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in selected else {} + forwarded = {"input": "hi [frame-pick]", "store": False, **kept_metadata} + assert json.loads(selected_turn.frame) == ( + {"type": "response.create", "response": forwarded} if nested else {"type": "response.create", **forwarded} + ) + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "again"} + assert key_required.inputs == ["hi", "again"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + assert admitted["litellm_metadata"] == {"guardrails": ["key-required"]} + + @pytest.mark.asyncio + async def test_frame_guard_forwards_string_metadata_named_like_a_pick_without_running_it(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + metadata = {"guardrails": "frame-pick", "guardrail_config": "strict", "trace": "t1"} + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick) 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, + ) + guarded = await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata})) + + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": metadata} + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == [] + + @pytest.mark.asyncio + async def test_frame_guard_hands_a_selected_guardrail_its_config_with_the_top_level_winning(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + frame = { + "type": "response.create", + "input": "hi", + "metadata": {"guardrail_config": {"frame-pick": {"threshold": 1}, "metadata-only": {"threshold": 3}}}, + "guardrails": ["frame-pick"], + "guardrail_config": {"frame-pick": {"threshold": 2}}, + } + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrail_config": {"key-required": {"threshold": 0}}}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps(frame)) + + assert frame_pick.configs == [ + {"key-required": {"threshold": 0}, "frame-pick": {"threshold": 2}, "metadata-only": {"threshold": 3}} + ] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": {}} + + @pytest.mark.asyncio + async def test_frame_guard_runs_streaming_scoped_guardrails_and_skips_non_streaming_ones_on_a_socket_turn(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + streaming_only = _FrameGuardrail("streaming-only") + streaming_only.apply_stream_scope("streaming") + streaming_during_call = _FrameGuardrail("streaming-during-call", event_hook=GuardrailEventHooks.during_call) + streaming_during_call.apply_stream_scope("streaming") + non_streaming_only = _FrameGuardrail("non-streaming-only") + non_streaming_only.apply_stream_scope("non_streaming") + frame = { + "type": "response.create", + "input": "hi", + "guardrails": ["streaming-only", "streaming-during-call", "non-streaming-only"], + } + + with _guardrails_on(streaming_only, streaming_during_call, non_streaming_only) 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(frame)) + + assert streaming_only.inputs == ["hi"] + assert streaming_during_call.inputs == ["hi"] + assert non_streaming_only.inputs == [] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + + @pytest.mark.asyncio + async def test_frame_guard_blocks_a_turn_with_the_message_a_selected_during_call_guardrail_raises(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, blocks=True) + + with _guardrails_on(during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + unselected_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Content blocked by during-pick$"): + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["during-pick"]})) + + 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"), + [ + ( + {"guardrails": "frame-pick"}, + "response.create guardrails: 'str' instances are not allowed as a Sequence value", + ), + ({"guardrails": [123]}, "response.create guardrails.0.str: Input should be a valid string"), + ({"guardrail_config": []}, "response.create guardrail_config: Input should be a valid dictionary"), + ( + {"policies": "frame-policy"}, + "response.create policies: 'str' instances are not allowed as a Sequence value", + ), + ({"policies": [1]}, "response.create policies.0: Input should be a valid string"), + ( + {"litellm_metadata": {"guardrails": "frame-pick"}}, + "response.create litellm_metadata.guardrails: 'str' instances are not allowed as a Sequence value", + ), + ( + {"litellm_metadata": {"guardrail_config": []}}, + "response.create litellm_metadata.guardrail_config: Input should be a valid dictionary", + ), + ], + ) + async def test_frame_guard_refuses_and_bills_a_frame_whose_guardrail_selection_is_malformed( + self, selection: dict[str, object], message: str + ): + 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, 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(), + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + 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): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + key_required = _FrameGuardrail("key-required") + policy_pick = _FrameGuardrail("policy-pick") + policy_registry = get_policy_registry() + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + try: + policy_registry.add_policy("frame-policy", Policy(guardrails=PolicyGuardrails(add=["policy-pick"]))) + with _guardrails_on(key_required, policy_pick) 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", "input": "hi", "policies": ["frame-policy"]}) + ) + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + + assert key_required.inputs == ["hi"] + assert policy_pick.inputs == ["hi"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + @pytest.mark.asyncio async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): import threading from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) span = threading.RLock() guard = _responses_ws_frame_guard( @@ -2510,7 +2815,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] return {**data, "model": "onprem-llama"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2543,7 +2848,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2571,7 +2876,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {key: value for key, value in data.items() if key != "instructions"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2596,7 +2901,7 @@ class TestResponsesWSFirstFrameModelAuth: data["metadata"]["pii_tokens"] = {"": "alice@example.com"} return {**data, "input": "email please"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2617,19 +2922,279 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) + async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field: str): + 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() + 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", [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: dict[str, object]): + 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", "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", + "string-global-toggle", + "nested", + "beside-nested", + "policies", + "nested-policies", + ], + ) + async def test_frame_guard_refuses_guardrail_and_policy_picks_from_a_team_without_permission(self, frame: dict[str, object]): + 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 + @pytest.mark.parametrize( + ("modify_guardrails", "picked_inputs"), + [(True, ["hi"]), (False, [])], + ids=["open", "locked"], + ) + @pytest.mark.parametrize( + "team_lookup", + [ + lambda: patch( + "litellm.proxy.response_api_endpoints.endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})), + ), + lambda: patch("litellm.proxy.proxy_server.prisma_client", None), + ], + ids=["team-row-gone", "no-database"], + ) + async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_cannot_be_loaded( + self, + modify_guardrails: bool, + picked_inputs: list[str], + team_lookup: Callable[[], contextlib.AbstractContextManager[object]], + ): + 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, team_lookup(): + 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 + + 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): + 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"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() ) - with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): - await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "include_guardrail_response": True, + "guardrail_config": {"frame-pick": {"threshold": 1}}, + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + ) + ) - proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert json.loads(guarded.frame) == { + "type": "response.create", + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + + @pytest.mark.asyncio + async def test_a_blocked_socket_turn_answers_with_the_guardrail_fields_http_returns(self): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + frame_pick = _FrameGuardrail("frame-pick", blocks=True) + websocket = MagicMock(send_text=AsyncMock()) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + handler = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-haiku-5-5", + logging_obj=Logging( + model="anthropic/claude-haiku-5-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=_responses_ws_frame_guard( + {"model": "claude-managed", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ), + ) + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + ) + + assert json.loads(websocket.send_text.await_args.args[0]) == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "Content blocked by frame-pick", + "provider_specific_fields": { + "error": "Content blocked by frame-pick", + "guardrail_name": "frame-pick", + "guardrail_mode": "pre_call", + }, + }, + } @pytest.mark.asyncio async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): @@ -2637,7 +3202,7 @@ class TestResponsesWSFirstFrameModelAuth: 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() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=ValueError("blocked")) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down")) @@ -2679,7 +3244,7 @@ class TestResponsesWSFirstFrameModelAuth: async def fake_llm_call(): return failure if provider_rejected else None - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() @@ -2768,7 +3333,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() row_written = asyncio.Event() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=lambda **kwargs: row_written.set()) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) @@ -2858,7 +3423,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -2948,7 +3513,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -3022,7 +3587,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: { **data, @@ -3083,7 +3648,7 @@ class TestResponsesWSFirstFrameModelAuth: rejection = litellm.RateLimitError( message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai" ) - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c5136ff6e4a..3194124ad1c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5360,6 +5360,88 @@ async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_ assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and_skips_ones_already_selected(): + 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 + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"guardrails": [configured, "audit"]}, + } + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="baseline", scope="*")] + attachment_registry._initialized = True + + try: + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + +def test_add_guardrails_from_auth_metadata_lists_each_key_and_team_guardrail_once_keeping_the_key_s_settings(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {}} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + metadata={"guardrails": [configured, "audit"]}, + team_metadata={"guardrails": ["audit", "pii_blocker", "toxic"]}, + ) + + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name="metadata" + ) + + assert data["metadata"]["guardrails"] == [configured, "audit", "toxic"] + + +def test_add_guardrails_from_auth_metadata_keeps_a_selected_guardrail_s_settings_when_a_key_policy_adds_it(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {"guardrails": [configured, "audit"]}} + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + + try: + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", metadata={"policies": ["baseline"]}), + data=data, + metadata_variable_name="metadata", + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_applies_inherited_parent_guardrail_when_child_condition_misses(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry diff --git a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py index e2c1a7c2c09..dd933c508cb 100644 --- a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py +++ b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py @@ -362,7 +362,17 @@ def test_merge_guardrails_with_existing_creates_metadata_when_missing(): } -def test_merge_guardrails_with_existing_raises_on_unhashable_guardrail(): - data = {"metadata": {"guardrails": [{"unhashable": True}]}} - with pytest.raises(TypeError): - _merge_guardrails_with_existing(data, ["g1"]) +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_over_the_model_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + 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"]