From 5aab7f5b9dfe2f68dbaea881f0ff21f5cd460ea9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 06:52:36 +0000 Subject: [PATCH] fix(proxy): keep non-bedrock passthrough streams streaming under post-call guardrails --- .../guardrail_translation/handler.py | 17 ++- litellm/proxy/common_request_processing.py | 21 +++- .../guardrail_translation/test_handler.py | 22 ++++ .../proxy/test_common_request_processing.py | 102 ++++++++++++++++++ 4 files changed, 159 insertions(+), 3 deletions(-) diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 42b708d112..a67228fd2f 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -278,6 +278,18 @@ class LlmPassthroughRouteHandler(BaseTranslation): request_data=request_data, ) + @staticmethod + def _resolve_event_stream_de_anonymizer(provider: Optional[str]): + handler_cls = _get_provider_handlers().get(provider or "") + return getattr(handler_cls, "de_anonymize_event_stream", None) + + @staticmethod + def supports_event_stream_de_anonymization(provider: Optional[str]) -> bool: + return ( + LlmPassthroughRouteHandler._resolve_event_stream_de_anonymizer(provider) + is not None + ) + @staticmethod async def de_anonymize_event_stream( body_bytes: bytes, @@ -286,8 +298,9 @@ class LlmPassthroughRouteHandler(BaseTranslation): data: dict, ) -> bytes: provider = data.get("custom_llm_provider") - handler_cls = _get_provider_handlers().get(provider or "") - de_anonymize = getattr(handler_cls, "de_anonymize_event_stream", None) + de_anonymize = LlmPassthroughRouteHandler._resolve_event_stream_de_anonymizer( + provider + ) if de_anonymize is None: verbose_proxy_logger.debug( "LlmPassthroughRouteHandler: no event-stream handler for provider=%s, " diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cc3ada76c5..c4744d61d4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1338,7 +1338,10 @@ class ProxyBaseLLMRequestProcessing: else: generator = response - if self._has_post_call_guardrails_for_passthrough(): + if ( + self._has_post_call_guardrails_for_passthrough() + and self._passthrough_provider_has_stream_guardrail_handler() + ): body_bytes = b"".join( [chunk async for chunk in generator] # type: ignore[union-attr] ) @@ -1765,6 +1768,22 @@ class ProxyBaseLLMRequestProcessing: return True return False + def _passthrough_provider_has_stream_guardrail_handler(self) -> bool: + """ + True when the resolved passthrough provider has an event-stream guardrail + handler that can rewrite buffered frames. Only such providers may have + their stream buffered for post-call guardrails; every other provider must + keep streaming so the response is not silently turned into a non-streaming + body when an unrelated post-call guardrail is registered. + """ + from litellm.llms.pass_through.guardrail_translation.handler import ( + LlmPassthroughRouteHandler, + ) + + return LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( + self.data.get("custom_llm_provider") + ) + async def _handle_non_streaming_allm_passthrough_route( self, response: Any, diff --git a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py index cc9c6eda20..1f6cafe5f3 100644 --- a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py @@ -187,3 +187,25 @@ class TestDeAnonymizeEventStream: ) assert result is body + + +class TestSupportsEventStreamDeAnonymization: + def test_bedrock_is_supported(self): + assert ( + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization("bedrock") + is True + ) + + def test_unknown_provider_is_not_supported(self): + assert ( + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( + "anthropic" + ) + is False + ) + + def test_missing_provider_is_not_supported(self): + assert ( + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization(None) + is False + ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 31235a8df0..1ee14be635 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2644,3 +2644,105 @@ class TestEventStreamAllmPassthroughRoute: assert result.headers.get("x-litellm-model-id") == "bedrock/claude" # content-length from custom_headers is filtered; Starlette sets the correct value from body assert result.headers.get("content-length") != "99" + + +class TestAllmPassthroughStreamingProviderGate: + """ + Regression: the streaming-buffer gate for allm_passthrough_route must only + fire for providers that have an event-stream guardrail handler (Bedrock). + + A non-Bedrock streaming passthrough response must keep streaming even when a + post-call guardrail is registered globally, instead of being silently + buffered into a non-streaming Response. Bedrock must still be buffered so the + converse-stream de-anonymization handler can rewrite frames. + """ + + def _build_processing_obj(self, custom_llm_provider: str) -> ProxyBaseLLMRequestProcessing: + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-123" + logging_obj.cost_breakdown = None + data = { + "custom_llm_provider": custom_llm_provider, + "litellm_logging_obj": logging_obj, + } + return ProxyBaseLLMRequestProcessing(data=data) + + async def _run(self, processing_obj, monkeypatch, chunks): + import litellm.proxy.common_request_processing as crp + from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth + + async def streaming_response(): + for chunk in chunks: + yield chunk + + async def fake_route_request(**kwargs): + async def _llm_call(): + return streaming_response() + + return _llm_call() + + monkeypatch.setattr(crp, "route_request", fake_route_request) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock() + + return await processing_obj.base_process_llm_request( + request=MagicMock(spec=Request, headers={}), + fastapi_response=Response(), + user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"), + route_type="allm_passthrough_route", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=None, + llm_router=None, + skip_pre_call_logic=True, + ) + + @pytest.mark.asyncio + async def test_non_bedrock_stream_is_not_buffered(self, monkeypatch): + processing_obj = self._build_processing_obj("anthropic") + chunks = [b"chunk-1", b"chunk-2"] + + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ): + result = await self._run(processing_obj, monkeypatch, chunks) + + assert isinstance(result, StreamingResponse) + streamed = [chunk async for chunk in result.body_iterator] + assert streamed == chunks + + @pytest.mark.asyncio + async def test_bedrock_stream_is_buffered_through_handler(self, monkeypatch): + processing_obj = self._build_processing_obj("bedrock") + chunks = [b"raw-1", b"raw-2"] + + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: + result = await self._run(processing_obj, monkeypatch, chunks) + + assert isinstance(result, Response) + assert not isinstance(result, StreamingResponse) + assert result.body == b"modified-body" + mock_handler.assert_awaited_once()