fix(policy_engine): skip pipeline-managed guardrails in the response-path guardrail loop

This commit is contained in:
mateo-berri 2026-08-28 17:59:15 -07:00
parent e6edd62f5d
commit aeac6a412c
2 changed files with 34 additions and 1 deletions

View file

@ -2825,6 +2825,7 @@ class ProxyLogging:
if pipeline_response is not None:
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
pipeline_managed: Final = _pipeline_managed_guardrail_names(data)
guardrail_callbacks: Final[list[CustomGuardrail]] = []
other_callbacks: Final[list[CustomLogger]] = []
try:
@ -2849,12 +2850,18 @@ class ProxyLogging:
guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False)
callback
for callback in guardrail_callbacks
if getattr(callback, "run_in_parallel", False)
and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed)
)
for callback in guardrail_callbacks:
# Main - V2 Guardrails implementation
if callback.guardrail_name and callback.guardrail_name in pipeline_managed:
continue
if getattr(callback, "run_in_parallel", False):
continue

View file

@ -951,6 +951,32 @@ async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouch
assert "guardrails" not in data["metadata"]
@pytest.mark.asyncio
async def test_post_call_pipeline_managed_default_on_guardrail_runs_exactly_once(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {"count": 0}
class CountingGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
seen["count"] += 1
return None
monkeypatch.setattr(
litellm,
"callbacks",
[CountingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)],
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data()
await proxy_logging.post_call_success_hook(
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
)
assert seen["count"] == 1
@pytest.mark.asyncio
async def test_post_call_pipeline_replacement_response_reaches_caller(
proxy_logging, make_user_api_key_auth, monkeypatch