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