From 6475443efbfecb52ef693bfd4e7149908b4283a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:00:25 -0700 Subject: [PATCH] fix(policy_engine): run per-chunk hook guardrails whose post_call pipeline cannot stream The per-chunk streaming hook skipped every guardrail stepped by a post_call pipeline, even when the pipeline is left out of the stream for lacking the unified apply_guardrail interface, so a default_on guardrail that only implements async_post_call_streaming_hook stopped governing streams it governed on the merge base. The skip set now comes from the pipelines that gate the stream, the same way the iterator hook already computes it --- litellm/proxy/utils.py | 29 ++++++++++++--- .../proxy_logging/test_guardrail_pipeline.py | 37 ++++++++++++++++++- 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b36061c6be..d40cc34476a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -555,6 +555,24 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> return False +def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool: + return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None + + +def _stream_gated_guardrail_names( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> frozenset[str]: + if not _route_supports_streaming_pipelines(user_api_key_dict): + return frozenset() + return _pipeline_step_guardrail_names( + tuple( + (policy_name, pipeline) + for policy_name, pipeline in _post_call_pipelines(request_data) + if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps) + ) + ) + + def _streamable_post_call_pipelines( request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> tuple[tuple[str, "GuardrailPipeline"], ...]: @@ -571,13 +589,12 @@ def _streamable_post_call_pipelines( post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: return () - route: Final = user_api_key_dict.request_route - if route and resolve_endpoint_translation(user_api_key_dict, None) is None: + if not _route_supports_streaming_pipelines(user_api_key_dict): verbose_proxy_logger.warning( "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " "on their own: %s", - route, + user_api_key_dict.request_route, ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines), ) return () @@ -3271,15 +3288,15 @@ class ProxyLogging: # dict lookups + llm_router.get_deployment() per callback per chunk. _cached_guardrail_data: dict | None = None _guardrail_data_computed = False - pipeline_managed: Final = ( - _pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() + pipeline_gated: Final = ( + _stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() ) for callback in litellm.callbacks: try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): - if callback.guardrail_name in pipeline_managed: + if callback.guardrail_name in pipeline_gated: continue # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 73dd6746b29..5cb595840fc 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -2128,7 +2128,11 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 return None - managed = RecordingGuardrail( + class UnifiedRecordingGuardrail(RecordingGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + managed = UnifiedRecordingGuardrail( guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True ) free = RecordingGuardrail( @@ -2147,3 +2151,34 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( assert result is not None assert seen.get("gr-post") is None assert seen["gr-free"] == 1 + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_stream( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class ChunkHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + seen["response"] = response + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [ChunkHookGuardrail(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(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + ) + + assert result is not None + assert seen["count"] == 1 + assert seen["response"] == "hello "