Merge pull request #38734 from BerriAI/litellm_fix_bedrock_buffered_responses_stream

fix(bedrock): route streamed responses-API output through the unified guardrail

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -82,6 +82,10 @@ class ResponsesStreamChunk(TypedDict, total=False):
 
     type: ReadOnly[str]
     text: ReadOnly[str]
+    delta: ReadOnly[str]
+    item_id: ReadOnly[str]
+    output_index: ReadOnly[int]
+    content_index: ReadOnly[int]
 
 
 def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
@@ -658,8 +662,32 @@ class OpenAIResponsesHandler(BaseTranslation):
     def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
         """
         Get the string so far from the responses so far.
+
+        ``response.output_text.done`` events carry the whole part in ``text``, while
+        ``response.output_text.delta`` events carry fragments in ``delta``. A stream
+        that dies before its done event (``response.failed`` / ``response.incomplete``)
+        has text only in deltas, so per content part the done text wins when present
+        and the joined deltas fill in otherwise, never both.
         """
-        return "".join([response.get("text", "") for response in responses_so_far])
+        keyed_events: Final = tuple(
+            (
+                (event.get("item_id"), event.get("output_index"), event.get("content_index")),
+                event.get("text"),
+                event.get("delta"),
+            )
+            for event in responses_so_far
+            if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str)
+        )
+
+        def part_text(part_key: tuple[object, object, object]) -> str:
+            done_texts: Final = tuple(
+                text for key, text, _ in keyed_events if key == part_key and isinstance(text, str)
+            )
+            if done_texts:
+                return done_texts[-1]
+            return "".join(delta for key, _, delta in keyed_events if key == part_key and isinstance(delta, str))
+
+        return "".join(part_text(key) for key in dict.fromkeys(key for key, _, _ in keyed_events))
 
     def _has_text_content(self, response: "ResponsesAPIResponse") -> bool:
         """

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -31,6 +31,7 @@ from litellm.caching import DualCache
 from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
 from litellm.exceptions import ModifyResponseException
 from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
 from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
 from litellm.litellm_core_utils.litellm_logging import (
     _get_masked_values,  # pyright: ignore[reportPrivateUsage]  # the shared header-masking helper has no public name
@@ -215,6 +216,16 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]:
     return redacted if isinstance(redacted, list) else assessments
 
 
+_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
+
+
+def _is_responses_api_route(request_route: str | None) -> bool:
+    if request_route is None:
+        return False
+    call_types: Final = get_call_types_for_route(request_route)
+    return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types)
+
+
 class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
     # During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
     # OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
@@ -2709,6 +2720,24 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                 yield streamed_chunk
             return
 
+        # Responses-API events are neither chat-completions chunks nor raw
+        # Anthropic SSE, so the assembly below cannot scan them; the unified
+        # guardrail's translation layer can, with buffering semantics kept.
+        if _is_responses_api_route(user_api_key_dict.request_route):
+            from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
+                UnifiedLLMGuardrails,
+            )
+
+            async for translated_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
+                user_api_key_dict=user_api_key_dict,
+                response=response,
+                request_data=request_data,
+                guardrail_to_apply=self,
+                buffer_until_moderated_default=True,
+            ):
+                yield translated_chunk
+            return
+
         # Import here to avoid circular imports
         from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
         from litellm.main import stream_chunk_builder

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -828,6 +828,24 @@ class MockPassThroughGuardrail(CustomGuardrail):
         return inputs
 
 
+class MockRecordingGuardrail(MockPassThroughGuardrail):
+    """Pass-through guardrail that records every apply_guardrail inputs payload"""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.seen_inputs: List[GenericGuardrailAPIInputs] = []
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        self.seen_inputs.append(inputs)
+        return inputs
+
+
 class TestOpenAIResponsesHandlerStreamingOutputProcessing:
     """Test streaming output processing functionality"""
 
@@ -1104,6 +1122,80 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
         output_text = result[-1]["response"]["output"][0]["content"][0]["text"]
         assert output_text == original_text
 
+    @pytest.mark.asyncio
+    async def test_failed_stream_scans_delta_text(self):
+        """A stream ending in response.failed has text only in delta events; the
+        fallback scan must assemble and scan it instead of skipping on an empty string."""
+        handler = OpenAIResponsesHandler()
+        guardrail = MockRecordingGuardrail(guardrail_name="test")
+
+        responses_so_far = [
+            {"type": "response.created", "response": {"id": "resp_123"}},
+            {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_123"}},
+            {
+                "type": "response.output_text.delta",
+                "item_id": "msg_123",
+                "output_index": 0,
+                "content_index": 0,
+                "delta": "Hello",
+            },
+            {
+                "type": "response.output_text.delta",
+                "item_id": "msg_123",
+                "output_index": 0,
+                "content_index": 0,
+                "delta": " world",
+            },
+            {"type": "response.failed", "response": {"id": "resp_123", "status": "failed"}},
+        ]
+
+        result = await handler.process_output_streaming_response(
+            responses_so_far=responses_so_far,
+            guardrail_to_apply=guardrail,
+            litellm_logging_obj=None,
+        )
+
+        assert result == responses_so_far
+        assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["Hello world"]]
+
+    def test_get_streaming_string_so_far_prefers_done_text_over_deltas(self):
+        """The done event repeats the whole part, so deltas must not be double counted;
+        a part with no done event yet still contributes its joined deltas."""
+        handler = OpenAIResponsesHandler()
+
+        events = [
+            {
+                "type": "response.output_text.delta",
+                "item_id": "msg_1",
+                "output_index": 0,
+                "content_index": 0,
+                "delta": "Hello",
+            },
+            {
+                "type": "response.output_text.delta",
+                "item_id": "msg_1",
+                "output_index": 0,
+                "content_index": 0,
+                "delta": " world",
+            },
+            {
+                "type": "response.output_text.done",
+                "item_id": "msg_1",
+                "output_index": 0,
+                "content_index": 0,
+                "text": "Hello world",
+            },
+            {
+                "type": "response.output_text.delta",
+                "item_id": "msg_2",
+                "output_index": 1,
+                "content_index": 0,
+                "delta": "; unfinished",
+            },
+        ]
+
+        assert handler.get_streaming_string_so_far(events) == "Hello world; unfinished"
+
 
 class TestGetStructuredMessages:
     """Test the get_structured_messages method for Responses API handler."""

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -5592,6 +5592,156 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
     assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail"
 
 
+def _responses_stream_events() -> list:
+    from litellm.types.llms.openai import (
+        OutputTextDeltaEvent,
+        ResponseCompletedEvent,
+        ResponsesAPIResponse,
+        ResponsesAPIStreamEvents,
+    )
+
+    deltas = [
+        OutputTextDeltaEvent(
+            type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
+            item_id="msg_lit6457",
+            output_index=0,
+            content_index=0,
+            delta=part,
+        )
+        for part in ("Hello", " world")
+    ]
+    completed = ResponseCompletedEvent(
+        type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
+        response=ResponsesAPIResponse(
+            id="resp_lit6457",
+            created_at=1234567890,
+            model="gpt-4o",
+            object="response",
+            status="completed",
+            output=[
+                {
+                    "type": "message",
+                    "id": "msg_lit6457",
+                    "status": "completed",
+                    "role": "assistant",
+                    "content": [{"type": "output_text", "text": "Hello world"}],
+                }
+            ],
+        ),
+    )
+    return [*deltas, completed]
+
+
+@pytest.mark.asyncio
+async def test_responses_api_stream_scans_output_and_replays_buffered_events():
+    """Streamed /v1/responses events must be scanned via the unified translation
+    layer, not fed to stream_chunk_builder (which raises APIError on them)."""
+    guardrail = BedrockGuardrail(
+        guardrail_name="bedrock-responses-stream",
+        guardrailIdentifier="test-id",
+        guardrailVersion="DRAFT",
+        event_hook=GuardrailEventHooks.post_call,
+        default_on=True,
+    )
+    stream_events = _responses_stream_events()
+    order = []
+    yielded = []
+
+    async def record_scan(*args, **kwargs):
+        order.append("scan")
+        return {"action": "NONE", "assessments": [], "outputs": []}
+
+    async def mock_stream():
+        for event in stream_events:
+            yield event
+
+    with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)):
+        async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"),
+            response=mock_stream(),
+            request_data={"model": "gpt-4o", "input": "hi"},
+        ):
+            order.append("chunk")
+            yielded.append(chunk)
+
+    assert order == ["scan", "chunk", "chunk", "chunk"]
+    assert len(yielded) == len(stream_events)
+    assert all(emitted is original for emitted, original in zip(yielded, stream_events))
+
+
+def _responses_failed_stream_events() -> list:
+    from litellm.types.llms.openai import (
+        OutputTextDeltaEvent,
+        ResponseFailedEvent,
+        ResponsesAPIResponse,
+        ResponsesAPIStreamEvents,
+    )
+
+    deltas = [
+        OutputTextDeltaEvent(
+            type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
+            item_id="msg_lit6457_failed",
+            output_index=0,
+            content_index=0,
+            delta=part,
+        )
+        for part in ("Hello", " world")
+    ]
+    failed = ResponseFailedEvent(
+        type=ResponsesAPIStreamEvents.RESPONSE_FAILED,
+        response=ResponsesAPIResponse(
+            id="resp_lit6457_failed",
+            created_at=1234567890,
+            model="gpt-4o",
+            object="response",
+            status="failed",
+            output=[],
+        ),
+    )
+    return [*deltas, failed]
+
+
+@pytest.mark.asyncio
+async def test_responses_api_failed_stream_scans_delta_text_before_replay():
+    """A responses stream that dies mid-generation carries its text only in delta
+    events; the end-of-stream scan must still see that text instead of skipping
+    on an empty assembled string and replaying the buffer unmoderated."""
+    guardrail = BedrockGuardrail(
+        guardrail_name="bedrock-responses-failed-stream",
+        guardrailIdentifier="test-id",
+        guardrailVersion="DRAFT",
+        event_hook=GuardrailEventHooks.post_call,
+        default_on=True,
+    )
+    stream_events = _responses_failed_stream_events()
+    order = []
+    scan_payloads = []
+    yielded = []
+
+    async def record_scan(*args, **kwargs):
+        order.append("scan")
+        scan_payloads.append(str(args) + str(kwargs))
+        return {"action": "NONE", "assessments": [], "outputs": []}
+
+    async def mock_stream():
+        for event in stream_events:
+            yield event
+
+    with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)):
+        async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"),
+            response=mock_stream(),
+            request_data={"model": "gpt-4o", "input": "hi"},
+        ):
+            order.append("chunk")
+            yielded.append(chunk)
+
+    assert order == ["scan", "chunk", "chunk", "chunk"]
+    assert "Hello world" in scan_payloads[0]
+    assert len(yielded) == len(stream_events)
+    assert all(emitted is original for emitted, original in zip(yielded, stream_events))
+
+
 @pytest.mark.asyncio
 async def test_apply_guardrail_debug_log_masks_signed_request_headers():
     import logging