From 5011b1263f3d322d551666b5e7ae75ca1f950e4c Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 30 Oct 2026 12:34:52 -0700 Subject: [PATCH 01/29] feat(proxy)!: run key, team and project guardrails on built-in provider pass-through routes The /openai, /azure, /anthropic, /vertex_ai and /gemini pass-through routes now merge the caller's key, team, project and policy guardrails with any named in the request body, the same way the native routes and /bedrock do. Requests without a JSON body are not scanned, and custom pass-through endpoints keep their opt-in guardrails config --- .../llm_passthrough_endpoints.py | 4 + .../pass_through_endpoints.py | 44 +++- .../test_passthrough_auth_guardrails.py | 193 ++++++++++++++++++ .../test_llm_pass_through_endpoints.py | 3 + .../test_pass_through_endpoints.py | 119 +++++++++++ 5 files changed, 361 insertions(+), 2 deletions(-) create mode 100644 tests/integration/authorization/test_passthrough_auth_guardrails.py diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 33b18c3d939..287475c868b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -414,6 +414,7 @@ async def gemini_proxy_route( custom_llm_provider="gemini", is_streaming_request=is_streaming_request, query_params=merged_params, + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( request, @@ -1079,6 +1080,7 @@ async def anthropic_proxy_route( target=str(updated_url), custom_headers=_upstream_headers_for_anthropic_route(request, user_api_key_dict, auth_header), is_streaming_request=is_streaming_request, + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( request, @@ -3192,6 +3194,7 @@ async def _base_vertex_proxy_route( upstream_headers_filter=partial( _without_dropped_anthropic_betas, _vertex_dropped_anthropic_betas(general_settings, request.headers) ), + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path try: @@ -3670,6 +3673,7 @@ class BaseOpenAIPassThroughHandler: if custom_llm_provider else None ), + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( request, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cfac28db1a9..23c1202ea07 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -60,6 +60,7 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -126,6 +127,7 @@ from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy modul _key_or_team_allows_client_pricing_override, # pyright: ignore[reportPrivateUsage] # reuse the proxy's pricing trust policy _strip_client_pricing_overrides, # pyright: ignore[reportPrivateUsage] # sanitize before trusted hooks add guardrail costs get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above + move_guardrails_to_metadata, # pyright: ignore[reportUnknownVariableType] # untyped data param; the same merge the native routes run redact_credential_headers, strip_client_guardrail_information, ) @@ -173,6 +175,7 @@ pass_through_endpoint_logging: Final = PassThroughEndpointLogging() _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) +_GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -854,6 +857,31 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _update_stream_param_based_on_request_body = update_stream_param_based_on_request_body +async def _inherit_auth_guardrails( + parsed_body: dict[str, object] | None, # mutable-ok: merges the guardrails into the live request body + user_api_key_dict: UserAPIKeyAuth, + configured_guardrails: Mapping[str, bool] | None, + inherit: bool, +) -> Mapping[str, bool] | Sequence[object] | None: + """Return the guardrails a pass-through request runs. + + The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, + project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` and returned, or None + when there are none. + """ + if configured_guardrails or not inherit or not parsed_body: + return configured_guardrails + if not isinstance(parsed_body.get("metadata"), dict): + parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata + await move_guardrails_to_metadata( + data=parsed_body, + _metadata_variable_name="metadata", + user_api_key_dict=user_api_key_dict, + ) + merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") + return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) if isinstance(merged, list) and merged else None + + def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | None) -> None: """Copy guardrail logging entries from ``guardrail_data`` onto ``request_data``. @@ -1212,6 +1240,7 @@ async def pass_through_request( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + inherit_auth_guardrails: bool = False, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1234,6 +1263,8 @@ async def pass_through_request( general_settings.pass_through_request_timeout, then 600s. upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion operation: Optional GenAI operation the route reports, unless the upstream names its own + inherit_auth_guardrails: Run the key, team, project and policy guardrails on a JSON body, as the + native routes do, when the pass-through guardrails config resolves to none """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1342,6 +1373,13 @@ async def pass_through_request( _parsed_body["metadata"]["guardrails"] = guardrails_to_run verbose_proxy_logger.debug("Added guardrails to passthrough request metadata: %s", guardrails_to_run) + post_call_guardrails: Final = await _inherit_auth_guardrails( + parsed_body=cast(dict[str, object] | None, _parsed_body), # cast-ok: json + user_api_key_dict=user_api_key_dict, + configured_guardrails=guardrails_to_run, + inherit=inherit_auth_guardrails, + ) + ## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it # Surface the requested model (when the body carries one) so logging/spans # read e.g. ``chat gpt-4o`` instead of ``chat unknown``. @@ -1876,7 +1914,7 @@ async def pass_through_request( logging_obj=logging_obj, ) - if response.status_code < 400 and response_body is not None and guardrails_to_run: + if response.status_code < 400 and response_body is not None and post_call_guardrails: # Build an enriched data dict: _parsed_body has been stripped of # `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint, # so we re-attach the configured guardrails here so should_run_guardrail @@ -1887,7 +1925,7 @@ async def pass_through_request( existing_metadata = {} hook_data["metadata"] = { **existing_metadata, - "guardrails": guardrails_to_run, + "guardrails": post_call_guardrails, } post_call_guardrail_data = hook_data response_body = await proxy_logging_obj.post_call_success_hook( @@ -2271,6 +2309,7 @@ def create_pass_through_route( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + inherit_auth_guardrails: bool = False, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2420,6 +2459,7 @@ def create_pass_through_route( timeout=cast(float | None, param_timeout), upstream_headers_filter=upstream_headers_filter, operation=operation, + inherit_auth_guardrails=inherit_auth_guardrails, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py new file mode 100644 index 00000000000..e92ebcb5594 --- /dev/null +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -0,0 +1,193 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.anthropic_sse import message_json +from integration._support.client import Gateway, Scenario +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +GUARDRAIL: Final = "passthrough-auth-guardrail" +BLOCK_MARKER: Final = "block-this-prompt" +OPENAI_ROUTE: Final = "/openai/v1/chat/completions" +ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" +GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) + +Level = Literal["key", "team", "project"] + + +def _route_body(route: str, prompt: str) -> dict[str, JsonValue]: + match route: + case "/anthropic/v1/messages": + return {"model": "claude-sonnet-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + case "/gemini/v1beta/models/gemini-2.5-flash:generateContent": + return {"contents": [{"role": "user", "parts": [{"text": prompt}]}]} + case _: + return {"model": "gpt-5.6", "messages": [{"role": "user", "content": prompt}]} + + +def _verdict(request: Request) -> Reply: + action: Final = "BLOCKED" if BLOCK_MARKER.encode() in request.body else "NONE" + return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode()) + + +def _openai_reply(_request: Request) -> Reply: + return chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False) + + +def _anthropic_reply(_request: Request) -> Reply: + return Reply(body=message_json("msg_passthrough_auth", "claude-sonnet-5-5", "Hi")) + + +def _gemini_reply(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 1, "totalTokenCount": 5}, + } + ).encode() + ) + + +@dataclass(frozen=True, slots=True) +class AuthGuardrailRig: + owned: OwnedProxy + vendor: Wire + upstreams: Mapping[str, Wire] + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[AuthGuardrailRig]: + directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails") + with ExitStack() as stack: + vendor: Final = stack.enter_context(wire_server(_verdict)) + upstreams: Final = MappingProxyType( + { + OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), + ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)), + GEMINI_ROUTE: stack.enter_context(wire_server(_gemini_reply)), + } + ) + config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + path: Final = directory / "passthrough-auth-guardrails.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "guardrails": [ + { + "guardrail_name": GUARDRAIL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": vendor.url, + "api_key": "synthetic-guardrail-key", + }, + } + ], + } + ) + ) + environment: Final = { + "OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url, + "OPENAI_API_KEY": "synthetic-openai-key", + "ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + "GEMINI_API_BASE": upstreams[GEMINI_ROUTE].url, + "GEMINI_API_KEY": "synthetic-gemini-key", + } + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) + yield AuthGuardrailRig(owned, vendor, upstreams) + + +def _key_with_guardrail(scenario: Scenario, level: Level) -> str: + attached: Final[dict[str, JsonValue]] = {"metadata": {"guardrails": [GUARDRAIL]}} + team: Final = scenario.team(**(attached if level == "team" else {})) + match level: + case "project": + return scenario.key(team_id=team, project_id=scenario.project(team, **attached)) + case "key": + return scenario.key(team_id=team, **attached) + case _: + return scenario.key(team_id=team) + + +def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str) -> httpx.Response: + return gateway.request("POST", route, body, key=key, headers={"x-goog-api-key": key}) + + +def _drain(rig: AuthGuardrailRig) -> None: + rig.vendor.drain() + for upstream in rig.upstreams.values(): + upstream.drain() + + +@pytest.mark.parametrize("level", ["key", "team", "project"]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_auth_guardrail_blocks_the_prompt_before_the_provider(rig: AuthGuardrailRig, route: str, level: Level) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, level) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, BLOCK_MARKER), key) + + assert response.status_code == 400, response.text + assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True] + assert rig.upstreams[route].drain() == () + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_auth_guardrail_lets_a_clean_prompt_reach_the_provider(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + marker: Final = f"clean-{uuid.uuid4().hex}" + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, marker), key) + + assert response.status_code == 200, response.text + assert [marker.encode() in request.body for request in rig.vendor.drain()] == [True] + assert [marker.encode() in request.body for request in rig.upstreams[route].drain()] == [True] + + +def test_empty_request_guardrails_list_does_not_drop_the_key_guardrail(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + _drain(rig) + + response: Final = gateway.request( + "POST", OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE, BLOCK_MARKER), "guardrails": []}, key=key + ) + + assert response.status_code == 400, response.text + assert rig.upstreams[OPENAI_ROUTE].drain() == () + + +def test_request_without_a_body_is_not_sent_to_the_guardrail(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + _drain(rig) + + response: Final = gateway.request("GET", "/openai/v1/models", key=key) + + assert response.status_code == 200, response.text + assert rig.vendor.drain() == () + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index ff27ac54a5b..4decee140e4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -344,6 +344,7 @@ class TestVertexAIPassThroughHandler: custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, upstream_headers_filter=mock.ANY, + inherit_auth_guardrails=True, ) @pytest.mark.asyncio @@ -444,6 +445,7 @@ class TestVertexAIPassThroughHandler: custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, upstream_headers_filter=mock.ANY, + inherit_auth_guardrails=True, ) @pytest.mark.parametrize( @@ -540,6 +542,7 @@ class TestVertexAIPassThroughHandler: custom_headers={"Authorization": f"Bearer {default_credentials}"}, is_streaming_request=False, upstream_headers_filter=mock.ANY, + inherit_auth_guardrails=True, ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index a9098ac3d59..294ad45d795 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11353,3 +11353,122 @@ async def test_passthrough_header_and_callback_share_finalized_cost( finally: cleanup() await client.aclose() + + +class TestInheritAuthGuardrails: + """Built-in provider routes run the caller's key, team and project guardrails on pass-through.""" + + @staticmethod + async def _run( + body: dict[str, object], *, inherit: bool, user_api_key_dict: UserAPIKeyAuth + ) -> tuple[list[dict[str, object]], AsyncMock]: + pre_call_data: Final[list[dict[str, object]]] = [] + + async def _record_pre_call(**kwargs: object) -> object: + pre_call_data.append(json.loads(json.dumps(kwargs["data"], default=str))) + return kwargs["data"] + + proxy_logging: Final = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=_record_pre_call) + proxy_logging.post_call_success_hook = AsyncMock(return_value={"id": "chatcmpl-1"}) + proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + upstream: Final = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id": "chatcmpl-1"}', + request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"), + ) + request: Final = MagicMock() + request.method = "POST" + request.query_params = {} + request.headers = MagicMock() + request.headers.copy.return_value = {} + module: Final = "litellm.proxy.pass_through_endpoints.pass_through_endpoints" + with ExitStack() as stack: + stack.enter_context( + patch( + f"{module}.HttpPassThroughEndpointHelpers.non_streaming_http_request_handler", + new_callable=AsyncMock, + return_value=upstream, + ) + ) + stack.enter_context(patch(f"{module}._is_streaming_response", return_value=False)) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging)) + stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) + stack.enter_context( + patch( + f"{module}.pass_through_endpoint_logging", + MagicMock(pass_through_async_success_handler=AsyncMock()), + ) + ) + stack.enter_context(patch(f"{module}.read_request_body", new_callable=AsyncMock, return_value=body)) + stack.enter_context(patch(f"{module}.safe_get_request_headers", return_value={})) + await pass_through_request( + request=request, + target="https://api.openai.com/v1/chat/completions", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + stream=False, + inherit_auth_guardrails=inherit, + ) + return pre_call_data, proxy_logging.post_call_success_hook + + @staticmethod + def _chat_body(**extra: object) -> dict[str, object]: + return {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], **extra} + + @pytest.mark.asyncio + async def test_key_team_and_project_guardrails_reach_pre_and_post_call(self): + auth: Final = UserAPIKeyAuth( + api_key="sk-test", + metadata={"guardrails": ["key-guard"]}, + team_metadata={"guardrails": ["team-guard"]}, + project_metadata={"guardrails": ["project-guard"]}, + ) + + pre_call, post_call = await self._run(self._chat_body(), inherit=True, user_api_key_dict=auth) + + assert sorted(pre_call[0]["metadata"]["guardrails"]) == ["key-guard", "project-guard", "team-guard"] + post_call.assert_awaited_once() + assert sorted(post_call.call_args.kwargs["data"]["metadata"]["guardrails"]) == [ + "key-guard", + "project-guard", + "team-guard", + ] + + @pytest.mark.asyncio + async def test_request_body_guardrails_are_added_to_the_key_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, _ = await self._run( + self._chat_body(guardrails=["body-guard"]), inherit=True, user_api_key_dict=auth + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"] + assert "guardrails" not in pre_call[0] + + @pytest.mark.asyncio + async def test_empty_request_body_guardrails_keep_the_key_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, _ = await self._run(self._chat_body(guardrails=[]), inherit=True, user_api_key_dict=auth) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard"] + + @pytest.mark.asyncio + async def test_routes_without_inheritance_keep_key_guardrails_off(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, post_call = await self._run(self._chat_body(), inherit=False, user_api_key_dict=auth) + + assert "guardrails" not in pre_call[0].get("metadata", {}) + post_call.assert_not_awaited() + + @pytest.mark.asyncio + async def test_request_without_a_body_gets_no_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, post_call = await self._run({}, inherit=True, user_api_key_dict=auth) + + assert "guardrails" not in pre_call[0].get("metadata", {}) + post_call.assert_not_awaited() From 205026dbe4fb486b2b529288285e055650687d5c Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 00:50:37 -0700 Subject: [PATCH 02/29] fix(proxy): relay pass-through responses untouched when no inherited guardrail runs post_call --- .../pass_through_endpoints.py | 20 +++++-- .../test_passthrough_auth_guardrails.py | 43 +++++++++++++- .../test_pass_through_endpoints.py | 58 +++++++++++++++++-- 3 files changed, 109 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 23c1202ea07..61651659556 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -141,6 +141,7 @@ from litellm.proxy.utils import normalize_route_for_root_path from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str from litellm.types import utils as types_utils +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.litellm_params import ProxyRequestState, wire_names from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -863,11 +864,11 @@ async def _inherit_auth_guardrails( configured_guardrails: Mapping[str, bool] | None, inherit: bool, ) -> Mapping[str, bool] | Sequence[object] | None: - """Return the guardrails a pass-through request runs. + """Return the guardrails a pass-through request runs after the provider responds. The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, - project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` and returned, or None - when there are none. + project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook, + and returned only when at least one of them runs post_call, so the response is relayed untouched otherwise. """ if configured_guardrails or not inherit or not parsed_body: return configured_guardrails @@ -879,7 +880,18 @@ async def _inherit_auth_guardrails( user_api_key_dict=user_api_key_dict, ) merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") - return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) if isinstance(merged, list) and merged else None + if not isinstance(merged, list) or not merged or not _named_guardrail_runs_post_call(parsed_body): + return None + return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) + + +def _named_guardrail_runs_post_call(data: dict[str, object]) -> bool: + return any( + guardrail.default_on is not True + and guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) + for guardrail in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + if isinstance(guardrail, CustomGuardrail) + ) def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | None) -> None: diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index e92ebcb5594..09b20bf6653 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -18,6 +18,7 @@ from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter GUARDRAIL: Final = "passthrough-auth-guardrail" +GLOBAL_POST_CALL: Final = "passthrough-global-post-call" BLOCK_MARKER: Final = "block-this-prompt" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" @@ -43,8 +44,17 @@ def _verdict(request: Request) -> Reply: return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode()) +_OPENAI_BODY: Final = json.dumps( + _JSON.validate_json(chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False).body), indent=2 +).encode() + + def _openai_reply(_request: Request) -> Reply: - return chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False) + return Reply(body=_OPENAI_BODY) + + +def _vendor_down(_request: Request) -> Reply: + return Reply(status=500, body=b'{"error": "guardrail vendor down"}') def _anthropic_reply(_request: Request) -> Reply: @@ -66,6 +76,7 @@ def _gemini_reply(_request: Request) -> Reply: class AuthGuardrailRig: owned: OwnedProxy vendor: Wire + post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -74,6 +85,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails") with ExitStack() as stack: vendor: Final = stack.enter_context(wire_server(_verdict)) + post_call_vendor: Final = stack.enter_context(wire_server(_vendor_down)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -97,7 +109,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_base": vendor.url, "api_key": "synthetic-guardrail-key", }, - } + }, + { + "guardrail_name": GLOBAL_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": True, + "api_base": post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, ], } ) @@ -111,7 +133,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "GEMINI_API_KEY": "synthetic-gemini-key", } owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) - yield AuthGuardrailRig(owned, vendor, upstreams) + yield AuthGuardrailRig(owned, vendor, post_call_vendor, upstreams) def _key_with_guardrail(scenario: Scenario, level: Level) -> str: @@ -132,6 +154,7 @@ def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str) def _drain(rig: AuthGuardrailRig) -> None: rig.vendor.drain() + rig.post_call_vendor.drain() for upstream in rig.upstreams.values(): upstream.drain() @@ -191,3 +214,17 @@ def test_request_without_a_body_is_not_sent_to_the_guardrail(rig: AuthGuardrailR assert response.status_code == 200, response.text assert rig.vendor.drain() == () assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +def test_pre_call_guardrail_relays_the_provider_reply_untouched(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert response.content == _OPENAI_BODY + assert len(rig.vendor.drain()) == 1 + assert rig.post_call_vendor.drain() == () diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 294ad45d795..c70e66329ad 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11360,8 +11360,25 @@ class TestInheritAuthGuardrails: @staticmethod async def _run( - body: dict[str, object], *, inherit: bool, user_api_key_dict: UserAPIKeyAuth + body: dict[str, object], + *, + inherit: bool, + user_api_key_dict: UserAPIKeyAuth, + guardrails: tuple[CustomGuardrail, ...] = (), ) -> tuple[list[dict[str, object]], AsyncMock]: + _, pre_call_data, post_call = await TestInheritAuthGuardrails._run_with_response( + body, inherit=inherit, user_api_key_dict=user_api_key_dict, guardrails=guardrails + ) + return pre_call_data, post_call + + @staticmethod + async def _run_with_response( + body: dict[str, object], + *, + inherit: bool, + user_api_key_dict: UserAPIKeyAuth, + guardrails: tuple[CustomGuardrail, ...] = (), + ) -> tuple[Response, list[dict[str, object]], AsyncMock]: pre_call_data: Final[list[dict[str, object]]] = [] async def _record_pre_call(**kwargs: object) -> object: @@ -11375,7 +11392,7 @@ class TestInheritAuthGuardrails: upstream: Final = httpx.Response( status_code=200, headers={"content-type": "application/json"}, - content=b'{"id": "chatcmpl-1"}', + content=b'{\n "id":"chatcmpl-1"\n}', request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"), ) request: Final = MagicMock() @@ -11403,7 +11420,8 @@ class TestInheritAuthGuardrails: ) stack.enter_context(patch(f"{module}.read_request_body", new_callable=AsyncMock, return_value=body)) stack.enter_context(patch(f"{module}.safe_get_request_headers", return_value={})) - await pass_through_request( + stack.enter_context(patch.object(litellm, "callbacks", list(guardrails))) + response: Final = await pass_through_request( request=request, target="https://api.openai.com/v1/chat/completions", custom_headers={}, @@ -11411,7 +11429,11 @@ class TestInheritAuthGuardrails: stream=False, inherit_auth_guardrails=inherit, ) - return pre_call_data, proxy_logging.post_call_success_hook + return response, pre_call_data, proxy_logging.post_call_success_hook + + @staticmethod + def _guardrail(name: str, event_hook: GuardrailEventHooks, *, default_on: bool = False) -> CustomGuardrail: + return CustomGuardrail(guardrail_name=name, event_hook=event_hook, default_on=default_on) @staticmethod def _chat_body(**extra: object) -> dict[str, object]: @@ -11426,7 +11448,15 @@ class TestInheritAuthGuardrails: project_metadata={"guardrails": ["project-guard"]}, ) - pre_call, post_call = await self._run(self._chat_body(), inherit=True, user_api_key_dict=auth) + pre_call, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("team-guard", GuardrailEventHooks.post_call), + ), + ) assert sorted(pre_call[0]["metadata"]["guardrails"]) == ["key-guard", "project-guard", "team-guard"] post_call.assert_awaited_once() @@ -11436,6 +11466,24 @@ class TestInheritAuthGuardrails: "team-guard", ] + @pytest.mark.asyncio + async def test_pre_call_only_guardrails_relay_the_response_untouched(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + response, pre_call, post_call = await self._run_with_response( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True), + ), + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard"] + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + @pytest.mark.asyncio async def test_request_body_guardrails_are_added_to_the_key_guardrails(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) From e1f8fb6c4687d3d264358465f9b8d8cd535f56c9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 01:10:32 -0700 Subject: [PATCH 03/29] fix(proxy): type the post-call gate input as a Mapping --- .../proxy/pass_through_endpoints/pass_through_endpoints.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 61651659556..70d677dfe4c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -885,10 +885,12 @@ async def _inherit_auth_guardrails( return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) -def _named_guardrail_runs_post_call(data: dict[str, object]) -> bool: +def _named_guardrail_runs_post_call(data: Mapping[str, object]) -> bool: return any( guardrail.default_on is not True - and guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) + and guardrail.should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + data=data, event_type=GuardrailEventHooks.post_call + ) for guardrail in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) if isinstance(guardrail, CustomGuardrail) ) From 9df7aa5cdb60fe435d8b17ae3f843813eb9475ae Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 01:40:19 -0700 Subject: [PATCH 04/29] fix(proxy): keep inherited post-call guardrails flat so model guardrails merge with them --- .../pass_through_endpoints.py | 2 +- .../test_passthrough_auth_guardrails.py | 67 +++++++++++++++++-- .../test_pass_through_endpoints.py | 16 ++++- 3 files changed, 78 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 07353793b86..cb37c6f30c3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -888,7 +888,7 @@ async def _inherit_auth_guardrails( merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") if not isinstance(merged, list) or not merged or not _named_guardrail_runs_post_call(parsed_body): return None - return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) + return _GUARDRAIL_ENTRIES.validate_python(merged) def _named_guardrail_runs_post_call(data: Mapping[str, object]) -> bool: diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 09b20bf6653..221b9b4c885 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -19,12 +19,15 @@ from pydantic import JsonValue, TypeAdapter GUARDRAIL: Final = "passthrough-auth-guardrail" GLOBAL_POST_CALL: Final = "passthrough-global-post-call" +KEY_POST_CALL: Final = "passthrough-key-post-call" +GUARDED_MODEL: Final = "passthrough-guarded-model" BLOCK_MARKER: Final = "block-this-prompt" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) Level = Literal["key", "team", "project"] @@ -53,8 +56,20 @@ def _openai_reply(_request: Request) -> Reply: return Reply(body=_OPENAI_BODY) -def _vendor_down(_request: Request) -> Reply: - return Reply(status=500, body=b'{"error": "guardrail vendor down"}') +def _lakera_clean(_request: Request) -> Reply: + return Reply(body=json.dumps({"flagged": False, "breakdown": [], "payload": []}).encode()) + + +def _lakera_flagged(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "flagged": True, + "breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 0}], + "payload": [], + } + ).encode() + ) def _anthropic_reply(_request: Request) -> Reply: @@ -77,6 +92,7 @@ class AuthGuardrailRig: owned: OwnedProxy vendor: Wire post_call_vendor: Wire + key_post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -85,7 +101,8 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails") with ExitStack() as stack: vendor: Final = stack.enter_context(wire_server(_verdict)) - post_call_vendor: Final = stack.enter_context(wire_server(_vendor_down)) + post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) + key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -99,6 +116,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I yaml.safe_dump( { **config, + "model_list": [ + *_JSON_LIST.validate_python(config.get("model_list") or []), + { + "model_name": GUARDED_MODEL, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "synthetic-openai-key", + "guardrails": [GUARDRAIL], + }, + }, + ], "guardrails": [ { "guardrail_name": GUARDRAIL, @@ -110,10 +138,20 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": KEY_POST_CALL, + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "post_call", + "default_on": False, + "api_base": key_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GLOBAL_POST_CALL, "litellm_params": { - "guardrail": "generic_guardrail_api", + "guardrail": "lakera_v2", "mode": "post_call", "default_on": True, "api_base": post_call_vendor.url, @@ -133,7 +171,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "GEMINI_API_KEY": "synthetic-gemini-key", } owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) - yield AuthGuardrailRig(owned, vendor, post_call_vendor, upstreams) + yield AuthGuardrailRig(owned, vendor, post_call_vendor, key_post_call_vendor, upstreams) def _key_with_guardrail(scenario: Scenario, level: Level) -> str: @@ -155,6 +193,7 @@ def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str) def _drain(rig: AuthGuardrailRig) -> None: rig.vendor.drain() rig.post_call_vendor.drain() + rig.key_post_call_vendor.drain() for upstream in rig.upstreams.values(): upstream.drain() @@ -228,3 +267,21 @@ def test_pre_call_guardrail_relays_the_provider_reply_untouched(rig: AuthGuardra assert response.content == _OPENAI_BODY assert len(rig.vendor.drain()) == 1 assert rig.post_call_vendor.drain() == () + + +def test_key_post_call_guardrail_runs_on_a_model_with_its_own_guardrails(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [KEY_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": GUARDED_MODEL}, + key, + ) + + assert response.status_code == 400, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 3ebd2f32953..2f2955d063f 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -68,7 +68,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, _merge_guardrails_with_existing from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes @@ -11829,6 +11829,20 @@ class TestInheritAuthGuardrails: "team-guard", ] + @pytest.mark.asyncio + async def test_post_call_guardrails_stay_named_after_a_model_guardrail_merge(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + merged: Final = _merge_guardrails_with_existing(post_call.call_args.kwargs["data"], ["model-guard"]) + assert merged["metadata"]["guardrails"] == ["key-guard", "model-guard"] + @pytest.mark.asyncio async def test_pre_call_only_guardrails_relay_the_response_untouched(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) From 9a870cfab730b0c709ddab42f5b2a12ef79f8219 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 01:51:21 -0700 Subject: [PATCH 05/29] fix(proxy): type configured pass-through guardrails as entries --- .../proxy/pass_through_endpoints/pass_through_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cb37c6f30c3..8ec0a89e77d 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -867,9 +867,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): async def _inherit_auth_guardrails( parsed_body: dict[str, object] | None, # mutable-ok: merges the guardrails into the live request body user_api_key_dict: UserAPIKeyAuth, - configured_guardrails: Mapping[str, bool] | None, + configured_guardrails: Sequence[object] | None, inherit: bool, -) -> Mapping[str, bool] | Sequence[object] | None: +) -> Sequence[object] | None: """Return the guardrails a pass-through request runs after the provider responds. The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, From 96ddbda50739c5757883a287c0ec93d9b4d05daa Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 10:09:55 -0700 Subject: [PATCH 06/29] fix(proxy): run the post-call pass when the key names a default_on post_call guardrail --- .../pass_through_endpoints.py | 11 ++++++----- .../test_pass_through_endpoints.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 8ec0a89e77d..1e701c9c0d5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -54,6 +54,7 @@ from litellm.constants import ( ) from litellm.integrations.custom_guardrail import ( CustomGuardrail, + guardrail_entry_names, guardrail_request_data_with_streaming, without_server_streaming_classification, ) @@ -886,14 +887,14 @@ async def _inherit_auth_guardrails( user_api_key_dict=user_api_key_dict, ) merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") - if not isinstance(merged, list) or not merged or not _named_guardrail_runs_post_call(parsed_body): - return None - return _GUARDRAIL_ENTRIES.validate_python(merged) + entries: Final = _GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else [] + return entries if entries and _named_guardrail_runs_post_call(parsed_body, entries) else None -def _named_guardrail_runs_post_call(data: Mapping[str, object]) -> bool: +def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequence[object]) -> bool: + names: Final = frozenset(chain.from_iterable(map(guardrail_entry_names, entries))) return any( - guardrail.default_on is not True + guardrail.guardrail_name in names and guardrail.should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict data=data, event_type=GuardrailEventHooks.post_call ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 2f2955d063f..3cac8a87c9f 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11861,6 +11861,20 @@ class TestInheritAuthGuardrails: post_call.assert_not_awaited() assert response.body == b'{\n "id":"chatcmpl-1"\n}' + @pytest.mark.asyncio + async def test_a_named_default_on_post_call_guardrail_still_checks_the_response(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["global-post"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_awaited_once() + assert post_call.call_args.kwargs["data"]["metadata"]["guardrails"] == ["global-post"] + @pytest.mark.asyncio async def test_request_body_guardrails_are_added_to_the_key_guardrails(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) From c0116078b939112f5aecb6577483bb2f6119dc06 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 12:13:16 -0700 Subject: [PATCH 07/29] fix(proxy): pass the logging object to pass-through post-call guardrails so apply_guardrail ones resolve the call type --- .../pass_through_endpoints.py | 3 + .../test_passthrough_auth_guardrails.py | 100 ++++++++++++++++-- .../test_pass_through_endpoints.py | 14 +++ 3 files changed, 107 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 1e701c9c0d5..a98b5eda401 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1984,6 +1984,9 @@ async def pass_through_request( **existing_metadata, "guardrails": post_call_guardrails, } + # The kwargs split popped the logging object off _parsed_body; unified (apply_guardrail) guardrails + # resolve a pass-through response's call type from it when the route has no API_ROUTE_TO_CALL_TYPES entry. + hook_data["litellm_logging_obj"] = logging_obj post_call_guardrail_data = hook_data response_body = await proxy_logging_obj.post_call_success_hook( data=hook_data, diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 221b9b4c885..c0fa3a233b8 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -20,11 +20,13 @@ from pydantic import JsonValue, TypeAdapter GUARDRAIL: Final = "passthrough-auth-guardrail" GLOBAL_POST_CALL: Final = "passthrough-global-post-call" KEY_POST_CALL: Final = "passthrough-key-post-call" +GENERIC_POST_CALL: Final = "passthrough-generic-post-call" GUARDED_MODEL: Final = "passthrough-guarded-model" BLOCK_MARKER: Final = "block-this-prompt" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" +CUSTOM_ROUTE: Final = "/passthrough-auth-custom" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) @@ -72,19 +74,27 @@ def _lakera_flagged(_request: Request) -> Reply: ) +_ANTHROPIC_BODY: Final = json.dumps( + _JSON.validate_json(message_json("msg_passthrough_auth", "claude-sonnet-5-5", "Hi")), indent=2 +).encode() +_GEMINI_BODY: Final = json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 1, "totalTokenCount": 5}, + }, + indent=2, +).encode() +_UPSTREAM_BODIES: Final = MappingProxyType( + {OPENAI_ROUTE: _OPENAI_BODY, ANTHROPIC_ROUTE: _ANTHROPIC_BODY, GEMINI_ROUTE: _GEMINI_BODY} +) + + def _anthropic_reply(_request: Request) -> Reply: - return Reply(body=message_json("msg_passthrough_auth", "claude-sonnet-5-5", "Hi")) + return Reply(body=_ANTHROPIC_BODY) def _gemini_reply(_request: Request) -> Reply: - return Reply( - body=json.dumps( - { - "candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}, "finishReason": "STOP"}], - "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 1, "totalTokenCount": 5}, - } - ).encode() - ) + return Reply(body=_GEMINI_BODY) @dataclass(frozen=True, slots=True) @@ -93,6 +103,7 @@ class AuthGuardrailRig: vendor: Wire post_call_vendor: Wire key_post_call_vendor: Wire + generic_post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -103,6 +114,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I vendor: Final = stack.enter_context(wire_server(_verdict)) post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) + generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -111,11 +123,22 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I } ) config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + general_settings: Final = _JSON.validate_python(config.get("general_settings") or {}) path: Final = directory / "passthrough-auth-guardrails.yaml" path.write_text( yaml.safe_dump( { **config, + "general_settings": { + **general_settings, + "pass_through_endpoints": [ + { + "path": CUSTOM_ROUTE, + "target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions", + "auth": True, + } + ], + }, "model_list": [ *_JSON_LIST.validate_python(config.get("model_list") or []), { @@ -148,6 +171,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": GENERIC_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": generic_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GLOBAL_POST_CALL, "litellm_params": { @@ -171,7 +204,9 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "GEMINI_API_KEY": "synthetic-gemini-key", } owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) - yield AuthGuardrailRig(owned, vendor, post_call_vendor, key_post_call_vendor, upstreams) + yield AuthGuardrailRig( + owned, vendor, post_call_vendor, key_post_call_vendor, generic_post_call_vendor, upstreams + ) def _key_with_guardrail(scenario: Scenario, level: Level) -> str: @@ -194,6 +229,7 @@ def _drain(rig: AuthGuardrailRig) -> None: rig.vendor.drain() rig.post_call_vendor.drain() rig.key_post_call_vendor.drain() + rig.generic_post_call_vendor.drain() for upstream in rig.upstreams.values(): upstream.drain() @@ -285,3 +321,47 @@ def test_key_post_call_guardrail_runs_on_a_model_with_its_own_guardrails(rig: Au assert response.status_code == 400, response.text assert len(rig.key_post_call_vendor.drain()) == 1 assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +@pytest.mark.parametrize("route", [ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_pre_call_guardrail_relays_every_provider_reply_untouched(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "team") + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert response.content == _UPSTREAM_BODIES[route] + assert len(rig.vendor.drain()) == 1 + assert rig.post_call_vendor.drain() == () + + +def test_custom_pass_through_endpoint_keeps_its_opt_in(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key( + team_id=scenario.team(), + metadata={"guardrails": [GUARDRAIL], "allowed_passthrough_routes": [CUSTOM_ROUTE]}, + ) + _drain(rig) + + response: Final = gateway.request("POST", CUSTOM_ROUTE, _route_body(OPENAI_ROUTE, BLOCK_MARKER), key=key) + + assert response.status_code == 200, response.text + assert rig.vendor.drain() == () + assert [BLOCK_MARKER.encode() in request.body for request in rig.upstreams[OPENAI_ROUTE].drain()] == [True] + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_key_apply_guardrail_post_call_scans_the_provider_reply(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert [b'"Hi' in request.body for request in rig.generic_post_call_vendor.drain()] == [True] diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 3cac8a87c9f..552847447cf 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11875,6 +11875,20 @@ class TestInheritAuthGuardrails: post_call.assert_awaited_once() assert post_call.call_args.kwargs["data"]["metadata"]["guardrails"] == ["global-post"] + @pytest.mark.asyncio + async def test_post_call_guardrails_get_the_request_logging_object(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-post"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-post", GuardrailEventHooks.post_call),), + ) + + post_call.assert_awaited_once() + assert isinstance(post_call.call_args.kwargs["data"]["litellm_logging_obj"], LiteLLMLoggingObj) + @pytest.mark.asyncio async def test_request_body_guardrails_are_added_to_the_key_guardrails(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) From 377124588edf3b8e11bb39b9f100d9adefc3b7cd Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 13:32:56 -0700 Subject: [PATCH 08/29] fix(proxy): keep the guardrails a client names in metadata on guarded pass-through keys and merge guardrail lists in linear time --- litellm/integrations/custom_guardrail.py | 27 ++++++++++++------- .../pass_through_endpoints.py | 15 ++++++++--- .../test_passthrough_auth_guardrails.py | 16 +++++++++++ .../test_pass_through_endpoints.py | 22 +++++++++++++++ 4 files changed, 68 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index dfe508f3630..3b444519fc6 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -5,7 +5,7 @@ import os import secrets from collections.abc import Iterable, Mapping, Sequence from datetime import datetime -from itertools import chain +from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -219,11 +219,16 @@ def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: any plain name in either order, else its first plain name stays""" entries: Final = tuple(chain.from_iterable(entry_lists)) settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries) - named_with_settings: Final = frozenset[str]().union(*settings_names) + first_plain_index: Final = _first_indexes( + (entry, index) for index, entry in enumerate(entries) if isinstance(entry, str) + ) + first_settings_index: Final = _first_indexes( + chain.from_iterable(zip(names, repeat(index)) for index, names in enumerate(settings_names)) + ) return [ entry for index, entry in enumerate(entries) - if not _superseded_guardrail_entry(entries, settings_names, named_with_settings, index) + if not _superseded_guardrail_entry(entry, index, settings_names[index], first_plain_index, first_settings_index) ] @@ -239,16 +244,20 @@ def _settings_entry_names(entry: object) -> frozenset[str]: return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset() +def _first_indexes(pairs: Iterable[tuple[str, int]]) -> Mapping[str, int]: + return MappingProxyType(dict(reversed(tuple(pairs)))) + + def _superseded_guardrail_entry( - entries: tuple[object, ...], - settings_names: tuple[frozenset[str], ...], - named_with_settings: frozenset[str], + entry: object, index: int, + settings_names: frozenset[str], + first_plain_index: Mapping[str, int], + first_settings_index: Mapping[str, 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]) + return entry in first_settings_index or first_plain_index[entry] < index + return bool(settings_names) and all(first_settings_index[name] < index for name in settings_names) def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a98b5eda401..8a6bd0fe7c5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -56,6 +56,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, guardrail_entry_names, guardrail_request_data_with_streaming, + merged_guardrail_entries, without_server_streaming_classification, ) from litellm.integrations.custom_logger import CustomLogger @@ -875,19 +876,27 @@ async def _inherit_auth_guardrails( The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook, - and returned only when at least one of them runs post_call, so the response is relayed untouched otherwise. + with the ones the client named in ``metadata.guardrails`` kept after them, and returned only when at least one + of them runs post_call, so the response is relayed untouched otherwise. """ if configured_guardrails or not inherit or not parsed_body: return configured_guardrails if not isinstance(parsed_body.get("metadata"), dict): parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata + client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") await move_guardrails_to_metadata( data=parsed_body, _metadata_variable_name="metadata", user_api_key_dict=user_api_key_dict, ) - merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") - entries: Final = _GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else [] + metadata: Final = parsed_body["metadata"] + merged: Final = (as_str_mapping(metadata) or {}).get("guardrails") + entries: Final = merged_guardrail_entries( + _GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else (), + _GUARDRAIL_ENTRIES.validate_python(client_named) if isinstance(client_named, list) else (), + ) + if entries and isinstance(metadata, dict): + metadata["guardrails"] = entries return entries if entries and _named_guardrail_runs_post_call(parsed_body, entries) else None diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index c0fa3a233b8..85c31a5e2bc 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -365,3 +365,19 @@ def test_key_apply_guardrail_post_call_scans_the_provider_reply(rig: AuthGuardra assert response.status_code == 200, response.text assert [b'"Hi' in request.body for request in rig.generic_post_call_vendor.drain()] == [True] + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, route, {**_route_body(route, BLOCK_MARKER), "metadata": {"guardrails": [GUARDRAIL]}}, key + ) + + assert response.status_code == 400, response.text + assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True] + assert rig.upstreams[route].drain() == () diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 552847447cf..fe3ba70d127 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11900,6 +11900,28 @@ class TestInheritAuthGuardrails: assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"] assert "guardrails" not in pre_call[0] + @pytest.mark.asyncio + async def test_guardrails_the_client_names_in_metadata_still_run_for_a_guarded_key(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, post_call = await self._run( + self._chat_body(metadata={"guardrails": ["client-pre", "client-post"]}), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("client-post", GuardrailEventHooks.post_call), + ), + ) + + assert pre_call[0]["metadata"] == {"guardrails": ["key-guard", "client-pre", "client-post"]} + post_call.assert_awaited_once() + assert post_call.call_args.kwargs["data"]["metadata"]["guardrails"] == [ + "key-guard", + "client-pre", + "client-post", + ] + @pytest.mark.asyncio async def test_empty_request_body_guardrails_keep_the_key_guardrails(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) From 8ec5c7ddf22f0490a7776de2496df5cc11da701c Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 14:12:32 -0700 Subject: [PATCH 09/29] perf(proxy): record applied guardrails in time linear in the names a request carries _process_guardrail_metadata checked each requested name against the growing applied_guardrails list, so a body naming 50k guardrails spent seconds of event loop time. Names are now filtered against sets and appended in one pass, with the same order, dedup and metadata bucket as before. --- litellm/proxy/utils.py | 52 ++++++++++---- .../proxy_logging/test_guardrail_pipeline.py | 70 ++++++++----------- 2 files changed, 67 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c7efbf9dacf..6e091391842 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -815,6 +815,29 @@ def _drop_response_steps_that_did_not_run( _drop_applied_guardrails_in_place(data.get(bucket_key), not_run) +def _add_applied_guardrails_in_place( + data: dict[str, object], # mutable-ok: the names are recorded in the live request body's metadata + names: Sequence[str], +) -> None: + """Adds each name the way add_guardrail_to_applied_guardrails_header does, in time linear in the names""" + from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header + + if not names: + return + metadata: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str-keyed bucket + if "applied_guardrails" not in metadata: + metadata["applied_guardrails"] = list(dict.fromkeys(names)) + return + applied: Final = metadata["applied_guardrails"] + if not isinstance(applied, list): + for name in names: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=name) + return + entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names + present: Final = frozenset(name for name in entries if isinstance(name, str)) + entries.extend(dict.fromkeys(name for name in names if name not in present)) + + def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None: if not isinstance(bucket, dict): return @@ -2831,10 +2854,6 @@ class ProxyLogging: policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on another request, so every step is added. """ - from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, - ) - metadata_standard: Final = data.get("metadata") or {} metadata_litellm: Final = data.get("litellm_metadata") or {} @@ -2854,18 +2873,21 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] - not_run: Final = ( - frozenset[str]() - if skip_guardrails - else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload + request: Final = cast("dict[str, object]", data) # cast-ok: request payload + not_run: Final = frozenset[str]() if skip_guardrails else _request_steps_that_have_not_run(request) + already_applied: Final = frozenset( + name + for name in cast("Sequence[object]", applied_guardrails) # cast-ok: checked to be a list above + if isinstance(name, str) + ) + _add_applied_guardrails_in_place( + request, + [ + name + for name in guardrails_in_metadata + if isinstance(name, str) and name not in already_applied and name not in not_run + ], ) - for guardrail_name in guardrails_in_metadata: - if ( - isinstance(guardrail_name, str) - and guardrail_name not in applied_guardrails - and guardrail_name not in not_run - ): - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( self, diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 5fd40696224..a46b56294bf 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import json +import time from collections.abc import AsyncGenerator, Iterator from contextlib import AbstractContextManager, nullcontext from copy import deepcopy @@ -275,54 +276,43 @@ async def test_process_guardrail_callback_enriches_and_reraises_http_exception( # --------------------------------------------------------------------------- -def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypatch): - calls: List[Dict[str, Any]] = [] - - def fake_add(request_data, guardrail_name): - calls.append({"data": request_data, "name": guardrail_name}) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) - data = {"metadata": {"guardrails": ["g1", "g2"]}} +def test_process_guardrail_metadata_adds_each_guardrail_once_in_order(proxy_logging): + data = {"metadata": {"guardrails": ["g1", "g2", "g1"]}} proxy_logging._process_guardrail_metadata(data) - snapshot = { - "call_count": len(calls), - "first_name": calls[0]["name"], - "second_name": calls[1]["name"], - "data_passed_is_input": all(c["data"] is data for c in calls), - } - assert snapshot == { - "call_count": 2, - "first_name": "g1", - "second_name": "g2", - "data_passed_is_input": True, - } + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeypatch): - calls: List[str] = [] - - def fake_add(request_data, guardrail_name): - calls.append(guardrail_name) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) +def test_process_guardrail_metadata_skips_already_applied(proxy_logging): data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} proxy_logging._process_guardrail_metadata(data) - assert calls == ["g2"] + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypatch): - from litellm.proxy.common_utils import callback_utils +def test_process_guardrail_metadata_records_in_the_litellm_metadata_bucket(proxy_logging): + data = {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2"]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2", "g1"]}} - monkeypatch.setattr( - callback_utils, - "add_guardrail_to_applied_guardrails_header", - MagicMock(side_effect=AssertionError("should not be called")), - ) - proxy_logging._process_guardrail_metadata({}) + +def test_process_guardrail_metadata_handles_many_guardrail_names_quickly(proxy_logging): + names = [f"guard-{i}" for i in range(50_000)] + data = {"metadata": {"guardrails": names, "applied_guardrails": []}} + started = time.perf_counter() + proxy_logging._process_guardrail_metadata(data) + assert time.perf_counter() - started < 2 + assert data["metadata"]["applied_guardrails"] == names + + +def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging): + data: Dict[str, Any] = {} + proxy_logging._process_guardrail_metadata(data) + assert data == {} + + +def test_process_guardrail_metadata_with_only_settings_entries_leaves_data_unchanged(proxy_logging): + data = {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): From ec52601580f3625def991113de7b0175274bca07 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 10:51:51 -0800 Subject: [PATCH 10/29] fix(proxy): keep pass-through post-call parity with what main and native run - A dict in the client's metadata.guardrails names guardrails by its keys, so a guarded key keeps them instead of dropping them - The inherited post-call pass reads key and team global-guardrail opt-outs, as native does - Router model guardrails no longer run on pass-through traffic that only shares the model name - merged_guardrail_entries walks the entries once, so a large list costs a fraction of before --- litellm/integrations/custom_guardrail.py | 49 +++----- .../pass_through_endpoints.py | 68 +++++++++--- litellm/proxy/utils.py | 10 +- .../test_passthrough_auth_guardrails.py | 105 +++++++++++++++++- .../integrations/test_custom_guardrail.py | 20 ++++ .../test_pass_through_endpoints.py | 73 ++++++++++++ .../test_passthrough_post_call_guardrails.py | 25 ++++- .../test_post_call_success_hook.py | 21 ++++ 8 files changed, 314 insertions(+), 57 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 3b444519fc6..965cfbb9ee8 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -5,7 +5,7 @@ import os import secrets from collections.abc import Iterable, Mapping, Sequence from datetime import datetime -from itertools import chain, repeat +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -211,55 +211,32 @@ def guardrail_request_data_with_streaming( } -_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) - - def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over any plain name in either order, else its first plain name stays""" entries: Final = tuple(chain.from_iterable(entry_lists)) - settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries) - first_plain_index: Final = _first_indexes( - (entry, index) for index, entry in enumerate(entries) if isinstance(entry, str) + named_with_settings: Final = frozenset( + chain.from_iterable(guardrail_entry_names(entry) for entry in entries if isinstance(entry, dict)) ) - first_settings_index: Final = _first_indexes( - chain.from_iterable(zip(names, repeat(index)) for index, names in enumerate(settings_names)) - ) - return [ - entry - for index, entry in enumerate(entries) - if not _superseded_guardrail_entry(entry, index, settings_names[index], first_plain_index, first_settings_index) - ] + seen: Final[set[str]] = set() # mutable-ok: local to this walk + merged: Final[list[object]] = [] # mutable-ok: local to this walk + for entry in entries: + names = guardrail_entry_names(entry) + if (names and seen.issuperset(names)) or (isinstance(entry, str) and entry in named_with_settings): + continue + merged.append(entry) + seen.update(names) + return merged def guardrail_entry_names(entry: object) -> tuple[str, ...]: if isinstance(entry, str): return (entry,) if isinstance(entry, dict): - return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return tuple(cast("Mapping[str, object]", entry)) # cast-ok: a guardrail settings entry is keyed by name return () -def _settings_entry_names(entry: object) -> frozenset[str]: - return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset() - - -def _first_indexes(pairs: Iterable[tuple[str, int]]) -> Mapping[str, int]: - return MappingProxyType(dict(reversed(tuple(pairs)))) - - -def _superseded_guardrail_entry( - entry: object, - index: int, - settings_names: frozenset[str], - first_plain_index: Mapping[str, int], - first_settings_index: Mapping[str, int], -) -> bool: - if isinstance(entry, str): - return entry in first_settings_index or first_plain_index[entry] < index - return bool(settings_names) and all(first_settings_index[name] < index for name in settings_names) - - 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/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a554cd13a09..8cf8a2ab40e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -115,6 +115,7 @@ from litellm.proxy.common_request_processing import ( open_sse_before_first_byte, resolve_litellm_call_id, ) +from litellm.proxy.common_utils.callback_utils import strip_callback_config from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -897,13 +898,52 @@ async def _inherit_auth_guardrails( ) metadata: Final = parsed_body["metadata"] merged: Final = (as_str_mapping(metadata) or {}).get("guardrails") + if not isinstance(merged, list): + return None entries: Final = merged_guardrail_entries( - _GUARDRAIL_ENTRIES.validate_python(merged) if isinstance(merged, list) else (), - _GUARDRAIL_ENTRIES.validate_python(client_named) if isinstance(client_named, list) else (), + _GUARDRAIL_ENTRIES.validate_python(merged), _client_named_guardrail_entries(client_named) ) if entries and isinstance(metadata, dict): metadata["guardrails"] = entries - return entries if entries and _named_guardrail_runs_post_call(parsed_body, entries) else None + admin_scoped_body: Final = { + **parsed_body, + "metadata": {**(as_str_mapping(metadata) or {}), **_admin_guardrail_metadata(user_api_key_dict)}, + } + return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None + + +def _client_named_guardrail_entries(client_named: object) -> Sequence[object]: + """The guardrails ``metadata.guardrails`` names: its entries, or the keys of the dict form the guardrails read""" + if isinstance(client_named, list): + return _GUARDRAIL_ENTRIES.validate_python(client_named) + return tuple(as_str_mapping(client_named) or ()) + + +def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: + """The key and team metadata ``should_run_guardrail`` reads a global guardrail opt-out from, as the native routes set it""" + return MappingProxyType( + { + "user_api_key_metadata": strip_callback_config(user_api_key_dict.metadata), + "user_api_key_team_metadata": strip_callback_config(user_api_key_dict.team_metadata), + } + ) + + +def _post_call_hook_data( + parsed_body: Mapping[str, object] | None, + guardrails: Sequence[object], + logging_obj: LiteLLMLoggingObj, + admin_metadata: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: the guardrails record their results in this request payload + """The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails, + which ``should_run_guardrail`` checks, and the logging object, which unified (apply_guardrail) guardrails resolve a + pass-through response's call type from when the route has no API_ROUTE_TO_CALL_TYPES entry""" + body: Final = parsed_body or {} + return { + **body, + "metadata": {**(as_str_mapping(body.get("metadata")) or {}), "guardrails": guardrails, **admin_metadata}, + "litellm_logging_obj": logging_obj, + } def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequence[object]) -> bool: @@ -1987,26 +2027,18 @@ async def pass_through_request( ) if response.status_code < 400 and response_body is not None and post_call_guardrails: - # Build an enriched data dict: _parsed_body has been stripped of - # `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint, - # so we re-attach the configured guardrails here so should_run_guardrail - # sees them. - hook_data: Final = dict(_parsed_body or {}) - existing_metadata = hook_data.get("metadata") - if not isinstance(existing_metadata, dict): - existing_metadata = {} - hook_data["metadata"] = { - **existing_metadata, - "guardrails": post_call_guardrails, - } - # The kwargs split popped the logging object off _parsed_body; unified (apply_guardrail) guardrails - # resolve a pass-through response's call type from it when the route has no API_ROUTE_TO_CALL_TYPES entry. - hook_data["litellm_logging_obj"] = logging_obj + hook_data: Final = _post_call_hook_data( + cast(dict[str, object] | None, _parsed_body), # cast-ok: json + post_call_guardrails, + logging_obj, + _admin_guardrail_metadata(user_api_key_dict) if inherit_auth_guardrails else MappingProxyType({}), + ) post_call_guardrail_data = hook_data response_body = await proxy_logging_obj.post_call_success_hook( data=hook_data, user_api_key_dict=user_api_key_dict, response=response_body, + merge_model_level_guardrails=not inherit_auth_guardrails, ) if isinstance(response_body, dict): content = json.dumps(response_body).encode("utf-8") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2fe2c46f2ad..516f5eac4d4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4343,6 +4343,7 @@ class ProxyLogging: data: dict, response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth, + merge_model_level_guardrails: bool = True, ): """ Allow user to modify outgoing data @@ -4352,6 +4353,9 @@ class ProxyLogging: 2. /embeddings 3. /image/generation 4. /files + + ``merge_model_level_guardrails`` is off for a request that never reaches a router deployment, so the guardrails + of a model that merely shares the body's ``model`` name don't run on it. """ from litellm.proxy.proxy_server import llm_router @@ -4369,7 +4373,11 @@ class ProxyLogging: guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run - guardrail_data: Final = check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + guardrail_data: Final = ( + check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + if merge_model_level_guardrails + else data + ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( callback diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 85c31a5e2bc..6a1f86ee018 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -21,7 +21,9 @@ GUARDRAIL: Final = "passthrough-auth-guardrail" GLOBAL_POST_CALL: Final = "passthrough-global-post-call" KEY_POST_CALL: Final = "passthrough-key-post-call" GENERIC_POST_CALL: Final = "passthrough-generic-post-call" +MODEL_POST_CALL: Final = "passthrough-model-post-call" GUARDED_MODEL: Final = "passthrough-guarded-model" +POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model" BLOCK_MARKER: Final = "block-this-prompt" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" @@ -104,6 +106,7 @@ class AuthGuardrailRig: post_call_vendor: Wire key_post_call_vendor: Wire generic_post_call_vendor: Wire + model_post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -115,6 +118,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) + model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -149,6 +153,14 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "guardrails": [GUARDRAIL], }, }, + { + "model_name": POST_CALL_GUARDED_MODEL, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "synthetic-openai-key", + "guardrails": [MODEL_POST_CALL], + }, + }, ], "guardrails": [ { @@ -181,6 +193,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": MODEL_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": model_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GLOBAL_POST_CALL, "litellm_params": { @@ -205,7 +227,13 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I } owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) yield AuthGuardrailRig( - owned, vendor, post_call_vendor, key_post_call_vendor, generic_post_call_vendor, upstreams + owned, + vendor, + post_call_vendor, + key_post_call_vendor, + generic_post_call_vendor, + model_post_call_vendor, + upstreams, ) @@ -230,6 +258,7 @@ def _drain(rig: AuthGuardrailRig) -> None: rig.post_call_vendor.drain() rig.key_post_call_vendor.drain() rig.generic_post_call_vendor.drain() + rig.model_post_call_vendor.drain() for upstream in rig.upstreams.values(): upstream.drain() @@ -381,3 +410,77 @@ def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig assert response.status_code == 400, response.text assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True] assert rig.upstreams[route].drain() == () + + +@pytest.mark.parametrize("settings", [{}, True]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_guardrail_the_client_names_in_the_dict_form_still_runs_for_a_guarded_key( + rig: AuthGuardrailRig, route: str, settings: JsonValue +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, route, {**_route_body(route, BLOCK_MARKER), "metadata": {"guardrails": {GUARDRAIL: settings}}}, key + ) + + assert response.status_code == 400, response.text + assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True] + assert rig.upstreams[route].drain() == () + + +def _key_with_opt_out(scenario: Scenario, opt_out: str) -> str: + guarded: Final[dict[str, JsonValue]] = {"guardrails": [GENERIC_POST_CALL]} + match opt_out: + case "key_disable": + return scenario.key(metadata=guarded, disable_global_guardrails=True) + case "key_opted_out": + return scenario.key(metadata={**guarded, "opted_out_global_guardrails": [GLOBAL_POST_CALL]}) + case "team_disable": + return scenario.key(team_id=scenario.team(disable_global_guardrails=True), metadata=guarded) + case "team_opted_out": + return scenario.key( + team_id=scenario.team(metadata={"opted_out_global_guardrails": [GLOBAL_POST_CALL]}), metadata=guarded + ) + case _: + return scenario.key(metadata=guarded) + + +@pytest.mark.parametrize( + ("opt_out", "global_vendor_calls"), + [("key_disable", 0), ("key_opted_out", 0), ("team_disable", 0), ("team_opted_out", 0), ("none", 1)], +) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, "/v1/chat/completions"]) +def test_post_call_pass_honors_a_global_guardrail_opt_out_as_the_native_route_does( + rig: AuthGuardrailRig, route: str, opt_out: str, global_vendor_calls: int +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_opt_out(scenario, opt_out) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert len(rig.post_call_vendor.drain()) == global_vendor_calls + + +def test_model_guardrail_does_not_run_on_pass_through_traffic_that_names_the_model(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": POST_CALL_GUARDED_MODEL}, + key, + ) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert rig.model_post_call_vendor.drain() == () diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index c32c1b71f75..302a41a16ff 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -4,6 +4,8 @@ import datetime as dt import json import pickle import threading +import time +from collections.abc import Callable from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock @@ -3802,6 +3804,24 @@ def test_default_on_guardrail() -> None: ) +def test_merged_guardrail_entries_stays_close_to_a_plain_dedup_on_a_large_list(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries + + names: Final = [f"g{index}" for index in range(300_000)] + + def _best_cpu_seconds(run: Callable[[], object]) -> float: + def _timed() -> float: + started: Final = time.process_time() + run() + return time.process_time() - started + + return min(_timed() for _ in range(3)) + + assert _best_cpu_seconds(lambda: merged_guardrail_entries(names)) < 25 * _best_cpu_seconds( + lambda: list(dict.fromkeys(names)) + ) + + def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): from litellm.integrations.custom_guardrail import merged_guardrail_entries diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 5700335bae8..b82cb8c7f91 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -13368,6 +13368,79 @@ class TestInheritAuthGuardrails: merged: Final = _merge_guardrails_with_existing(post_call.call_args.kwargs["data"], ["model-guard"]) assert merged["metadata"]["guardrails"] == ["key-guard", "model-guard"] + @pytest.mark.asyncio + @pytest.mark.parametrize("settings", [{}, True]) + async def test_guardrail_a_client_names_in_the_dict_form_stays_after_the_keys(self, settings: object): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, _ = await self._run( + {**self._chat_body(), "metadata": {"guardrails": {"client-guard": settings}}}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.pre_call),), + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "client-guard"] + + @pytest.mark.asyncio + async def test_dict_form_stays_as_sent_for_a_key_without_guardrails(self): + pre_call, post_call = await self._run( + {**self._chat_body(), "metadata": {"guardrails": {"client-guard": {}}}}, + inherit=True, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + guardrails=(self._guardrail("client-guard", GuardrailEventHooks.post_call),), + ) + + assert pre_call[0]["metadata"]["guardrails"] == {"client-guard": {}} + post_call.assert_not_awaited() + + @pytest.mark.asyncio + async def test_post_call_pass_reads_opt_outs_from_the_key_and_team_not_the_body(self): + auth: Final = UserAPIKeyAuth( + api_key="sk-test", + metadata={"guardrails": ["key-guard"], "disable_global_guardrails": True}, + team_metadata={"opted_out_global_guardrails": ["global-post"]}, + ) + forged: Final = {"user_api_key_metadata": {"disable_global_guardrails": False}} + + _, post_call = await self._run( + {**self._chat_body(), "metadata": forged}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + metadata: Final = post_call.call_args.kwargs["data"]["metadata"] + assert metadata["user_api_key_metadata"]["disable_global_guardrails"] is True + assert metadata["user_api_key_team_metadata"] == {"opted_out_global_guardrails": ["global-post"]} + + @pytest.mark.asyncio + async def test_inherited_post_call_pass_skips_model_level_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + assert post_call.call_args.kwargs["merge_model_level_guardrails"] is False + + @pytest.mark.asyncio + async def test_named_global_guardrail_the_key_opted_out_of_does_not_start_the_post_call_pass(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"opted_out_global_guardrails": ["global-post"]}) + + response, _, post_call = await self._run_with_response( + {**self._chat_body(), "metadata": {"guardrails": ["global-post"]}}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + @pytest.mark.asyncio async def test_pre_call_only_guardrails_relay_the_response_untouched(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 7b06905c8aa..7b1ce703a4c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -130,6 +130,29 @@ class TestPassthroughPostCallGuardrails: call_kwargs = mock_proxy_logging.post_call_success_hook.call_args assert call_kwargs.kwargs["response"] == _GEMINI_RESPONSE + @patch(_COLLECT, return_value=["rubrik"]) + async def test_configured_endpoint_keeps_model_guardrails_and_reads_no_key_opt_outs( + self, + mock_collect, + ): + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=_GEMINI_RESPONSE) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + + with _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)): + await pass_through_request( + request=_make_mock_request(), + target="https://example.com/v1/generateContent", + custom_headers={"Content-Type": "application/json"}, + user_api_key_dict=_make_user_api_key_dict(), + stream=False, + ) + + call_kwargs = mock_proxy_logging.post_call_success_hook.call_args.kwargs + assert call_kwargs["merge_model_level_guardrails"] is True + assert call_kwargs["data"]["metadata"].keys() == {"guardrails"} + @patch(_COLLECT, return_value=[]) async def test_post_call_success_hook_skipped_when_no_guardrails( self, @@ -216,7 +239,7 @@ class TestPassthroughPostCallGuardrails: block path dropping it.""" mock_response = _make_httpx_response(_GEMINI_RESPONSE) - def _block(*, data, user_api_key_dict, response): + def _block(*, data, user_api_key_dict, response, merge_model_level_guardrails): metadata = data.setdefault("metadata", {}) metadata.setdefault("standard_logging_guardrail_information", []).append( {"guardrail_name": "rubrik", "guardrail_status": "guardrail_intervened"} diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py index 53d8948869f..7f6429750c8 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -59,6 +59,27 @@ async def test_post_call_success_hook_runs_other_callback_and_replaces_response( assert out == new_response +@pytest.mark.asyncio +@pytest.mark.parametrize(("merge", "guardrails_seen"), [(True, ["key-guard", "model-guard"]), (False, ["key-guard"])]) +async def test_post_call_success_hook_merges_model_level_guardrails_unless_told_not_to( + proxy_logging, make_user_api_key_auth, monkeypatch, merge, guardrails_seen +): + router = MagicMock() + router.registry.return_value.deployments.return_value = [{"litellm_params": {"guardrails": ["model-guard"]}}] + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + g = _make_guardrail(should_run=False) + monkeypatch.setattr(litellm, "callbacks", [g]) + + await proxy_logging.post_call_success_hook( + data={"model": "alias", "metadata": {"guardrails": ["key-guard"]}}, + response=MagicMock(), + user_api_key_dict=make_user_api_key_auth(), + merge_model_level_guardrails=merge, + ) + + assert g.should_run_guardrail.call_args.kwargs["data"]["metadata"]["guardrails"] == guardrails_seen + + @pytest.mark.asyncio async def test_post_call_success_hook_guardrail_should_not_run_skipped( proxy_logging, make_user_api_key_auth, monkeypatch From 44adfe43cbe11190292a87791caf36b0c94aee12 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 10:59:58 -0800 Subject: [PATCH 11/29] test(proxy): send the opt-out cells a model the native route can serve --- .../authorization/test_passthrough_auth_guardrails.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 6a1f86ee018..c00109f3132 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -461,7 +461,12 @@ def test_post_call_pass_honors_a_global_guardrail_opt_out_as_the_native_route_do key: Final = _key_with_opt_out(scenario, opt_out) _drain(rig) - response: Final = _send(gateway, route, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + response: Final = _send( + gateway, + route, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": GUARDED_MODEL}, + key, + ) assert response.status_code == 200, response.text assert len(rig.generic_post_call_vendor.drain()) == 1 From b427237942bb5add84b9e375f76794d6f24b04a0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 11:12:31 -0800 Subject: [PATCH 12/29] perf(proxy): dedupe an all-plain guardrail list without walking it entry by entry --- litellm/integrations/custom_guardrail.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 965cfbb9ee8..0aaadc033e5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -215,6 +215,8 @@ def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over any plain name in either order, else its first plain name stays""" entries: Final = tuple(chain.from_iterable(entry_lists)) + if all(isinstance(entry, str) for entry in entries): + return list(dict.fromkeys(entries)) named_with_settings: Final = frozenset( chain.from_iterable(guardrail_entry_names(entry) for entry in entries if isinstance(entry, dict)) ) From c3c165a683121271642423974554d058dd53b3d9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 11:14:45 -0800 Subject: [PATCH 13/29] fix(proxy): run only guardrails in the inherited pass-through post-call pass Other callback hooks now stay off relayed responses, as on main: a hook written for native responses no longer turns a provider's 200 into a 500, and the Langfuse root-span stamp no longer exports prompts and replies past the no-log and redaction settings. --- .../pass_through_endpoints.py | 2 +- litellm/proxy/utils.py | 13 ++++--- .../test_passthrough_auth_guardrails.py | 36 +++++++++++++++++++ .../test_pass_through_endpoints.py | 4 +-- .../test_passthrough_post_call_guardrails.py | 6 ++-- .../test_post_call_success_hook.py | 28 ++++++++++++--- 6 files changed, 72 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 8cf8a2ab40e..d641e6f598e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2038,7 +2038,7 @@ async def pass_through_request( data=hook_data, user_api_key_dict=user_api_key_dict, response=response_body, - merge_model_level_guardrails=not inherit_auth_guardrails, + guardrails_only=inherit_auth_guardrails, ) if isinstance(response_body, dict): content = json.dumps(response_body).encode("utf-8") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 516f5eac4d4..eb36d1ae361 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4343,7 +4343,7 @@ class ProxyLogging: data: dict, response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth, - merge_model_level_guardrails: bool = True, + guardrails_only: bool = False, ): """ Allow user to modify outgoing data @@ -4354,8 +4354,9 @@ class ProxyLogging: 3. /image/generation 4. /files - ``merge_model_level_guardrails`` is off for a request that never reaches a router deployment, so the guardrails - of a model that merely shares the body's ``model`` name don't run on it. + ``guardrails_only`` is for a relayed response that never reached a router deployment and isn't in the shape the + other hooks expect: only the guardrails the request selected run, not the guardrails of a model that merely + shares the body's ``model`` name, nor the callbacks that aren't guardrails. """ from litellm.proxy.proxy_server import llm_router @@ -4374,9 +4375,7 @@ class ProxyLogging: try: # Merge model-level guardrails before checking which guardrails to run guardrail_data: Final = ( - check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) - if merge_model_level_guardrails - else data + data if guardrails_only else check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( @@ -4445,7 +4444,7 @@ class ProxyLogging: ############ Handle CustomLogger ############################### ################################################################# - for callback in other_callbacks: + for callback in () if guardrails_only else other_callbacks: callback_response: LLMResponseTypes | None = await callback.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, response=response ) diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index c00109f3132..4ace2af3462 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -25,6 +25,19 @@ MODEL_POST_CALL: Final = "passthrough-model-post-call" GUARDED_MODEL: Final = "passthrough-guarded-model" POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model" BLOCK_MARKER: Final = "block-this-prompt" +RAISING_HOOK_MODEL: Final = "passthrough-model-a-logger-hook-cannot-read" +_RAISING_HOOK_MODULE: Final = f""" +from litellm.integrations.custom_logger import CustomLogger + + +class RaisingHook(CustomLogger): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + if data.get("model") == "{RAISING_HOOK_MODEL}": + raise ValueError("this hook only reads native responses") + + +proxy_handler_instance = RaisingHook() +""" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" @@ -128,11 +141,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I ) config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) general_settings: Final = _JSON.validate_python(config.get("general_settings") or {}) + (directory / "passthrough_raising_hook.py").write_text(_RAISING_HOOK_MODULE) path: Final = directory / "passthrough-auth-guardrails.yaml" path.write_text( yaml.safe_dump( { **config, + "litellm_settings": { + **_JSON.validate_python(config.get("litellm_settings") or {}), + "callbacks": ["passthrough_raising_hook.proxy_handler_instance"], + }, "general_settings": { **general_settings, "pass_through_endpoints": [ @@ -489,3 +507,21 @@ def test_model_guardrail_does_not_run_on_pass_through_traffic_that_names_the_mod assert response.status_code == 200, response.text assert len(rig.generic_post_call_vendor.drain()) == 1 assert rig.model_post_call_vendor.drain() == () + + +def test_post_call_pass_leaves_a_logger_hook_for_native_responses_out(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": RAISING_HOOK_MODEL}, + key, + ) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index b82cb8c7f91..0667a6d47d6 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -13415,7 +13415,7 @@ class TestInheritAuthGuardrails: assert metadata["user_api_key_team_metadata"] == {"opted_out_global_guardrails": ["global-post"]} @pytest.mark.asyncio - async def test_inherited_post_call_pass_skips_model_level_guardrails(self): + async def test_inherited_post_call_pass_runs_only_the_guardrails_the_request_selected(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) _, post_call = await self._run( @@ -13425,7 +13425,7 @@ class TestInheritAuthGuardrails: guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), ) - assert post_call.call_args.kwargs["merge_model_level_guardrails"] is False + assert post_call.call_args.kwargs["guardrails_only"] is True @pytest.mark.asyncio async def test_named_global_guardrail_the_key_opted_out_of_does_not_start_the_post_call_pass(self): diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 7b1ce703a4c..20529384c33 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -131,7 +131,7 @@ class TestPassthroughPostCallGuardrails: assert call_kwargs.kwargs["response"] == _GEMINI_RESPONSE @patch(_COLLECT, return_value=["rubrik"]) - async def test_configured_endpoint_keeps_model_guardrails_and_reads_no_key_opt_outs( + async def test_configured_endpoint_keeps_every_post_call_hook_and_reads_no_key_opt_outs( self, mock_collect, ): @@ -150,7 +150,7 @@ class TestPassthroughPostCallGuardrails: ) call_kwargs = mock_proxy_logging.post_call_success_hook.call_args.kwargs - assert call_kwargs["merge_model_level_guardrails"] is True + assert call_kwargs["guardrails_only"] is False assert call_kwargs["data"]["metadata"].keys() == {"guardrails"} @patch(_COLLECT, return_value=[]) @@ -239,7 +239,7 @@ class TestPassthroughPostCallGuardrails: block path dropping it.""" mock_response = _make_httpx_response(_GEMINI_RESPONSE) - def _block(*, data, user_api_key_dict, response, merge_model_level_guardrails): + def _block(*, data, user_api_key_dict, response, guardrails_only): metadata = data.setdefault("metadata", {}) metadata.setdefault("standard_logging_guardrail_information", []).append( {"guardrail_name": "rubrik", "guardrail_status": "guardrail_intervened"} diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py index 7f6429750c8..61e7e47f7b8 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -60,9 +60,11 @@ async def test_post_call_success_hook_runs_other_callback_and_replaces_response( @pytest.mark.asyncio -@pytest.mark.parametrize(("merge", "guardrails_seen"), [(True, ["key-guard", "model-guard"]), (False, ["key-guard"])]) -async def test_post_call_success_hook_merges_model_level_guardrails_unless_told_not_to( - proxy_logging, make_user_api_key_auth, monkeypatch, merge, guardrails_seen +@pytest.mark.parametrize( + ("guardrails_only", "guardrails_seen"), [(False, ["key-guard", "model-guard"]), (True, ["key-guard"])] +) +async def test_post_call_success_hook_merges_model_level_guardrails_unless_guardrails_only( + proxy_logging, make_user_api_key_auth, monkeypatch, guardrails_only, guardrails_seen ): router = MagicMock() router.registry.return_value.deployments.return_value = [{"litellm_params": {"guardrails": ["model-guard"]}}] @@ -74,12 +76,30 @@ async def test_post_call_success_hook_merges_model_level_guardrails_unless_told_ data={"model": "alias", "metadata": {"guardrails": ["key-guard"]}}, response=MagicMock(), user_api_key_dict=make_user_api_key_auth(), - merge_model_level_guardrails=merge, + guardrails_only=guardrails_only, ) assert g.should_run_guardrail.call_args.kwargs["data"]["metadata"]["guardrails"] == guardrails_seen +@pytest.mark.asyncio +async def test_post_call_success_hook_guardrails_only_leaves_other_callbacks_out( + proxy_logging, make_user_api_key_auth, monkeypatch +): + other = MagicMock(spec=CustomLogger) + other.async_post_call_success_hook = AsyncMock(return_value={"replaced": True}) + g = _make_guardrail(override={"guarded": True}) + monkeypatch.setattr(litellm, "callbacks", [other, g]) + + out = await proxy_logging.post_call_success_hook( + data={}, response={"original": True}, user_api_key_dict=make_user_api_key_auth(), guardrails_only=True + ) + + other.async_post_call_success_hook.assert_not_called() + g.async_post_call_success_hook.assert_awaited_once() + assert out == {"guarded": True} + + @pytest.mark.asyncio async def test_post_call_success_hook_guardrail_should_not_run_skipped( proxy_logging, make_user_api_key_auth, monkeypatch From 83a95188869d10d6db91ced56423f5123acdf2dc Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 11:27:46 -0800 Subject: [PATCH 14/29] test(proxy): run the global guardrail first so the opt-out cells see the post-call data as built --- .../test_passthrough_auth_guardrails.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 4ace2af3462..dbdab6b76a6 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -181,6 +181,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I }, ], "guardrails": [ + { + "guardrail_name": GLOBAL_POST_CALL, + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "post_call", + "default_on": True, + "api_base": post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GUARDRAIL, "litellm_params": { @@ -221,16 +231,6 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, - { - "guardrail_name": GLOBAL_POST_CALL, - "litellm_params": { - "guardrail": "lakera_v2", - "mode": "post_call", - "default_on": True, - "api_base": post_call_vendor.url, - "api_key": "synthetic-guardrail-key", - }, - }, ], } ) From 03290d59026a9ad78c2543f0d043a69318785fe2 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 11:46:10 -0800 Subject: [PATCH 15/29] fix(proxy): keep the provider charge when a pass-through post-call guardrail blocks the response A post-call block or a guardrail failure on a pass-through route happens after the provider answered and billed the call. Settle the cost on the logging object before the failure path writes its row, as the native routes do, so budget and TPM limits count the call. --- .../pass_through_endpoints.py | 103 +++++++++++++++--- .../pass_through_endpoints/success_handler.py | 12 ++ .../test_passthrough_auth_guardrails.py | 33 +++++- .../test_passthrough_post_call_guardrails.py | 35 ++++++ 4 files changed, 166 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d641e6f598e..2ec3fa50fdd 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,10 +6,19 @@ import posixpath import re import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Iterable, + Mapping, + MutableMapping, + Sequence, +) from dataclasses import dataclass from datetime import datetime -from functools import reduce +from functools import partial, reduce from itertools import chain, count, groupby, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast @@ -174,7 +183,7 @@ from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( ) from .llm_provider_handlers.vertex_passthrough_logging_handler import VertexPassthroughLoggingHandler from .streaming_handler import PassThroughStreamingHandler -from .success_handler import PassThroughEndpointLogging +from .success_handler import PassThroughEndpointLogging, finalize_cost_without_success_logs from .upstream_usage_headers import ( UpstreamReportedUsage, apply_upstream_reported_model, @@ -185,9 +194,12 @@ from .upstream_usage_headers import ( if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig +_GuardedResponse = TypeVar("_GuardedResponse") + router: Final = APIRouter() pass_through_endpoint_logging: Final = PassThroughEndpointLogging() +_blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch=finalize_cost_without_success_logs) _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) @@ -933,12 +945,14 @@ def _post_call_hook_data( parsed_body: Mapping[str, object] | None, guardrails: Sequence[object], logging_obj: LiteLLMLoggingObj, - admin_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + inherited: bool, ) -> dict[str, object]: # mutable-ok: the guardrails record their results in this request payload """The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails, which ``should_run_guardrail`` checks, and the logging object, which unified (apply_guardrail) guardrails resolve a pass-through response's call type from when the route has no API_ROUTE_TO_CALL_TYPES entry""" body: Final = parsed_body or {} + admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) if inherited else MappingProxyType({}) return { **body, "metadata": {**(as_str_mapping(body.get("metadata")) or {}), "guardrails": guardrails, **admin_metadata}, @@ -946,6 +960,49 @@ def _post_call_hook_data( } +def _provider_charge_settlement( + response: httpx.Response, + response_body: Mapping[str, object] | None, + url_route: str, + start_time: datetime, + logging_obj: LiteLLMLoggingObj, + parsed_body: Mapping[str, object] | None, + custom_llm_provider: str | None, + handler_kwargs: Mapping[str, object], +) -> Callable[[], Awaitable[None]]: + return partial( + _blocked_response_cost_logging.pass_through_async_success_handler, + finalize_before_dispatch=True, + httpx_response=response, + response_body=response_body, + url_route=url_route, + result="", + start_time=start_time, + end_time=datetime.now(start_time.tzinfo), + logging_obj=logging_obj, + cache_hit=False, + request_body=parsed_body or {}, + custom_llm_provider=custom_llm_provider, + **handler_kwargs, + ) + + +async def _run_post_call_guardrails_keeping_the_charge( + run_guardrails: Callable[[], Awaitable[_GuardedResponse]], + settle_provider_charge: Callable[[], Awaitable[None]], +) -> _GuardedResponse: + """A guardrail that blocks the response, or fails, does so after the provider answered and billed the call, so the + charge is settled before the failure path writes its row, as on the native routes.""" + try: + return await run_guardrails() + except Exception: + try: + await settle_provider_charge() + except Exception: # noqa: BLE001 - the guardrail's outcome still reaches the client + verbose_proxy_logger.exception("pass_through_endpoint: could not settle a blocked response's cost") + raise + + def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequence[object]) -> bool: names: Final = frozenset(chain.from_iterable(map(guardrail_entry_names, entries))) return any( @@ -958,7 +1015,7 @@ def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequenc ) -def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | None) -> None: +def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict[str, object] | None) -> None: """Copy guardrail logging entries from ``guardrail_data`` onto ``request_data``. Post-call guardrails run against a throwaway ``hook_data`` dict (its @@ -972,13 +1029,16 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non if guardrail_data is None: return source_key: Final = get_metadata_variable_name_from_kwargs(guardrail_data) - source_metadata: Final = guardrail_data.get(source_key) or {} + source_metadata: Final = as_str_mapping(guardrail_data.get(source_key)) or {} entries: Final = source_metadata.get("standard_logging_guardrail_information") if not entries: return _, metadata = get_or_create_metadata_bucket(request_data) - metadata.setdefault("standard_logging_guardrail_information", list(entries)) + metadata.setdefault( + "standard_logging_guardrail_information", + list(cast("Iterable[object]", entries)), # cast-ok: guardrails record a list of entries here + ) _GOOGLE_PASSTHROUGH_PROVIDERS: Final = frozenset( @@ -2027,18 +2087,29 @@ async def pass_through_request( ) if response.status_code < 400 and response_body is not None and post_call_guardrails: + json_body: Final = cast(dict[str, object] | None, _parsed_body) # cast-ok: json hook_data: Final = _post_call_hook_data( - cast(dict[str, object] | None, _parsed_body), # cast-ok: json - post_call_guardrails, - logging_obj, - _admin_guardrail_metadata(user_api_key_dict) if inherit_auth_guardrails else MappingProxyType({}), + json_body, post_call_guardrails, logging_obj, user_api_key_dict, inherit_auth_guardrails ) post_call_guardrail_data = hook_data - response_body = await proxy_logging_obj.post_call_success_hook( - data=hook_data, - user_api_key_dict=user_api_key_dict, - response=response_body, - guardrails_only=inherit_auth_guardrails, + response_body = await _run_post_call_guardrails_keeping_the_charge( + partial( + proxy_logging_obj.post_call_success_hook, + data=hook_data, + user_api_key_dict=user_api_key_dict, + response=response_body, + guardrails_only=inherit_auth_guardrails, + ), + _provider_charge_settlement( + response, + cast(dict[str, object] | None, response_body), # cast-ok: json + str(url), + start_time, + logging_obj, + json_body, + custom_llm_provider, + cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above + ), ) if isinstance(response_body, dict): content = json.dumps(response_body).encode("utf-8") diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 8f55906498b..eb0ece24c0e 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -79,6 +79,18 @@ def _websocket_frames( ) +async def finalize_cost_without_success_logs( + *, + logging_obj: LiteLLMLoggingObj, + standard_logging_response_object: object, + **_success_log_fields: object, # kwargs-ok: accepts every field a log dispatch receives; none are needed here +) -> None: + """Settles the provider's charge on the logging object for a response that won't be logged as a success""" + from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost + + await finalize_proxy_response_cost(logging_obj, standard_logging_response_object) + + class PassThroughEndpointLogging: def __init__( self, diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index dbdab6b76a6..3d7bbe18a24 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -11,9 +11,10 @@ import httpx import pytest import yaml from integration._support.anthropic_sse import message_json -from integration._support.client import Gateway, Scenario +from integration._support.client import Gateway, Scenario, eventually from integration._support.openai_wire import chat_reply from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.safeguards_rig import key_hash from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -525,3 +526,33 @@ def test_post_call_pass_leaves_a_logger_hook_for_native_responses_out(rig: AuthG assert response.status_code == 200, response.text assert len(rig.generic_post_call_vendor.drain()) == 1 assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +def _charged_rows(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ...]: + return tuple( + eventually( + lambda: gateway.rows( + 'SELECT status, spend, total_tokens FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (key_hash(key),) + ), + bool, + ) + ) + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( + rig: AuthGuardrailRig, route: str +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [KEY_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[route].drain()) == 1 + rows: Final = _charged_rows(gateway, key) + assert [row["status"] for row in rows] == ["failure"], rows + assert float(str(rows[0]["spend"])) > 0, rows + assert int(str(rows[0]["total_tokens"])) > 0, rows diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 20529384c33..0a750335fd6 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -18,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.proxy._types import ProxyException +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import _run_post_call_guardrails_keeping_the_charge _PT_MOD = "litellm.proxy.pass_through_endpoints.pass_through_endpoints" _COLLECT = "litellm.proxy.pass_through_endpoints.passthrough_guardrails.PassthroughGuardrailHandler.collect_guardrails" @@ -335,3 +336,37 @@ def test_modify_response_exception_importable_from_both_paths(): ) assert FromExceptions is FromGuardrail + + +@pytest.mark.asyncio +async def test_a_blocked_response_settles_the_provider_charge_and_keeps_the_block(): + blocked = HTTPException(status_code=400, detail="blocked") + settle = AsyncMock() + + with pytest.raises(HTTPException) as caught: + await _run_post_call_guardrails_keeping_the_charge(AsyncMock(side_effect=blocked), settle) + + assert caught.value is blocked + settle.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_a_failed_settlement_does_not_replace_the_block(): + blocked = HTTPException(status_code=400, detail="blocked") + + with pytest.raises(HTTPException) as caught: + await _run_post_call_guardrails_keeping_the_charge( + AsyncMock(side_effect=blocked), AsyncMock(side_effect=ValueError("no price")) + ) + + assert caught.value is blocked + + +@pytest.mark.asyncio +async def test_an_allowed_response_leaves_the_charge_to_the_success_path(): + settle = AsyncMock() + + response = await _run_post_call_guardrails_keeping_the_charge(AsyncMock(return_value={"id": "kept"}), settle) + + assert response == {"id": "kept"} + settle.assert_not_awaited() From 0aaa521e15be135f33c894c329418cfd6f3ffcc4 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 12:09:44 -0800 Subject: [PATCH 16/29] fix(proxy): scan a pass-through reply's text, not its JSON envelope, in post-call guardrails A post-call guardrail on a route with no typed call type read the serialized response, so a unix timestamp or an embedding float tripped a phone or card pattern and blocked a clean reply. Without response_fields it now reads the strings the reply carries, minus ids and encoded data. --- .../guardrail_translation/README.md | 2 +- .../guardrail_translation/handler.py | 50 ++++++++++++++++++- .../test_passthrough_auth_guardrails.py | 32 +++++++++++- .../guardrail_translation/test_handler.py | 40 +++++++++++++++ 4 files changed, 119 insertions(+), 5 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/README.md b/litellm/llms/pass_through/guardrail_translation/README.md index db4c0704e75..79d96e3049e 100644 --- a/litellm/llms/pass_through/guardrail_translation/README.md +++ b/litellm/llms/pass_through/guardrail_translation/README.md @@ -24,7 +24,7 @@ litellm/proxy/pass_through_endpoints/ The `PassThroughEndpointHandler` enables guardrails to run on passthrough endpoint requests by: 1. **Field Targeting**: Extracts specific fields from the request/response using JSONPath expressions configured in `request_fields` / `response_fields` -2. **Full Payload Fallback**: If no field targeting is configured, processes the entire payload +2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload, or the text fields of the response 3. **Config Access**: Uses `get_passthrough_guardrails_config()` / `set_passthrough_guardrails_config()` helpers to access the passthrough guardrails configuration stored in request metadata ## Example Config diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index dea8bcb73a4..817367b8844 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -13,7 +13,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import without_server_streaming_classification -from litellm.litellm_core_utils.core_helpers import has_attributes +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, has_attributes from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -29,6 +29,47 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +# Strings in a provider's JSON response that name or encode something rather than carry the reply's text. +_NON_TEXT_RESPONSE_FIELDS: Final = frozenset( + ( + "b64_json", + "bytesBase64Encoded", + "embedding", + "finishReason", + "finish_reason", + "id", + "model", + "modelVersion", + "object", + "responseId", + "role", + "service_tier", + "status", + "system_fingerprint", + "thoughtSignature", + "type", + ) +) + + +def _response_text(response: Mapping[str, object]) -> str: + """The text a JSON response carries, without its timestamps, usage counts, ids and vectors, which a guardrail + would otherwise read as phone or card numbers.""" + pending: Final[list[object]] = [response] # mutable-ok: local to this walk + texts: Final[list[str]] = [] # mutable-ok: local to this walk + while pending: + value = pending.pop() + fields = as_str_mapping(value) + items = as_object_list(value) + if isinstance(value, str): + texts.append(value) + elif fields is not None: + pending.extend(reversed([item for name, item in fields.items() if name not in _NON_TEXT_RESPONSE_FIELDS])) + elif items is not None: + pending.extend(reversed(items)) + return "\n".join(text for text in texts if text) + + class PassThroughEndpointHandler(BaseTranslation): """ Handler for processing passthrough endpoint requests with guardrails. @@ -153,6 +194,7 @@ class PassThroughEndpointHandler(BaseTranslation): litellm_logging_obj: Optional logging object user_api_key_dict: User API key metadata to pass to guardrails """ + response_fields: Final = as_str_mapping(response) if not isinstance(response, dict): verbose_proxy_logger.debug("PassThroughEndpointHandler: Response is not a dict, skipping") return response @@ -168,7 +210,11 @@ class PassThroughEndpointHandler(BaseTranslation): field_expressions: Final = settings.response_fields if settings else None # Extract text to check - text_to_check: Final = self._extract_text_for_guardrail(response, field_expressions) + text_to_check: Final = ( + self._extract_text_for_guardrail(response, field_expressions) + if field_expressions + else _response_text(response_fields or {}) + ) if not text_to_check: return response diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 3d7bbe18a24..02030cfb5cb 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -43,6 +43,7 @@ OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" CUSTOM_ROUTE: Final = "/passthrough-auth-custom" +PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) @@ -66,7 +67,11 @@ def _verdict(request: Request) -> Reply: _OPENAI_BODY: Final = json.dumps( - _JSON.validate_json(chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False).body), indent=2 + { + **_JSON.validate_json(chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False).body), + "created": 1793600123, + }, + indent=2, ).encode() @@ -222,6 +227,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": PHONE_POST_CALL, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "patterns": [ + {"pattern_type": "prebuilt", "pattern_name": "us_phone", "action": "BLOCK"} + ], + }, + }, { "guardrail_name": MODEL_POST_CALL, "litellm_params": { @@ -415,6 +431,18 @@ def test_key_apply_guardrail_post_call_scans_the_provider_reply(rig: AuthGuardra assert [b'"Hi' in request.body for request in rig.generic_post_call_vendor.drain()] == [True] +def test_post_call_guardrail_does_not_read_a_reply_timestamp_as_a_phone_number(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [PHONE_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert response.json()["created"] == 1793600123 + + @pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig: AuthGuardrailRig, route: str) -> None: gateway: Final = rig.owned.gateway @@ -539,7 +567,7 @@ def _charged_rows(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ... ) -@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +@pytest.mark.parametrize("route", [ANTHROPIC_ROUTE, GEMINI_ROUTE]) def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( rig: AuthGuardrailRig, route: str ) -> None: diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 7cbabb95931..d2a439557fe 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -1,5 +1,6 @@ import json from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock import pytest from typing_extensions import ReadOnly, TypedDict @@ -226,3 +227,42 @@ def test_passthrough_route_block_before_any_invoke_frame_is_a_whole_message_in_c "message_stop", ] assert events[2]["delta"]["text"] == "blocked" + + +async def _scanned_response_texts(response: dict[str, object]) -> tuple[str, ...]: + guardrail: Final = MagicMock(guardrail_name="post-call", apply_guardrail=AsyncMock()) + await PassThroughEndpointHandler().process_output_response(response, guardrail, request_data={}) + return tuple(text for call in guardrail.apply_guardrail.await_args_list for text in call.kwargs["inputs"]["texts"]) + + +@pytest.mark.asyncio +async def test_post_call_scan_reads_the_reply_text_and_not_the_numbers_around_it(): + response: Final = { + "id": "chatcmpl-1793600123", + "object": "chat.completion", + "created": 1793600123, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Paris."}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4111111111111111, "total_tokens": 5551234567}, + } + + assert await _scanned_response_texts(response) == ("Paris.",) + + +@pytest.mark.asyncio +async def test_post_call_scan_still_reads_every_string_a_reply_carries(): + response: Final = { + "choices": [ + {"message": {"content": "call 555-123-4567", "tool_calls": [{"function": {"arguments": '{"a": "b"}'}}]}} + ], + "candidates": [{"content": {"parts": [{"text": "card 4111 1111 1111 1111"}]}}], + } + + assert await _scanned_response_texts(response) == ('call 555-123-4567\n{"a": "b"}\ncard 4111 1111 1111 1111',) + + +@pytest.mark.asyncio +async def test_post_call_scan_skips_a_reply_with_no_text(): + response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]} + + assert await _scanned_response_texts(response) == () From 95ffe42ad1ad7b9ef40485bf2fd4ef1b6bf30987 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 12:57:16 -0800 Subject: [PATCH 17/29] fix(proxy): skip reply metadata by position, not by field name at any depth A tool call's arguments are the model's own output, so a field named id or status in them stays in the post-call scan, as do those fields on a custom endpoint's own JSON. Timing asserts in the merge and metadata tests are replaced with deterministic checks. --- .../guardrail_translation/handler.py | 49 +++++++++++++------ .../integrations/test_custom_guardrail.py | 38 ++++++++------ .../guardrail_translation/test_handler.py | 15 ++++++ .../proxy_logging/test_guardrail_pipeline.py | 13 ++--- 4 files changed, 78 insertions(+), 37 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 817367b8844..3c712a213d4 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -7,6 +7,7 @@ to extract specific fields for guardrail processing. """ from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, Protocol, runtime_checkable from pydantic import TypeAdapter @@ -29,44 +30,60 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth -# Strings in a provider's JSON response that name or encode something rather than carry the reply's text. -_NON_TEXT_RESPONSE_FIELDS: Final = frozenset( +# The fields of a provider's JSON reply that name or encode something rather than carry its text, by the field they sit +# under. Anything outside these positions, such as a tool call's arguments, is the model's own output and is read whole. +_REPLY_TOP_LEVEL_FIELDS: Final = frozenset( ( - "b64_json", - "bytesBase64Encoded", - "embedding", - "finishReason", - "finish_reason", "id", "model", "modelVersion", "object", "responseId", - "role", "service_tier", "status", "system_fingerprint", - "thoughtSignature", - "type", ) ) +_REPLY_NESTED_FIELDS: Final = MappingProxyType( + { + "candidates": frozenset(("finishReason",)), + "choices": frozenset(("finish_reason",)), + "content": frozenset(("role", "type")), + "data": frozenset(("b64_json", "embedding", "object")), + "message": frozenset(("role",)), + "output": frozenset(("id", "role", "status", "type")), + "parts": frozenset(("thoughtSignature",)), + "predictions": frozenset(("bytesBase64Encoded",)), + "tool_calls": frozenset(("id", "type")), + } +) +_REPLY_BODY_FIELDS: Final = frozenset(("candidates", "choices", "content", "data", "output", "predictions")) def _response_text(response: Mapping[str, object]) -> str: - """The text a JSON response carries, without its timestamps, usage counts, ids and vectors, which a guardrail - would otherwise read as phone or card numbers.""" - pending: Final[list[object]] = [response] # mutable-ok: local to this walk + """The strings a JSON response carries, without the numbers around them or a provider reply's ids and encoded data, + which a guardrail would otherwise read as phone or card numbers.""" + top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _REPLY_BODY_FIELDS & response.keys() else None + pending: Final[list[tuple[object, frozenset[str] | None]]] = [(response, top_level)] # mutable-ok: stack texts: Final[list[str]] = [] # mutable-ok: local to this walk while pending: - value = pending.pop() + value, skipped = pending.pop() fields = as_str_mapping(value) items = as_object_list(value) if isinstance(value, str): texts.append(value) elif fields is not None: - pending.extend(reversed([item for name, item in fields.items() if name not in _NON_TEXT_RESPONSE_FIELDS])) + pending.extend( + reversed( + [ + (item, None if skipped is None else _REPLY_NESTED_FIELDS.get(name)) + for name, item in fields.items() + if skipped is None or name not in skipped + ] + ) + ) elif items is not None: - pending.extend(reversed(items)) + pending.extend((item, skipped) for item in reversed(items)) return "\n".join(text for text in texts if text) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 302a41a16ff..e6a11f92a0c 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -4,10 +4,8 @@ import datetime as dt import json import pickle import threading -import time -from collections.abc import Callable from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch import pytest @@ -3804,22 +3802,32 @@ def test_default_on_guardrail() -> None: ) -def test_merged_guardrail_entries_stays_close_to_a_plain_dedup_on_a_large_list(): - from litellm.integrations.custom_guardrail import merged_guardrail_entries +def test_merged_guardrail_entries_dedupes_an_all_plain_list_without_walking_it(): + from litellm.integrations import custom_guardrail - names: Final = [f"g{index}" for index in range(300_000)] + names: Final = [f"g{index % 500}" for index in range(1000)] - def _best_cpu_seconds(run: Callable[[], object]) -> float: - def _timed() -> float: - started: Final = time.process_time() - run() - return time.process_time() - started + with patch.object(custom_guardrail, "guardrail_entry_names") as entry_names: + merged: Final = custom_guardrail.merged_guardrail_entries(names) - return min(_timed() for _ in range(3)) + assert merged == [f"g{index}" for index in range(500)] + entry_names.assert_not_called() - assert _best_cpu_seconds(lambda: merged_guardrail_entries(names)) < 25 * _best_cpu_seconds( - lambda: list(dict.fromkeys(names)) - ) + +def test_merged_guardrail_entries_reads_each_entry_of_a_mixed_list_a_fixed_number_of_times(): + from litellm.integrations import custom_guardrail + + entries: Final = [*(f"g{index}" for index in range(1000)), {"g1": {"threshold": 1}}] + + with patch.object( + custom_guardrail, "guardrail_entry_names", wraps=custom_guardrail.guardrail_entry_names + ) as entry_names: + merged: Final = custom_guardrail.merged_guardrail_entries(entries) + + assert "g1" not in merged + assert merged[-1] == {"g1": {"threshold": 1}} + assert len(merged) == 1000 + assert entry_names.call_count <= 2 * len(entries) def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index d2a439557fe..62f49b0047b 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -261,6 +261,21 @@ async def test_post_call_scan_still_reads_every_string_a_reply_carries(): assert await _scanned_response_texts(response) == ('call 555-123-4567\n{"a": "b"}\ncard 4111 1111 1111 1111',) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"id": "alice@example.com"}}}]}}]}, + {"content": [{"input": {"id": "alice@example.com", "message": {"role": "alice@example.com"}}}]}, + {"choices": [{"message": {"tool_calls": [{"function": {"nested": {"type": "alice@example.com"}}}]}}]}, + {"id": "alice@example.com", "result": {"status": "alice@example.com"}}, + ], +) +async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_model_or_api_wrote(response): + assert all("alice@example.com" in text for text in await _scanned_response_texts(response)) + assert await _scanned_response_texts(response) != () + + @pytest.mark.asyncio async def test_post_call_scan_skips_a_reply_with_no_text(): response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]} diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index a46b56294bf..cc84ff9f247 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,7 +11,6 @@ from __future__ import annotations import asyncio import json -import time from collections.abc import AsyncGenerator, Iterator from contextlib import AbstractContextManager, nullcontext from copy import deepcopy @@ -294,12 +293,14 @@ def test_process_guardrail_metadata_records_in_the_litellm_metadata_bucket(proxy assert data == {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2", "g1"]}} -def test_process_guardrail_metadata_handles_many_guardrail_names_quickly(proxy_logging): - names = [f"guard-{i}" for i in range(50_000)] - data = {"metadata": {"guardrails": names, "applied_guardrails": []}} - started = time.perf_counter() +def test_process_guardrail_metadata_does_not_search_the_applied_list_per_name(proxy_logging): + class _NoMembershipChecks(list): + def __contains__(self, item): + raise AssertionError("searched the applied list") + + names = [f"guard-{i}" for i in range(100)] + data = {"metadata": {"guardrails": names, "applied_guardrails": _NoMembershipChecks()}} proxy_logging._process_guardrail_metadata(data) - assert time.perf_counter() - started < 2 assert data["metadata"]["applied_guardrails"] == names From fe159ac944b6fb8917b3347f1d1b865bc3056afc Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 13:32:52 -0800 Subject: [PATCH 18/29] fix(proxy): run policy pipelines in the pass-through post-call pass and read reply integers outside provider replies A policy's post_call pipeline now runs on built-in pass-through responses, with the pipeline state request setup moved out of the body carried into the post-call data, so the pipeline's on_fail decides instead of the step running as a plain guardrail. A non-provider JSON response is read whole, integers included, and tool-call integers in a provider reply are read, so a number in what the model or API wrote still reaches the guardrail while a reply's timestamps and usage counts stay out. --- .../guardrail_translation/README.md | 2 +- .../guardrail_translation/handler.py | 28 +++++++-- .../pass_through_endpoints.py | 51 ++++++++++++--- .../test_passthrough_auth_guardrails.py | 63 +++++++++++++++++++ .../integrations/test_custom_guardrail.py | 10 +-- .../guardrail_translation/test_handler.py | 31 +++++++++ .../test_pass_through_endpoints.py | 37 +++++++++++ 7 files changed, 200 insertions(+), 22 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/README.md b/litellm/llms/pass_through/guardrail_translation/README.md index 79d96e3049e..8a3501dd972 100644 --- a/litellm/llms/pass_through/guardrail_translation/README.md +++ b/litellm/llms/pass_through/guardrail_translation/README.md @@ -24,7 +24,7 @@ litellm/proxy/pass_through_endpoints/ The `PassThroughEndpointHandler` enables guardrails to run on passthrough endpoint requests by: 1. **Field Targeting**: Extracts specific fields from the request/response using JSONPath expressions configured in `request_fields` / `response_fields` -2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload, or the text fields of the response +2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload. A response is read whole, integers included, except a provider's reply, where the ids, timestamps and usage counts around the text are left out 3. **Config Access**: Uses `get_passthrough_guardrails_config()` / `set_passthrough_guardrails_config()` helpers to access the passthrough guardrails configuration stored in request metadata ## Example Config diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 3c712a213d4..80d03527acb 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: # The fields of a provider's JSON reply that name or encode something rather than carry its text, by the field they sit -# under. Anything outside these positions, such as a tool call's arguments, is the model's own output and is read whole. +# under, and the fields that hold what the model wrote, which are read whole. _REPLY_TOP_LEVEL_FIELDS: Final = frozenset( ( "id", @@ -57,13 +57,27 @@ _REPLY_NESTED_FIELDS: Final = MappingProxyType( "tool_calls": frozenset(("id", "type")), } ) -_REPLY_BODY_FIELDS: Final = frozenset(("candidates", "choices", "content", "data", "output", "predictions")) +_REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters")) +_REPLY_MARKER_FIELDS: Final = frozenset(("candidates", "choices", "object", "output", "predictions")) +_IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data")) +_NO_FIELDS: Final[frozenset[str]] = frozenset() + + +def _is_provider_reply(response: Mapping[str, object]) -> bool: + return bool(_REPLY_MARKER_FIELDS & response.keys()) or _IMAGE_REPLY_FIELDS.issubset(response) + + +def _reply_field_position(name: str, skipped: frozenset[str] | None) -> frozenset[str] | None: + if skipped is None or name in _REPLY_MODEL_OUTPUT_FIELDS: + return None + return _REPLY_NESTED_FIELDS.get(name, _NO_FIELDS) def _response_text(response: Mapping[str, object]) -> str: - """The strings a JSON response carries, without the numbers around them or a provider reply's ids and encoded data, - which a guardrail would otherwise read as phone or card numbers.""" - top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _REPLY_BODY_FIELDS & response.keys() else None + """The text a guardrail reads from a JSON response. In a provider's reply that is its strings without the ids and + encoded data around them, so a timestamp or a usage count doesn't read as a phone or card number. What the model + wrote into a tool call, and any other API's JSON, is read whole, integers included. Floats are scores and vectors.""" + top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _is_provider_reply(response) else None pending: Final[list[tuple[object, frozenset[str] | None]]] = [(response, top_level)] # mutable-ok: stack texts: Final[list[str]] = [] # mutable-ok: local to this walk while pending: @@ -72,11 +86,13 @@ def _response_text(response: Mapping[str, object]) -> str: items = as_object_list(value) if isinstance(value, str): texts.append(value) + elif skipped is None and isinstance(value, int) and not isinstance(value, bool): + texts.append(str(value)) elif fields is not None: pending.extend( reversed( [ - (item, None if skipped is None else _REPLY_NESTED_FIELDS.get(name)) + (item, _reply_field_position(name, skipped)) for name, item in fields.items() if skipped is None or name not in skipped ] diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 2ec3fa50fdd..f377756b75a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -158,7 +158,7 @@ from litellm.proxy.route_priority import ( configured_pass_through_routes_first, shadowed_by_builtin_claim, ) -from litellm.proxy.utils import normalize_route_for_root_path +from litellm.proxy.utils import normalize_route_for_root_path, pipeline_managed_guardrail_names from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str from litellm.types import utils as types_utils @@ -204,6 +204,7 @@ _blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch= _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) _GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) +_POLICY_PIPELINE_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails") _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -896,7 +897,7 @@ async def _inherit_auth_guardrails( The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook, with the ones the client named in ``metadata.guardrails`` kept after them, and returned only when at least one - of them runs post_call, so the response is relayed untouched otherwise. + of them runs post_call. """ if configured_guardrails or not inherit or not parsed_body: return configured_guardrails @@ -941,21 +942,42 @@ def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, ) +def _policy_pipeline_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: + """The policy pipelines resolved for a request that inherits auth guardrails, from the logging metadata that + request setup moved the body's ``metadata`` into""" + if not inherited: + return MappingProxyType({}) + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} + return {field: metadata[field] for field in _POLICY_PIPELINE_FIELDS if field in metadata} + + +def _starts_post_call_pass(guardrails: Sequence[object] | None, policy_pipelines: Mapping[str, object]) -> bool: + return bool(guardrails) or bool(pipeline_managed_guardrail_names({"metadata": dict(policy_pipelines)}, "post_call")) + + def _post_call_hook_data( parsed_body: Mapping[str, object] | None, - guardrails: Sequence[object], + guardrails: Sequence[object] | None, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth, + policy_pipelines: Mapping[str, object], inherited: bool, ) -> dict[str, object]: # mutable-ok: the guardrails record their results in this request payload """The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails, - which ``should_run_guardrail`` checks, and the logging object, which unified (apply_guardrail) guardrails resolve a - pass-through response's call type from when the route has no API_ROUTE_TO_CALL_TYPES entry""" + which ``should_run_guardrail`` checks, the policy pipelines to run, and the logging object, which unified + (apply_guardrail) guardrails resolve a pass-through response's call type from when the route has no + API_ROUTE_TO_CALL_TYPES entry""" body: Final = parsed_body or {} admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) if inherited else MappingProxyType({}) return { **body, - "metadata": {**(as_str_mapping(body.get("metadata")) or {}), "guardrails": guardrails, **admin_metadata}, + "metadata": { + **(as_str_mapping(body.get("metadata")) or {}), + **policy_pipelines, + "guardrails": guardrails or (), + **admin_metadata, + }, "litellm_logging_obj": logging_obj, } @@ -2086,10 +2108,23 @@ async def pass_through_request( logging_obj=logging_obj, ) - if response.status_code < 400 and response_body is not None and post_call_guardrails: + policy_pipelines: Final = _policy_pipeline_metadata( + cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above + inherit_auth_guardrails, + ) + if ( + response.status_code < 400 + and response_body is not None + and _starts_post_call_pass(post_call_guardrails, policy_pipelines) + ): json_body: Final = cast(dict[str, object] | None, _parsed_body) # cast-ok: json hook_data: Final = _post_call_hook_data( - json_body, post_call_guardrails, logging_obj, user_api_key_dict, inherit_auth_guardrails + json_body, + post_call_guardrails, + logging_obj, + user_api_key_dict, + policy_pipelines, + inherit_auth_guardrails, ) post_call_guardrail_data = hook_data response_body = await _run_post_call_guardrails_keeping_the_charge( diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 02030cfb5cb..d292fe41059 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -44,6 +44,9 @@ ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" CUSTOM_ROUTE: Final = "/passthrough-auth-custom" PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call" +PIPELINE_BLOCKS: Final = "passthrough-pipeline-blocks-the-response" +PIPELINE_ALLOWS: Final = "passthrough-pipeline-allows-the-response" +PIPELINE_BEFORE_THE_CALL: Final = "passthrough-pipeline-before-the-call" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) @@ -61,6 +64,13 @@ def _route_body(route: str, prompt: str) -> dict[str, JsonValue]: return {"model": "gpt-5.6", "messages": [{"role": "user", "content": prompt}]} +def _pipeline_policy(mode: str, guardrail: str, on_fail: str) -> dict[str, JsonValue]: + return { + "guardrails": {"add": [guardrail]}, + "pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": on_fail}]}, + } + + def _verdict(request: Request) -> Reply: action: Final = "BLOCKED" if BLOCK_MARKER.encode() in request.body else "NONE" return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode()) @@ -167,6 +177,15 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I } ], }, + "policies": { + PIPELINE_BLOCKS: _pipeline_policy("post_call", KEY_POST_CALL, "block"), + PIPELINE_ALLOWS: _pipeline_policy("post_call", KEY_POST_CALL, "allow"), + PIPELINE_BEFORE_THE_CALL: _pipeline_policy("pre_call", GUARDRAIL, "block"), + }, + "policy_attachments": [ + {"policy": policy, "tags": [policy]} + for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL) + ], "model_list": [ *_JSON_LIST.validate_python(config.get("model_list") or []), { @@ -567,6 +586,50 @@ def _charged_rows(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ... ) +def _tagged_body(route: str, tag: str) -> dict[str, JsonValue]: + return {**_route_body(route, "What is the capital of France?"), "metadata": {"tags": [tag]}} + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_post_call_policy_pipeline_blocks_a_flagged_pass_through_response(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, route, _tagged_body(route, PIPELINE_BLOCKS), key) + + assert response.status_code == 400, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[route].drain()) == 1 + assert [row["status"] for row in _charged_rows(gateway, key)] == ["failure"] + + +def test_post_call_policy_pipeline_decides_what_a_flagged_step_does(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_ALLOWS), key) + + assert response.status_code == 200, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + + +def test_pre_call_policy_pipeline_relays_the_provider_bytes(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_BEFORE_THE_CALL), key) + + assert response.status_code == 200, response.text + assert response.content == _OPENAI_BODY + assert len(rig.vendor.drain()) == 1 + + @pytest.mark.parametrize("route", [ANTHROPIC_ROUTE, GEMINI_ROUTE]) def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( rig: AuthGuardrailRig, route: str diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index e6a11f92a0c..1e93db413b1 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3802,16 +3802,12 @@ def test_default_on_guardrail() -> None: ) -def test_merged_guardrail_entries_dedupes_an_all_plain_list_without_walking_it(): - from litellm.integrations import custom_guardrail +def test_merged_guardrail_entries_dedupes_an_all_plain_list(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries names: Final = [f"g{index % 500}" for index in range(1000)] - with patch.object(custom_guardrail, "guardrail_entry_names") as entry_names: - merged: Final = custom_guardrail.merged_guardrail_entries(names) - - assert merged == [f"g{index}" for index in range(500)] - entry_names.assert_not_called() + assert merged_guardrail_entries(names) == [f"g{index}" for index in range(500)] def test_merged_guardrail_entries_reads_each_entry_of_a_mixed_list_a_fixed_number_of_times(): diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 62f49b0047b..b712f6c471c 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -276,6 +276,37 @@ async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_ assert await _scanned_response_texts(response) != () +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response, expected", + [ + ({"result": {"credit_card": 4111111111111111, "score": 0.4111111111111111, "ok": True}}, "4111111111111111"), + ( + {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"card": 4111111111111111}}}]}}]}, + "4111111111111111", + ), + ({"choices": [{"message": {"tool_calls": [{"function": {"arguments": {"n": 5551234567}}}]}}]}, "5551234567"), + ], +) +async def test_post_call_scan_reads_an_integer_in_what_the_model_or_api_wrote(response, expected): + assert await _scanned_response_texts(response) == (expected,) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"created": 1793600123, "data": [{"url": "https://x.example/a.png"}]}, + {"id": "batch_1", "object": "batch", "created_at": 1793600123, "request_counts": {"total": 5551234567}}, + {"predictions": [{"embeddings": {"values": [0.0123456789]}, "total": 5551234567}]}, + ], +) +async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(response): + assert all( + "1793600123" not in text and "5551234567" not in text for text in await _scanned_response_texts(response) + ) + + @pytest.mark.asyncio async def test_post_call_scan_skips_a_reply_with_no_text(): response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]} diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0667a6d47d6..f9a3f04e9e7 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -85,6 +85,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id from litellm.proxy.utils import InternalUsageCache, ProxyLogging, _merge_guardrails_with_existing, hash_token from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.utils import CallTypes from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, @@ -13382,6 +13383,42 @@ class TestInheritAuthGuardrails: assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "client-guard"] + @staticmethod + def _body_with_policy_pipeline(mode: str) -> dict[str, object]: + pipeline: Final = GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="step-guard")]) + return { + **TestInheritAuthGuardrails._chat_body(), + "metadata": { + "_guardrail_pipelines": [("policy", pipeline)], + "_pipeline_managed_guardrails": {"step-guard"}, + }, + } + + @pytest.mark.asyncio + async def test_post_call_policy_pipeline_starts_the_post_call_pass_with_its_state(self): + _, post_call = await self._run( + self._body_with_policy_pipeline("post_call"), + inherit=True, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + post_call.assert_awaited_once() + metadata: Final = post_call.await_args.kwargs["data"]["metadata"] + assert metadata["_pipeline_managed_guardrails"] == {"step-guard"} + assert [policy for policy, _ in metadata["_guardrail_pipelines"]] == ["policy"] + + @pytest.mark.asyncio + @pytest.mark.parametrize(("mode", "inherit"), [("pre_call", True), ("post_call", False)]) + async def test_policy_pipeline_leaves_the_response_alone_without_a_post_call_run(self, mode: str, inherit: bool): + response, _, post_call = await self._run_with_response( + self._body_with_policy_pipeline(mode), + inherit=inherit, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + @pytest.mark.asyncio async def test_dict_form_stays_as_sent_for_a_key_without_guardrails(self): pre_call, post_call = await self._run( From 83eb02cb31985d0344eae7a1fe93a0e3b6e4c7b0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 13:43:45 -0800 Subject: [PATCH 19/29] fix(proxy): read JSON field names in what a post-call guardrail scans outside a provider's envelope A field name can carry the text a guardrail looks for, as in a custom reply keyed by email address. Custom JSON and tool-call arguments now read their names too, and a provider reply's envelope stays out. --- .../guardrail_translation/README.md | 2 +- .../guardrail_translation/handler.py | 25 +++++++++++-------- .../guardrail_translation/test_handler.py | 14 ++++++++++- 3 files changed, 28 insertions(+), 13 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/README.md b/litellm/llms/pass_through/guardrail_translation/README.md index 8a3501dd972..74bde8e809c 100644 --- a/litellm/llms/pass_through/guardrail_translation/README.md +++ b/litellm/llms/pass_through/guardrail_translation/README.md @@ -24,7 +24,7 @@ litellm/proxy/pass_through_endpoints/ The `PassThroughEndpointHandler` enables guardrails to run on passthrough endpoint requests by: 1. **Field Targeting**: Extracts specific fields from the request/response using JSONPath expressions configured in `request_fields` / `response_fields` -2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload. A response is read whole, integers included, except a provider's reply, where the ids, timestamps and usage counts around the text are left out +2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload. A response is read whole, field names and integers included, except a provider's reply, where the ids, timestamps and usage counts around the text are left out 3. **Config Access**: Uses `get_passthrough_guardrails_config()` / `set_passthrough_guardrails_config()` helpers to access the passthrough guardrails configuration stored in request metadata ## Example Config diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 80d03527acb..31e92201a41 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -6,7 +6,7 @@ It uses the field targeting configuration from litellm_logging_obj to extract specific fields for guardrail processing. """ -from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, Protocol, runtime_checkable @@ -73,10 +73,21 @@ def _reply_field_position(name: str, skipped: frozenset[str] | None) -> frozense return _REPLY_NESTED_FIELDS.get(name, _NO_FIELDS) +def _reply_entries( + fields: Mapping[str, object], skipped: frozenset[str] | None +) -> Iterator[tuple[object, frozenset[str] | None]]: + for name, item in fields.items(): + if skipped is None: + yield name, None + if skipped is None or name not in skipped: + yield item, _reply_field_position(name, skipped) + + def _response_text(response: Mapping[str, object]) -> str: """The text a guardrail reads from a JSON response. In a provider's reply that is its strings without the ids and encoded data around them, so a timestamp or a usage count doesn't read as a phone or card number. What the model - wrote into a tool call, and any other API's JSON, is read whole, integers included. Floats are scores and vectors.""" + wrote into a tool call, and any other API's JSON, is read whole, field names and integers included. Floats are + scores and vectors.""" top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _is_provider_reply(response) else None pending: Final[list[tuple[object, frozenset[str] | None]]] = [(response, top_level)] # mutable-ok: stack texts: Final[list[str]] = [] # mutable-ok: local to this walk @@ -89,15 +100,7 @@ def _response_text(response: Mapping[str, object]) -> str: elif skipped is None and isinstance(value, int) and not isinstance(value, bool): texts.append(str(value)) elif fields is not None: - pending.extend( - reversed( - [ - (item, _reply_field_position(name, skipped)) - for name, item in fields.items() - if skipped is None or name not in skipped - ] - ) - ) + pending.extend(reversed(tuple(_reply_entries(fields, skipped)))) elif items is not None: pending.extend((item, skipped) for item in reversed(items)) return "\n".join(text for text in texts if text) diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index b712f6c471c..e19939ddd3e 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -289,7 +289,19 @@ async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_ ], ) async def test_post_call_scan_reads_an_integer_in_what_the_model_or_api_wrote(response, expected): - assert await _scanned_response_texts(response) == (expected,) + assert any(expected in text for text in await _scanned_response_texts(response)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"results": {"alice@example.com": "active"}}, + {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"alice@example.com": "active"}}}]}}]}, + ], +) +async def test_post_call_scan_reads_a_field_name_in_what_the_model_or_api_wrote(response): + assert any("alice@example.com" in text for text in await _scanned_response_texts(response)) @pytest.mark.asyncio From c7928956d14ddef133a2e6746707f393a6f3fa81 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 13:47:05 -0800 Subject: [PATCH 20/29] fix(proxy): read an Anthropic message as a provider reply in the post-call scan An Anthropic message has no choices or candidates, so its ids and usage counts were read whole. Its stop_reason and usage now mark it as a provider reply, and its id, role, type and stop fields stay out. --- .../guardrail_translation/handler.py | 21 +++++++++++++++++-- .../guardrail_translation/test_handler.py | 19 +++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 31e92201a41..57044b96f44 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -34,21 +34,26 @@ if TYPE_CHECKING: # under, and the fields that hold what the model wrote, which are read whole. _REPLY_TOP_LEVEL_FIELDS: Final = frozenset( ( + "finish_reason", "id", "model", "modelVersion", "object", "responseId", + "role", "service_tier", "status", + "stop_reason", + "stop_sequence", "system_fingerprint", + "type", ) ) _REPLY_NESTED_FIELDS: Final = MappingProxyType( { "candidates": frozenset(("finishReason",)), "choices": frozenset(("finish_reason",)), - "content": frozenset(("role", "type")), + "content": frozenset(("id", "role", "type")), "data": frozenset(("b64_json", "embedding", "object")), "message": frozenset(("role",)), "output": frozenset(("id", "role", "status", "type")), @@ -58,7 +63,19 @@ _REPLY_NESTED_FIELDS: Final = MappingProxyType( } ) _REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters")) -_REPLY_MARKER_FIELDS: Final = frozenset(("candidates", "choices", "object", "output", "predictions")) +_REPLY_MARKER_FIELDS: Final = frozenset( + ( + "candidates", + "choices", + "finish_reason", + "object", + "output", + "predictions", + "stop_reason", + "usage", + "usageMetadata", + ) +) _IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data")) _NO_FIELDS: Final[frozenset[str]] = frozenset() diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index e19939ddd3e..43c8ce9cc1a 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -319,6 +319,25 @@ async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(respons ) +@pytest.mark.asyncio +async def test_post_call_scan_reads_only_the_text_of_an_anthropic_message(): + response: Final = { + "id": "msg_5551234567", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [ + {"type": "text", "text": "Paris."}, + {"type": "tool_use", "id": "toolu_5551234567", "name": "lookup", "input": {"card": 4111111111111111}}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5551234567, "output_tokens": 1}, + } + + assert await _scanned_response_texts(response) == ("Paris.\nlookup\ncard\n4111111111111111",) + + @pytest.mark.asyncio async def test_post_call_scan_skips_a_reply_with_no_text(): response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]} From 1d956d84d16f29e244961330ec149d0f8a559269 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 13:54:09 -0800 Subject: [PATCH 21/29] fix(proxy): tell a provider reply from a custom API's JSON by the type of its fields A field name such as usage or output also appears in custom JSON, which then lost its integers. A reply now counts as a provider's by a list under choices, candidates, output or predictions, a string object, or an Anthropic or Cohere message. --- .../guardrail_translation/handler.py | 24 ++++++++----------- .../guardrail_translation/test_handler.py | 4 ++++ 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 57044b96f44..ac91b7bd260 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -63,25 +63,21 @@ _REPLY_NESTED_FIELDS: Final = MappingProxyType( } ) _REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters")) -_REPLY_MARKER_FIELDS: Final = frozenset( - ( - "candidates", - "choices", - "finish_reason", - "object", - "output", - "predictions", - "stop_reason", - "usage", - "usageMetadata", - ) -) +_REPLY_LIST_FIELDS: Final = ("candidates", "choices", "output", "predictions") _IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data")) _NO_FIELDS: Final[frozenset[str]] = frozenset() def _is_provider_reply(response: Mapping[str, object]) -> bool: - return bool(_REPLY_MARKER_FIELDS & response.keys()) or _IMAGE_REPLY_FIELDS.issubset(response) + """Whether a JSON response has the shape of a provider's reply, by the type of the field that names it, since a + custom API's JSON can use the same field names.""" + return ( + any(isinstance(response.get(field), list) for field in _REPLY_LIST_FIELDS) + or isinstance(response.get("object"), str) + or (response.get("type") == "message" and "stop_reason" in response) + or ("finish_reason" in response and "message" in response) + or _IMAGE_REPLY_FIELDS.issubset(response) + ) def _reply_field_position(name: str, skipped: frozenset[str] | None) -> frozenset[str] | None: diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 43c8ce9cc1a..e028ff2a7bd 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -286,6 +286,9 @@ async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_ "4111111111111111", ), ({"choices": [{"message": {"tool_calls": [{"function": {"arguments": {"n": 5551234567}}}]}}]}, "5551234567"), + ({"usage": {}, "customer": {"phone": 5551234567}}, "5551234567"), + ({"output": {"card": 4111111111111111}}, "4111111111111111"), + ({"type": "message", "phone": 5551234567}, "5551234567"), ], ) async def test_post_call_scan_reads_an_integer_in_what_the_model_or_api_wrote(response, expected): @@ -311,6 +314,7 @@ async def test_post_call_scan_reads_a_field_name_in_what_the_model_or_api_wrote( {"created": 1793600123, "data": [{"url": "https://x.example/a.png"}]}, {"id": "batch_1", "object": "batch", "created_at": 1793600123, "request_counts": {"total": 5551234567}}, {"predictions": [{"embeddings": {"values": [0.0123456789]}, "total": 5551234567}]}, + {"finish_reason": "COMPLETE", "message": {"role": "assistant"}, "usage": {"total": 5551234567}}, ], ) async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(response): From de63702f0835ecc606a72a6d4b761d48053b3c2b Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 17:27:38 -0800 Subject: [PATCH 22/29] fix(guardrails): tolerate pass-through bodies request setup hasn't normalized --- .../generic_guardrail_api.py | 5 +- .../guardrails/guardrail_hooks/lakera_ai.py | 5 +- .../guardrail_hooks/lakera_ai_v2.py | 3 + .../guardrail_hooks/noma/noma_v2.py | 10 ++- litellm/proxy/litellm_pre_call_utils.py | 3 +- ...guardrails_on_unexpected_request_fields.py | 83 +++++++++++++++++++ .../guardrail_hooks/test_lakera_ai_v2.py | 29 +++++++ .../unit/proxy/test_litellm_pre_call_utils.py | 27 ++++++ 8 files changed, 160 insertions(+), 5 deletions(-) create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 7e052796545..f3309b98462 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -365,7 +365,10 @@ class GenericGuardrailAPI(CustomGuardrail): top_level_metadata: Final = request_data.get("metadata", {}) # Merge both sources, preferring litellm_metadata if both exist - metadata_dict: Final = {**top_level_metadata, **litellm_metadata} + metadata_dict: Final = { + **(top_level_metadata if isinstance(top_level_metadata, dict) else {}), + **(litellm_metadata if isinstance(litellm_metadata, dict) else {}), + } if not metadata_dict: return result_metadata diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..f7eb05e2ab6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -219,7 +219,10 @@ class lakeraAI_Moderation(CustomGuardrail): } ) elif "input" in data and isinstance(data["input"], list): - text = "\n".join(data["input"]) + text = "\n".join(item for item in data["input"] if isinstance(item, str)) + if not text: + verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text") + return _json_data = json.dumps( { "input": text, diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 3e5c54b4501..d127e10e685 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -781,6 +781,9 @@ class LakeraAIGuardrail(CustomGuardrail): # Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"] post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages + if not post_call_messages: + verbose_proxy_logger.debug("Lakera AI: not running guardrail. No inspectable text in data") + return response # Call Lakera guardrail lakera_guardrail_response, _ = await self.call_v2_guard( diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 43eb54d1464..71772ad6c46 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( @@ -63,6 +64,11 @@ class _Action(str, enum.Enum): GUARDRAIL_INTERVENED = "GUARDRAIL_INTERVENED" +def _user_api_key_alias(request_data: Mapping[str, object], metadata_key: str) -> object: + """The key alias recorded under ``metadata_key``, which a pass-through client can set to anything.""" + return (as_str_mapping(request_data.get(metadata_key)) or {}).get("user_api_key_alias") + + class NomaV2Guardrail(CustomGuardrail): reads_stream_marks: ClassVar[bool] = True """Noma keeps one record per streamed reply when every scan carries ``stream`` and only the last one carries the @@ -336,8 +342,8 @@ class NomaV2Guardrail(CustomGuardrail): # (ports v1 fallback from PR #16832). if application_id is None: application_id = self._get_non_empty_str( - request_data.get("litellm_metadata", {}).get("user_api_key_alias") - ) or self._get_non_empty_str(request_data.get("metadata", {}).get("user_api_key_alias")) + _user_api_key_alias(request_data, "litellm_metadata") + ) or self._get_non_empty_str(_user_api_key_alias(request_data, "metadata")) try: payload: Final = self._build_scan_payload( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..501c8352b10 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4377,11 +4377,12 @@ 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.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"), + "model": _model if isinstance(_model, str) else None, "tags": all_tags, } ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py new file mode 100644 index 00000000000..f1e1dca04c3 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py @@ -0,0 +1,83 @@ +"""Guardrails on built-in pass-through routes read client JSON that request setup hasn't normalized.""" + +from typing import Final +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( + GenericGuardrailAPI, +) +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation +from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail + +NOT_AN_OBJECT: Final = (None, 0, False, "", [], "text", 7) + + +@pytest.mark.parametrize("value", NOT_AN_OBJECT) +def test_generic_guardrail_reads_metadata_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + + metadata: Final = guardrail._extract_user_api_key_metadata( + {"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value} + ) + + assert metadata["user_api_key_hash"] == "hash" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", NOT_AN_OBJECT) +async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None: + guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma") + verdict: Final = httpx.Response( + 200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid") + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)): + result: Final = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}, "litellm_metadata": value}, + input_type="request", + ) + + assert result["texts"] == ["hello"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]]) +async def test_legacy_lakera_has_no_text_to_scan_in_token_array_embeddings(token_arrays: list[object]) -> None: + guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid") + vendor_error: Final = httpx.HTTPStatusError( + "vendor called", + request=httpx.Request("POST", "http://vendor.invalid"), + response=httpx.Response(500, text="vendor called"), + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)): + await guardrail._check( + data={"input": token_arrays}, + user_api_key_dict=UserAPIKeyAuth(api_key="key"), + call_type="pass_through_endpoint", + ) + + +@pytest.mark.asyncio +async def test_legacy_lakera_still_scans_string_inputs() -> None: + guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid") + vendor_error: Final = httpx.HTTPStatusError( + "vendor called", + request=httpx.Request("POST", "http://vendor.invalid"), + response=httpx.Response(500, text="vendor called"), + ) + + with ( + patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)), + pytest.raises(Exception, match="vendor called"), + ): + await guardrail._check( + data={"input": ["hello"]}, + user_api_key_dict=UserAPIKeyAuth(api_key="key"), + call_type="pass_through_endpoint", + ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 23288a6136f..52c3c9eb92d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -2413,6 +2413,35 @@ async def test_lakera_post_call_allows_clean_content(): assert result is llm_response +@pytest.mark.asyncio +async def test_post_call_with_no_messages_and_no_reply_text_returns_the_response(): + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + response: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}}]} + rejected: Final = HTTPException(status_code=400, detail="either content, messages, or tools must be provided") + with patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)): + result: Final = await guardrail.async_post_call_success_hook( + data={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response=response, + ) + assert result == response + + +@pytest.mark.asyncio +async def test_post_call_with_messages_still_reaches_lakera(): + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + rejected: Final = HTTPException(status_code=400, detail="rejected") + with ( + patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)), + pytest.raises(HTTPException), + ): + await guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hello"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response={}, + ) + + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_lakera_post_call_masks_pii_and_allows(): diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..eeffba68969 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5327,6 +5327,33 @@ async def test_embedding_header_forwarding_without_model_group_config(): litellm.model_group_settings = original_model_group_settings +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [0, False, [], {}, 7, True]) +async def test_policy_engine_leaves_a_model_that_is_not_a_string_alone( + model: object, monkeypatch: pytest.MonkeyPatch +) -> None: + 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 + + policy_registry = get_policy_registry() + monkeypatch.setattr( + policy_registry, "_policies", {"baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"]))} + ) + monkeypatch.setattr(policy_registry, "_initialized", True) + attachment_registry = get_attachment_registry() + monkeypatch.setattr(attachment_registry, "_attachments", [PolicyAttachment(policy="baseline", scope="*")]) + monkeypatch.setattr(attachment_registry, "_initialized", True) + data = {"model": model, "metadata": {}} + + await add_guardrails_from_policy_engine( + data=data, metadata_variable_name="metadata", user_api_key_dict=UserAPIKeyAuth(api_key="test-key") + ) + + assert data["model"] == model + assert data["metadata"]["guardrails"] == ["pii_blocker"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine(): """ From 1a5b56c9a279488e46b0affb8b33de8617c7ca1f Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 17:58:15 -0800 Subject: [PATCH 23/29] fix(guardrails): read Responses messages in legacy Lakera and skip bodies it can't read on pass-through A list input of Responses messages now reaches legacy Lakera instead of being dropped, and a pass-through body with no messages or input skips the empty vendor call. Lakera post_call on /gemini has nothing to read, so the rig's Gemini block cells use a generic post_call guardrail. Moves the regression cells into each guardrail's own test file. --- .../guardrails/guardrail_hooks/lakera_ai.py | 22 ++++- .../test_passthrough_auth_guardrails.py | 52 ++++++++++-- .../test_generic_guardrail_api.py | 11 +++ ...guardrails_on_unexpected_request_fields.py | 83 ------------------- .../guardrail_hooks/test_lakera_ai.py | 57 +++++++++++++ .../guardrail_hooks/test_noma_v2.py | 20 +++++ 6 files changed, 156 insertions(+), 89 deletions(-) delete mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index f7eb05e2ab6..228dc417142 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -15,6 +15,7 @@ from typing import Final, Literal import httpx from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -22,6 +23,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -45,6 +47,20 @@ INPUT_POSITIONING_MAP: Final = { Role.ASSISTANT.value: 2, } +_CONTENT_PARTS: Final = TypeAdapter(list[object]) + + +def _input_item_text(item: object) -> str: + """The text of one entry of a list ``input``: a string, or a Responses message's string or parts ``content``.""" + if isinstance(item, str): + return item + content: Final = (as_str_mapping(item) or {}).get("content") + if isinstance(content, str): + return content + parts: Final = _CONTENT_PARTS.validate_python(content) if isinstance(content, list) else () + part_texts: Final = ((as_str_mapping(part) or {}).get("text") for part in parts) + return "\n".join(text for text in part_texts if isinstance(text, str)) + class lakeraAI_Moderation(CustomGuardrail): @classmethod @@ -219,7 +235,7 @@ class lakeraAI_Moderation(CustomGuardrail): } ) elif "input" in data and isinstance(data["input"], list): - text = "\n".join(item for item in data["input"] if isinstance(item, str)) + text = "\n".join(filter(None, map(_input_item_text, data["input"]))) if not text: verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text") return @@ -230,6 +246,10 @@ class lakeraAI_Moderation(CustomGuardrail): } ) + if not _json_data and call_type == "pass_through_endpoint": + verbose_proxy_logger.debug("Skipping lakera prompt injection, pass-through body has no messages or input") + return + verbose_proxy_logger.debug("Lakera AI Request Args %s", _json_data) # https://platform.lakera.ai/account/api-keys diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index d292fe41059..c8b91077d91 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -22,6 +22,7 @@ GUARDRAIL: Final = "passthrough-auth-guardrail" GLOBAL_POST_CALL: Final = "passthrough-global-post-call" KEY_POST_CALL: Final = "passthrough-key-post-call" GENERIC_POST_CALL: Final = "passthrough-generic-post-call" +BLOCKING_POST_CALL: Final = "passthrough-post-call-blocks-every-reply" MODEL_POST_CALL: Final = "passthrough-model-post-call" GUARDED_MODEL: Final = "passthrough-guarded-model" POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model" @@ -76,6 +77,10 @@ def _verdict(request: Request) -> Reply: return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode()) +def _block_every_reply(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode()) + + _OPENAI_BODY: Final = json.dumps( { **_JSON.validate_json(chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False).body), @@ -147,6 +152,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) + blocking_post_call_vendor: Final = stack.enter_context(wire_server(_block_every_reply)) model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) upstreams: Final = MappingProxyType( { @@ -236,6 +242,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": BLOCKING_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": blocking_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GENERIC_POST_CALL, "litellm_params": { @@ -590,7 +606,7 @@ def _tagged_body(route: str, tag: str) -> dict[str, JsonValue]: return {**_route_body(route, "What is the capital of France?"), "metadata": {"tags": [tag]}} -@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE]) def test_post_call_policy_pipeline_blocks_a_flagged_pass_through_response(rig: AuthGuardrailRig, route: str) -> None: gateway: Final = rig.owned.gateway with gateway.scenario() as scenario: @@ -630,15 +646,41 @@ def test_pre_call_policy_pipeline_relays_the_provider_bytes(rig: AuthGuardrailRi assert len(rig.vendor.drain()) == 1 -@pytest.mark.parametrize("route", [ANTHROPIC_ROUTE, GEMINI_ROUTE]) -def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( - rig: AuthGuardrailRig, route: str -) -> None: +def test_lakera_post_call_has_nothing_to_read_on_a_gemini_body(rig: AuthGuardrailRig) -> None: gateway: Final = rig.owned.gateway with gateway.scenario() as scenario: key: Final = scenario.key(metadata={"guardrails": [KEY_POST_CALL]}) _drain(rig) + response: Final = _send(gateway, GEMINI_ROUTE, _route_body(GEMINI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert rig.key_post_call_vendor.drain() == () + assert len(rig.upstreams[GEMINI_ROUTE].drain()) == 1 + + +def test_post_call_policy_pipeline_lakera_step_has_nothing_to_read_on_a_gemini_body(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, GEMINI_ROUTE, _tagged_body(GEMINI_ROUTE, PIPELINE_BLOCKS), key) + + assert response.status_code == 200, response.text + assert rig.key_post_call_vendor.drain() == () + assert len(rig.upstreams[GEMINI_ROUTE].drain()) == 1 + + +@pytest.mark.parametrize(("route", "guardrail"), [(ANTHROPIC_ROUTE, KEY_POST_CALL), (GEMINI_ROUTE, BLOCKING_POST_CALL)]) +def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( + rig: AuthGuardrailRig, route: str, guardrail: str +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [guardrail]}) + _drain(rig) + response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key) assert response.status_code == 400, response.text diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index b5f794cb136..daa9945e406 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -2692,3 +2692,14 @@ async def test_during_mcp_call_blocked_raises() -> None: ) assert exc_info.value.message == "credential detected in tool input" + + +@pytest.mark.parametrize("value", (None, 0, False, "", [], "text", 7)) +def test_extract_user_api_key_metadata_reads_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + + metadata: Final = guardrail._extract_user_api_key_metadata( + {"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value} + ) + + assert metadata["user_api_key_hash"] == "hash" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py deleted file mode 100644 index f1e1dca04c3..00000000000 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_guardrails_on_unexpected_request_fields.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Guardrails on built-in pass-through routes read client JSON that request setup hasn't normalized.""" - -from typing import Final -from unittest.mock import AsyncMock, patch - -import httpx -import pytest - -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( - GenericGuardrailAPI, -) -from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation -from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail - -NOT_AN_OBJECT: Final = (None, 0, False, "", [], "text", 7) - - -@pytest.mark.parametrize("value", NOT_AN_OBJECT) -def test_generic_guardrail_reads_metadata_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None: - guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") - - metadata: Final = guardrail._extract_user_api_key_metadata( - {"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value} - ) - - assert metadata["user_api_key_hash"] == "hash" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("value", NOT_AN_OBJECT) -async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None: - guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma") - verdict: Final = httpx.Response( - 200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid") - ) - - with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)): - result: Final = await guardrail.apply_guardrail( - inputs={"texts": ["hello"]}, - request_data={"metadata": {}, "litellm_metadata": value}, - input_type="request", - ) - - assert result["texts"] == ["hello"] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]]) -async def test_legacy_lakera_has_no_text_to_scan_in_token_array_embeddings(token_arrays: list[object]) -> None: - guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid") - vendor_error: Final = httpx.HTTPStatusError( - "vendor called", - request=httpx.Request("POST", "http://vendor.invalid"), - response=httpx.Response(500, text="vendor called"), - ) - - with patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)): - await guardrail._check( - data={"input": token_arrays}, - user_api_key_dict=UserAPIKeyAuth(api_key="key"), - call_type="pass_through_endpoint", - ) - - -@pytest.mark.asyncio -async def test_legacy_lakera_still_scans_string_inputs() -> None: - guardrail: Final = lakeraAI_Moderation(api_key="key", api_base="http://vendor.invalid") - vendor_error: Final = httpx.HTTPStatusError( - "vendor called", - request=httpx.Request("POST", "http://vendor.invalid"), - response=httpx.Response(500, text="vendor called"), - ) - - with ( - patch.object(guardrail.async_handler, "post", new=AsyncMock(side_effect=vendor_error)), - pytest.raises(Exception, match="vendor called"), - ): - await guardrail._check( - data={"input": ["hello"]}, - user_api_key_dict=UserAPIKeyAuth(api_key="key"), - call_type="pass_through_endpoint", - ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py new file mode 100644 index 00000000000..afba8f1b39d --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -0,0 +1,57 @@ +"""The legacy Lakera guardrail reads client JSON that request setup hasn't normalized on pass-through routes.""" + +import json +from typing import Final +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation + +VENDOR: Final = "http://vendor.invalid" +CLEAN: Final = httpx.Response(200, json={"results": [{"flagged": False}]}, request=httpx.Request("POST", VENDOR)) + + +async def _sent_to_lakera(data: dict[str, object], call_type: str = "pass_through_endpoint") -> tuple[str, ...]: + guardrail: Final = lakeraAI_Moderation(api_key="key", api_base=VENDOR) + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=CLEAN)) as post: + await guardrail._check(data=data, user_api_key_dict=UserAPIKeyAuth(api_key="key"), call_type=call_type) + return tuple(call.kwargs["data"] for call in post.await_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]]) +async def test_token_arrays_have_no_text_to_send(token_arrays: list[object]) -> None: + assert await _sent_to_lakera({"input": token_arrays}) == () + + +@pytest.mark.asyncio +async def test_string_entries_of_a_list_input_are_sent() -> None: + sent: Final = await _sent_to_lakera({"input": ["hello", "there"]}) + + assert [json.loads(body)["input"] for body in sent] == ["hello\nthere"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", + ["card 4111", [{"type": "input_text", "text": "card 4111"}, {"type": "input_image", "image_url": "http://x"}]], +) +async def test_a_responses_message_in_a_list_input_is_sent(content: object) -> None: + sent: Final = await _sent_to_lakera({"input": [{"role": "user", "content": content}]}) + + assert [json.loads(body)["input"] for body in sent] == ["card 4111"] + + +@pytest.mark.asyncio +async def test_a_pass_through_body_with_no_messages_or_input_sends_nothing() -> None: + contents: Final = {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]} + + assert await _sent_to_lakera(contents) == () + + +@pytest.mark.asyncio +async def test_a_native_body_with_no_messages_or_input_still_reaches_lakera() -> None: + assert await _sent_to_lakera({"prompt": "hello"}, call_type="text_completion") == ("",) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py index b145b990539..8d430d83787 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py @@ -1,7 +1,9 @@ import asyncio import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from litellm.proxy.guardrails.guardrail_hooks.noma import NomaV2Guardrail @@ -747,3 +749,21 @@ class TestNomaV2ApplicationIdResolution: payload = call_mock.call_args.kwargs["payload"] assert "application_id" not in payload + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", (None, 0, False, "", [], "text", 7)) +async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None: + guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma") + verdict: Final = httpx.Response( + 200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid") + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)): + result: Final = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}, "litellm_metadata": value}, + input_type="request", + ) + + assert result["texts"] == ["hello"] From 880f28451f4a2c37d769ca9c2524a2fd4963e94f Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 04:06:39 -0800 Subject: [PATCH 24/29] fix(proxy): keep main's outcome where pass-through guardrails reached too far - Drop a top-level guardrails or guardrail_config the merge can't combine, which returned 500 where main stripped the field. - Count only a pre_call policy pipeline when withholding the Anthropic user_id. - Read /bedrock invoke replies whole for generic guardrails, as on main. - Keep the provider's body for logging when a hook returns a non-dict. - Settle a no-cost blocked response at 0 so its failure row keeps its usage. - Leave custom endpoints uncharged on a failed post-call hook, as on main. - Keep a legacy Lakera string list's empty entries. --- litellm/litellm_core_utils/litellm_logging.py | 11 +- .../guardrail_translation/handler.py | 1 + .../guardrail_translation/handler.py | 4 +- .../guardrails/guardrail_hooks/lakera_ai.py | 3 +- .../pass_through_endpoints.py | 28 +++- .../pass_through_endpoints/success_handler.py | 10 +- .../proxy/spend_tracking/cost_finalization.py | 6 +- litellm/proxy/utils.py | 29 +++- .../test_passthrough_auth_guardrails.py | 151 +++++++++++++++++- .../guardrail_translation/test_handler.py | 22 +++ .../guardrail_hooks/test_lakera_ai.py | 8 + .../test_pass_through_endpoints.py | 66 ++++++++ .../test_passthrough_post_call_guardrails.py | 48 ++++++ .../spend_tracking/test_cost_finalization.py | 17 ++ .../utils/proxy_logging/test_pre_call_hook.py | 39 ++++- 15 files changed, 422 insertions(+), 21 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 639003c9a6a..551e08f870e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2566,7 +2566,10 @@ class Logging(LiteLLMLoggingBaseClass): **kwargs, ) - def prepare_response_cost(self, response: object, now: datetime.datetime) -> PreparedCost | None: + def prepare_response_cost( + self, response: object, now: datetime.datetime, settle_unknown_cost_at_zero: bool = False + ) -> PreparedCost | None: + """The response's cost, ready to finalize, or ``None`` when it has no cost result, unless that settles at 0""" if self.prepared_cost is not None: return self.prepared_cost recorded_cache_hit: Final[object] = self.model_call_details.get("cache_hit") @@ -2592,8 +2595,10 @@ class Logging(LiteLLMLoggingBaseClass): publish_logging_payload=False, ) payload: Final = prepared[3] - raw_cost: Final[object] = self.model_call_details.get("response_cost") - if payload is None or not isinstance(raw_cost, (int, float)): + recorded_cost: Final[object] = self.model_call_details.get("response_cost") + unknown_cost: Final = 0.0 if settle_unknown_cost_at_zero else None + raw_cost: Final = recorded_cost if isinstance(recorded_cost, (int, float)) else unknown_cost + if payload is None or raw_cost is None: return None recorded_usage: Final[object] = self.model_call_details.get("combined_usage_object") response_usage: Final[object] = getattr(prepared[2], "usage", None) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 3a65a8928bf..8cca405cd97 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -1416,6 +1416,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + read_whole_response=True, ) if not isinstance(response, dict): diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index ac91b7bd260..712e51f7afc 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -233,6 +233,7 @@ class PassThroughEndpointHandler(BaseTranslation): litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, + read_whole_response: bool = False, ) -> object: """ Process output response by applying guardrails to targeted fields. @@ -242,6 +243,7 @@ class PassThroughEndpointHandler(BaseTranslation): guardrail_to_apply: The guardrail instance to apply litellm_logging_obj: Optional logging object user_api_key_dict: User API key metadata to pass to guardrails + read_whole_response: Without targeted fields, read the response as JSON, ids and usage included """ response_fields: Final = as_str_mapping(response) if not isinstance(response, dict): @@ -261,7 +263,7 @@ class PassThroughEndpointHandler(BaseTranslation): # Extract text to check text_to_check: Final = ( self._extract_text_for_guardrail(response, field_expressions) - if field_expressions + if field_expressions or read_whole_response else _response_text(response_fields or {}) ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index 228dc417142..f93204f1bc9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -235,7 +235,8 @@ class lakeraAI_Moderation(CustomGuardrail): } ) elif "input" in data and isinstance(data["input"], list): - text = "\n".join(filter(None, map(_input_item_text, data["input"]))) + item_texts: Final = tuple(map(_input_item_text, _CONTENT_PARTS.validate_python(data["input"]))) + text = "\n".join(item_texts) if any(item_texts) else "" if not text: verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text") return diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f377756b75a..7696d9fc344 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -886,6 +886,17 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _update_stream_param_based_on_request_body = update_stream_param_based_on_request_body +def _drop_unmergeable_guardrail_fields( + parsed_body: dict[str, object], # mutable-ok: strips the request body in place before the guardrail merge +) -> None: + """Drop a ``guardrails`` that isn't a list or an object and a ``guardrail_config`` that isn't an object, which the + merge can't combine with the metadata's own. The provider receives neither field.""" + if "guardrails" in parsed_body and not isinstance(parsed_body["guardrails"], list | dict): + parsed_body.pop("guardrails") + if "guardrail_config" in parsed_body and not isinstance(parsed_body["guardrail_config"], dict): + parsed_body.pop("guardrail_config") + + async def _inherit_auth_guardrails( parsed_body: dict[str, object] | None, # mutable-ok: merges the guardrails into the live request body user_api_key_dict: UserAPIKeyAuth, @@ -901,6 +912,7 @@ async def _inherit_auth_guardrails( """ if configured_guardrails or not inherit or not parsed_body: return configured_guardrails + _drop_unmergeable_guardrail_fields(parsed_body) if not isinstance(parsed_body.get("metadata"), dict): parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") @@ -1012,12 +1024,16 @@ def _provider_charge_settlement( async def _run_post_call_guardrails_keeping_the_charge( run_guardrails: Callable[[], Awaitable[_GuardedResponse]], settle_provider_charge: Callable[[], Awaitable[None]], + settles_the_charge: bool = True, ) -> _GuardedResponse: """A guardrail that blocks the response, or fails, does so after the provider answered and billed the call, so the - charge is settled before the failure path writes its row, as on the native routes.""" + charge is settled before the failure path writes its row, as on the native routes. A custom endpoint runs every + hook here, not only guardrails, and doesn't settle, so its failure row stays uncharged.""" try: return await run_guardrails() except Exception: + if not settles_the_charge: + raise try: await settle_provider_charge() except Exception: # noqa: BLE001 - the guardrail's outcome still reaches the client @@ -1649,7 +1665,7 @@ async def pass_through_request( ) hooked_body: Final = _json_object(_parsed_body) client_user_id: Final = _anthropic_user_id_to_forward( - hooked_body, str(url), proxy_logging_obj.request_has_pre_call_guardrails + hooked_body, str(url), proxy_logging_obj.request_is_scanned_before_the_provider ) kwargs = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( # rebind-ok: pre-existing rebinding on a rename-only line user_api_key_dict=user_api_key_dict, @@ -2127,7 +2143,7 @@ async def pass_through_request( inherit_auth_guardrails, ) post_call_guardrail_data = hook_data - response_body = await _run_post_call_guardrails_keeping_the_charge( + guarded_body: Final = await _run_post_call_guardrails_keeping_the_charge( partial( proxy_logging_obj.post_call_success_hook, data=hook_data, @@ -2145,14 +2161,16 @@ async def pass_through_request( custom_llm_provider, cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above ), + settles_the_charge=inherit_auth_guardrails, ) - if isinstance(response_body, dict): + if isinstance(guarded_body, dict): + response_body = guarded_body content = json.dumps(response_body).encode("utf-8") _content_modified = True else: verbose_proxy_logger.debug( "pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response", - type(response_body).__name__, + type(guarded_body).__name__, ) elif response_body is None: verbose_proxy_logger.debug( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index eb0ece24c0e..335b4889037 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -2,7 +2,7 @@ import asyncio import json from collections.abc import Mapping, Sequence from contextlib import suppress -from datetime import datetime +from datetime import datetime, timezone from types import MappingProxyType from typing import Any, Final from urllib.parse import urlparse @@ -85,9 +85,13 @@ async def finalize_cost_without_success_logs( standard_logging_response_object: object, **_success_log_fields: object, # kwargs-ok: accepts every field a log dispatch receives; none are needed here ) -> None: - """Settles the provider's charge on the logging object for a response that won't be logged as a success""" - from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost + """Settles the provider's charge on the logging object for a response that won't be logged as a success. A response + with no cost result settles at 0, the spend its success row would carry, so the failure row keeps its usage.""" + from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost + prepare_response_cost( + logging_obj, standard_logging_response_object, datetime.now(timezone.utc), settle_unknown_cost_at_zero=True + ) await finalize_proxy_response_cost(logging_obj, standard_logging_response_object) diff --git a/litellm/proxy/spend_tracking/cost_finalization.py b/litellm/proxy/spend_tracking/cost_finalization.py index c42579714e8..2db4b05f1e6 100644 --- a/litellm/proxy/spend_tracking/cost_finalization.py +++ b/litellm/proxy/spend_tracking/cost_finalization.py @@ -11,9 +11,11 @@ from litellm.proxy.hooks.search_free_tier import SharedCounter, allocate_search_ from litellm.types.utils import CostBreakdown, StandardLoggingPayload -def prepare_response_cost(logging_obj: Logging, response: object, now: datetime) -> PreparedCost | None: +def prepare_response_cost( + logging_obj: Logging, response: object, now: datetime, settle_unknown_cost_at_zero: bool = False +) -> PreparedCost | None: try: - return logging_obj.prepare_response_cost(response, now) + return logging_obj.prepare_response_cost(response, now, settle_unknown_cost_at_zero) except Exception: # noqa: BLE001 # preparation must not prevent serving an already completed response verbose_proxy_logger.exception("Unable to prepare completed response cost for %s", logging_obj.litellm_call_id) return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index eb36d1ae361..d4a2b6299e2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -155,6 +155,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import ( + as_object_list, coerce_token_limit, get_or_create_metadata_bucket, independent_snapshot, @@ -719,6 +720,20 @@ def pipeline_managed_guardrail_names( ) +def _runs_after_the_provider(pipeline_entry: object) -> bool: + if not isinstance(pipeline_entry, tuple): + return False + pipeline: Final = cast("tuple[object, ...]", pipeline_entry)[-1] # cast-ok: the policy engine's (name, pipeline) + return getattr(pipeline, "mode", None) == "post_call" + + +def _holds_pipeline_before_the_provider(pipelines: object) -> bool: + if not pipelines: + return False + entries: Final = as_object_list(pipelines) + return entries is None or not all(map(_runs_after_the_provider, entries)) + + def _ran_or_runs_outside_pipelines( callback: CustomGuardrail | None, data: Mapping[str, object], @@ -3108,8 +3123,18 @@ class ProxyLogging: request_data: Mapping[str, object], event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call, ) -> bool: - if _policy_state_metadata(request_data).get("_guardrail_pipelines"): - return True + return bool(_policy_state_metadata(request_data).get("_guardrail_pipelines")) or ( + self._request_has_guardrail_callback(request_data, event_type) + ) + + def request_is_scanned_before_the_provider(self, request_data: Mapping[str, object]) -> bool: + """``request_has_pre_call_guardrails``, except that a resolved policy pipeline running post_call doesn't count""" + return _holds_pipeline_before_the_provider( + _policy_state_metadata(request_data).get("_guardrail_pipelines") + ) or self._request_has_guardrail_callback(request_data, GuardrailEventHooks.pre_call) + + @staticmethod + def _request_has_guardrail_callback(request_data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool: caps: Final = ProxyLogging.callback_capabilities() if caps.has_content_enforcer: return True diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index c8b91077d91..690b018d8f3 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -44,6 +44,11 @@ OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" CUSTOM_ROUTE: Final = "/passthrough-auth-custom" +CUSTOM_GUARDED_ROUTE: Final = "/passthrough-auth-custom-guarded" +UNPRICED_MODEL: Final = "passthrough-auth-unpriced-model" +MONITOR_POST_CALL: Final = "passthrough-auth-monitor-post-call" +MONITORED_PROMPT: Final = "reply-with-the-monitored-id" +MONITORED_REPLY_ID: Final = "msg_passthrough_auth_monitored" PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call" PIPELINE_BLOCKS: Final = "passthrough-pipeline-blocks-the-response" PIPELINE_ALLOWS: Final = "passthrough-pipeline-allows-the-response" @@ -94,6 +99,18 @@ def _openai_reply(_request: Request) -> Reply: return Reply(body=_OPENAI_BODY) +def _lakera_pii_only(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "flagged": True, + "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}], + "payload": [], + } + ).encode() + ) + + def _lakera_clean(_request: Request) -> Reply: return Reply(body=json.dumps({"flagged": False, "breakdown": [], "payload": []}).encode()) @@ -125,8 +142,20 @@ _UPSTREAM_BODIES: Final = MappingProxyType( ) -def _anthropic_reply(_request: Request) -> Reply: - return Reply(body=_ANTHROPIC_BODY) +_MONITORED_ANTHROPIC_BODY: Final = json.dumps( + _JSON.validate_json(message_json(MONITORED_REPLY_ID, "claude-sonnet-5-5", "Hi")) +).encode() + + +_UNPRICED_ANTHROPIC_BODY: Final = json.dumps( + _JSON.validate_json(message_json("msg_passthrough_auth_unpriced", UNPRICED_MODEL, "Hi")) +).encode() + + +def _anthropic_reply(request: Request) -> Reply: + if UNPRICED_MODEL.encode() in request.body: + return Reply(body=_UNPRICED_ANTHROPIC_BODY) + return Reply(body=_MONITORED_ANTHROPIC_BODY if MONITORED_PROMPT.encode() in request.body else _ANTHROPIC_BODY) def _gemini_reply(_request: Request) -> Reply: @@ -153,6 +182,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) blocking_post_call_vendor: Final = stack.enter_context(wire_server(_block_every_reply)) + monitor_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_pii_only)) model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) upstreams: Final = MappingProxyType( { @@ -180,7 +210,14 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "path": CUSTOM_ROUTE, "target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions", "auth": True, - } + }, + { + "path": CUSTOM_GUARDED_ROUTE, + "target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions", + "auth": True, + "cost_per_request": 0.25, + "guardrails": {BLOCKING_POST_CALL: None}, + }, ], }, "policies": { @@ -252,6 +289,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": MONITOR_POST_CALL, + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "post_call", + "default_on": False, + "on_flagged": "monitor", + "api_base": monitor_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GENERIC_POST_CALL, "litellm_params": { @@ -689,3 +737,100 @@ def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( assert [row["status"] for row in rows] == ["failure"], rows assert float(str(rows[0]["spend"])) > 0, rows assert int(str(rows[0]["total_tokens"])) > 0, rows + + +def _row_fields(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ...]: + return tuple( + eventually( + lambda: gateway.rows( + 'SELECT status, spend, total_tokens, api_base, request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (key_hash(key),), + ), + bool, + ) + ) + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +@pytest.mark.parametrize("value", [None, False, 7, "not-a-list"]) +def test_a_body_guardrails_that_is_not_a_list_is_dropped_for_a_guarded_key( + rig: AuthGuardrailRig, route: str, value: JsonValue +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GUARDRAIL]}) + _drain(rig) + + response: Final = _send(gateway, route, {**_route_body(route, "Hello"), "guardrails": value}, key) + + assert response.status_code == 200, response.text + assert len(rig.vendor.drain()) == 1 + assert [b'"guardrails"' in request.body for request in rig.upstreams[route].drain()] == [False] + + +def test_a_blocked_response_with_no_cost_result_keeps_its_usage_on_the_failure_row(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [BLOCKING_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, ANTHROPIC_ROUTE, {**_route_body(ANTHROPIC_ROUTE, "Hello"), "model": UNPRICED_MODEL}, key + ) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[ANTHROPIC_ROUTE].drain()) == 1 + rows: Final = _row_fields(gateway, key) + assert [row["status"] for row in rows] == ["failure"], rows + assert float(str(rows[0]["spend"])) == 0, rows + assert int(str(rows[0]["total_tokens"])) > 0, rows + assert rows[0]["api_base"], rows + + +def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail_blocks(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key( + team_id=scenario.team(), metadata={"allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE]} + ) + _drain(rig) + + response: Final = gateway.request("POST", CUSTOM_GUARDED_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=key) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + assert [float(str(row["spend"])) for row in _row_fields(gateway, key)] == [0], response.text + + +def test_a_monitored_lakera_flag_keeps_the_provider_id_on_the_row(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [MONITOR_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, ANTHROPIC_ROUTE, _route_body(ANTHROPIC_ROUTE, MONITORED_PROMPT), key) + + assert response.status_code == 200, response.text + assert response.content == _MONITORED_ANTHROPIC_BODY + assert [row["request_id"] for row in _row_fields(gateway, key)] == [MONITORED_REPLY_ID] + + +@pytest.mark.parametrize( + ("policy", "forwarded"), [(PIPELINE_ALLOWS, True), (PIPELINE_BEFORE_THE_CALL, False)], ids=["post-call", "pre-call"] +) +def test_only_a_pipeline_that_runs_before_the_provider_withholds_the_client_user_id( + rig: AuthGuardrailRig, policy: str, forwarded: bool +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + body: Final = { + **_route_body(ANTHROPIC_ROUTE, "Hello"), + "metadata": {"tags": [policy], "user_id": "client-user-1"}, + } + + response: Final = _send(gateway, ANTHROPIC_ROUTE, body, key) + + assert response.status_code == 200, response.text + assert [b"client-user-1" in request.body for request in rig.upstreams[ANTHROPIC_ROUTE].drain()] == [forwarded] diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index aec1672b335..02bc6e0d938 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -832,6 +832,28 @@ class TestBedrockPassthroughGuardrailHandlerOutput: sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] assert "blocked model output" in sent_texts[0] + @pytest.mark.asyncio + async def test_non_converse_response_is_read_whole_with_its_envelope(self): + response = { + "id": "msg_marker", + "type": "message", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": "stop-marker", + "content": [{"type": "text", "text": "model output"}], + "usage": {"input_tokens": 5551234}, + } + guardrail = _make_guardrail({"texts": []}) + + await BedrockPassthroughGuardrailHandler().process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data={"endpoint": "model/anthropic.claude-3-sonnet/invoke"}, + ) + + sent_text = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"][0] + assert all(part in sent_text for part in ("msg_marker", "stop-marker", "5551234", "model output")) + def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: import json diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py index afba8f1b39d..bacb576d6c9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -55,3 +55,11 @@ async def test_a_pass_through_body_with_no_messages_or_input_sends_nothing() -> @pytest.mark.asyncio async def test_a_native_body_with_no_messages_or_input_still_reaches_lakera() -> None: assert await _sent_to_lakera({"prompt": "hello"}, call_type="text_completion") == ("",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("strings", "text"), [(["a", "", "b"], "a\n\nb"), (["a", ""], "a\n")]) +async def test_empty_strings_of_a_list_input_keep_their_place(strings: list[object], text: str) -> None: + sent: Final = await _sent_to_lakera({"input": strings}) + + assert [json.loads(body)["input"] for body in sent] == [text] diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f9a3f04e9e7..1fd7ee4d70c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -56,6 +56,8 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + _anthropic_user_id_to_forward, + _inherit_auth_guardrails, WebsocketFrameModelAllowed, WebsocketFrameModelGate, WebsocketFrameModelRefused, @@ -13328,6 +13330,70 @@ class TestInheritAuthGuardrails: def _chat_body(**extra: object) -> dict[str, object]: return {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], **extra} + @pytest.mark.asyncio + @pytest.mark.parametrize("value", [None, False, True, 7, "name"]) + @pytest.mark.parametrize( + ("key_metadata", "body_metadata"), + [({"guardrails": ["key-guard"]}, None), ({}, {"guardrails": ["body-guard"]}), ({}, None)], + ids=["key-guardrails", "metadata-guardrails", "no-guardrails"], + ) + async def test_a_top_level_guardrails_the_merge_cant_add_to_is_dropped(self, value, key_metadata, body_metadata): + body: Final[dict[str, object]] = { + **self._chat_body(guardrails=value), + **({"metadata": body_metadata} if body_metadata is not None else {}), + } + + with patch("litellm.proxy.proxy_server.premium_user", True): + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test", metadata=key_metadata), None, True) + + assert "guardrails" not in body + assert (body["metadata"].get("guardrails") or []) == [ + *key_metadata.get("guardrails", ()), + *(body_metadata or {}).get("guardrails", ()), + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", [None, False, 7, "config", ["config"]]) + async def test_a_guardrail_config_the_merge_cant_add_to_is_dropped(self, value): + body: Final[dict[str, object]] = self._chat_body( + guardrail_config=value, metadata={"guardrail_config": {"kept": True}} + ) + + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) + + assert "guardrail_config" not in body + assert body["metadata"]["guardrail_config"] == {"kept": True} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("body_guardrails", "merged"), + [(["body-guard"], ["key-guard", "body-guard"]), ({"body-guard": {}}, ["key-guard", "body-guard"])], + ids=["list", "object"], + ) + async def test_a_top_level_guardrails_list_or_object_still_merges(self, body_guardrails, merged): + body: Final[dict[str, object]] = self._chat_body(guardrails=body_guardrails) + + with patch("litellm.proxy.proxy_server.premium_user", True): + await _inherit_auth_guardrails( + body, UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}), None, True + ) + + assert body["metadata"]["guardrails"] == merged + + @pytest.mark.parametrize(("mode", "forwarded"), [("post_call", "user-1"), ("pre_call", None)]) + def test_only_a_pipeline_that_runs_before_the_provider_withholds_the_user_id(self, mode, forwarded): + pipeline: Final = GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="guard")]) + body: Final = {"model": "claude", "metadata": {"user_id": "user-1", "_guardrail_pipelines": [("p", pipeline)]}} + + with patch.object(litellm, "callbacks", []): + user_id: Final = _anthropic_user_id_to_forward( + body, + "https://api.anthropic.com/v1/messages", + ProxyLogging(user_api_key_cache=DualCache()).request_is_scanned_before_the_provider, + ) + + assert user_id == forwarded + @pytest.mark.asyncio async def test_key_team_and_project_guardrails_reach_pre_and_post_call(self): auth: Final = UserAPIKeyAuth( diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 0a750335fd6..667625cdb7d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -228,6 +228,54 @@ class TestPassthroughPostCallGuardrails: assert body["error"]["guardrail_name"] == "rubrik" assert body["error"]["model"] == "gemini-2.0-flash" + @patch(_COLLECT, return_value=["rubrik"]) + async def test_configured_endpoint_writes_no_charge_for_a_failed_post_call_hook(self, mock_collect): + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock(side_effect=HTTPException(status_code=418, detail="no")) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + with ( + _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)), + patch(f"{_PT_MOD}._blocked_response_cost_logging") as settlement, + pytest.raises(ProxyException), + ): + await pass_through_request( + request=_make_mock_request(), + target="https://example.com/v1/generateContent", + custom_headers={"Content-Type": "application/json"}, + user_api_key_dict=_make_user_api_key_dict(), + stream=False, + ) + + settlement.pass_through_async_success_handler.assert_not_called() + + @patch(_COLLECT, return_value=["rubrik"]) + async def test_a_hook_result_that_is_not_a_dict_leaves_the_body_the_logs_read(self, mock_collect): + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=MagicMock(id="rebuilt")) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + success_handler = AsyncMock() + + with ( + _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)), + patch( + f"{_PT_MOD}.pass_through_endpoint_logging", + MagicMock(pass_through_async_success_handler=success_handler), + ), + ): + result = await pass_through_request( + request=_make_mock_request(), + target="https://example.com/v1/generateContent", + custom_headers={"Content-Type": "application/json"}, + user_api_key_dict=_make_user_api_key_dict(), + stream=False, + ) + + assert success_handler.await_args.kwargs["response_body"] == _GEMINI_RESPONSE + assert json.loads(result.body) == _GEMINI_RESPONSE + @patch(_COLLECT, return_value=["rubrik"]) async def test_deny_forwards_guardrail_logging_info_to_failure_hook( self, diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index c9515c1232e..3dffc02e7b3 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -1,6 +1,7 @@ import asyncio from datetime import datetime, timezone from typing import Final +from unittest.mock import patch import pytest @@ -633,3 +634,19 @@ async def test_guardrail_json_rejection_retains_completed_provider_usage( assert (payload["prompt_tokens"], payload["completion_tokens"]) == (10, 5) assert payload["response_cost"] == 0.5 ProxyLogging._callback_capabilities_cache.clear() + + +def test_a_response_with_no_cost_result_settles_at_zero_only_when_asked(logging_obj: Logging) -> None: + from litellm.proxy.spend_tracking.cost_finalization import prepare_response_cost + + unpriced: Final = response() + del unpriced._hidden_params["response_cost"] + + with patch.object(logging_obj, "response_cost_calculator", return_value=None): + assert prepare_response_cost(logging_obj, unpriced, NOW) is None + prepared: Final = prepare_response_cost(logging_obj, unpriced, NOW, settle_unknown_cost_at_zero=True) + + assert prepared is not None + assert prepared.amount == 0.0 + assert prepared.payload["response_cost"] == 0.0 + assert (prepared.usage.prompt_tokens, prepared.usage.completion_tokens) == (10, 5) diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 26ee9f04d01..77542f1c3f4 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -16,6 +16,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.utils import CallTypesLiteral @@ -1014,3 +1015,39 @@ def test_request_has_pre_call_guardrails_reads_pipelines_from_both_metadata_buck monkeypatch.setattr(litellm, "callbacks", []) assert proxy_logging.request_has_pre_call_guardrails(request_data) is expected + + +@pytest.mark.parametrize( + ("pipelines", "scanned"), + [ + ([("policy", GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g")]))], False), + ([("policy", GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g")]))], True), + ( + [ + ("after", GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g")])), + ("before", GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g")])), + ], + True, + ), + (["p1"], True), + (5, True), + ([], False), + ], + ids=["post-call", "pre-call", "both", "unresolved-entry", "non-list", "empty"], +) +def test_only_a_pipeline_running_before_the_provider_scans_the_request(proxy_logging, monkeypatch, pipelines, scanned): + monkeypatch.setattr(litellm, "callbacks", []) + request_data: Final = {"metadata": {"_guardrail_pipelines": pipelines}} + + assert proxy_logging.request_is_scanned_before_the_provider(request_data) is scanned + assert proxy_logging.request_has_pre_call_guardrails(request_data) is bool(pipelines) + + +def test_a_pre_call_guardrail_scans_the_request_beside_a_post_call_pipeline(proxy_logging, monkeypatch): + guardrail: Final = CustomGuardrail(guardrail_name="g", default_on=True, event_hook=GuardrailEventHooks.pre_call) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="other")]) + + assert proxy_logging.request_is_scanned_before_the_provider( + {"metadata": {"_guardrail_pipelines": [("p", pipeline)]}} + ) From 5d8d7b811f716041b5175ed2f0b45d38d84ff817 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 04:15:39 -0800 Subject: [PATCH 25/29] refactor(proxy): move the post-call result branch out of pass_through_request --- .../pass_through_endpoints.py | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7696d9fc344..121bf0ed2c8 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1041,6 +1041,23 @@ async def _run_post_call_guardrails_keeping_the_charge( raise +def _response_after_post_call_guardrails( + guarded_body: object, + response_body: dict[str, object] | None, # mutable-ok: pass_through_request's own response dict, returned as is + content: bytes, +) -> tuple[dict[str, object] | None, bytes, bool]: # mutable-ok: pass_through_request keeps a dict for the id rewrite + """The response body, its bytes and whether they changed, once the guardrails returned ``guarded_body``. A hook that + returns anything but a JSON object leaves the response as the provider sent it.""" + if isinstance(guarded_body, dict): + guarded: Final = cast("dict[str, object]", guarded_body) # cast-ok: the hook's JSON response + return guarded, json.dumps(guarded).encode("utf-8"), True + verbose_proxy_logger.debug( + "pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response", + type(guarded_body).__name__, + ) + return response_body, content, False + + def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequence[object]) -> bool: names: Final = frozenset(chain.from_iterable(map(guardrail_entry_names, entries))) return any( @@ -2163,15 +2180,9 @@ async def pass_through_request( ), settles_the_charge=inherit_auth_guardrails, ) - if isinstance(guarded_body, dict): - response_body = guarded_body - content = json.dumps(response_body).encode("utf-8") - _content_modified = True - else: - verbose_proxy_logger.debug( - "pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response", - type(guarded_body).__name__, - ) + response_body, content, _content_modified = _response_after_post_call_guardrails( + guarded_body, response_body, content + ) elif response_body is None: verbose_proxy_logger.debug( "pass_through_endpoint: response body not JSON-parseable, skipping post-call guardrails" From 68ecee226a0d73b39a4cfe1fb7ca2b6070d51d33 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 04:43:54 -0800 Subject: [PATCH 26/29] fix(proxy): read tag and opt-out metadata the post-call guardrail selection trusts --- .../pass_through_endpoints.py | 15 +++--- .../test_passthrough_auth_guardrails.py | 52 ++++++++++++++++++- .../test_pass_through_endpoints.py | 16 ++++++ 3 files changed, 74 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 121bf0ed2c8..6ae21b02aff 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -204,7 +204,7 @@ _blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch= _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) _GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) -_POLICY_PIPELINE_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails") +_POST_CALL_SELECTION_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails", "tags") _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -930,9 +930,10 @@ async def _inherit_auth_guardrails( ) if entries and isinstance(metadata, dict): metadata["guardrails"] = entries + admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) admin_scoped_body: Final = { **parsed_body, - "metadata": {**(as_str_mapping(metadata) or {}), **_admin_guardrail_metadata(user_api_key_dict)}, + **{key: {**(as_str_mapping(parsed_body.get(key)) or {}), **admin_metadata} for key in sorted(_METADATA_KEYS)}, } return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None @@ -954,14 +955,14 @@ def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, ) -def _policy_pipeline_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: - """The policy pipelines resolved for a request that inherits auth guardrails, from the logging metadata that - request setup moved the body's ``metadata`` into""" +def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: + """The policy pipelines resolved for a request that inherits auth guardrails, and the tags a tag-scoped guardrail + mode reads, from the logging metadata that request setup moved the body's ``metadata`` into""" if not inherited: return MappingProxyType({}) litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} - return {field: metadata[field] for field in _POLICY_PIPELINE_FIELDS if field in metadata} + return {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata} def _starts_post_call_pass(guardrails: Sequence[object] | None, policy_pipelines: Mapping[str, object]) -> bool: @@ -2141,7 +2142,7 @@ async def pass_through_request( logging_obj=logging_obj, ) - policy_pipelines: Final = _policy_pipeline_metadata( + policy_pipelines: Final = _post_call_selection_metadata( cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above inherit_auth_guardrails, ) diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 690b018d8f3..d73e3db7d4c 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -11,7 +11,7 @@ import httpx import pytest import yaml from integration._support.anthropic_sse import message_json -from integration._support.client import Gateway, Scenario, eventually +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually from integration._support.openai_wire import chat_reply from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.safeguards_rig import key_hash @@ -23,6 +23,8 @@ GLOBAL_POST_CALL: Final = "passthrough-global-post-call" KEY_POST_CALL: Final = "passthrough-key-post-call" GENERIC_POST_CALL: Final = "passthrough-generic-post-call" BLOCKING_POST_CALL: Final = "passthrough-post-call-blocks-every-reply" +TAG_SCOPED_POST_CALL: Final = "passthrough-tag-scoped-post-call" +POST_CALL_TAG: Final = "scan-the-reply" MODEL_POST_CALL: Final = "passthrough-model-post-call" GUARDED_MODEL: Final = "passthrough-guarded-model" POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model" @@ -148,7 +150,10 @@ _MONITORED_ANTHROPIC_BODY: Final = json.dumps( _UNPRICED_ANTHROPIC_BODY: Final = json.dumps( - _JSON.validate_json(message_json("msg_passthrough_auth_unpriced", UNPRICED_MODEL, "Hi")) + { + **JSON_OBJECT.validate_json(message_json("msg_passthrough_auth_unpriced", UNPRICED_MODEL, "Hi")), + "usage": {"input_tokens": 300, "output_tokens": 200}, + } ).encode() @@ -289,6 +294,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": TAG_SCOPED_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": {"tags": {POST_CALL_TAG: "post_call"}, "default": "logging_only"}, + "default_on": False, + "api_base": blocking_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": MONITOR_POST_CALL, "litellm_params": { @@ -787,6 +802,39 @@ def test_a_blocked_response_with_no_cost_result_keeps_its_usage_on_the_failure_r assert rows[0]["api_base"], rows +def test_a_tpm_limit_counts_blocked_responses_with_no_cost_result(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [BLOCKING_POST_CALL]}, tpm_limit=600) + _drain(rig) + body: Final = {**_route_body(ANTHROPIC_ROUTE, "Hello"), "model": UNPRICED_MODEL} + + limited: Final = eventually( + lambda: _send(gateway, ANTHROPIC_ROUTE, body, key), + lambda response: response.status_code != 400, + seconds=70, + ) + + assert limited.status_code == 429, limited.text + + +def test_a_tag_scoped_post_call_guardrail_scans_a_reply_for_a_request_with_its_tag(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [TAG_SCOPED_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "Hello"), "metadata": {"tags": [POST_CALL_TAG]}}, + key, + ) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail_blocks(rig: AuthGuardrailRig) -> None: gateway: Final = rig.owned.gateway with gateway.scenario() as scenario: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 1fd7ee4d70c..3d89e1f21e3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -13517,6 +13517,22 @@ class TestInheritAuthGuardrails: assert metadata["user_api_key_metadata"]["disable_global_guardrails"] is True assert metadata["user_api_key_team_metadata"] == {"opted_out_global_guardrails": ["global-post"]} + @pytest.mark.asyncio + async def test_a_body_opt_out_in_litellm_metadata_does_not_stop_the_post_call_pass(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["global-post"]}) + + _, post_call = await self._run( + { + **self._chat_body(), + "litellm_metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, + }, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_awaited_once() + @pytest.mark.asyncio async def test_inherited_post_call_pass_runs_only_the_guardrails_the_request_selected(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) From 7e6a82d90ac9b288818ba02d30a20a95c1a722d1 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 10:39:23 -0800 Subject: [PATCH 27/29] fix(proxy): keep main's outcome for malformed body fields and custom endpoints on the post-call pass --- .../integrations/custom_guardrail.py | 7 +- .../guardrail_translation/handler.py | 5 +- .../generic_guardrail_api.py | 23 +++-- .../pass_through_endpoints.py | 40 ++++++-- .../test_passthrough_auth_guardrails.py | 93 +++++++++++++++++++ .../integrations/test_custom_guardrail.py | 26 ++++++ .../guardrail_translation/test_handler.py | 14 ++- .../test_generic_guardrail_api.py | 49 ++++++++++ .../test_pass_through_endpoints.py | 77 ++++++++++++++- 9 files changed, 312 insertions(+), 22 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/custom_guardrail.py b/enterprise/litellm_enterprise/integrations/custom_guardrail.py index 58c24219b47..56452bf3e85 100644 --- a/enterprise/litellm_enterprise/integrations/custom_guardrail.py +++ b/enterprise/litellm_enterprise/integrations/custom_guardrail.py @@ -3,6 +3,11 @@ from typing import List, Optional, Union from litellm.types.guardrails import GuardrailEventHooks, Mode +def _names_a_mode_tag(tag: object, mode: Mode) -> bool: + """Whether a request tag, which a client's body can make anything, is one of ``mode``'s tags""" + return isinstance(tag, str) and tag in mode.tags + + class EnterpriseCustomGuardrailHelper: @staticmethod def _should_run_if_mode_by_tag( @@ -55,7 +60,7 @@ class EnterpriseCustomGuardrailHelper: matched_mode = None if request_tags: for tag in request_tags: - if tag in event_hook.tags: + if _names_a_mode_tag(tag, event_hook): matched_mode = event_hook.tags[tag] break diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 712e51f7afc..0af14c23654 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -243,7 +243,8 @@ class PassThroughEndpointHandler(BaseTranslation): guardrail_to_apply: The guardrail instance to apply litellm_logging_obj: Optional logging object user_api_key_dict: User API key metadata to pass to guardrails - read_whole_response: Without targeted fields, read the response as JSON, ids and usage included + read_whole_response: Without targeted fields, read the response as JSON, ids and usage included, as a call + with no logging object does, such as a logging_only scan """ response_fields: Final = as_str_mapping(response) if not isinstance(response, dict): @@ -263,7 +264,7 @@ class PassThroughEndpointHandler(BaseTranslation): # Extract text to check text_to_check: Final = ( self._extract_text_for_guardrail(response, field_expressions) - if field_expressions or read_whole_response + if field_expressions or read_whole_response or litellm_logging_obj is None else _response_text(response_fields or {}) ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index f3309b98462..5e47ec88a3e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -378,12 +378,13 @@ class GenericGuardrailAPI(CustomGuardrail): # Fields in metadata are already prefixed with 'user_api_key_' for field_name in GenericGuardrailAPIMetadata.__annotations__: value = metadata_dict.get(field_name) - if value is not None: + if isinstance(value, str): result_metadata[field_name] = value # handle user_api_key_token = user_api_key_hash - if metadata_dict.get("user_api_key_token") is not None: - result_metadata["user_api_key_hash"] = metadata_dict.get("user_api_key_token") + token: Final = metadata_dict.get("user_api_key_token") + if isinstance(token, str): + result_metadata["user_api_key_hash"] = token verbose_proxy_logger.debug( "Generic Guardrail API: Extracted user metadata: %s", @@ -392,6 +393,14 @@ class GenericGuardrailAPI(CustomGuardrail): return result_metadata + def _dynamic_request_body_params(self, request_body: object) -> Mapping[str, object]: + """The ``extra_body`` a request's ``body`` asks this guardrail to send. A custom endpoint's client can send a + ``body`` field of any shape, which asks for nothing unless it names this guardrail's params in the usual form.""" + try: + return as_str_mapping(self.get_guardrail_dynamic_request_body_params(request_body)) or {} # pyright: ignore[reportArgumentType, reportUnknownArgumentType] # a client's JSON, validated by the call + except (TypeError, ValidationError): + return {} + def _fail_open_passthrough( self, *, @@ -561,16 +570,16 @@ class GenericGuardrailAPI(CustomGuardrail): tools: Final = inputs.get("tools") structured_messages: Final = inputs.get("structured_messages") tool_calls: Final = inputs.get("tool_calls") - model: Final = inputs.get("model") + model: Final[object] = inputs.get("model") # Use provided request_data or create an empty dict if request_data is None: request_data = {} - request_body: Final = request_data.get("body") or {} + request_body: Final[object] = request_data.get("body") # Merge additional provider specific params from config and dynamic params - dynamic_params: Final[Mapping[str, object]] = self.get_guardrail_dynamic_request_body_params(request_body) or {} + dynamic_params: Final = self._dynamic_request_body_params(request_body) additional_params: Final = {**self.additional_provider_specific_params, **dynamic_params} # Extract user API key metadata @@ -597,7 +606,7 @@ class GenericGuardrailAPI(CustomGuardrail): tool_calls=tool_calls, additional_provider_specific_params=additional_params, input_type=input_type, - model=model, + model=model if isinstance(model, str) else None, ) headers: Final = self._build_request_headers() diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6ae21b02aff..d90b7301a4d 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -204,6 +204,11 @@ _blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch= _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) _GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) +_FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = ( + ("guardrails", (list, dict)), + ("guardrail_config", (dict,)), + ("proxy_server_request", (dict,)), +) _POST_CALL_SELECTION_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails", "tags") _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -889,12 +894,11 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): def _drop_unmergeable_guardrail_fields( parsed_body: dict[str, object], # mutable-ok: strips the request body in place before the guardrail merge ) -> None: - """Drop a ``guardrails`` that isn't a list or an object and a ``guardrail_config`` that isn't an object, which the - merge can't combine with the metadata's own. The provider receives neither field.""" - if "guardrails" in parsed_body and not isinstance(parsed_body["guardrails"], list | dict): - parsed_body.pop("guardrails") - if "guardrail_config" in parsed_body and not isinstance(parsed_body["guardrail_config"], dict): - parsed_body.pop("guardrail_config") + """Drop a ``guardrails`` that isn't a list or an object, and a ``guardrail_config`` or ``proxy_server_request`` that + isn't an object, which the merge and the guardrails can't read. The provider receives none of these fields.""" + for field, shape in _FIELDS_GUARDRAILS_READ_BY_SHAPE: + if field in parsed_body and not isinstance(parsed_body[field], shape): + parsed_body.pop(field) async def _inherit_auth_guardrails( @@ -913,6 +917,7 @@ async def _inherit_auth_guardrails( if configured_guardrails or not inherit or not parsed_body: return configured_guardrails _drop_unmergeable_guardrail_fields(parsed_body) + client_sent_metadata: Final = "metadata" in parsed_body if not isinstance(parsed_body.get("metadata"), dict): parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") @@ -922,6 +927,9 @@ async def _inherit_auth_guardrails( user_api_key_dict=user_api_key_dict, ) metadata: Final = parsed_body["metadata"] + if not metadata and not client_sent_metadata: + parsed_body.pop("metadata") + return None merged: Final = (as_str_mapping(metadata) or {}).get("guardrails") if not isinstance(merged, list): return None @@ -955,6 +963,12 @@ def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, ) +def _holds_resolved_pipelines(pipelines: object) -> bool: + """Whether a pipeline slot holds the policy engine's (name, pipeline) tuples, which a JSON body can't""" + entries: Final = as_object_list(pipelines) + return entries is not None and all(isinstance(entry, tuple) for entry in entries) + + def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: """The policy pipelines resolved for a request that inherits auth guardrails, and the tags a tag-scoped guardrail mode reads, from the logging metadata that request setup moved the body's ``metadata`` into""" @@ -962,7 +976,10 @@ def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) return MappingProxyType({}) litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} - return {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata} + selection: Final = {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata} + if "_guardrail_pipelines" in selection and not _holds_resolved_pipelines(selection["_guardrail_pipelines"]): + selection.pop("_guardrail_pipelines") + return selection def _starts_post_call_pass(guardrails: Sequence[object] | None, policy_pipelines: Mapping[str, object]) -> bool: @@ -980,10 +997,11 @@ def _post_call_hook_data( """The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails, which ``should_run_guardrail`` checks, the policy pipelines to run, and the logging object, which unified (apply_guardrail) guardrails resolve a pass-through response's call type from when the route has no - API_ROUTE_TO_CALL_TYPES entry""" + API_ROUTE_TO_CALL_TYPES entry. A custom endpoint whose guardrails don't run post_call gets no logging object, so a + global unified guardrail still can't resolve its response.""" body: Final = parsed_body or {} admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) if inherited else MappingProxyType({}) - return { + data: Final = { **body, "metadata": { **(as_str_mapping(body.get("metadata")) or {}), @@ -991,8 +1009,10 @@ def _post_call_hook_data( "guardrails": guardrails or (), **admin_metadata, }, - "litellm_logging_obj": logging_obj, } + if not inherited and not _named_guardrail_runs_post_call(data, guardrails or ()): + return data + return {**data, "litellm_logging_obj": logging_obj} def _provider_charge_settlement( diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index d73e3db7d4c..207ba6fff75 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -47,6 +47,8 @@ ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" CUSTOM_ROUTE: Final = "/passthrough-auth-custom" CUSTOM_GUARDED_ROUTE: Final = "/passthrough-auth-custom-guarded" +CUSTOM_PRE_CALL_ROUTE: Final = "/passthrough-auth-custom-pre-call" +GLOBAL_UNIFIED_POST_CALL: Final = "passthrough-global-unified-post-call" UNPRICED_MODEL: Final = "passthrough-auth-unpriced-model" MONITOR_POST_CALL: Final = "passthrough-auth-monitor-post-call" MONITORED_PROMPT: Final = "reply-with-the-monitored-id" @@ -174,6 +176,7 @@ class AuthGuardrailRig: post_call_vendor: Wire key_post_call_vendor: Wire generic_post_call_vendor: Wire + global_unified_vendor: Wire model_post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -186,6 +189,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) + global_unified_vendor: Final = stack.enter_context(wire_server(_verdict)) blocking_post_call_vendor: Final = stack.enter_context(wire_server(_block_every_reply)) monitor_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_pii_only)) model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) @@ -223,6 +227,12 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "cost_per_request": 0.25, "guardrails": {BLOCKING_POST_CALL: None}, }, + { + "path": CUSTOM_PRE_CALL_ROUTE, + "target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions", + "auth": True, + "guardrails": {GUARDRAIL: None}, + }, ], }, "policies": { @@ -315,6 +325,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_key": "synthetic-guardrail-key", }, }, + { + "guardrail_name": GLOBAL_UNIFIED_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": True, + "api_base": global_unified_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, { "guardrail_name": GENERIC_POST_CALL, "litellm_params": { @@ -365,6 +385,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I post_call_vendor, key_post_call_vendor, generic_post_call_vendor, + global_unified_vendor, model_post_call_vendor, upstreams, ) @@ -850,6 +871,78 @@ def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail assert [float(str(row["spend"])) for row in _row_fields(gateway, key)] == [0], response.text +def test_a_custom_endpoint_with_only_a_pre_call_guardrail_never_calls_a_global_unified_guardrail( + rig: AuthGuardrailRig, +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + custom_key: Final = scenario.key(metadata={"allowed_passthrough_routes": [CUSTOM_PRE_CALL_ROUTE]}) + built_in_key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + rig.global_unified_vendor.drain() + + custom_response: Final = gateway.request( + "POST", CUSTOM_PRE_CALL_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=custom_key + ) + + assert custom_response.status_code == 200, custom_response.text + assert len(rig.vendor.drain()) == 1 + assert rig.global_unified_vendor.drain() == () + + built_in_response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), built_in_key) + + assert built_in_response.status_code == 200, built_in_response.text + assert len(rig.global_unified_vendor.drain()) == 1 + + +@pytest.mark.parametrize( + "field", ["Hello there", 7, ["a", "b"], {"guardrails": 7}], ids=["text", "int", "list", "dict"] +) +def test_a_custom_endpoints_own_body_field_still_gets_its_reply_scanned( + rig: AuthGuardrailRig, field: JsonValue +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE]}) + + response: Final = gateway.request( + "POST", CUSTOM_GUARDED_ROUTE, {**_route_body(OPENAI_ROUTE, "Hello"), "body": field}, key=key + ) + + assert response.status_code == 400, response.text + assert "synthetic block" in response.text + + +@pytest.mark.parametrize( + ("guardrail", "fields"), + [ + (GUARDRAIL, {"proxy_server_request": 5}), + (GUARDRAIL, {"metadata": {"user_api_key_alias": 5}}), + (TAG_SCOPED_POST_CALL, {"metadata": {"tags": [1, {"a": 1}]}}), + ( + None, + { + "metadata": {"_pipeline_managed_guardrails": []}, + "litellm_metadata": {"_guardrail_pipelines": [["p", "q"]]}, + }, + ), + ], + ids=["proxy-server-request", "key-alias", "tags", "pipeline-slots"], +) +def test_a_malformed_proxy_field_in_a_body_still_gets_the_provider_reply( + rig: AuthGuardrailRig, guardrail: str | None, fields: Mapping[str, JsonValue] +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [guardrail]} if guardrail else {}) + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE, "Hello"), **fields}, key) + + assert response.status_code == 200, response.text + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + def test_a_monitored_lakera_flag_keeps_the_provider_id_on_the_row(rig: AuthGuardrailRig) -> None: gateway: Final = rig.owned.gateway with gateway.scenario() as scenario: diff --git a/tests/unit/enterprise/integrations/test_custom_guardrail.py b/tests/unit/enterprise/integrations/test_custom_guardrail.py index 84f34144d99..6c6fb13ff59 100644 --- a/tests/unit/enterprise/integrations/test_custom_guardrail.py +++ b/tests/unit/enterprise/integrations/test_custom_guardrail.py @@ -330,3 +330,29 @@ async def test_baggage_never_changes_the_guardrail_mode_key_team_and_deployment_ assert await runs_post_call_guardrail(()) is runs_post_call assert await runs_post_call_guardrail(("usecase:dev",)) is runs_post_call + + +@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"]) +def test_a_request_tag_that_cant_match_falls_back_to_the_default_mode(monkeypatch, bucket): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cg = CustomGuardrail( + guardrail_name="test_guardrail", + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + event_hook=Mode(tags={"test_tag": "post_call"}, default="pre_call"), + default_on=True, + ) + + assert ( + cg.should_run_guardrail( + data={bucket: {"tags": [1, {"a": 1}, ["b"]]}}, + event_type=GuardrailEventHooks.pre_call, + ) + is True + ) + assert ( + cg.should_run_guardrail( + data={bucket: {"tags": [{"a": 1}, "test_tag"]}}, + event_type=GuardrailEventHooks.pre_call, + ) + is False + ) diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index e028ff2a7bd..1ba79b704da 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -231,10 +231,22 @@ def test_passthrough_route_block_before_any_invoke_frame_is_a_whole_message_in_c async def _scanned_response_texts(response: dict[str, object]) -> tuple[str, ...]: guardrail: Final = MagicMock(guardrail_name="post-call", apply_guardrail=AsyncMock()) - await PassThroughEndpointHandler().process_output_response(response, guardrail, request_data={}) + await PassThroughEndpointHandler().process_output_response( + response, guardrail, MagicMock(passthrough_guardrails_config=None), request_data={} + ) return tuple(text for call in guardrail.apply_guardrail.await_args_list for text in call.kwargs["inputs"]["texts"]) +@pytest.mark.asyncio +async def test_a_scan_with_no_logging_object_reads_the_response_as_json(): + response: Final = {"response": {"id": "chatcmpl-1", "object": "chat.completion", "choices": [{"text": "Paris."}]}} + guardrail: Final = MagicMock(guardrail_name="logging-only", apply_guardrail=AsyncMock()) + + await PassThroughEndpointHandler().process_output_response(response, guardrail, request_data={}) + + assert json.loads(guardrail.apply_guardrail.await_args.kwargs["inputs"]["texts"][0]) == response + + @pytest.mark.asyncio async def test_post_call_scan_reads_the_reply_text_and_not_the_numbers_around_it(): response: Final = { diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index daa9945e406..c09eb4ed5c5 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -2703,3 +2703,52 @@ def test_extract_user_api_key_metadata_reads_past_a_litellm_metadata_that_is_not ) assert metadata["user_api_key_hash"] == "hash" + + +def test_extract_user_api_key_metadata_leaves_out_a_value_that_isnt_a_string() -> None: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + + metadata: Final = guardrail._extract_user_api_key_metadata( + {"metadata": {"user_api_key_alias": 5, "user_api_key_token": 7, "user_api_key_team_id": "team"}} + ) + + assert metadata == {"user_api_key_team_id": "team"} + + +async def _vendor_payload(request_data: dict[str, object], inputs: dict[str, object]) -> dict[str, object]: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + vendor_reply: Final = MagicMock(json=MagicMock(return_value={"action": "NONE"})) + + with patch.object(guardrail.async_handler, "post", return_value=vendor_reply) as post: + await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="response") + + return post.call_args.kwargs["json"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + ( + "Hello there", + 7, + True, + ["a", "b"], + {"guardrails": 7}, + {"metadata": {"guardrails": 7}}, + {"guardrails": [{"generic": 7}]}, + ), +) +async def test_a_body_field_a_custom_endpoints_client_sent_asks_the_vendor_for_nothing(body: object) -> None: + payload: Final = await _vendor_payload({"body": body}, {"texts": ["reply"]}) + + assert payload["texts"] == ["reply"] + assert not payload.get("additional_provider_specific_params") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ({"id": 1}, ["gpt"], 7)) +async def test_a_reply_model_that_isnt_a_string_is_left_out_of_the_vendor_request(model: object) -> None: + payload: Final = await _vendor_payload({}, {"texts": ["reply"], "model": model}) + + assert payload["texts"] == ["reply"] + assert payload.get("model") is None diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 3d89e1f21e3..da74d6c9cf4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -58,6 +58,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( SafeRouteAdder, _anthropic_user_id_to_forward, _inherit_auth_guardrails, + _post_call_hook_data, + _post_call_selection_metadata, + _starts_post_call_pass, WebsocketFrameModelAllowed, WebsocketFrameModelGate, WebsocketFrameModelRefused, @@ -13347,11 +13350,83 @@ class TestInheritAuthGuardrails: await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test", metadata=key_metadata), None, True) assert "guardrails" not in body - assert (body["metadata"].get("guardrails") or []) == [ + assert (body.get("metadata", {}).get("guardrails") or []) == [ *key_metadata.get("guardrails", ()), *(body_metadata or {}).get("guardrails", ()), ] + @pytest.mark.asyncio + async def test_a_body_with_nothing_to_merge_gets_no_metadata(self): + body: Final[dict[str, object]] = self._chat_body() + + assert await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) is None + + assert body == self._chat_body() + + @pytest.mark.asyncio + async def test_a_body_with_nothing_to_merge_keeps_the_metadata_it_sent(self): + body: Final[dict[str, object]] = self._chat_body(metadata={}) + + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) + + assert body["metadata"] == {} + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", [5, "text", ["list"], None]) + async def test_a_proxy_server_request_that_isnt_an_object_never_reaches_the_guardrails(self, value): + body: Final[dict[str, object]] = self._chat_body(proxy_server_request=value) + + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) + + assert "proxy_server_request" not in body + + @pytest.mark.parametrize("slot", [[["p", "q"]], [7], 7, "text", [{"mode": "post_call"}]]) + def test_a_pipeline_slot_a_body_sent_does_not_start_the_post_call_pass(self, slot): + kwargs: Final = { + "litellm_params": { + "metadata": {"_guardrail_pipelines": slot, "_pipeline_managed_guardrails": [], "tags": ["t"]} + } + } + + selection: Final = _post_call_selection_metadata(kwargs, inherited=True) + + assert selection == {"_pipeline_managed_guardrails": [], "tags": ["t"]} + assert _starts_post_call_pass(None, selection) is False + + def test_a_pipeline_the_policy_engine_resolved_starts_the_post_call_pass(self): + pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="guard")]) + kwargs: Final = {"litellm_params": {"metadata": {"_guardrail_pipelines": [("p", pipeline)]}}} + + selection: Final = _post_call_selection_metadata(kwargs, inherited=True) + + assert selection == {"_guardrail_pipelines": [("p", pipeline)]} + assert _starts_post_call_pass(None, selection) is True + + @pytest.mark.parametrize( + ("event_hook", "inherited", "gets_the_logging_object"), + [ + (GuardrailEventHooks.pre_call, False, False), + (GuardrailEventHooks.post_call, False, True), + (GuardrailEventHooks.pre_call, True, True), + ], + ids=["custom-pre-call", "custom-post-call", "built-in-route"], + ) + def test_a_custom_endpoint_gets_the_logging_object_only_for_a_post_call_guardrail( + self, event_hook, inherited, gets_the_logging_object + ): + logging_obj: Final = MagicMock() + + with patch.object( + litellm.logging_callback_manager, + "get_custom_loggers_for_type", + return_value=[self._guardrail("guard", event_hook)], + ): + data: Final = _post_call_hook_data( + self._chat_body(), ["guard"], logging_obj, UserAPIKeyAuth(api_key="sk-test"), {}, inherited + ) + + assert data.get("litellm_logging_obj") is (logging_obj if gets_the_logging_object else None) + @pytest.mark.asyncio @pytest.mark.parametrize("value", [None, False, 7, "config", ["config"]]) async def test_a_guardrail_config_the_merge_cant_add_to_is_dropped(self, value): From b35faae71e23db8ae8addd5fb3aa3cf875e9056a Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 12:38:25 -0800 Subject: [PATCH 28/29] fix(proxy): keep spend_logs_metadata on the row of a blocked pass-through response --- .../pass_through_endpoints.py | 23 +++++-- .../test_passthrough_auth_guardrails.py | 69 ++++++++++++++++++- 2 files changed, 87 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d90b7301a4d..a2b77b3168e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -209,7 +209,12 @@ _FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = ( ("guardrail_config", (dict,)), ("proxy_server_request", (dict,)), ) -_POST_CALL_SELECTION_FIELDS: Final = ("_guardrail_pipelines", "_pipeline_managed_guardrails", "tags") +_POST_CALL_SELECTION_FIELDS: Final = ( + "_guardrail_pipelines", + "_pipeline_managed_guardrails", + "spend_logs_metadata", + "tags", +) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -970,8 +975,9 @@ def _holds_resolved_pipelines(pipelines: object) -> bool: def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: - """The policy pipelines resolved for a request that inherits auth guardrails, and the tags a tag-scoped guardrail - mode reads, from the logging metadata that request setup moved the body's ``metadata`` into""" + """The policy pipelines resolved for a request that inherits auth guardrails, the tags a tag-scoped guardrail mode + reads and the spend_logs_metadata a blocked response's row keeps, from the logging metadata that request setup moved + the body's ``metadata`` into""" if not inherited: return MappingProxyType({}) litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} @@ -1100,10 +1106,19 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict[str, ``standard_logging_guardrail_information`` there and not on the dict the failure handler forwards to ``post_call_failure_hook``. Without this the otel guardrail span is emitted on allow but missing on block. Carry the - entries over so the failure path matches the unified path. + entries over so the failure path matches the unified path. The failure row + reads its ``metadata`` here too, so it keeps the ``spend_logs_metadata`` a + built-in route's post-call pass was given, as its success row would. """ if guardrail_data is None: return + spend_logs_metadata: Final = (as_str_mapping(guardrail_data.get("metadata")) or {}).get("spend_logs_metadata") + if spend_logs_metadata is not None: + request_fields: Final = cast("Mapping[str, object]", request_data) # cast-ok: the JSON body plus logging kwargs + request_data["metadata"] = { # rebind-ok: the failure path hands this dict to post_call_failure_hook + **(as_str_mapping(request_fields.get("metadata")) or {}), + "spend_logs_metadata": spend_logs_metadata, + } source_key: Final = get_metadata_variable_name_from_kwargs(guardrail_data) source_metadata: Final = as_str_mapping(guardrail_data.get(source_key)) or {} entries: Final = source_metadata.get("standard_logging_guardrail_information") diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 207ba6fff75..222fffb3e7f 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -57,6 +57,7 @@ PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call" PIPELINE_BLOCKS: Final = "passthrough-pipeline-blocks-the-response" PIPELINE_ALLOWS: Final = "passthrough-pipeline-allows-the-response" PIPELINE_BEFORE_THE_CALL: Final = "passthrough-pipeline-before-the-call" +PIPELINE_REPLACES: Final = "passthrough-pipeline-replaces-the-response" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) @@ -239,10 +240,11 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I PIPELINE_BLOCKS: _pipeline_policy("post_call", KEY_POST_CALL, "block"), PIPELINE_ALLOWS: _pipeline_policy("post_call", KEY_POST_CALL, "allow"), PIPELINE_BEFORE_THE_CALL: _pipeline_policy("pre_call", GUARDRAIL, "block"), + PIPELINE_REPLACES: _pipeline_policy("post_call", KEY_POST_CALL, "modify_response"), }, "policy_attachments": [ {"policy": policy, "tags": [policy]} - for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL) + for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL, PIPELINE_REPLACES) ], "model_list": [ *_JSON_LIST.validate_python(config.get("model_list") or []), @@ -775,6 +777,71 @@ def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( assert int(str(rows[0]["total_tokens"])) > 0, rows +def _spend_logs_metadata_of_the_failure_row(gateway: Gateway, key: str) -> JsonValue: + rows: Final = eventually( + lambda: gateway.rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (key_hash(key),)), + bool, + ) + assert [row["status"] for row in rows] == ["failure"], rows + return _JSON.validate_python(rows[0]["metadata"]).get("spend_logs_metadata") + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_a_blocked_response_keeps_the_spend_logs_metadata_on_its_row(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"spend_logs_metadata": {"team_field": "team"}}) + key: Final = scenario.key( + team_id=team, + metadata={"guardrails": [BLOCKING_POST_CALL], "spend_logs_metadata": {"key_field": "key"}}, + ) + _drain(rig) + body: Final = { + **_route_body(route, "What is the capital of France?"), + "metadata": {"spend_logs_metadata": {"body_field": "body"}}, + } + + response: Final = _send(gateway, route, body, key) + + assert response.status_code == 400, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) == { + "body_field": "body", + "key_field": "key", + "team_field": "team", + } + + +def test_a_response_a_pipeline_replaces_keeps_the_spend_logs_metadata_on_its_row(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"spend_logs_metadata": {"key_field": "key"}}) + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_REPLACES), key) + + assert response.status_code == 200, response.text + assert b"content_filter" in response.content, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) == {"key_field": "key"} + + +def test_a_custom_endpoints_blocked_response_row_gets_no_spend_logs_metadata(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key( + team_id=scenario.team(), + metadata={ + "allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE], + "spend_logs_metadata": {"key_field": "key"}, + }, + ) + _drain(rig) + + response: Final = gateway.request("POST", CUSTOM_GUARDED_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=key) + + assert response.status_code == 400, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) is None + + def _row_fields(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ...]: return tuple( eventually( From 4d45e22c375cf8e1efdc388a5466cfad27c56107 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 12:45:06 -0800 Subject: [PATCH 29/29] test(proxy): name the modify_response pipeline as the stacked header tests do --- .../authorization/test_passthrough_auth_guardrails.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 222fffb3e7f..b3d711c3e9d 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -57,7 +57,7 @@ PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call" PIPELINE_BLOCKS: Final = "passthrough-pipeline-blocks-the-response" PIPELINE_ALLOWS: Final = "passthrough-pipeline-allows-the-response" PIPELINE_BEFORE_THE_CALL: Final = "passthrough-pipeline-before-the-call" -PIPELINE_REPLACES: Final = "passthrough-pipeline-replaces-the-response" +PIPELINE_MODIFIES: Final = "passthrough-pipeline-modifies-the-response" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) @@ -240,11 +240,11 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I PIPELINE_BLOCKS: _pipeline_policy("post_call", KEY_POST_CALL, "block"), PIPELINE_ALLOWS: _pipeline_policy("post_call", KEY_POST_CALL, "allow"), PIPELINE_BEFORE_THE_CALL: _pipeline_policy("pre_call", GUARDRAIL, "block"), - PIPELINE_REPLACES: _pipeline_policy("post_call", KEY_POST_CALL, "modify_response"), + PIPELINE_MODIFIES: _pipeline_policy("post_call", KEY_POST_CALL, "modify_response"), }, "policy_attachments": [ {"policy": policy, "tags": [policy]} - for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL, PIPELINE_REPLACES) + for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL, PIPELINE_MODIFIES) ], "model_list": [ *_JSON_LIST.validate_python(config.get("model_list") or []), @@ -817,7 +817,7 @@ def test_a_response_a_pipeline_replaces_keeps_the_spend_logs_metadata_on_its_row key: Final = scenario.key(metadata={"spend_logs_metadata": {"key_field": "key"}}) _drain(rig) - response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_REPLACES), key) + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_MODIFIES), key) assert response.status_code == 200, response.text assert b"content_filter" in response.content, response.text