From 774a015aad0fb87ca50ed96f8ab46fe7abced885 Mon Sep 17 00:00:00 2001 From: akraines Date: Wed, 4 Feb 2026 08:44:30 +0200 Subject: [PATCH] fix: Preserve streaming content on guardrail-sampled chunks (#20027) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The unified guardrail's streaming iterator hook processes every Nth chunk (sampling_rate, default 5). On each sampled chunk it calls process_output_streaming_response, which combines all accumulated text into the first chunk and clears all subsequent chunks to "". The hook then yielded `processed_items[-1]` — the last item, whose content had been cleared to "". This permanently lost every Nth chunk's content, causing random missing words/tokens in the client output (observed in Roo Code, Open WebUI, etc.). Fix: deep-copy the current chunk before guardrail processing runs, then yield the original (unmodified) chunk. The guardrail validation still executes and can block if it detects a problem, but the stream content is preserved. Co-authored-by: Claude Opus 4.5 --- .../unified_guardrail/unified_guardrail.py | 25 +++-- .../test_unified_guardrail.py | 100 +++++++++++++++++- 2 files changed, 114 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 80f9860bdff..f07f65d10f5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -6,6 +6,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint 3. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_post_call_streaming_iterator_hook """ +import copy from typing import Any, AsyncGenerator, List, Optional, Union from litellm._logging import verbose_proxy_logger @@ -349,22 +350,26 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply.guardrail_name, ) + # Deep-copy the current chunk before guardrail processing. + # process_output_streaming_response modifies responses_so_far + # in-place: it puts the combined guardrailed text in the first + # chunk and clears all subsequent chunks to "". Without this + # copy, yielding processed_items[-1] would yield an empty + # string, permanently losing this chunk's content. + original_item = copy.deepcopy(item) + endpoint_translation = endpoint_guardrail_translation_mappings[ CallTypes(call_type) ]() - processed_items = ( - await endpoint_translation.process_output_streaming_response( - responses_so_far=responses_so_far, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=request_data.get("litellm_logging_obj"), - user_api_key_dict=user_api_key_dict, - ) + await endpoint_translation.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, ) - last_item = processed_items[-1] - - yield last_item + yield original_item else: yield item diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 2a33c56b56a..b41cded1d0a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -8,12 +8,13 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTra from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypes +from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices class RecordingGuardrail(CustomGuardrail): @@ -131,3 +132,100 @@ class TestUnifiedLLMGuardrails: ) assert guardrail.event_history == [GuardrailEventHooks.during_call] + + class TestAsyncPostCallStreamingIteratorHook: + @pytest.mark.asyncio + async def test_streaming_content_not_lost_on_sampled_chunks(self): + """ + Verify that every chunk's content is preserved in the output stream. + + The bug: process_output_streaming_response puts the combined + guardrailed text in the first chunk and clears all subsequent + chunks to "". The hook then yielded processed_items[-1] (the + cleared last item), permanently losing every Nth chunk's content. + """ + + class _ContentClearingTranslation(BaseTranslation): + """Simulates the real OpenAI handler behavior that triggers the bug.""" + + async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj=None): # type: ignore[override] + return data + + async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None): # type: ignore[override] + return response + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + ): + # Simulate what the real handler does: + # put combined text in first chunk, clear the rest + combined = "" + for resp in responses_so_far: + for choice in resp.choices: + if choice.delta and choice.delta.content: + combined += choice.delta.content + + first_set = False + for resp in responses_so_far: + for choice in resp.choices: + if not first_set: + choice.delta.content = combined + first_set = True + else: + choice.delta.content = "" + + return responses_so_far + + # Override the mapping to use our content-clearing translation + unified_module.endpoint_guardrail_translation_mappings = { + CallTypes.acompletion: _ContentClearingTranslation, + } + + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + + # Create 10 streaming chunks with distinct content + chunks = [] + for i in range(10): + chunk = ModelResponseStream( + choices=[StreamingChoices( + delta=Delta(content=f"word{i} ", role="assistant"), + finish_reason=None, + )], + ) + chunks.append(chunk) + + async def mock_stream(): + for chunk in chunks: + yield chunk + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + request_route="/v1/chat/completions", + ) + + request_data = { + "guardrail_to_apply": guardrail, + "model": "gpt-4", + } + + # Collect all yielded chunks + yielded_contents = [] + async for item in handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + content = item.choices[0].delta.content if item.choices[0].delta else None + yielded_contents.append(content) + + # Every chunk should have non-empty content + for i, content in enumerate(yielded_contents): + assert content is not None and content != "", ( + f"Chunk {i} lost its content (got {content!r}). " + f"Expected non-empty content for every streamed chunk." + )