diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 6700700ff7c..e5183ac29d4 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/litellm/proxy/common_utils/sse_keepalive.py @@ -21,6 +21,17 @@ def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None: return interval +def keepalive_ping_has_fired(elapsed_seconds: float, ping_interval_seconds: float | str | None) -> bool: + """Whether a keepalive ping has already gone out, which flushes the response headers. + + A caller that discovers a failure after that point cannot raise its way to the client, since + the status line is already on the wire. With pings disabled nothing flushes early, so a raise + still carries its real status. + """ + interval: Final = _coerce_interval(ping_interval_seconds) + return interval is not None and elapsed_seconds >= interval + + def wrap_sse_stream_with_keepalive_pings( stream: AsyncGenerator[str, None], ping_interval_seconds: float | str | None, diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 293bb5ae8a0..e8c6eba581c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -14,6 +14,7 @@ import copy import json import re import sys +import time from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import datetime, timezone from itertools import accumulate, groupby @@ -41,6 +42,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import _serialize_http_exception_detail +from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, anthropic_sse_error_frames, @@ -2587,6 +2589,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): from litellm.types.utils import TextCompletionResponse # Collect all chunks to process them together + started_at: Final = time.monotonic() all_chunks: Final[list[ModelResponseStream]] = [] async for chunk in response: all_chunks.append(chunk) @@ -2627,14 +2630,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): except HTTPException as block_exc: block_detail: Final = block_exc.detail # A policy block is the only 400 carrying a structured detail; a service failure - # either details a plain string or reports a non-400 status, and re-raising those - # keeps their real status instead of reporting an outage as a guardrail decision - if not raw_sse or block_exc.status_code != 400 or not isinstance(block_detail, Mapping): + # either details a plain string or reports a non-400 status. Re-raising a service + # failure keeps its real status, but only while the headers are unflushed: past the + # first keepalive ping the raise reaches nobody, so it has to travel as a frame too + is_block: Final = raw_sse and block_exc.status_code == 400 and isinstance(block_detail, Mapping) + headers_flushed: Final = keepalive_ping_has_fired( + time.monotonic() - started_at, litellm.anthropic_sse_ping_interval_seconds + ) + if not raw_sse or (not is_block and not headers_flushed): 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): + for error_frame in anthropic_sse_error_frames( + block_message if is_block else f"{block_exc.status_code}: {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 4d287dbcb9b..f4f4003d5ee 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 @@ -2972,6 +2972,27 @@ async def test_streaming_hook_reraises_guardrail_service_failures(): assert exc.value.status_code == 500 +@pytest.mark.asyncio +async def test_streaming_hook_frames_a_service_failure_once_a_keepalive_ping_flushed_the_headers(): + """Past the ping the status line is already on the wire, so a raise reaches the client as nothing. + + The failure has to travel as a frame instead, carrying its real status in the message. + """ + guardrail = _sse_guardrail() + + with ( + patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api, + patch.object(litellm, "anthropic_sse_ping_interval_seconds", 0.0001), + ): + mock_api.side_effect = HTTPException(status_code=503, detail="Bedrock is unavailable") + delivered = await _drain_streaming_hook(guardrail) + + body = b"".join(delivered).decode() + frame = next(line for line in body.splitlines() if line.startswith("data: ")) + message = json.loads(frame[6:])["error"]["message"] + assert message == "503: Bedrock is unavailable" + + @pytest.mark.asyncio async def test_streaming_hook_reraises_a_service_failure_that_details_a_mapping(): """InvokeGuardrailChecks details a Mapping on its 500, so detail shape alone cannot mean "block"."""