From a8dc662b933fce540c99b1d9007adc212cc5fe5c Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 14:27:44 -0800 Subject: [PATCH] fix(proxy): keep post-call guardrail headers from the pass before a successful reply --- .../pass_through_endpoints.py | 7 +- .../test_passthrough_auth_guardrails.py | 17 ++- .../test_pass_through_endpoints.py | 103 +++++++++--------- 3 files changed, 70 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 793d07bfdde..7b28b25ffc6 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1444,8 +1444,11 @@ async def _log_passthrough_upstream_failure_keeping_guardrail_headers( kwargs: Mapping[str, object] | None, ) -> Mapping[str, str]: """Log an upstream error like ``_log_passthrough_upstream_failure`` and return the guardrail response headers, - read before the failure hooks rebuild ``kwargs``' metadata without the guardrail fields.""" - guardrail_headers: Final = _guardrail_response_headers(parsed_body, kwargs) + read before the failure hooks rebuild ``kwargs``' metadata without the guardrail fields. A successful response + has no failure hooks to read around, so it returns none.""" + guardrail_headers: Final = ( + _guardrail_response_headers(parsed_body, kwargs) if response.status_code >= 400 else MappingProxyType({}) + ) await _log_passthrough_upstream_failure( response=response, user_api_key_dict=user_api_key_dict, diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index e8d38689547..819a2653f6b 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -47,6 +47,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_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]) @@ -181,10 +182,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_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) + for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL, PIPELINE_MODIFIES) ], "model_list": [ *_JSON_LIST.validate_python(config.get("model_list") or []), @@ -618,6 +620,19 @@ def test_post_call_policy_pipeline_decides_what_a_flagged_step_does(rig: AuthGua assert KEY_POST_CALL in response.headers["x-litellm-applied-guardrails"].split(",") +def test_response_a_policy_pipeline_modifies_names_the_step_that_ran(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_MODIFIES), key) + + assert response.status_code == 200, response.text + assert response.json()["error"]["type"] == "content_filter" + assert KEY_POST_CALL in response.headers["x-litellm-applied-guardrails"].split(",") + + def test_pre_call_policy_pipeline_relays_the_provider_bytes(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 1ee3ee6c0a9..6d22cb8f739 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 @@ -8,7 +8,7 @@ import os import sys import uuid import zlib -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime @@ -64,6 +64,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( _build_passthrough_failure_request_payload, _registered_pass_through_routes, _guardrail_response_headers, + _log_passthrough_upstream_failure_keeping_guardrail_headers, _strip_client_guardrail_fields, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -503,38 +504,61 @@ async def test_pass_through_request_preserves_proxy_exception_status(): assert exc.value.param == "limit" -@pytest.mark.asyncio -async def test_guardrail_proxy_exception_carries_the_guardrails_that_ran(): - blocked: Final = ProxyException(message="blocked by guardrail", type="None", param="None", code=400) - - async def record_then_block( - user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str, endpoint_type: object +class _RecordsThenRaisesProxyException(CustomGuardrail): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, ) -> None: data["metadata"] = {"applied_guardrails": ["passed-guard"]} # mutable-ok: a guardrail records on the live body - raise blocked + raise ProxyException(message="blocked by guardrail", type="None", param="None", code=400) - with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: - mock_proxy_logging.post_call_failure_hook = AsyncMock() - mock_proxy_logging.pre_call_hook = AsyncMock(side_effect=record_then_block) - mock_request = MagicMock(spec=Request) - mock_request.method = "POST" - mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "messages": []}') - mock_request.headers = Headers({"content-type": "application/json"}) - mock_request.query_params = QueryParams({}) +@pytest.mark.asyncio +async def test_guardrail_proxy_exception_carries_the_guardrails_that_ran(): + guardrail: Final = _RecordsThenRaisesProxyException( + guardrail_name="passed-guard", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + mock_request: Final = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "messages": []}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) - with pytest.raises(ProxyException) as exc: - await pass_through_request( - request=mock_request, - target="http://test.com/v1/chat/completions", - custom_headers={}, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), - ) + with ( + pytest.MonkeyPatch.context() as litellm_patch, + patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ), + pytest.raises(ProxyException) as exc, + ): + litellm_patch.setattr(litellm, "callbacks", [guardrail]) + await pass_through_request( + request=mock_request, + target="http://test.com/v1/chat/completions", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) - assert exc.value is blocked assert exc.value.headers["x-litellm-applied-guardrails"] == "passed-guard" +@pytest.mark.asyncio +async def test_successful_response_keeps_no_guardrail_headers_from_before_the_post_call_pass(): + headers: Final = await _log_passthrough_upstream_failure_keeping_guardrail_headers( + response=httpx.Response(status_code=200), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + request_payload={}, + logging_obj=MagicMock(), + parsed_body={"metadata": {"applied_guardrails": ["pre-call-guard"]}}, + kwargs=None, + ) + + assert headers == {} + + def test_is_langfuse_route(): """ Test that the is_langfuse_route method correctly identifies Langfuse routes @@ -13312,8 +13336,6 @@ class TestInheritAuthGuardrails: inherit: bool, user_api_key_dict: UserAPIKeyAuth, guardrails: tuple[CustomGuardrail, ...] = (), - post_call_side_effect: Callable[..., Awaitable[None]] | None = None, - failure_side_effect: Callable[..., Awaitable[None]] | None = None, ) -> tuple[Response, list[dict[str, object]], AsyncMock]: pre_call_data: Final[list[dict[str, object]]] = [] @@ -13323,10 +13345,7 @@ class TestInheritAuthGuardrails: 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"}, side_effect=post_call_side_effect - ) - proxy_logging.post_call_failure_hook = AsyncMock(side_effect=failure_side_effect) + 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, @@ -13378,30 +13397,6 @@ class TestInheritAuthGuardrails: def _chat_body(**extra: object) -> dict[str, object]: return {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], **extra} - @pytest.mark.asyncio - async def test_modified_response_names_a_guardrail_the_failure_hook_drops(self): - async def _modify_response(**kwargs: object) -> None: - data: Final = cast(dict[str, dict[str, object]], kwargs["data"]) # cast-ok: the hook data the handler built - data["metadata"]["applied_guardrails"] = ["key-guard"] - raise ModifyResponseException( - message="replaced", model="gpt-4o", request_data=data, guardrail_name="key-guard" - ) - - async def _drop_guardrail_fields(**kwargs: object) -> None: - cast(dict[str, dict[str, object]], kwargs["request_data"])["metadata"].pop("applied_guardrails", None) # cast-ok: payload - - response, _, _ = await self._run_with_response( - self._chat_body(), - inherit=True, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}), - guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), - post_call_side_effect=_modify_response, - failure_side_effect=_drop_guardrail_fields, - ) - - assert response.status_code == 200 - assert response.headers["x-litellm-applied-guardrails"] == "key-guard" - @pytest.mark.asyncio async def test_key_team_and_project_guardrails_reach_pre_and_post_call(self): auth: Final = UserAPIKeyAuth(