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 fef984d7044..5f391767608 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 @@ -2164,3 +2164,132 @@ async def test_streaming_post_call_output_only_path_passes_request_data_to_make_ c = mock_make.call_args assert c.kwargs.get("source") == "OUTPUT" assert c.kwargs.get("request_data") is request_data + + +# --------------------------------------------------------------------------- +# SSE error-frame conversion: HTTPException must not escape the generator +# --------------------------------------------------------------------------- + + +def _make_stream_chunks() -> list: + return [ + litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="hello", role="assistant"), + finish_reason="stop", + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + ] + + +@pytest.mark.asyncio +async def test_streaming_hook_converts_http_exception_to_sse_error_frame_parallel_path(): + """HTTPException raised during parallel INPUT+OUTPUT validation must not escape + the async generator — it must be caught and emitted as in-band SSE error frames. + + Without the fix, the exception propagates after the HTTP 200 header has been + written, producing a broken/truncated stream the client cannot parse. + """ + import json as _json + + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + } + # post_call hook — no pre/during, so the parallel (INPUT+OUTPUT) path is taken + guardrail = BedrockGuardrail( + guardrail_name="bedrock-sse-parallel", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + + chunks = _make_stream_chunks() + + async def mock_stream(): + for c in chunks: + yield c + + hard_block = HTTPException(status_code=400, detail="Violated guardrail policy") + + with patch.object( + guardrail, "make_bedrock_api_request", AsyncMock(side_effect=hard_block) + ): + out = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data=request_data, + ): + out.append(chunk) + + # Must yield exactly 2 SSE frames: the error payload and [DONE] + assert len(out) == 2, f"Expected 2 SSE frames, got {len(out)}: {out}" + assert out[1] == "data: [DONE]\n\n" + + error_frame = out[0] + assert error_frame.startswith("data: ") + payload = _json.loads(error_frame[len("data: "):].rstrip()) + assert payload["error"]["code"] == 400 + assert payload["error"]["type"] == "guardrail_violation" + assert "Violated guardrail policy" in payload["error"]["message"] + + +@pytest.mark.asyncio +async def test_streaming_hook_converts_http_exception_to_sse_error_frame_output_only_path(): + """Same as above but for the output-only path (pre/during already validated INPUT). + + When event_hook is during_call, should_validate_input is False and the method + takes the single OUTPUT make_bedrock_api_request branch instead of asyncio.gather. + The HTTPException must still be caught and converted to SSE error frames. + """ + import json as _json + + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + } + guardrail = BedrockGuardrail( + guardrail_name="bedrock-sse-output-only", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.during_call, + default_on=True, + ) + + chunks = _make_stream_chunks() + + async def mock_stream(): + for c in chunks: + yield c + + hard_block = HTTPException(status_code=400, detail="Violated guardrail policy") + + with patch.object( + guardrail, "make_bedrock_api_request", AsyncMock(side_effect=hard_block) + ): + out = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data=request_data, + ): + out.append(chunk) + + assert len(out) == 2, f"Expected 2 SSE frames, got {len(out)}: {out}" + assert out[1] == "data: [DONE]\n\n" + + error_frame = out[0] + assert error_frame.startswith("data: ") + payload = _json.loads(error_frame[len("data: "):].rstrip()) + assert payload["error"]["code"] == 400 + assert payload["error"]["type"] == "guardrail_violation" + assert "Violated guardrail policy" in payload["error"]["message"]