From c6f5763443504501e88113e783d2e1ec88b9264d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:04:52 -0700 Subject: [PATCH 01/10] feat(guardrails): run legacy post-call hooks as streaming pipeline steps A post_call pipeline step whose guardrail only implements the older async_post_call_success_hook used to skip the stream entirely: PR #38721 fails that shape open with a warning. The streaming step now assembles the buffered stream into the response the hook expects, runs the hook, ends the stream with the hook's exception when it raises, and delivers the hook's rewrite through the same event write-back the unified guardrails use on chat, Responses, and Messages streams (Messages gets the Anthropic shape). A stream a pipeline manages no longer runs the same hook again after the stream ends. A guardrail with neither the unified interface nor a post-call hook keeps the fail-open, as does a rewrite the buffer cannot be patched with. --- .../chat/guardrail_translation/handler.py | 5 + .../guardrail_translation/base_translation.py | 7 + litellm/proxy/common_request_processing.py | 15 +- .../proxy/policy_engine/pipeline_executor.py | 113 ++++++++- litellm/proxy/utils.py | 30 +-- .../test_anthropic_guardrail_handler.py | 26 +++ .../policy_engine/test_pipeline_executor.py | 150 +++++++++++- .../test_proxy_logging_hook_detection.py | 26 +++ .../proxy_logging/test_guardrail_pipeline.py | 220 +++++++++++++----- 9 files changed, 505 insertions(+), 87 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e486be12fe2..f61cfe58e80 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -176,6 +176,11 @@ class AnthropicMessagesHandler(BaseTranslation): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() + def post_call_hook_response(self, response: object) -> object: + if not isinstance(response, ModelResponse): + return response + return self.adapter.translate_openai_response_to_anthropic(response) + @staticmethod def _build_streaming_usage_response( responses_so_far: Sequence[object], diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index afd8e0f67f7..bef472b2882 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.""" + 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 + ended stream. Chat and Responses scan that shape already; a translation that scans a + different one (Messages scans an OpenAI-shaped ModelResponse) overrides this.""" + return response + @staticmethod def transform_user_api_key_dict_to_metadata( user_api_key_dict: Any | None, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9720e4b1cf8..7e85b66682c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3328,9 +3328,10 @@ class ProxyBaseLLMRequestProcessing: has completed. Guardrails routed through unified_guardrail are skipped, since they already ran - via its streaming iterator. Guardrails that override - async_post_call_success_hook directly run here, including those that implement - apply_guardrail but keep their native lifecycle hooks. + via its streaming iterator, and so are guardrails a post_call policy pipeline + manages, since the pipeline ran them against the buffered stream. Guardrails + that override async_post_call_success_hook directly run here, including those + that implement apply_guardrail but keep their native lifecycle hooks. This is audit-only — content has already been delivered to the client. @@ -3340,12 +3341,18 @@ class ProxyBaseLLMRequestProcessing: _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import _check_and_merge_model_level_guardrails + from litellm.proxy.utils import ( + _check_and_merge_model_level_guardrails, + pipeline_managed_guardrail_names, + ) guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) + pipeline_managed: Final = pipeline_managed_guardrail_names(captured_data, "post_call") for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue + if cb.guardrail_name in pipeline_managed: + continue if not cb.should_run_guardrail( data=guardrail_data, event_type=GuardrailEventHooks.post_call, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9bc10949e9f..25c54d72f25 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -121,6 +121,80 @@ class _StreamRewriteObserver(CustomGuardrail): return outputs +class _ScannedTextRecorder(CustomGuardrail): + def __init__(self, guardrail_name: str) -> None: + super().__init__(guardrail_name=guardrail_name) + self.texts: tuple[str, ...] | None = None + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + self.texts = _text_snapshot(inputs.get("texts")) + return inputs + + +class _LegacyHookStreamAdapter(CustomGuardrail): + """Runs a guardrail that only implements the legacy post-call hook (no unified + ``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The + endpoint translation hands it the texts it scanned plus the assembled response under + ``request_data["response"]``; the hook gets that response in the shape its route gives + non-streaming hooks, an exception it raises ends the stream through the executor's + fail/error classification, and a replacement response is re-scanned by the same translation + so its texts reach the client through the translation's ended-stream write-back. A + replacement whose scanned texts do not line up with the originals is undeliverable, so the + executor releases the original chunks.""" + + def __init__( + self, + inner: CustomGuardrail, + endpoint_translation: "BaseTranslation", + user_api_key_dict: "UserAPIKeyAuth", + ) -> None: + super().__init__(guardrail_name=inner.guardrail_name) + self.inner: Final = inner + self.endpoint_translation: Final = endpoint_translation + self.user_api_key_dict: Final = user_api_key_dict + + def structured_messages_cover_full_request(self) -> bool: + return self.inner.structured_messages_cover_full_request() + + @_logged_by_inner_guardrail + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + replacement: Final = await self.inner.async_post_call_success_hook( + data=request_data, + user_api_key_dict=self.user_api_key_dict, + response=self.endpoint_translation.post_call_hook_response(request_data.get("response")), + ) + if replacement is None: + return inputs + scanned: Final = _text_snapshot(inputs.get("texts")) + rewritten: Final = await self._scanned_texts(replacement, logging_obj) + if scanned is None or rewritten is None or len(rewritten) != len(scanned): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + return {**inputs, "texts": list(rewritten)} + + async def _scanned_texts(self, response: object, logging_obj: "LiteLLMLoggingObj | None") -> tuple[str, ...] | None: + recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown") + await self.endpoint_translation.process_output_response( + response=response, + guardrail_to_apply=recorder, + litellm_logging_obj=logging_obj, + user_api_key_dict=self.user_api_key_dict, + ) + return recorder.texts + + def _prepare_hook_input( step: PipelineStep, callback: CustomGuardrail, @@ -286,16 +360,23 @@ class PipelineExecutor: endpoint_translation: "BaseTranslation", streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place hook_input: dict[str, object], # mutable-ok: same request-payload shape as data - user_api_key_dict: "UserAPIKeyAuth | None", + user_api_key_dict: "UserAPIKeyAuth", litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> None: """Run one streaming post_call step through the endpoint translation, delivering - text rewrites on translations that support ended-stream write-back. A rewrite that - cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation - without write-back, or one the translation refused with - ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the - originals and the step passes, so the client gets the stream the merge base sent.""" - observer: Final = _StreamRewriteObserver(callback) + text rewrites on translations that support ended-stream write-back. A guardrail + without the unified interface runs its legacy post-call hook against the assembled + response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the client + yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one + the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: + the buffered chunks go back to the originals and the step passes, so the client gets + the stream the merge base sent.""" + scanner: Final = ( + callback + if PipelineExecutor.supports_unified_execution(callback) + else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict) + ) + observer: Final = _StreamRewriteObserver(scanner) deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites originals: Final = copy.deepcopy(streaming_chunks) try: @@ -379,11 +460,11 @@ class PipelineExecutor: if isinstance(response, dict): callback.mark_pre_call_hook_ran(response) elif mode == "post_call" and streaming_chunks is not None: - if not use_unified or endpoint_translation is None: + if endpoint_translation is None: return ( "error", None, - f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", + f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation", None, ) await PipelineExecutor._run_streaming_step( @@ -433,10 +514,20 @@ class PipelineExecutor: @staticmethod def supports_unified_execution(callback: CustomGuardrail) -> bool: - """Whether this guardrail runs through the unified apply_guardrail path, - the interface streaming pipeline execution requires.""" + """Whether this guardrail runs through the unified apply_guardrail path.""" return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + @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 + ) + @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: """Look up an initialized guardrail callback by name from litellm.callbacks.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b36061c6be..bbec1b842c4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -451,7 +451,7 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps) -def _pipeline_managed_guardrail_names( +def pipeline_managed_guardrail_names( data: Mapping[str, object], mode: Literal["pre_call", "post_call"] ) -> frozenset[str]: return _pipeline_step_guardrail_names( @@ -514,9 +514,9 @@ def _merge_pipeline_metadata_writes( _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: +def _pipeline_step_supports_streaming(guardrail_name: str) -> bool: callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - return callback is not None and PipelineExecutor.supports_unified_execution(callback) + return callback is not None and PipelineExecutor.supports_streaming_execution(callback) def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: @@ -541,14 +541,15 @@ def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> No def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: unsupported: Final = tuple( dict.fromkeys( - step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail) + step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step.guardrail) ) ) if not unsupported: return True verbose_proxy_logger.warning( - "Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, " - "which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s", + "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_name, ", ".join(unsupported), ) @@ -562,11 +563,12 @@ def _streamable_post_call_pipelines( The post_call pipelines a streaming response can be gated through. Streaming pipelines scan the buffered stream through the endpoint guardrail - translation of the request route, so every step's guardrail needs the - unified apply_guardrail interface 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. + 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 + 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. """ post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: @@ -1968,7 +1970,7 @@ class ProxyLogging: ) # Get pipeline-managed guardrails to skip in normal loop - pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call") + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call") caps: Final = ProxyLogging._callback_capabilities() # Skip the per-request callback walk entirely when nothing in @@ -2956,7 +2958,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, "post_call") + pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run @@ -3272,7 +3274,7 @@ class ProxyLogging: _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_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() ) for callback in litellm.callbacks: diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index bf40f781fa3..a10a2aae0ee 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2156,3 +2156,29 @@ class TestAnthropicMessagesHandlerStreamingScanKey: assert open_key == StreamingScanKey(texts=("hi",)) assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] assert ended_key != open_key + + +class TestAnthropicMessagesHandlerPostCallHookResponse: + def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self): + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + assembled = ModelResponse( + id="msg_1", + model="claude", + choices=[Choices(message=Message(role="assistant", content="hello world"), finish_reason="stop")], + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + ) + + hook_response = AnthropicMessagesHandler().post_call_hook_response(assembled) + + assert hook_response["type"] == "message" + assert hook_response["role"] == "assistant" + assert hook_response["content"] == [{"type": "text", "text": "hello world"}] + assert hook_response["stop_reason"] == "end_turn" + assert hook_response["usage"]["input_tokens"] == 1 + assert hook_response["usage"]["output_tokens"] == 2 + + def test_anything_else_reaches_the_hook_untouched(self): + native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]} + + assert AnthropicMessagesHandler().post_call_hook_response(native) is native 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 54ef1f79f4d..db403a17c28 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -534,7 +534,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch): ], ) - monkeypatch.setattr(litellm, "callbacks", []) result = await PipelineExecutor.execute_steps( steps=pipeline.steps, @@ -1153,3 +1152,152 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] + + +class _LegacyHookGuardrail(CustomGuardrail): + """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" + + def __init__(self, replacement=None, raises=None): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + self.replacement = replacement + self.raises = raises + self.calls = [] + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response}) + if self.raises is not None: + raise self.raises + return self.replacement + + +class _NativeHooksGuardrail(_LegacyHookGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + +class _LegacyScanningTranslation: + """Stores the assembled response under request_data["response"] before scanning, like the + chat, Responses, and Messages handlers, hands hooks a route-native shape, and re-extracts one + text per entry of a replacement's "texts".""" + + delivers_ended_stream_text_rewrites = True + + def post_call_hook_response(self, response): + return {"native": True, "text": response["text"]} + + 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, + ): + request_data.setdefault("response", {"text": responses_so_far[0]["text"]}) + outputs = 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, + ) + responses_so_far[0]["text"] = outputs["texts"][0] + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": list(response["texts"])}, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="block", on_error="next"): + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + return await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail=on_fail, on_error=on_error)], + mode="post_call", + data={"model": "m"}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="p", + streaming_chunks=chunks, + endpoint_translation=_LegacyScanningTranslation(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail]) +async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class): + guardrail = guardrail_class(replacement={"texts": ["[REWRITTEN] hello world"]}) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in guardrail.calls] == [{"native": True, "text": "hello world"}] + assert guardrail.calls[0]["data"]["model"] == "m" + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert len(guardrail.calls) == 1 + assert chunks == [_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_streaming_step_blocks_with_the_legacy_hook_exception(monkeypatch): + exc = HTTPException(status_code=400, detail={"error": "output blocked"}) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, _LegacyHookGuardrail(raises=exc), chunks) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["fail"] + assert result.original_exception is exc + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatch): + chunks = [_chunk()] + + result = await _run_legacy_streaming_step( + monkeypatch, _LegacyHookGuardrail(raises=ValueError("boom")), chunks, on_error="block" + ) + + assert result.terminal_action == "block" + assert [step.outcome for step in result.step_results] == ["error"] + assert result.step_results[0].error_detail == "boom" + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement={"texts": ["split", "in two"]}) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 9f1321aec2c..f4165f98ba5 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -671,6 +671,32 @@ async def test_deferred_stream_guardrails_run_native_hook_when_opted_out(monkeyp assert routed.native_hooks_ran == [] +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monkeypatch): + """A post_call pipeline step already ran the opted-out guardrail's own hook against + the buffered stream, so the deferred audit must not run it a second time.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == [] + + @pytest.mark.asyncio async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch): """The realtime path calls apply_guardrail directly, so the opt-out has to be 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..d01f881c446 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 @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import json +from copy import deepcopy import logging from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -1497,29 +1498,78 @@ async def _async_chunk_iter(chunks: List[Any]): yield chunk -def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported( +def _legacy_hook_stream_guardrail( + seen: Dict[str, Any], + rewrite: Callable[[Any], Any] | None = None, + raises: Exception | None = None, + native_lifecycle: bool = False, +) -> CustomGuardrail: + class LegacyHookGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = native_lifecycle + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + seen["data"] = data + seen["user_api_key_dict"] = user_api_key_dict + seen["response"] = deepcopy(response) + if raises is not None: + raise raises + return None if rewrite is None else rewrite(response) + + if native_lifecycle: + + class NativeLifecycleGuardrail(LegacyHookGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail") + + return NativeLifecycleGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) + return LegacyHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail: + class IteratorHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["count"] = seen.get("count", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + return IteratorHookGuardrail(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"] + return litellm.ModelResponse(**payload) + + +def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only( make_user_api_key_auth, monkeypatch, caplog ): - class NativeOnlyGuardrail(CustomGuardrail): - pass - supported = _unified_stream_guardrail({}) - native_only = NativeOnlyGuardrail(guardrail_name="gr-native", event_hook=GuardrailEventHooks.post_call) - monkeypatch.setattr(litellm, "callbacks", [supported, native_only]) - governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + legacy = _legacy_hook_stream_guardrail({}) + legacy.guardrail_name = "gr-legacy" + iterator_only = _iterator_hook_only_guardrail("gr-iterator", {}) + monkeypatch.setattr(litellm, "callbacks", [supported, legacy, iterator_only]) + governed = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-legacy", on_fail="block")], + ) ungoverned = GuardrailPipeline( mode="post_call", - steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")], + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-iterator", on_fail="block")], ) - pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")]) + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-iterator", on_fail="block")]) data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}} with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions")) assert streamable == (("governed", governed),) - assert any("'ungoverned'" in message and "gr-native" in message for message in _warnings(caplog)) - assert not any("'governed'" in message for message in _warnings(caplog)) + assert any("'ungoverned'" in message and "gr-iterator" in message for message in _warnings(caplog)) + assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog)) def test_streamable_post_call_pipelines_is_empty_on_route_without_translation( @@ -1569,55 +1619,123 @@ async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_u @pytest.mark.asyncio @pytest.mark.parametrize("native_lifecycle", [False, True]) -async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_lacks_unified_support( +async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite( proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog ): seen: Dict[str, Any] = {} - if native_lifecycle: - - class NativeOnlyGuardrail(CustomGuardrail): - use_native_lifecycle_hooks = True - - async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): - seen["count"] = seen.get("count", 0) + 1 - return inputs - - else: - - class NativeOnlyGuardrail(CustomGuardrail): - async def async_post_call_success_hook(self, data, user_api_key_dict, response): - seen["count"] = seen.get("count", 0) + 1 - return response - - monkeypatch.setattr( - litellm, - "callbacks", - [NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], - ) + guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_model_response, native_lifecycle=native_lifecycle) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) chunks = _stream_chunks() - delivered: List[Any] = [] + auth = make_user_api_key_auth(request_route="/v1/chat/completions") with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), - data=data, - call_type="completion", - guardrails_only=True, + user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True ) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert out is not None and out.get("stream") is True + assert seen["count"] == 1 + assert isinstance(seen["response"], litellm.ModelResponse) + assert seen["response"].choices[0].message.content == "hello world" + assert seen["data"]["messages"] == data["messages"] + assert seen["user_api_key_dict"] is auth + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert delivered[0].choices[0].delta.content == "[REWRITTEN] hello world" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + assert _warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + 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(chunks), request_data=data, + ) + ] + + assert seen["count"] == 1 + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert [item.choices[0].delta.content for item in delivered] == ["hello ", "world"] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_ends_stream_with_legacy_hook_exception( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + blocked = HTTPException(status_code=400, detail={"error": "output blocked"}) + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, raises=blocked)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + 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, ): delivered.append(item) - assert out is not None - assert out.get("stream") is True - assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] - assert len(delivered) == 2 - assert seen.get("count") is None - assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + with pytest.raises(HTTPException) as info: + await _drain() + + assert seen["count"] == 1 + assert delivered == [] + assert info.value is blocked + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + def rewrite(response: Any) -> Dict[str, Any]: + return {**response, "content": [{"type": "text", "text": "[REWRITTEN] " + response["content"][0]["text"]}]} + + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=rewrite)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + 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/messages"), + response=_async_chunk_iter(_anthropic_sse_chunks()), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert seen["response"]["content"][0]["text"] == "hello world" + assert seen["response"]["role"] == "assistant" + raw = b"".join(delivered).decode() + assert "[REWRITTEN] hello world" in raw + assert raw.count("event: content_block_delta") == 1 + for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + assert f"event: {expected_event}" in raw @pytest.mark.asyncio @@ -1625,19 +1743,7 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} - - class IteratorHookGuardrail(CustomGuardrail): - async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): - seen["count"] = seen.get("count", 0) + 1 - async for item in response: - item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" - yield item - - monkeypatch.setattr( - litellm, - "callbacks", - [IteratorHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], - ) + monkeypatch.setattr(litellm, "callbacks", [_iterator_hook_only_guardrail("gr-post", seen)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) From 3f30c05493091f62d1edcfb0af58aa33a370a335 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:39:39 -0700 Subject: [PATCH 02/10] fix(guardrails): skip only stream-gated pipeline guardrails in the deferred post-call pass --- litellm/proxy/common_request_processing.py | 6 +- litellm/proxy/utils.py | 37 +++++++++-- .../test_proxy_logging_hook_detection.py | 63 +++++++++++++++++++ 3 files changed, 98 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7e85b66682c..2a6eb0a16ff 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3343,15 +3343,15 @@ class ProxyBaseLLMRequestProcessing: from litellm.proxy.proxy_server import llm_router as _global_llm_router from litellm.proxy.utils import ( _check_and_merge_model_level_guardrails, - pipeline_managed_guardrail_names, + stream_gated_guardrail_names, ) guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router) - pipeline_managed: Final = pipeline_managed_guardrail_names(captured_data, "post_call") + stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict) for cb in litellm.callbacks: if not isinstance(cb, CustomGuardrail): continue - if cb.guardrail_name in pipeline_managed: + if cb.guardrail_name in stream_gated: continue if not cb.should_run_guardrail( data=guardrail_data, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bbec1b842c4..106fab7af8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -538,12 +538,16 @@ def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> No ) -def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: - unsupported: Final = tuple( +def _pipeline_unsupported_streaming_guardrails(pipeline: "GuardrailPipeline") -> tuple[str, ...]: + return tuple( dict.fromkeys( step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step.guardrail) ) ) + + +def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline) if not unsupported: return True verbose_proxy_logger.warning( @@ -573,13 +577,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_has_endpoint_translation(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 () @@ -590,6 +593,30 @@ def _streamable_post_call_pipelines( ) +def _route_has_endpoint_translation(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]: + """ + The guardrails whose post_call pipelines gate a streaming response on this + route: the selection ``_streamable_post_call_pipelines`` makes, without its + warnings, so the post-call pass deferred to the end of the stream skips + exactly the guardrails the pipelines already ran and no others. + """ + if not _route_has_endpoint_translation(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 not _pipeline_unsupported_streaming_guardrails(pipeline) + ) + ) + + def _prompt_block_text(block: object) -> str: if isinstance(block, str): return block diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index f4165f98ba5..10de7e7fc3e 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -697,6 +697,69 @@ async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monk assert pipeline_managed.native_hooks_ran == [] +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch): + """A pipeline step with neither streaming interface keeps the whole pipeline off the + stream, so the deferred audit is the only place the opted-out guardrail's own hook + still runs, the way it did before pipelines ran on streams.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + class NeitherHookGuardrail(CustomGuardrail): + pass + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither]) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="keeps_native", on_fail="next"), + PipelineStep(guardrail="gr-neither", on_fail="block"), + ], + ) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + +@pytest.mark.asyncio +async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch): + """A route with no endpoint guardrail translation cannot gate the stream through its + pipelines, so the deferred audit still owes the opted-out guardrail its own hook.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.utils import Choices, Message, ModelResponse + + pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [pipeline_managed]) + pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")]) + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={ + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, + }, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"), + captured_logging_obj=_streaming_logging_obj(), + assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), + cache_hit=False, + ) + + assert pipeline_managed.native_hooks_ran == ["post_call"] + + @pytest.mark.asyncio async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch): """The realtime path calls apply_guardrail directly, so the opt-out has to be From e700e79cce9dac5d0fda0b5d31210283230c9de7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:52:38 -0700 Subject: [PATCH 03/10] refactor(guardrails): keep the rescanned inputs instead of rebuilding the texts list --- litellm/proxy/policy_engine/pipeline_executor.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index afc9a29afd9..089f0683823 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -128,7 +128,7 @@ class _StreamRewriteObserver(CustomGuardrail): class _ScannedTextRecorder(CustomGuardrail): def __init__(self, guardrail_name: str) -> None: super().__init__(guardrail_name=guardrail_name) - self.texts: tuple[str, ...] | None = None + self.inputs: GenericGuardrailAPIInputs | None = None @_logged_by_inner_guardrail async def apply_guardrail( @@ -138,7 +138,7 @@ class _ScannedTextRecorder(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: - self.texts = _text_snapshot(inputs.get("texts")) + self.inputs = inputs return inputs @@ -183,12 +183,16 @@ class _LegacyHookStreamAdapter(CustomGuardrail): if replacement is None: return inputs scanned: Final = _text_snapshot(inputs.get("texts")) - rewritten: Final = await self._scanned_texts(replacement, logging_obj) + rescanned: Final = await self._rescan(replacement, logging_obj) + rewritten: Final = None if rescanned is None else rescanned.get("texts") if scanned is None or rewritten is None or len(rewritten) != len(scanned): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") - return {**inputs, "texts": list(rewritten)} + rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} + return rewritten_inputs - async def _scanned_texts(self, response: object, logging_obj: "LiteLLMLoggingObj | None") -> tuple[str, ...] | None: + async def _rescan( + self, response: object, logging_obj: "LiteLLMLoggingObj | None" + ) -> GenericGuardrailAPIInputs | None: recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown") await self.endpoint_translation.process_output_response( response=response, @@ -196,7 +200,7 @@ class _LegacyHookStreamAdapter(CustomGuardrail): litellm_logging_obj=logging_obj, user_api_key_dict=self.user_api_key_dict, ) - return recorder.texts + return recorder.inputs def _prepare_hook_input( From 80782d478ddbdd2e145b0cb72eb3c653efbe0ed0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:17:36 -0700 Subject: [PATCH 04/10] fix(policy_engine): keep legacy stream steps in sync and reject tool-call rewrites Each streaming step drops the response an earlier step's translation stored under request_data["response"], so a later legacy hook sees the stream as the steps before it left it instead of the first step's snapshot. A legacy replacement whose tool calls differ from the scanned chunks is now undeliverable like a text mismatch, so the original stream is released with a warning instead of delivering the text while dropping the tool-call change --- .../proxy/policy_engine/pipeline_executor.py | 17 ++-- .../policy_engine/test_pipeline_executor.py | 78 ++++++++++++++++--- 2 files changed, 81 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 089f0683823..b4f6da18d1c 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -150,8 +150,8 @@ class _LegacyHookStreamAdapter(CustomGuardrail): non-streaming hooks, an exception it raises ends the stream through the executor's fail/error classification, and a replacement response is re-scanned by the same translation so its texts reach the client through the translation's ended-stream write-back. A - replacement whose scanned texts do not line up with the originals is undeliverable, so the - executor releases the original chunks.""" + replacement whose scanned texts do not line up with the originals, or whose tool calls + differ from them, is undeliverable, so the executor releases the original chunks.""" def __init__( self, @@ -184,8 +184,12 @@ class _LegacyHookStreamAdapter(CustomGuardrail): return inputs scanned: Final = _text_snapshot(inputs.get("texts")) rescanned: Final = await self._rescan(replacement, logging_obj) - rewritten: Final = None if rescanned is None else rescanned.get("texts") - if scanned is None or rewritten is None or len(rewritten) != len(scanned): + if scanned is None or rescanned is None: + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + rewritten: Final = rescanned.get("texts") + if rewritten is None or len(rewritten) != len(scanned): + raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} return rewritten_inputs @@ -380,7 +384,9 @@ class PipelineExecutor: yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the originals and the step passes, so the client gets - the stream the merge base sent.""" + the stream the merge base sent. The response an earlier step's translation stored under + ``request_data["response"]`` is dropped first, so this step's hook sees the stream as + the steps before it left it.""" scanner: Final = ( callback if PipelineExecutor.supports_unified_execution(callback) @@ -389,6 +395,7 @@ class PipelineExecutor: observer: Final = _StreamRewriteObserver(scanner) deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites originals: Final = copy.deepcopy(streaming_chunks) + hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored try: if deliver_rewrites: await endpoint_translation.process_output_streaming_response( 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 a44afc6ef96..860f83a06bb 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -675,7 +675,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch): ], ) - result = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, @@ -1298,8 +1297,8 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri class _LegacyHookGuardrail(CustomGuardrail): """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" - def __init__(self, replacement=None, raises=None): - super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + def __init__(self, replacement=None, raises=None, guardrail_name="masker"): + super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True) self.replacement = replacement self.raises = raises self.calls = [] @@ -1339,7 +1338,7 @@ class _LegacyScanningTranslation: ): request_data.setdefault("response", {"text": responses_so_far[0]["text"]}) outputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [responses_so_far[0]["text"]]}, + inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, @@ -1350,8 +1349,11 @@ class _LegacyScanningTranslation: async def process_output_response( self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None ): + inputs = {"texts": list(response["texts"])} + if response.get("tool_calls"): + inputs["tool_calls"] = list(response["tool_calls"]) await guardrail_to_apply.apply_guardrail( - inputs={"texts": list(response["texts"])}, + inputs=inputs, request_data={"response": response}, input_type="response", logging_obj=litellm_logging_obj, @@ -1359,10 +1361,26 @@ class _LegacyScanningTranslation: return response +def _legacy_replacement(*texts, tool_calls=None): + return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} + + async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="block", on_error="next"): - monkeypatch.setattr(litellm, "callbacks", [guardrail]) + return await _run_legacy_streaming_steps(monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error) + + +async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail="block", on_error="next"): + monkeypatch.setattr(litellm, "callbacks", list(guardrails)) return await PipelineExecutor.execute_steps( - steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail=on_fail, on_error=on_error)], + steps=[ + PipelineStep( + guardrail=guardrail.guardrail_name, + on_pass="next" if position + 1 < len(guardrails) else "allow", + on_fail=on_fail, + on_error=on_error, + ) + for position, guardrail in enumerate(guardrails) + ], mode="post_call", data={"model": "m"}, user_api_key_dict=MagicMock(), @@ -1376,7 +1394,7 @@ async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="bl @pytest.mark.asyncio @pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail]) async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class): - guardrail = guardrail_class(replacement={"texts": ["[REWRITTEN] hello world"]}) + guardrail = guardrail_class(replacement=_legacy_replacement("[REWRITTEN] hello world")) chunks = [_chunk()] with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -1434,7 +1452,7 @@ async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatc @pytest.mark.asyncio async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog): - guardrail = _LegacyHookGuardrail(replacement={"texts": ["split", "in two"]}) + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("split", "in two")) chunks = [_chunk()] with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): @@ -1442,3 +1460,45 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up _assert_passed_with_discard_warning(result, caplog) assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail( + replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call]) + ) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[])) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): + masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) + auditor = _LegacyHookGuardrail(replacement=None, guardrail_name="auditor") + chunks = [_chunk()] + + result = await _run_legacy_streaming_steps(monkeypatch, [masker, auditor], chunks, on_fail="next") + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass", "pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert [call["response"] for call in masker.calls] == [{"native": True, "text": "hello world"}] + assert [call["response"] for call in auditor.calls] == [{"native": True, "text": "[REWRITTEN] hello world"}] From ea427e33d8208034f934e52a7293f7f0bfa020ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:44:14 -0700 Subject: [PATCH 05/10] fix(policy_engine): deliver legacy stream rewrites made in place A legacy post-call hook that changes the response it was handed and returns None (the model armor guardrail masks choice content that way) used to have that rewrite dropped on the streaming pipeline path, since the adapter only re-scanned a returned replacement. The adapter now re-scans the response it handed the hook when the hook returns None, so an in-place rewrite reaches the client through the same ended-stream write-back --- .../proxy/policy_engine/pipeline_executor.py | 13 ++++--- .../policy_engine/test_pipeline_executor.py | 37 +++++++++++++++---- 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index b4f6da18d1c..45a113fcda2 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -148,8 +148,9 @@ class _LegacyHookStreamAdapter(CustomGuardrail): endpoint translation hands it the texts it scanned plus the assembled response under ``request_data["response"]``; the hook gets that response in the shape its route gives non-streaming hooks, an exception it raises ends the stream through the executor's - fail/error classification, and a replacement response is re-scanned by the same translation - so its texts reach the client through the translation's ended-stream write-back. A + fail/error classification, and the response it hands back, or the one it changed in place + and returned ``None`` for, is re-scanned by the same translation so its texts reach the + client through the translation's ended-stream write-back. A replacement whose scanned texts do not line up with the originals, or whose tool calls differ from them, is undeliverable, so the executor releases the original chunks.""" @@ -175,15 +176,17 @@ class _LegacyHookStreamAdapter(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: + hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response")) replacement: Final = await self.inner.async_post_call_success_hook( data=request_data, user_api_key_dict=self.user_api_key_dict, - response=self.endpoint_translation.post_call_hook_response(request_data.get("response")), + response=hooked, ) - if replacement is None: + rewrite: Final = hooked if replacement is None else replacement + if rewrite is None: return inputs scanned: Final = _text_snapshot(inputs.get("texts")) - rescanned: Final = await self._rescan(replacement, logging_obj) + rescanned: Final = await self._rescan(rewrite, logging_obj) if scanned is None or rescanned is None: raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") rewritten: Final = rescanned.get("texts") 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 860f83a06bb..ae7754e5821 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1297,16 +1297,19 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri class _LegacyHookGuardrail(CustomGuardrail): """A guardrail with only the legacy post-call hook: it never defines apply_guardrail.""" - def __init__(self, replacement=None, raises=None, guardrail_name="masker"): + def __init__(self, replacement=None, raises=None, guardrail_name="masker", rewrite_in_place=None): super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True) self.replacement = replacement self.raises = raises + self.rewrite_in_place = rewrite_in_place self.calls = [] async def async_post_call_success_hook(self, data, user_api_key_dict, response): self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response}) if self.raises is not None: raise self.raises + if self.rewrite_in_place is not None: + response["text"] = self.rewrite_in_place return self.replacement @@ -1325,7 +1328,7 @@ class _LegacyScanningTranslation: delivers_ended_stream_text_rewrites = True def post_call_hook_response(self, response): - return {"native": True, "text": response["text"]} + return {"native": True, "text": response["text"], "tool_calls": response["tool_calls"]} async def process_output_streaming_response( self, @@ -1336,7 +1339,9 @@ class _LegacyScanningTranslation: request_data=None, deliver_ended_stream_rewrites=False, ): - request_data.setdefault("response", {"text": responses_so_far[0]["text"]}) + request_data.setdefault( + "response", {"text": responses_so_far[0]["text"], "tool_calls": [dict(responses_so_far[0]["tool_call"])]} + ) outputs = await guardrail_to_apply.apply_guardrail( inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, request_data=request_data, @@ -1349,7 +1354,7 @@ class _LegacyScanningTranslation: async def process_output_response( self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None ): - inputs = {"texts": list(response["texts"])} + inputs = {"texts": [response["text"]] if "text" in response else list(response["texts"])} if response.get("tool_calls"): inputs["tool_calls"] = list(response["tool_calls"]) await guardrail_to_apply.apply_guardrail( @@ -1361,6 +1366,10 @@ class _LegacyScanningTranslation: return response +def _native(text): + return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]} + + def _legacy_replacement(*texts, tool_calls=None): return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} @@ -1403,12 +1412,26 @@ async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypa assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] assert chunks[0]["text"] == "[REWRITTEN] hello world" - assert [call["response"] for call in guardrail.calls] == [{"native": True, "text": "hello world"}] + assert [call["response"] for call in guardrail.calls] == [_native("hello world")] assert guardrail.calls[0]["data"]["model"] == "m" assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] assert not any("discarded" in record.getMessage() for record in caplog.records) +@pytest.mark.asyncio +async def test_streaming_step_delivers_a_legacy_rewrite_made_in_place(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(rewrite_in_place="[REWRITTEN] hello world") + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert chunks[0]["text"] == "[REWRITTEN] hello world" + assert not any("discarded" in record.getMessage() for record in caplog.records) + + @pytest.mark.asyncio async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog): guardrail = _LegacyHookGuardrail(replacement=None) @@ -1500,5 +1523,5 @@ async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(mon assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass", "pass"] assert chunks[0]["text"] == "[REWRITTEN] hello world" - assert [call["response"] for call in masker.calls] == [{"native": True, "text": "hello world"}] - assert [call["response"] for call in auditor.calls] == [{"native": True, "text": "[REWRITTEN] hello world"}] + assert [call["response"] for call in masker.calls] == [_native("hello world")] + assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")] From ce7cec1a3635bdb66292daa45ec9fe39f37bbffa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:48:22 -0700 Subject: [PATCH 06/10] fix(policy_engine): keep discarded stream rewrites out of the applied-guardrails header A streaming step whose rewrite the executor threw away (a tool-call rewrite, a text rewrite the translation cannot write back, or one the adapter refused) still marked its guardrail as applied, so the header claimed an output the client never received. The step now returns right after releasing the original chunks, which leaves the header as the merge base sent it --- litellm/proxy/policy_engine/pipeline_executor.py | 10 ++++++---- .../proxy/policy_engine/test_pipeline_executor.py | 1 + 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 45a113fcda2..9b0d1e60839 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -387,7 +387,8 @@ class PipelineExecutor: yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the originals and the step passes, so the client gets - the stream the merge base sent. The response an earlier step's translation stored under + the stream the merge base sent, and the guardrail stays out of the applied-guardrails + header since its output never reached the client. The response an earlier step's translation stored under ``request_data["response"]`` is dropped first, so this step's hook sees the stream as the steps before it left it.""" scanner: Final = ( @@ -419,9 +420,10 @@ class PipelineExecutor: ) except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) - else: - if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): - _release_original_chunks(step.guardrail, streaming_chunks, originals) + return + if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): + _release_original_chunks(step.guardrail, streaming_chunks, originals) + return if not callback.records_own_guardrail_information: add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) 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 ae7754e5821..2af318c88cd 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1146,6 +1146,7 @@ def _assert_passed_with_discard_warning(result, caplog): assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) @pytest.mark.asyncio From 9e67c083092da5e20b85ed87e935d0af168cc66d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:16:39 -0700 Subject: [PATCH 07/10] fix(policy_engine): keep tool-only legacy streams deliverable A Messages stream that ends with only tool_use blocks reaches the legacy step without a texts key, while the non-streaming rescan of the same response sends an empty list. Compare both as empty and keep the stream when the hook left the tool calls alone. --- .../proxy/policy_engine/pipeline_executor.py | 15 +++- .../policy_engine/test_pipeline_executor.py | 84 ++++++++++++++++++- 2 files changed, 91 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 9b0d1e60839..52985d49c7e 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None: return None if texts is None else tuple(texts) +def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]: + return tuple(texts or ()) + + def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None: return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls) @@ -152,7 +156,9 @@ class _LegacyHookStreamAdapter(CustomGuardrail): and returned ``None`` for, is re-scanned by the same translation so its texts reach the client through the translation's ended-stream write-back. A replacement whose scanned texts do not line up with the originals, or whose tool calls - differ from them, is undeliverable, so the executor releases the original chunks.""" + differ from them, is undeliverable, so the executor releases the original chunks. A stream + that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as + long as the hook left the tool calls alone.""" def __init__( self, @@ -185,15 +191,16 @@ class _LegacyHookStreamAdapter(CustomGuardrail): rewrite: Final = hooked if replacement is None else replacement if rewrite is None: return inputs - scanned: Final = _text_snapshot(inputs.get("texts")) rescanned: Final = await self._rescan(rewrite, logging_obj) - if scanned is None or rescanned is None: + if rescanned is None: raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") rewritten: Final = rescanned.get("texts") - if rewritten is None or len(rewritten) != len(scanned): + if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")): raise UndeliverableStreamRewrite(self.guardrail_name or "unknown") + if not rewritten: + return inputs rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten} return rewritten_inputs 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 2af318c88cd..98674843c73 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1367,6 +1367,44 @@ class _LegacyScanningTranslation: return response +class _ToolOnlyLegacyScanningTranslation(_LegacyScanningTranslation): + """Like the Messages handler on a tool-only message: the ended-stream scan omits "texts" from + the inputs, while the non-streaming scan of the same response sends an empty list.""" + + 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, + ): + request_data.setdefault("response", {"text": "", "tool_calls": [dict(responses_so_far[0]["tool_call"])]}) + await guardrail_to_apply.apply_guardrail( + inputs={"tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + await guardrail_to_apply.apply_guardrail( + inputs={"texts": [], "tool_calls": list(response.get("tool_calls") or [])}, + request_data={"response": response}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return response + + +def _tool_only_chunk(): + return {"text": "", "tool_call": _chunk()["tool_call"]} + + def _native(text): return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]} @@ -1375,11 +1413,17 @@ def _legacy_replacement(*texts, tool_calls=None): return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls} -async def _run_legacy_streaming_step(monkeypatch, guardrail, chunks, on_fail="block", on_error="next"): - return await _run_legacy_streaming_steps(monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error) +async def _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, on_fail="block", on_error="next", translation=None +): + return await _run_legacy_streaming_steps( + monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error, translation=translation + ) -async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail="block", on_error="next"): +async def _run_legacy_streaming_steps( + monkeypatch, guardrails, chunks, on_fail="block", on_error="next", translation=None +): monkeypatch.setattr(litellm, "callbacks", list(guardrails)) return await PipelineExecutor.execute_steps( steps=[ @@ -1397,7 +1441,7 @@ async def _run_legacy_streaming_steps(monkeypatch, guardrails, chunks, on_fail=" call_type="completion", policy_name="p", streaming_chunks=chunks, - endpoint_translation=_LegacyScanningTranslation(), + endpoint_translation=_LegacyScanningTranslation() if translation is None else translation, ) @@ -1513,6 +1557,38 @@ async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls( assert chunks == [_chunk()] +@pytest.mark.asyncio +async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alone(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=None) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + assert chunks == [_tool_only_chunk()] + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog): + masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}} + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call])) + chunks = [_tool_only_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_legacy_streaming_step( + monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation() + ) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_tool_only_chunk()] + + @pytest.mark.asyncio async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) From 359c26aa1d1403b0dc10c2d7a47b984d13c468c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:27:46 -0700 Subject: [PATCH 08/10] test(policy_engine): cover the legacy stream paths with no response and no rescan A translation that hands the hook no assembled response leaves the stream as it is, and a rewrite the translation cannot rescan is released as the original stream with a warning. --- .../policy_engine/test_pipeline_executor.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) 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 98674843c73..6de7f7b593d 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1589,6 +1589,68 @@ 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"].""" + + 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 _UnscannableRewriteTranslation(_LegacyScanningTranslation): + """Like the chat handler on a response whose choices are plain dicts: the non-streaming scan + never hands anything to the guardrail.""" + + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + 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) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_rescan(monkeypatch, caplog): + guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("hello [MASKED]")) + chunks = [_chunk()] + + result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation()) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + @pytest.mark.asyncio async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch): masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world")) 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 09/10] 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", From 354365eecaf6c70aba2a649fa324bf0369438f85 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:15:01 -0700 Subject: [PATCH 10/10] test(proxy): give the pipeline-managed native hook audit test its chat completions route --- tests/test_litellm/proxy/test_proxy_logging_hook_detection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 42888ca7d54..28ff4571b44 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -688,7 +688,7 @@ async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monk "messages": [{"role": "user", "content": "hi"}], "metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]}, }, - captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"), + captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), captured_logging_obj=_streaming_logging_obj(), assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]), cache_hit=False,