fix(proxy): relay pass-through responses untouched when no inherited guardrail runs post_call

This commit is contained in:
yucheng 2026-10-31 00:50:37 -07:00
parent 5011b1263f
commit 205026dbe4
3 changed files with 109 additions and 12 deletions

View file

@ -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:

View file

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

View file

@ -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"]})