From c0116078b939112f5aecb6577483bb2f6119dc06 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 12:13:16 -0700 Subject: [PATCH] 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"]})