From e8b8bf73c4b94c22bdfd5ba4d6d47c1c5d59e65b Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:55:31 -0700 Subject: [PATCH 01/10] fix(responses): run the guardrails a Responses WebSocket frame selects --- litellm/proxy/litellm_pre_call_utils.py | 5 +- .../proxy/response_api_endpoints/endpoints.py | 89 ++++++++++++-- .../response_api_endpoints/test_endpoints.py | 113 +++++++++++++++++- 3 files changed, 196 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1e7cdc44c70..9c2d8feb364 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4220,7 +4220,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: @@ -4263,10 +4263,11 @@ async def add_guardrails_from_policy_engine( all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias + _model: Final = data.get("model") context: Final = PolicyMatchContext( team_alias=_team_alias if isinstance(_team_alias, str) else None, key_alias=_key_alias if isinstance(_key_alias, str) else None, - model=data.get("model"), + model=_model if isinstance(_model, str) else None, tags=all_tags, ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 97bbf893e3e..0e78f7afbd9 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import copy import json import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence @@ -75,6 +76,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, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, @@ -2017,8 +2019,12 @@ _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") +_PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") +_NESTED_BODY_ONLY_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS) +_GUARDRAIL_NAMES: Final = TypeAdapter(list[str]) +_OBJECTS: Final = TypeAdapter(list[object]) +_TURN_COPIED_CONTAINERS: Final = (list, dict) class _PreCallGuardrailWalk(Protocol): @@ -2057,21 +2063,21 @@ def _responses_ws_frame_guard( 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): + if isinstance(nested, dict) and any(key in frame_obj for key in _NESTED_BODY_ONLY_FIELDS): raise ValueError( - "response.create carries input, instructions or tools both at the top level and under response" + "response.create carries input, instructions, tools or a guardrail field 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 + if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS }, "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}, + metadata_key: await _metadata_with_frame_guardrails(bucket, body, model, metadata_key, user_api_key_dict), } try: guarded_data: Final = await proxy_logging_obj.pre_call_hook( @@ -2089,7 +2095,11 @@ def _responses_ws_frame_guard( 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} + kept: Final = { + key: value + for key, value in body.items() + if key not in _GUARDED_FRAME_FIELDS and key not in _PROXY_ONLY_FRAME_FIELDS + } rewritten: Final = ( {**frame_obj, "response": {**kept, **guarded_fields}} if isinstance(nested, dict) @@ -2106,6 +2116,71 @@ def _responses_ws_frame_guard( return guard +async def _metadata_with_frame_guardrails( + bucket: Mapping[str, object], + body: 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""" + metadata: Final = ( + await _metadata_with_frame_policies(bucket, body["policies"], model, metadata_key, user_api_key_dict) + if "policies" in body + else bucket + ) + 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 ()), + *_GUARDRAIL_NAMES.validate_python(body["guardrails"]), + ] + } + if "guardrails" in body + else {} + ), + **( + { + "guardrail_config": { + **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), + **_WS_FRAME.validate_python(body["guardrail_config"]), + } + } + if "guardrail_config" in body + else {} + ), + } + + +async def _metadata_with_frame_policies( + bucket: Mapping[str, object], + policies: object, + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + selection: Final = { + "model": model, + "policies": 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/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 985aee061bc..1c0a04bdf90 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2226,6 +2226,112 @@ class TestResponsesWSFirstFrameModelAuth: assert "guardrails" not in (scanned.get("metadata") or {}) assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, nested: bool): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + scans: list[dict] = [] + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + scans.append(data) + return data + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + admitted = { + "model": "gpt-4o-mini", + "litellm_metadata": { + "guardrails": ["key-required"], + "guardrail_config": {"key-required": {"threshold": 1}}, + }, + } + guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj, MagicMock()) + selection = { + "guardrails": ["frame-pick"], + "guardrail_config": {"frame-pick": {"threshold": 2}}, + "include_guardrail_response": True, + } + body = {"input": "hi", "store": False, **selection} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + selected = await guard(json.dumps(frame)) + plain = await guard(json.dumps({"type": "response.create", "input": "again"})) + + assert scans[0]["litellm_metadata"]["guardrails"] == ["key-required", "frame-pick"] + assert scans[0]["litellm_metadata"]["guardrail_config"] == { + "key-required": {"threshold": 1}, + "frame-pick": {"threshold": 2}, + } + assert not set(selection) & set(scans[0]) + unselected = {"input": "hi", "store": False} + assert json.loads(selected.frame) == ( + {"type": "response.create", "response": unselected} if nested else {"type": "response.create", **unselected} + ) + assert scans[1]["litellm_metadata"]["guardrails"] == ["key-required"] + assert scans[1]["litellm_metadata"]["guardrail_config"] == {"key-required": {"threshold": 1}} + assert json.loads(plain.frame) == {"type": "response.create", "input": "again"} + assert admitted["litellm_metadata"] == { + "guardrails": ["key-required"], + "guardrail_config": {"key-required": {"threshold": 1}}, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "selection", [{"guardrails": "frame-pick"}, {"guardrails": [{"name": "frame-pick"}]}, {"guardrail_config": []}] + ) + async def test_frame_guard_refuses_a_frame_whose_guardrail_selection_is_malformed(self, selection): + from pydantic import ValidationError + + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + 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(), + proxy_logging_obj, + MagicMock(), + ) + + with pytest.raises(ValidationError): + await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self): + from litellm.proxy._types import UserAPIKeyAuth + 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 + + proxy_logging_obj = MagicMock() + 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"]}}, + UserAPIKeyAuth(api_key="hashed"), + proxy_logging_obj, + MagicMock(), + ) + 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"]))) + + 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 + + scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"] + assert scanned["litellm_metadata"]["guardrails"] == ["key-required", "policy-pick"] + assert "policies" not in scanned + 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 @@ -2367,7 +2473,10 @@ 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( + "top_level", [{"input": "forbidden"}, {"guardrails": ["frame-pick"]}, {"include_guardrail_response": True}] + ) + async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self, top_level): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard proxy_logging_obj = MagicMock() @@ -2377,7 +2486,7 @@ class TestResponsesWSFirstFrameModelAuth: ) with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): - await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + await guard(json.dumps({"type": "response.create", **top_level, "response": {}})) proxy_logging_obj.pre_call_hook.assert_not_awaited() From d62dd7ac9afee238845fb71ca71249631731caff Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:27:38 -0700 Subject: [PATCH 02/10] 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 --- .../proxy/response_api_endpoints/endpoints.py | 27 ++++++++-- .../response_api_endpoints/test_endpoints.py | 51 ++++++++++++++----- 2 files changed, 61 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 0e78f7afbd9..e295d953790 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2022,7 +2022,6 @@ _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") _PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") _NESTED_BODY_ONLY_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS) -_GUARDRAIL_NAMES: Final = TypeAdapter(list[str]) _OBJECTS: Final = TypeAdapter(list[object]) _TURN_COPIED_CONTAINERS: Final = (list, dict) @@ -2116,6 +2115,14 @@ def _responses_ws_frame_guard( return guard +class _FrameGuardrailSelection(TypedDict, total=False): + guardrails: ReadOnly[Sequence[str]] + guardrail_config: ReadOnly[Mapping[str, object]] + + +_FRAME_GUARDRAIL_SELECTION: Final = TypeAdapter(_FrameGuardrailSelection) + + async def _metadata_with_frame_guardrails( bucket: Mapping[str, object], body: Mapping[str, object], @@ -2125,6 +2132,7 @@ async def _metadata_with_frame_guardrails( ) -> 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""" + selection: Final = _frame_guardrail_selection(body) metadata: Final = ( await _metadata_with_frame_policies(bucket, body["policies"], model, metadata_key, user_api_key_dict) if "policies" in body @@ -2138,25 +2146,34 @@ async def _metadata_with_frame_guardrails( { "guardrails": [ *(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()), - *_GUARDRAIL_NAMES.validate_python(body["guardrails"]), + *selection["guardrails"], ] } - if "guardrails" in body + if "guardrails" in selection else {} ), **( { "guardrail_config": { **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), - **_WS_FRAME.validate_python(body["guardrail_config"]), + **selection["guardrail_config"], } } - if "guardrail_config" in body + if "guardrail_config" in selection else {} ), } +def _frame_guardrail_selection(body: Mapping[str, object]) -> _FrameGuardrailSelection: + try: + return _FRAME_GUARDRAIL_SELECTION.validate_python( + {key: body[key] for key in ("guardrails", "guardrail_config") if key in body} + ) + except ValidationError as e: + raise ValueError(f"response.create {first_validation_error(e)}") from e + + async def _metadata_with_frame_policies( bucket: Mapping[str, object], policies: object, diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 1c0a04bdf90..701731eaf65 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -6,6 +6,7 @@ import asyncio import contextlib import dataclasses import json +import re import unittest from collections.abc import Mapping from typing import Any, Final, Literal, cast @@ -1801,11 +1802,19 @@ class TestResponsesWSStreamingFirstMessage: @pytest.mark.parametrize( ("failure", "expected_status", "expected_type"), [ - (ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), 429, "rate_limit_exceeded"), + ( + ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), + 429, + "rate_limit_exceeded", + ), (litellm.BadRequestError(message="bad input", model="m", llm_provider="openai"), 400, "invalid_request_error"), (RuntimeError("boom"), 500, "server_error"), (ValueError(f"{RouterErrors.no_deployments_with_tag_routing.value}=['free']"), 401, "invalid_request_error"), - (ProxyException(message="Authentication Error", type="auth_error", param=None, code=401), 401, "invalid_request_error"), + ( + ProxyException(message="Authentication Error", type="auth_error", param=None, code=401), + 401, + "invalid_request_error", + ), (HTTPException(status_code=403, detail="key not allowed"), 403, "invalid_request_error"), ], ids=["crossed-budget", "bad-request", "unexpected", "tag-routing", "proxy-auth", "http-exception"], @@ -2278,11 +2287,17 @@ class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio @pytest.mark.parametrize( - "selection", [{"guardrails": "frame-pick"}, {"guardrails": [{"name": "frame-pick"}]}, {"guardrail_config": []}] + ("selection", "message"), + [ + ( + {"guardrails": "frame-pick"}, + "response.create guardrails: 'str' instances are not allowed as a Sequence value", + ), + ({"guardrails": [{"name": "frame-pick"}]}, "response.create guardrails.0: Input should be a valid string"), + ({"guardrail_config": []}, "response.create guardrail_config: Input should be a valid dictionary"), + ], ) - async def test_frame_guard_refuses_a_frame_whose_guardrail_selection_is_malformed(self, selection): - from pydantic import ValidationError - + async def test_frame_guard_refuses_a_frame_whose_guardrail_selection_is_malformed(self, selection, message): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard proxy_logging_obj = MagicMock() @@ -2294,7 +2309,7 @@ class TestResponsesWSFirstFrameModelAuth: MagicMock(), ) - with pytest.raises(ValidationError): + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) proxy_logging_obj.pre_call_hook.assert_not_awaited() @@ -2685,6 +2700,7 @@ class TestResponsesWSFirstFrameModelAuth: ws.receive_text = AsyncMock( return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []}) ) + async def close_over_the_network(**kwargs): await asyncio.sleep(0) @@ -2704,7 +2720,9 @@ class TestResponsesWSFirstFrameModelAuth: def route_request(data, **kwargs): async def relay(): - guarded = await data["responses_ws_frame_guard"](json.dumps({"type": "response.create", "input": "two"})) + guarded = await data["responses_ws_frame_guard"]( + json.dumps({"type": "response.create", "input": "two"}) + ) guarded.bill_refusal(refusal) if relay_outcome == "still_relaying": await asyncio.Event().wait() @@ -2722,6 +2740,7 @@ class TestResponsesWSFirstFrameModelAuth: side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=slow_failure_hook) + async def release_over_the_network(user_api_key_dict): await asyncio.sleep(0.01) @@ -2764,7 +2783,9 @@ class TestResponsesWSFirstFrameModelAuth: row_may_finish.set() await asyncio.wait_for(row_written.wait(), timeout=1) - assert proxy_logging_obj.post_call_failure_hook.await_args.kwargs["original_exception"] is (crash if crashed else refusal) + assert proxy_logging_obj.post_call_failure_hook.await_args.kwargs["original_exception"] is ( + crash if crashed else refusal + ) @pytest.mark.asyncio @pytest.mark.parametrize("ending", ["client_closed", "crash"]) @@ -2799,7 +2820,9 @@ class TestResponsesWSFirstFrameModelAuth: def route_request(data, **kwargs): async def relay(): - guarded = await data["responses_ws_frame_guard"](json.dumps({"type": "response.create", "input": "two"})) + guarded = await data["responses_ws_frame_guard"]( + json.dumps({"type": "response.create", "input": "two"}) + ) guarded.bill_refusal(refusal) if ending == "crash": raise crash @@ -2874,7 +2897,9 @@ class TestResponsesWSFirstFrameModelAuth: def route_request(data, **kwargs): async def relay(): - guarded = await data["responses_ws_frame_guard"](json.dumps({"type": "response.create", "input": "two"})) + guarded = await data["responses_ws_frame_guard"]( + json.dumps({"type": "response.create", "input": "two"}) + ) guarded.bill_refusal(refusal) data["litellm_metadata"]["standard_logging_guardrail_information"] = "written by a broken callback" raise RuntimeError("relay crashed") @@ -4576,7 +4601,9 @@ class TestResponsesInputTokens: assert response.status_code == 200, response.text sent = counter.await_args.kwargs["responses_request"] - assert {name: sent[name] for name in ("tools", "tool_choice", "parallel_tool_calls") if name in sent} == expected + assert { + name: sent[name] for name in ("tools", "tool_choice", "parallel_tool_calls") if name in sent + } == expected def test_previous_response_id_reaches_the_counter_authorized_next_to_the_id_the_client_held(self) -> None: from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD From b6d870d7a4faa6e390d33ba88123bde226f7e925 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:20:04 -0700 Subject: [PATCH 03/10] 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 --- .../proxy/response_api_endpoints/endpoints.py | 99 ++++-- .../response_api_endpoints/test_endpoints.py | 311 +++++++++++------- 2 files changed, 270 insertions(+), 140 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index e295d953790..fb15e38a3ff 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -7,6 +7,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, 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, @@ -56,6 +57,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + proxy_exception_from_route_error, resolve_counted_tool_fields, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id @@ -2020,13 +2022,15 @@ _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" _GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") _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, *_PROXY_ONLY_FRAME_FIELDS) +_UNFORWARDED_FRAME_FIELDS: Final = (*_NESTED_BODY_ONLY_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, *, @@ -2036,15 +2040,24 @@ 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"], + ) -> object: ... + 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], ) -> 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.""" metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) admitted_bucket: Final = admitted.get(metadata_key) bucket: Final = { @@ -2062,10 +2075,14 @@ def _responses_ws_frame_guard( 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 _NESTED_BODY_ONLY_FIELDS): + stray: Final = ( + next((key for key in _NESTED_BODY_ONLY_FIELDS if key in frame_obj), None) + if isinstance(nested, dict) + else None + ) + if stray is not None: raise ValueError( - "response.create carries input, instructions, tools or a guardrail field both at the top level " - "and under response" + 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 = { @@ -2082,9 +2099,12 @@ def _responses_ws_frame_guard( 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 ) + 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) - raise + raise ValueError(proxy_exception_from_route_error(e).message) 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" @@ -2095,9 +2115,9 @@ def _responses_ws_frame_guard( # 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 + key: _without_guardrail_selection(value) if key == "metadata" else value for key, value in body.items() - if key not in _GUARDED_FRAME_FIELDS and key not in _PROXY_ONLY_FRAME_FIELDS + if key not in _UNFORWARDED_FRAME_FIELDS } rewritten: Final = ( {**frame_obj, "response": {**kept, **guarded_fields}} @@ -2116,8 +2136,9 @@ def _responses_ws_frame_guard( class _FrameGuardrailSelection(TypedDict, total=False): - guardrails: ReadOnly[Sequence[str]] + 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) @@ -2132,12 +2153,19 @@ async def _metadata_with_frame_guardrails( ) -> 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""" - selection: Final = _frame_guardrail_selection(body) + selections: Final = _frame_guardrail_selections(body) + policies: Final = selections[-1].get("policies") metadata: Final = ( - await _metadata_with_frame_policies(bucket, body["policies"], model, metadata_key, user_api_key_dict) - if "policies" in body + await _metadata_with_frame_policies(bucket, policies, model, metadata_key, user_api_key_dict) + if policies is not None 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 { @@ -2146,44 +2174,67 @@ async def _metadata_with_frame_guardrails( { "guardrails": [ *(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()), - *selection["guardrails"], + *selected_guardrails, ] } - if "guardrails" in selection + if selected_guardrails else {} ), **( { "guardrail_config": { **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), - **selection["guardrail_config"], + **selected_config, } } - if "guardrail_config" in selection + if selected_config else {} ), } -def _frame_guardrail_selection(body: Mapping[str, object]) -> _FrameGuardrailSelection: +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,), _WS_FRAME.validate_python(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 _frame_guardrail_selection( + location: tuple[str, ...], source: Mapping[str, object], selectors: tuple[str, ...] +) -> _FrameGuardrailSelection: try: - return _FRAME_GUARDRAIL_SELECTION.validate_python( - {key: body[key] for key in ("guardrails", "guardrail_config") if key in body} - ) + 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)}") from 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 + } async def _metadata_with_frame_policies( bucket: Mapping[str, object], - policies: object, + policies: Sequence[str], model: object, metadata_key: str, user_api_key_dict: UserAPIKeyAuth, ) -> Mapping[str, object]: selection: Final = { "model": model, - "policies": policies, + "policies": list(policies), metadata_key: {key: _turn_copy(value) for key, value in bucket.items()}, } await add_guardrails_from_policy_engine( diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 701731eaf65..ce6727db5de 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -20,8 +20,12 @@ from fastapi.testclient import TestClient from httpx import Response import litellm -from litellm.proxy._types import ProxyException +from litellm.caching.caching import DualCache +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.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors from tests._master_key import MASTER_KEY @@ -1802,19 +1806,11 @@ class TestResponsesWSStreamingFirstMessage: @pytest.mark.parametrize( ("failure", "expected_status", "expected_type"), [ - ( - ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), - 429, - "rate_limit_exceeded", - ), + (ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), 429, "rate_limit_exceeded"), (litellm.BadRequestError(message="bad input", model="m", llm_provider="openai"), 400, "invalid_request_error"), (RuntimeError("boom"), 500, "server_error"), (ValueError(f"{RouterErrors.no_deployments_with_tag_routing.value}=['free']"), 401, "invalid_request_error"), - ( - ProxyException(message="Authentication Error", type="auth_error", param=None, code=401), - 401, - "invalid_request_error", - ), + (ProxyException(message="Authentication Error", type="auth_error", param=None, code=401), 401, "invalid_request_error"), (HTTPException(status_code=403, detail="key not allowed"), 403, "invalid_request_error"), ], ids=["crossed-budget", "bad-request", "unexpected", "tag-routing", "proxy-auth", "http-exception"], @@ -1972,6 +1968,48 @@ 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, cache, data, call_type): + return self._scan(data) + + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + self._scan(data) + + def _scan(self, data): + self.inputs.append(data["input"]) + self.configs.append(data["litellm_metadata"].get("guardrail_config")) + if self.blocks: + raise HTTPException(status_code=400, detail={"error": f"Content blocked by {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): @@ -2103,7 +2141,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", @@ -2134,7 +2172,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) @@ -2172,7 +2210,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() @@ -2191,7 +2229,7 @@ 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() @@ -2213,13 +2251,14 @@ 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"]}}, @@ -2236,55 +2275,88 @@ class TestResponsesWSFirstFrameModelAuth: assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] @pytest.mark.asyncio - @pytest.mark.parametrize("nested", [False, True]) - async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, nested: bool): + @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): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - scans: list[dict] = [] - - async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): - scans.append(data) - return data - - proxy_logging_obj = MagicMock() - proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) - admitted = { - "model": "gpt-4o-mini", - "litellm_metadata": { - "guardrails": ["key-required"], - "guardrail_config": {"key-required": {"threshold": 1}}, - }, - } - guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj, MagicMock()) - selection = { - "guardrails": ["frame-pick"], - "guardrail_config": {"frame-pick": {"threshold": 2}}, - "include_guardrail_response": True, - } - body = {"input": "hi", "store": False, **selection} + 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} - selected = await guard(json.dumps(frame)) - plain = await guard(json.dumps({"type": "response.create", "input": "again"})) + 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"})) - assert scans[0]["litellm_metadata"]["guardrails"] == ["key-required", "frame-pick"] - assert scans[0]["litellm_metadata"]["guardrail_config"] == { - "key-required": {"threshold": 1}, - "frame-pick": {"threshold": 2}, - } - assert not set(selection) & set(scans[0]) - unselected = {"input": "hi", "store": False} - assert json.loads(selected.frame) == ( - {"type": "response.create", "response": unselected} if nested else {"type": "response.create", **unselected} + 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 scans[1]["litellm_metadata"]["guardrails"] == ["key-required"] - assert scans[1]["litellm_metadata"]["guardrail_config"] == {"key-required": {"threshold": 1}} - assert json.loads(plain.frame) == {"type": "response.create", "input": "again"} - assert admitted["litellm_metadata"] == { - "guardrails": ["key-required"], - "guardrail_config": {"key-required": {"threshold": 1}}, + 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_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_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 @pytest.mark.parametrize( ("selection", "message"), @@ -2293,58 +2365,71 @@ class TestResponsesWSFirstFrameModelAuth: {"guardrails": "frame-pick"}, "response.create guardrails: 'str' instances are not allowed as a Sequence value", ), - ({"guardrails": [{"name": "frame-pick"}]}, "response.create guardrails.0: Input should be a valid string"), + ({"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"), + ( + {"metadata": {"guardrails": "frame-pick"}}, + "response.create 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_a_frame_whose_guardrail_selection_is_malformed(self, selection, message): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() - 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(), - proxy_logging_obj, - MagicMock(), - ) + key_required = _FrameGuardrail("key-required") - with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) + with _guardrails_on(key_required) 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})) - proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert key_required.inputs == [] @pytest.mark.asyncio async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self): - from litellm.proxy._types import UserAPIKeyAuth 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 - proxy_logging_obj = MagicMock() - 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"]}}, - UserAPIKeyAuth(api_key="hashed"), - proxy_logging_obj, - MagicMock(), - ) + 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"]))) - - guarded = await guard(json.dumps({"type": "response.create", "input": "hi", "policies": ["frame-policy"]})) + 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 - scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"] - assert scanned["litellm_metadata"]["guardrails"] == ["key-required", "policy-pick"] - assert "policies" not in scanned + assert key_required.inputs == ["hi"] + assert policy_pick.inputs == ["hi"] assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} @pytest.mark.asyncio @@ -2353,7 +2438,7 @@ class TestResponsesWSFirstFrameModelAuth: 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( @@ -2381,7 +2466,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() @@ -2414,7 +2499,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() @@ -2442,7 +2527,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() @@ -2467,7 +2552,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() @@ -2489,18 +2574,24 @@ class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio @pytest.mark.parametrize( - "top_level", [{"input": "forbidden"}, {"guardrails": ["frame-pick"]}, {"include_guardrail_response": True}] + ("top_level", "field"), + [ + ({"input": "forbidden"}, "input"), + ({"guardrails": ["frame-pick"]}, "guardrails"), + ({"include_guardrail_response": True}, "include_guardrail_response"), + ], ) - async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self, top_level): + async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self, top_level, field): 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() ) + 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="both at the top level"): + with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match=f"^{re.escape(message)}$"): await guard(json.dumps({"type": "response.create", **top_level, "response": {}})) proxy_logging_obj.pre_call_hook.assert_not_awaited() @@ -2511,7 +2602,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")) @@ -2553,7 +2644,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() @@ -2642,7 +2733,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) @@ -2700,7 +2791,6 @@ class TestResponsesWSFirstFrameModelAuth: ws.receive_text = AsyncMock( return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []}) ) - async def close_over_the_network(**kwargs): await asyncio.sleep(0) @@ -2720,9 +2810,7 @@ class TestResponsesWSFirstFrameModelAuth: def route_request(data, **kwargs): async def relay(): - guarded = await data["responses_ws_frame_guard"]( - json.dumps({"type": "response.create", "input": "two"}) - ) + guarded = await data["responses_ws_frame_guard"](json.dumps({"type": "response.create", "input": "two"})) guarded.bill_refusal(refusal) if relay_outcome == "still_relaying": await asyncio.Event().wait() @@ -2735,12 +2823,11 @@ 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 ) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=slow_failure_hook) - async def release_over_the_network(user_api_key_dict): await asyncio.sleep(0.01) @@ -2783,9 +2870,7 @@ class TestResponsesWSFirstFrameModelAuth: row_may_finish.set() await asyncio.wait_for(row_written.wait(), timeout=1) - assert proxy_logging_obj.post_call_failure_hook.await_args.kwargs["original_exception"] is ( - crash if crashed else refusal - ) + assert proxy_logging_obj.post_call_failure_hook.await_args.kwargs["original_exception"] is (crash if crashed else refusal) @pytest.mark.asyncio @pytest.mark.parametrize("ending", ["client_closed", "crash"]) @@ -2820,9 +2905,7 @@ class TestResponsesWSFirstFrameModelAuth: def route_request(data, **kwargs): async def relay(): - guarded = await data["responses_ws_frame_guard"]( - json.dumps({"type": "response.create", "input": "two"}) - ) + guarded = await data["responses_ws_frame_guard"](json.dumps({"type": "response.create", "input": "two"})) guarded.bill_refusal(refusal) if ending == "crash": raise crash @@ -2830,7 +2913,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 ) @@ -2897,16 +2980,14 @@ class TestResponsesWSFirstFrameModelAuth: def route_request(data, **kwargs): async def relay(): - guarded = await data["responses_ws_frame_guard"]( - json.dumps({"type": "response.create", "input": "two"}) - ) + guarded = await data["responses_ws_frame_guard"](json.dumps({"type": "response.create", "input": "two"})) guarded.bill_refusal(refusal) data["litellm_metadata"]["standard_logging_guardrail_information"] = "written by a broken callback" raise RuntimeError("relay crashed") 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, @@ -2967,7 +3048,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() @@ -4601,9 +4682,7 @@ class TestResponsesInputTokens: assert response.status_code == 200, response.text sent = counter.await_args.kwargs["responses_request"] - assert { - name: sent[name] for name in ("tools", "tool_choice", "parallel_tool_calls") if name in sent - } == expected + assert {name: sent[name] for name in ("tools", "tool_choice", "parallel_tool_calls") if name in sent} == expected def test_previous_response_id_reaches_the_counter_authorized_next_to_the_id_the_client_held(self) -> None: from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD From ad55e6eaaaf71cf88d9b30c2899c6e49015c6886 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:50:19 -0700 Subject: [PATCH 04/10] 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. --- litellm/proxy/litellm_pre_call_utils.py | 16 +++++++-- .../response_api_endpoints/test_endpoints.py | 17 ++++++--- .../unit/proxy/test_litellm_pre_call_utils.py | 36 +++++++++++++++++++ 3 files changed, 62 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index fc76779f634..b063670c346 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4255,10 +4255,20 @@ 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))) + selected_names: Final = frozenset( + itertools.chain.from_iterable( + _REQUEST_FIELDS.validate_python(entry).keys() if isinstance(entry, dict) else (entry,) + for entry in existing_guardrails_list + ) + ) + combined: Final = [ + *existing_guardrails_list, + *(name for name in dict.fromkeys(resolved_guardrails) if name not in selected_names), + ] data[metadata_variable_name]["guardrails"] = combined verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 11dbc1b0ca2..63e5315ecd1 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.utils import ProxyLogging 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 @@ -1983,15 +1984,23 @@ class _FrameGuardrail(CustomGuardrail): self.rewrites = rewrites self.blocks = blocks - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + 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, user_api_key_dict, call_type): + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: self._scan(data) - def _scan(self, data): + def _scan(self, data: dict[str, object]) -> dict[str, object]: self.inputs.append(data["input"]) - self.configs.append(data["litellm_metadata"].get("guardrail_config")) + 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}"}) return {**data, "input": f"{data['input']} [{self.guardrail_name}]"} if self.rewrites else data diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 81a196af676..e82d2d66e3c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5292,6 +5292,42 @@ 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"] + + @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 From a17ad50dd3c8893027411187cd697c4424c65627 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:05:58 -0700 Subject: [PATCH 05/10] fix(responses): run streaming-scoped guardrails on socket turns, keep guardrail fields on blocks, merge guardrail lists by name --- litellm/integrations/custom_guardrail.py | 27 ++++- litellm/proxy/litellm_pre_call_utils.py | 70 ++++-------- .../proxy/response_api_endpoints/endpoints.py | 59 ++++++---- litellm/proxy/utils.py | 4 +- litellm/responses/streaming_iterator.py | 36 ++++-- .../integrations/test_custom_guardrail.py | 10 ++ .../response_api_endpoints/test_endpoints.py | 108 +++++++++++++++++- .../unit/proxy/test_litellm_pre_call_utils.py | 46 ++++++++ .../utils/helpers/test_guardrail_merge.py | 10 +- 9 files changed, 282 insertions(+), 88 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..15a29ddfb7e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,12 +3,14 @@ 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 import httpx +from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -207,6 +209,29 @@ 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 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""" + 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)] + + +def _guardrail_entry_names(entry: object) -> frozenset[str]: + if isinstance(entry, str): + return frozenset((entry,)) + if isinstance(entry, dict): + return frozenset(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return frozenset() + + +def _named_earlier(names: tuple[frozenset[str], ...], index: int) -> bool: + return bool(names[index]) and names[index] <= frozenset[str]().union(*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/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b063670c346..0285bc270c7 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 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -3895,30 +3895,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 @@ -4013,14 +3997,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]: @@ -4259,16 +4239,7 @@ def _apply_resolved_guardrails_to_metadata( _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () ) - selected_names: Final = frozenset( - itertools.chain.from_iterable( - _REQUEST_FIELDS.validate_python(entry).keys() if isinstance(entry, dict) else (entry,) - for entry in existing_guardrails_list - ) - ) - combined: Final = [ - *existing_guardrails_list, - *(name for name in dict.fromkeys(resolved_guardrails) if name not in selected_names), - ] + 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,12 +4299,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 - _model: Final = data.get("model") - 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=_model if isinstance(_model, str) else None, - 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/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 1f6d3b5bf08..3ebf99d176a 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -36,7 +36,7 @@ from litellm._logging import verbose_proxy_logger from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES -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 @@ -89,7 +89,12 @@ from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_proxy 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.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 ( @@ -2057,8 +2062,8 @@ _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, *_PROXY_ONLY_FRAME_FIELDS) -_UNFORWARDED_FRAME_FIELDS: Final = (*_NESTED_BODY_ONLY_FIELDS, "litellm_metadata") +_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) @@ -2108,26 +2113,27 @@ def _responses_ws_frame_guard( async def guard(frame: str) -> GuardedResponsesWsFrame: frame_obj: Final = _WS_FRAME.validate_json(frame) nested: Final = frame_obj.get("response") - stray: Final = ( - next((key for key in _NESTED_BODY_ONLY_FIELDS if key in frame_obj), None) - if isinstance(nested, dict) - else None - ) + 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 = { - **{ - key: _without_untrusted_metadata(key, value) - for key, value in body.items() - if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS + frame_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: await _metadata_with_frame_guardrails( + bucket, body, model, metadata_key, user_api_key_dict + ), }, - "model": model, - "litellm_call_id": str(uuid.uuid4()), - metadata_key: await _metadata_with_frame_guardrails(bucket, body, model, metadata_key, user_api_key_dict), - } + is_streaming=True, + ) 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 @@ -2137,7 +2143,8 @@ def _responses_ws_frame_guard( ) except Exception as e: await _bill_rejected_frame(rejected_frame_auth, e, frame_data) - raise ValueError(proxy_exception_from_route_error(e).message) from e + proxy_exception: Final = proxy_exception_from_route_error(e) + raise ResponsesWsFrameRejection(proxy_exception.message, proxy_exception.provider_specific_fields) 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" @@ -2153,7 +2160,10 @@ def _responses_ws_frame_guard( if key not in _UNFORWARDED_FRAME_FIELDS } rewritten: Final = ( - {**frame_obj, "response": {**kept, **guarded_fields}} + { + **{key: value for key, value in frame_obj.items() if key not in _PROXY_ONLY_FRAME_FIELDS}, + "response": {**kept, **guarded_fields}, + } if isinstance(nested, dict) else {**kept, **guarded_fields} ) @@ -2168,6 +2178,15 @@ def _responses_ws_frame_guard( 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) + + class _FrameGuardrailSelection(TypedDict, total=False): guardrails: ReadOnly[Sequence[str | Mapping[str, Mapping[str, object]]]] guardrail_config: ReadOnly[Mapping[str, object]] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0ddfc5ce62f..0ce92bc21a1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -146,6 +146,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 @@ -9086,8 +9087,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 c1f422c7d7b..4969227f79a 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2172,6 +2172,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: @@ -2437,7 +2460,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( @@ -2865,14 +2889,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: @@ -3286,7 +3302,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 614ff97553a..de5f6995dcc 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3786,3 +3786,13 @@ def test_default_on_guardrail() -> None: ) is True ) + + +def test_merged_guardrail_entries_keeps_the_first_entry_per_guardrail_and_any_entry_that_names_a_new_one(): + 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", + ] diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 63e5315ecd1..5431c633623 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2002,7 +2002,10 @@ class _FrameGuardrail(CustomGuardrail): 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}"}) + 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 @@ -2350,6 +2353,33 @@ class TestResponsesWSFirstFrameModelAuth: ] 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 @@ -2588,7 +2618,9 @@ class TestResponsesWSFirstFrameModelAuth: [ ({"input": "forbidden"}, "input"), ({"guardrails": ["frame-pick"]}, "guardrails"), - ({"include_guardrail_response": True}, "include_guardrail_response"), + ({"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): @@ -2606,6 +2638,78 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj.pre_call_hook.assert_not_awaited() + @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", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "include_guardrail_response": True, + "guardrail_config": {"frame-pick": {"threshold": 1}}, + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + ) + ) + + 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): import litellm diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index e82d2d66e3c..2ba12b543de 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5328,6 +5328,52 @@ async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and 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 e505517eeff..4ae87dda343 100644 --- a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py +++ b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py @@ -360,7 +360,9 @@ 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"] 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 06/10] 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"] From 754225427c43c0c22ecdd7eaa04f1c629da5da92 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:12:15 -0700 Subject: [PATCH 07/10] style(responses): format the frame guard and give the pass-through guardrail test an enterprise license --- litellm/proxy/response_api_endpoints/endpoints.py | 9 +++++++-- .../test_passthrough_guardrails.py | 3 ++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 58272269235..9894ee1035b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2137,7 +2137,10 @@ 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 + 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) @@ -2211,7 +2214,9 @@ def _responses_ws_frame_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 + 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") 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 7ebb82386a8..ddf22b920b5 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py @@ -74,9 +74,10 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: class TestPassthroughGuardrailHandlerCollectGuardrails: - def test_inherits_a_key_guardrail_listed_with_settings_by_its_name(self): + def test_inherits_a_key_guardrail_listed_with_settings_by_its_name(self, monkeypatch): 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"]} ) From 7d5037cd79c8b314d3e8ec568ab7bddd563adfe9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:52:42 -0700 Subject: [PATCH 08/10] fix(proxy): keep pass-through guardrail settings, gate frame policy picks, skip moderation-only hooks on socket turns --- litellm/proxy/auth/auth_checks.py | 8 ++- .../passthrough_guardrails.py | 32 ++------- .../proxy/response_api_endpoints/endpoints.py | 29 ++++---- litellm/proxy/utils.py | 2 + .../test_passthrough_guardrails.py | 13 ++-- .../response_api_endpoints/test_endpoints.py | 71 ++++++++++++++++++- 6 files changed, 107 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3044caca42e..beb8f614f4f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -981,10 +981,12 @@ def guardrail_modification_check(request_body: Mapping[str, object], team_object failing loudly at the auth layer so operators see an explicit 403 instead of a confusing silent-ignore. """ - from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails + if requests_guardrail_modification(request_body): + require_guardrail_modification_permission(team_object) - if not requests_guardrail_modification(request_body): - return + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: + from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails if not can_modify_guardrails(team_object): raise HTTPException( diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index f35d4346d87..bd8390d704e 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -7,11 +7,10 @@ Handles guardrail execution for passthrough endpoints with: - Automatic inheritance from org/team/key levels when enabled """ -from itertools import chain from typing import Any, Final from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import guardrail_entry_names +from litellm.integrations.custom_guardrail import merged_guardrail_entries from litellm.proxy._types import ( PassThroughGuardrailsConfig, PassThroughGuardrailSettings, @@ -207,7 +206,7 @@ class PassthroughGuardrailHandler: def collect_guardrails( user_api_key_dict: UserAPIKeyAuth, passthrough_guardrails_config: PassThroughGuardrailsConfigInput | None, - ) -> dict[str, bool] | None: + ) -> list[object] | None: """ Collect guardrails for a passthrough endpoint. @@ -227,7 +226,7 @@ class PassthroughGuardrailHandler: passthrough_guardrails_config: List or Dict of guardrail names/settings Returns: - Dict of guardrail names to run (format: {guardrail_name: True}), or None + Guardrail entries to run, each a name or a ``{name: settings}`` entry, or None """ from litellm.proxy.litellm_pre_call_utils import ( add_guardrails_from_key_or_team_metadata, @@ -244,15 +243,6 @@ class PassthroughGuardrailHandler: verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None - # Passthrough is enabled - collect guardrails - guardrails_to_run: Final[dict[str, bool]] = {} - - # Add passthrough-specific guardrails - for guardrail_name in normalized_config: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added passthrough-specific guardrail: %s", guardrail_name) - - # Add org/team/key level guardrails using shared helper temp_data: Final[dict[str, Any]] = {"metadata": {}} add_guardrails_from_key_or_team_metadata( key_metadata=user_api_key_dict.metadata, @@ -260,20 +250,12 @@ class PassthroughGuardrailHandler: data=temp_data, metadata_variable_name="metadata", ) + inherited_guardrails: Final[list[object]] = temp_data["metadata"].get("guardrails", []) + guardrails_to_run: Final = merged_guardrail_entries(normalized_config, inherited_guardrails) - # Merge inherited guardrails into guardrails_to_run - inherited_guardrails: Final = temp_data["metadata"].get("guardrails", []) - for guardrail_name in chain.from_iterable(guardrail_entry_names(entry) for entry in inherited_guardrails): - if guardrail_name not in guardrails_to_run: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added inherited guardrail (key/team level): %s", guardrail_name) + verbose_proxy_logger.debug("Collected guardrails for passthrough endpoint: %s", guardrails_to_run) - verbose_proxy_logger.debug( - "Collected guardrails for passthrough endpoint: %s", - list(guardrails_to_run.keys()), - ) - - return guardrails_to_run if guardrails_to_run else None + return guardrails_to_run or None @staticmethod def get_field_targeted_text( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5bf55c71e88..685f1086eb2 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -50,8 +50,8 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( get_team_object, - guardrail_modification_check, requests_guardrail_modification, + require_guardrail_modification_permission, ) from litellm.proxy.auth.auth_utils import carries_client_side_overrides from litellm.proxy.auth.user_api_key_auth import ( @@ -2183,6 +2183,7 @@ class _GuardrailWalk(Protocol): data: dict[str, object], # mutable-ok: ProxyLogging.during_call_hook hands the guardrails a dict they edit user_api_key_dict: UserAPIKeyAuth | None, call_type: Literal["aresponses"], + guardrails_only: bool, ) -> object: ... @@ -2192,13 +2193,16 @@ async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTab team_id: Final = user_api_key_dict.team_id if team_id is None or team_id == UI_TEAM_ID: return None - return await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_dict.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + return LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata) def _responses_ws_frame_guard( @@ -2274,7 +2278,7 @@ def _responses_ws_frame_guard( } pii_tokens: Final = _recorded_pii_tokens(guarded_data.get("metadata")) await proxy_logging_obj.during_call_hook( - data=guarded_data, user_api_key_dict=user_api_key_dict, call_type="aresponses" + data=guarded_data, user_api_key_dict=user_api_key_dict, call_type="aresponses", guardrails_only=True ) except Exception as e: raise await rejection(e, frame_data) from e @@ -2316,10 +2320,9 @@ async def _check_frame_may_modify_guardrails( user_api_key_dict: UserAPIKeyAuth, team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]], ) -> None: - """A frame that picks guardrails needs the team permission an HTTP body that picks them needs""" - modifying_source: Final = next((source for source in sources if requests_guardrail_modification(source)), None) - if modifying_source is not None: - guardrail_modification_check(modifying_source, await team_object(user_api_key_dict)) + """A frame that picks guardrails or policies needs the team's permission to modify guardrails""" + if any("policies" in source or requests_guardrail_modification(source) for source in sources): + require_guardrail_modification_permission(await team_object(user_api_key_dict)) def _forwarded_frame_fields(source: Mapping[str, object]) -> dict[str, object]: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8f51d14f14a..47159c81e71 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3497,6 +3497,7 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ): caps: Final = ProxyLogging.callback_capabilities() if not caps.has_guardrail and not caps.has_moderation_override: @@ -3508,6 +3509,7 @@ class ProxyLogging: if ( isinstance(callback, CustomLogger) and not isinstance(callback, CustomGuardrail) + and (not guardrails_only or callback.enforces_request_content) and _overrides_moderation_hook(callback) and user_api_key_dict is not None ): diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py index ddf22b920b5..ef19482b30a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py @@ -74,7 +74,8 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: class TestPassthroughGuardrailHandlerCollectGuardrails: - def test_inherits_a_key_guardrail_listed_with_settings_by_its_name(self, monkeypatch): + def test_inherits_a_key_guardrail_with_its_settings(self, monkeypatch: pytest.MonkeyPatch): + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) @@ -82,10 +83,12 @@ class TestPassthroughGuardrailHandlerCollectGuardrails: api_key="hashed", metadata={"guardrails": [{"pii": {"extra_body": {"threshold": 1}}}, "toxic"]} ) - assert PassthroughGuardrailHandler.collect_guardrails(key, ["passthrough-guard"]) == { - "passthrough-guard": True, - "pii": True, - "toxic": True, + guardrails = PassthroughGuardrailHandler.collect_guardrails(key, ["passthrough-guard", "toxic"]) + + assert guardrails == ["passthrough-guard", "toxic", {"pii": {"extra_body": {"threshold": 1}}}] + request_data = {"metadata": {"guardrails": guardrails}} + assert CustomGuardrail(guardrail_name="pii").get_guardrail_dynamic_request_body_params(request_data) == { + "threshold": 1 } diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 0fc88fc7f3b..9dea7fa0af9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2628,6 +2628,34 @@ class TestResponsesWSFirstFrameModelAuth: assert json.loads(unselected_turn.frame) == {"type": "response.create", "input": "plain"} assert during_pick.inputs == ["hi"] + @pytest.mark.asyncio + async def test_frame_guard_runs_only_the_moderation_hooks_that_enforce_request_content(self): + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + class _ModerationHook(CustomLogger): + def __init__(self, enforces_request_content: bool) -> None: + super().__init__() + self.enforces_request_content = enforces_request_content + self.inputs: list[object] = [] + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + self.inputs.append(data.get("input")) + + vendor_moderation = _ModerationHook(enforces_request_content=False) + content_enforcer = _ModerationHook(enforces_request_content=True) + + with _guardrails_on(vendor_moderation, content_enforcer) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert vendor_moderation.inputs == [] + assert content_enforcer.inputs == ["hi"] + @pytest.mark.asyncio @pytest.mark.parametrize( ("selection", "message"), @@ -2918,10 +2946,12 @@ class TestResponsesWSFirstFrameModelAuth: {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": True}}, {"type": "response.create", "response": {"input": "hi", "guardrails": ["frame-pick"]}}, {"type": "response.create", "litellm_metadata": {"guardrails": ["frame-pick"]}, "response": {"input": "hi"}}, + {"type": "response.create", "input": "hi", "policies": ["frame-policy"]}, + {"type": "response.create", "response": {"input": "hi", "policies": ["frame-policy"]}}, ], - ids=["top-level", "metadata", "global-toggle", "nested", "beside-nested"], + ids=["top-level", "metadata", "global-toggle", "nested", "beside-nested", "policies", "nested-policies"], ) - async def test_frame_guard_refuses_guardrail_changes_from_a_team_without_permission_like_http(self, frame): + async def test_frame_guard_refuses_guardrail_and_policy_picks_from_a_team_without_permission(self, frame): from litellm.proxy._types import LiteLLM_TeamTable from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard @@ -2977,6 +3007,43 @@ class TestResponsesWSFirstFrameModelAuth: assert frame_pick.inputs == ["hi"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("modify_guardrails", "picked_inputs"), + [(True, ["hi"]), (False, [])], + ids=["open", "locked"], + ) + async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_row_is_gone( + self, modify_guardrails: bool, picked_inputs: list[str] + ): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + key = UserAPIKeyAuth( + api_key="hashed", + team_id="deleted-team", + team_metadata={"guardrails": {"modify_guardrails": modify_guardrails}}, + ) + frame = json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + + with ( + _guardrails_on(frame_pick) as proxy_logging_obj, + patch( + "litellm.proxy.response_api_endpoints.endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})), + ), + ): + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, key, proxy_logging_obj, MagicMock() + ) + if modify_guardrails: + await guard(frame) + else: + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(frame) + + assert frame_pick.inputs == picked_inputs + @pytest.mark.asyncio async def test_frame_guard_forwards_the_pre_call_output_and_not_a_during_call_rewrite_like_http(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard From 607c59e3c57dfe5c2cba0a91688d5bd473dc8449 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:30:11 -0700 Subject: [PATCH 09/10] fix(responses): keep string metadata guardrail values as client metadata on socket turns --- .../proxy/response_api_endpoints/endpoints.py | 12 +++++-- .../response_api_endpoints/test_endpoints.py | 35 +++++++++++++++---- 2 files changed, 38 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 685f1086eb2..0578de15f3d 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2399,7 +2399,7 @@ def _frame_guardrail_selections(body: Mapping[str, object]) -> tuple[_FrameGuard ``litellm_metadata``, the top level last so its ``guardrail_config`` wins""" return ( *( - _frame_guardrail_selection((key,), _WS_FRAME.validate_python(value), _METADATA_GUARDRAIL_SELECTORS) + _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) ), @@ -2407,6 +2407,14 @@ def _frame_guardrail_selections(body: Mapping[str, object]) -> tuple[_FrameGuard ) +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 not isinstance(item, str)} + + def _frame_guardrail_selection( location: tuple[str, ...], source: Mapping[str, object], selectors: tuple[str, ...] ) -> _FrameGuardrailSelection: @@ -2422,7 +2430,7 @@ def _without_guardrail_selection(metadata: object) -> object: return { name: item for name, item in _WS_FRAME.validate_python(metadata).items() - if name not in _METADATA_GUARDRAIL_SELECTORS + if name not in _METADATA_GUARDRAIL_SELECTORS or isinstance(item, str) } diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 9dea7fa0af9..74f9162c992 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2530,7 +2530,7 @@ class TestResponsesWSFirstFrameModelAuth: ], 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): + 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") @@ -2557,6 +2557,27 @@ class TestResponsesWSFirstFrameModelAuth: 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.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"} + + 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(), + ) + 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 @@ -2672,8 +2693,8 @@ class TestResponsesWSFirstFrameModelAuth: ), ({"policies": [1]}, "response.create policies.0: Input should be a valid string"), ( - {"metadata": {"guardrails": "frame-pick"}}, - "response.create metadata.guardrails: 'str' instances are not allowed as a Sequence value", + {"litellm_metadata": {"guardrails": "frame-pick"}}, + "response.create litellm_metadata.guardrails: 'str' instances are not allowed as a Sequence value", ), ( {"litellm_metadata": {"guardrail_config": []}}, @@ -2682,7 +2703,7 @@ class TestResponsesWSFirstFrameModelAuth: ], ) async def test_frame_guard_refuses_and_bills_a_frame_whose_guardrail_selection_is_malformed( - self, selection, message + 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 @@ -2879,7 +2900,7 @@ class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) - async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field): + 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 @@ -2911,7 +2932,7 @@ class TestResponsesWSFirstFrameModelAuth: ], ids=["top-level", "metadata", "litellm-metadata"], ) - async def test_frame_guard_runs_the_guardrails_a_nested_frame_selects_beside_its_response(self, top_level): + 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") @@ -2951,7 +2972,7 @@ class TestResponsesWSFirstFrameModelAuth: ], ids=["top-level", "metadata", "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): + 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 From 18d96358d62c116d0de7bcb05c48824c07616043 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:31:48 -0700 Subject: [PATCH 10/10] fix(responses): ask no team permission for string metadata named like a guardrail pick on socket turns --- .../proxy/response_api_endpoints/endpoints.py | 15 +++++++++++++-- .../response_api_endpoints/test_endpoints.py | 17 ++++++++++++++++- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 110b627f36c..66db9b84dc8 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -2331,10 +2331,17 @@ async def _check_frame_may_modify_guardrails( 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(source) for source in sources): + 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 @@ -2422,7 +2429,11 @@ def _metadata_picks(key: str, metadata: object) -> Mapping[str, object]: validated: Final = _WS_FRAME.validate_python(metadata) if key != "metadata": return validated - return {name: item for name, item in validated.items() if not isinstance(item, str)} + 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( diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index fda037071ba..4d59c47aa3d 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2577,18 +2577,23 @@ class TestResponsesWSFirstFrameModelAuth: @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})) @@ -2983,12 +2988,22 @@ class TestResponsesWSFirstFrameModelAuth: {"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", "nested", "beside-nested", "policies", "nested-policies"], + 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