From ab9bb303cfb3bd2d60f3685663ecae26cce3d621 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Wed, 12 Aug 2026 19:24:51 -0700 Subject: [PATCH] fix(guardrails): let bedrock service failures keep their status instead of framing them as blocks --- .../guardrail_hooks/bedrock_guardrails.py | 17 +++++++----- .../test_bedrock_guardrails.py | 27 +++++++++++++++++-- .../guardrail_hooks/test_tool_permission.py | 25 ++++++++++++++--- 3 files changed, 57 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 345c6b524b1..84fa3473e4a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -40,6 +40,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import _serialize_http_exception_detail from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, anthropic_sse_error_frames, @@ -2624,14 +2625,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): logging_event_type=GuardrailEventHooks.post_call, ) except HTTPException as block_exc: - if not raw_sse: - raise - # A keepalive ping may already have flushed the response headers, so a raise - # cannot reach the client; deliver the block as an error frame instead block_detail: Final = block_exc.detail - for error_frame in anthropic_sse_error_frames( - block_detail if isinstance(block_detail, Mapping) else str(block_detail) - ): + # Only a policy block carries a dict detail; every API/transport failure details a + # plain string (see _is_input_too_large_error). Re-raising those keeps their real + # status instead of reporting an outage as a guardrail decision + if not raw_sse or not isinstance(block_detail, Mapping): + raise + # past the keepalive ping the headers are already flushed, so a raise cannot reach + # the client; the block travels as a frame + block_message, _ = _serialize_http_exception_detail(block_detail) + for error_frame in anthropic_sse_error_frames(block_message): yield error_frame return except ModifyResponseException as e: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 2022429eeb2..a8f39e0a0b1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2954,9 +2954,27 @@ async def test_streaming_hook_block_stream_keeps_upstream_identity(): assert b"my-proxy-alias" not in body +@pytest.mark.asyncio +async def test_streaming_hook_reraises_guardrail_service_failures(): + """A Bedrock outage must keep its status, not be reported to the caller as a guardrail decision. + + Only a policy block details a Mapping; every API/transport failure details a plain string. + """ + guardrail = _sse_guardrail() + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = HTTPException( + status_code=500, detail="Bedrock guardrail throttle retries exhausted" + ) + with pytest.raises(HTTPException) as exc: + await _drain_streaming_hook(guardrail) + + assert exc.value.status_code == 500 + + @pytest.mark.asyncio async def test_streaming_block_error_frame_message_is_a_string(): - """AnthropicErrorDetail.message is typed str, so a structured guardrail detail is serialized.""" + """AnthropicErrorDetail.message is typed str, built by the proxy's own detail serializer.""" guardrail = _sse_guardrail() with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: @@ -2967,8 +2985,10 @@ async def test_streaming_block_error_frame_message_is_a_string(): frame = next(line for line in b"".join(delivered).decode().splitlines() if line.startswith("data: ")) message = json.loads(frame[6:])["error"]["message"] + # AnthropicErrorDetail.message is typed str, and the proxy's own serializer produces the + # readable message rather than a repr of the detail dict assert isinstance(message, str) - assert json.loads(message)["error"] == "Violated guardrail policy" + assert message == "Violated guardrail policy" @pytest.mark.asyncio @@ -3071,6 +3091,9 @@ async def test_streaming_hook_yields_synthetic_block_stream_for_raw_anthropic_ss body = b"".join(delivered) assert b"Sorry, the model cannot answer this question." in body assert b"123-45-6789" not in body + # the upstream call was already paid for, so the block frame must still report its usage + assert b'"input_tokens": 5' in body + assert b'"output_tokens": 9' in body @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 60ba197317a..4b381b67f0e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -1217,10 +1217,10 @@ class TestToolPermissionGuardrailAnthropicMessages: @pytest.mark.asyncio async def test_rewrite_mode_keeps_the_stream_identity_it_had_before_the_shared_helper(self): - """The SSE helpers moved to litellm/proxy/guardrails/anthropic_sse.py must be a pure move. + """Well-formed SSE must round-trip exactly as it did before the helpers were shared. - That module can stamp the upstream message id and model onto the assembled response for - callers that ask for it; this path never did, and a client reads those bytes. + The shared module can stamp the upstream message id and model onto the assembled response + for callers that ask for it; this path never did, and a client reads those bytes. """ with patch.object(self.rewriting, "should_run_guardrail", return_value=True): out = await self._drain(self.rewriting, self._sse_chunks("Read")) @@ -1234,6 +1234,25 @@ class TestToolPermissionGuardrailAnthropicMessages: assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id" assert message_start["model"] == "unknown-model", "the rewritten stream must not adopt the upstream model" + @pytest.mark.asyncio + async def test_message_start_without_a_dict_message_fails_closed(self): + """Malformed SSE must not be forwarded unscanned. + + The shared assembler requires message_start.message to be a dict; the private helper it + replaced accepted anything, and assembled a response from it. + """ + events = [ + {"type": "message_start", "message": "not-a-dict"}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + {"type": "message_stop"}, + ] + chunks = [f"event: {e['type']}\ndata: {json.dumps(e)}\n\n".encode() for e in events] + + with patch.object(self.rewriting, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self._drain(self.rewriting, chunks) + def _resplit(self, chunks, size=7): joined = b"".join(chunks) return [joined[i : i + size] for i in range(0, len(joined), size)]