From 8d040d89e68c075b105a6369e5d95fdc7c1ba8f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:39:56 -0700 Subject: [PATCH] fix(policy_engine): keep legacy hooks off streams their route cannot assemble and off guardrails with their own iterator hook The streaming pipeline step only takes a post-call hook on routes whose translation assembles the streamed response (chat completions, Responses, Messages). On /v1/completions, the Gemini streamGenerateContent route, and A2A streams the pipeline is skipped with the merge-base warning and the hook runs on its own afterwards, instead of getting a None response while the header says the guardrail ran. A guardrail that overrides async_post_call_streaming_iterator_hook next to its post-call hook keeps its native per-chunk path rather than running buffered through the adapter --- .../chat/guardrail_translation/handler.py | 1 + .../guardrail_translation/base_translation.py | 7 ++ .../chat/guardrail_translation/handler.py | 1 + .../guardrail_translation/handler.py | 1 + .../proxy/policy_engine/pipeline_executor.py | 14 ++-- litellm/proxy/utils.py | 60 +++++++++++---- .../policy_engine/test_pipeline_executor.py | 46 +++-------- .../proxy_logging/test_guardrail_pipeline.py | 77 ++++++++++++++++++- 8 files changed, 148 insertions(+), 59 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index f61cfe58e80..b83235f660f 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -171,6 +171,7 @@ class AnthropicMessagesHandler(BaseTranslation): """ delivers_ended_stream_text_rewrites = True + assembles_streamed_response = True def __init__(self): super().__init__() diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index bef472b2882..375b22ff5f6 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -60,6 +60,13 @@ class BaseTranslation(ABC): text rewrites on every other translation, are undeliverable: the pipeline executor discards them and releases the original chunks.""" + assembles_streamed_response: ClassVar[bool] = False + """Whether ``process_output_streaming_response`` stores the assembled response of an + ended stream under ``request_data["response"]`` before scanning it, the way the chat, + Responses, and Messages translations do. A streaming pipeline runs a guardrail that only + has the legacy post-call hook against that response, so on a translation without it such + a guardrail keeps running on its own.""" + def post_call_hook_response(self, response: object) -> object: """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from the object the translation stores under ``request_data["response"]`` while scanning an diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 80292aef2cf..b9b88061fec 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -79,6 +79,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ delivers_ended_stream_text_rewrites = True + assembles_streamed_response = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b0f79552bc5..e0ba80706df 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -341,6 +341,7 @@ class OpenAIResponsesHandler(BaseTranslation): """ delivers_ended_stream_text_rewrites = True + assembles_streamed_response = True def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 52985d49c7e..ab5ad01fda9 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -556,12 +556,14 @@ class PipelineExecutor: @staticmethod def supports_streaming_execution(callback: CustomGuardrail) -> bool: """Whether a streaming pipeline step can run this guardrail against the buffered - stream: through the unified path, or through its own post-call hook on the - assembled response. A guardrail with neither (one that only rewrites the stream - through its iterator hook) has to keep running on its own.""" - return ( - PipelineExecutor.supports_unified_execution(callback) - or type(callback).async_post_call_success_hook is not CustomLogger.async_post_call_success_hook + stream: through the unified path, or through its post-call hook on the assembled + response when that hook is its only streaming path. A guardrail with its own + streaming iterator hook, or with neither hook, keeps running on its own.""" + callback_type: Final = type(callback) + return PipelineExecutor.supports_unified_execution(callback) or ( + callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook + and callback_type.async_post_call_streaming_iterator_hook + is CustomLogger.async_post_call_streaming_iterator_hook ) @staticmethod diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 23260785b97..67a579c4d15 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -192,6 +192,7 @@ if TYPE_CHECKING: from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -517,9 +518,17 @@ def _merge_pipeline_metadata_writes( _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _pipeline_step_supports_streaming(guardrail_name: str) -> bool: +def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool: callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - return callback is not None and PipelineExecutor.supports_streaming_execution(callback) + if callback is None: + return False + if PipelineExecutor.supports_unified_execution(callback): + return True + return ( + translation is not None + and type(translation).assembles_streamed_response + and PipelineExecutor.supports_streaming_execution(callback) + ) def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: @@ -541,42 +550,57 @@ def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> No ) -def _pipeline_unsupported_streaming_guardrails(pipeline: "GuardrailPipeline") -> tuple[str, ...]: +def _pipeline_unsupported_streaming_guardrails( + pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> tuple[str, ...]: return tuple( dict.fromkeys( - step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step.guardrail) + step.guardrail + for step in pipeline.steps + if not _pipeline_step_supports_streaming(step.guardrail, translation) ) ) -def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: - unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline) +def _pipeline_is_streamable( + policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" +) -> bool: + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation) if not unsupported: return True verbose_proxy_logger.warning( - "Policy '%s' has post_call pipeline guardrails with neither the unified apply_guardrail interface nor a " - "post-call hook, one of which streaming pipelines need; the stream skips the pipeline and its guardrails " - "run on their own: %s", + "Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they " + "need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a " + "route whose translation assembles the streamed response. The stream skips the pipeline and its " + "guardrails run on their own: %s", policy_name, ", ".join(unsupported), ) 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 _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None": + resolved: Final = resolve_endpoint_translation(user_api_key_dict, None) + return None if resolved is None else resolved[1] + + +def _route_supports_streaming_pipelines( + user_api_key_dict: UserAPIKeyAuth, translation: "BaseTranslation | None" +) -> bool: + return not user_api_key_dict.request_route or translation 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): + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if not _route_supports_streaming_pipelines(user_api_key_dict, translation): return frozenset() return _pipeline_step_guardrail_names( tuple( (policy_name, pipeline) for policy_name, pipeline in _post_call_pipelines(request_data) - if not _pipeline_unsupported_streaming_guardrails(pipeline) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) ) ) @@ -589,8 +613,9 @@ def _streamable_post_call_pipelines( Streaming pipelines scan the buffered stream through the endpoint guardrail translation of the request route, so every step's guardrail needs either the - unified apply_guardrail interface or a post-call hook to run against the - assembled response, and the route needs a translation. A pipeline that + unified apply_guardrail interface or, on a route whose translation assembles + the streamed response, a post-call hook that is its only streaming path, and + the route needs a translation. A pipeline that cannot be run that way yet is left out and its guardrails run on the stream on their own, the way they did before pipelines ran on streams at all, with a warning naming the pipeline. @@ -598,7 +623,8 @@ def _streamable_post_call_pipelines( post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: return () - if not _route_supports_streaming_pipelines(user_api_key_dict): + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if not _route_supports_streaming_pipelines(user_api_key_dict, translation): 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 " @@ -610,7 +636,7 @@ def _streamable_post_call_pipelines( return tuple( (policy_name, pipeline) for policy_name, pipeline in post_call_pipelines - if _pipeline_is_streamable(policy_name, pipeline) + if _pipeline_is_streamable(policy_name, pipeline, translation) ) diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 6de7f7b593d..e2523567065 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1589,28 +1589,14 @@ async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only assert chunks == [_tool_only_chunk()] -class _ResponselessLegacyScanningTranslation(_LegacyScanningTranslation): - """Like a handler that never stores the assembled response under request_data["response"].""" +class _NoHooksGuardrail(CustomGuardrail): + pass - def post_call_hook_response(self, response): - return response - async def process_output_streaming_response( - self, - responses_so_far, - guardrail_to_apply, - litellm_logging_obj=None, - user_api_key_dict=None, - request_data=None, - deliver_ended_stream_rewrites=False, - ): - await guardrail_to_apply.apply_guardrail( - inputs={"texts": [responses_so_far[0]["text"]]}, - request_data=request_data, - input_type="response", - logging_obj=litellm_logging_obj, - ) - return responses_so_far +class _IteratorAndLegacyHookGuardrail(_LegacyHookGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for item in response: + yield item class _UnscannableRewriteTranslation(_LegacyScanningTranslation): @@ -1623,21 +1609,11 @@ class _UnscannableRewriteTranslation(_LegacyScanningTranslation): return response -@pytest.mark.asyncio -async def test_streaming_step_leaves_the_stream_alone_when_the_hook_gets_no_response(monkeypatch, caplog): - guardrail = _LegacyHookGuardrail() - chunks = [_chunk()] - - result = await _run_legacy_streaming_step( - monkeypatch, guardrail, chunks, translation=_ResponselessLegacyScanningTranslation() - ) - - assert result.terminal_action == "allow" - assert [step.outcome for step in result.step_results] == ["pass"] - assert guardrail.calls[0]["response"] is None - assert chunks == [_chunk()] - assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] - assert not any("discarded" in record.getMessage() for record in caplog.records) +def test_streaming_execution_runs_legacy_hooks_only_when_that_hook_is_their_only_streaming_path(): + assert PipelineExecutor.supports_streaming_execution(_LegacyHookGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_NativeHooksGuardrail()) is True + assert PipelineExecutor.supports_streaming_execution(_IteratorAndLegacyHookGuardrail()) is False + assert PipelineExecutor.supports_streaming_execution(_NoHooksGuardrail(guardrail_name="neither")) is False @pytest.mark.asyncio 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 46d6f2bb7a9..ac303e8a027 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 @@ -28,7 +28,7 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header -from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines +from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( @@ -1539,6 +1539,21 @@ def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuar return IteratorHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) +def _iterator_and_legacy_hook_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorAndLegacyHookGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["success_hook_calls"] = seen.get("success_hook_calls", 0) + 1 + return None + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["iterator_hook_calls"] = seen.get("iterator_hook_calls", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorAndLegacyHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + def _rewritten_model_response(response: Any) -> litellm.ModelResponse: payload = response.model_dump() payload["choices"][0]["message"]["content"] = "[REWRITTEN] " + payload["choices"][0]["message"]["content"] @@ -1572,6 +1587,42 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog)) +@pytest.mark.parametrize( + "request_route", + [None, "/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], +) +def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response( + make_user_api_key_auth, monkeypatch, caplog, request_route +): + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail({})]) + legacy = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("legacy-governance", legacy)]}} + auth = make_user_api_key_auth(request_route=request_route) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'legacy-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_keeps_guardrails_with_their_own_iterator_hook_on_their_own_path( + make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", {})]) + both_hooks = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("both-hooks", both_hooks)]}} + auth = make_user_api_key_auth(request_route="/v1/chat/completions") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, auth) + + assert streamable == () + assert stream_gated_guardrail_names(data, auth) == frozenset() + assert any("'both-hooks'" in message and "gr-post" in message for message in _warnings(caplog)) + + def test_streamable_post_call_pipelines_is_empty_on_route_without_translation( make_user_api_key_auth, monkeypatch, caplog ): @@ -1762,6 +1813,30 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_runs_the_iterator_hook_of_a_guardrail_that_also_has_a_post_call_hook( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ) + ] + + assert seen == {"iterator_hook_calls": 1} + assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"] + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "rewrite_attribute, value",