diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 386b9d2772..09200c9684 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -221,16 +221,21 @@ async def create_response( f"Error consuming first chunk from generator: {e}" ) - # Fallback to a generic error stream + # Preserve status code from HTTPException (e.g., guardrail blocks) + error_status = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) + error_detail = getattr(e, "detail", "Error processing stream start") + if not isinstance(error_detail, str): + error_detail = str(error_detail) + async def error_gen_message() -> AsyncGenerator[str, None]: - yield f"data: {json.dumps({'error': {'message': 'Error processing stream start', 'code': status.HTTP_500_INTERNAL_SERVER_ERROR}})}\n\n" + yield f"data: {json.dumps({'error': {'message': error_detail, 'code': error_status}})}\n\n" yield "data: [DONE]\n\n" return StreamingResponse( error_gen_message(), media_type=media_type, headers=headers, - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + status_code=error_status, ) async def combined_generator() -> AsyncGenerator[str, None]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 2d3f048f81..3ee8e6a27b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -14,6 +14,8 @@ from fastapi import HTTPException if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +import json + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -203,8 +205,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): response.text, ) raise HTTPException( - status_code=response.status_code, - detail=f"Model Armor API error: {response.text}", + status_code=400, + detail=f"Model Armor API error (upstream {response.status_code}): {response.text}", ) json_response = response.json() @@ -746,8 +748,21 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): yield chunk return - except HTTPException: - raise + except HTTPException as e: + # Yield error as SSE event so create_response() detects it and + # returns a proper JSON error response with the correct status code. + # (Raising from a generator hits create_response's generic except → 500.) + detail = ( + e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} + ) + error_value = detail.get("error", detail) + if isinstance(error_value, dict): + error_obj = dict(error_value) + else: + error_obj = {"message": str(error_value)} + error_obj["code"] = e.status_code + yield f"data: {json.dumps({'error': error_obj})}\n\n" + return except Exception as e: verbose_proxy_logger.error( "Model Armor streaming error: %s", str(e), exc_info=True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 8080491f66..97ca55073d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -380,8 +380,9 @@ async def test_model_armor_api_error_handling(): call_type="completion" ) - assert exc_info.value.status_code == 500 + assert exc_info.value.status_code == 400 assert "Model Armor API error" in str(exc_info.value.detail) + assert "upstream 500" in str(exc_info.value.detail) @pytest.mark.asyncio @@ -485,6 +486,128 @@ async def test_model_armor_streaming_response(): assert len(result_chunks) > 0 mock_post.assert_called() +@pytest.mark.asyncio +async def test_model_armor_streaming_block_yields_sse_error(): + """Test that streaming content block yields SSE error event instead of raising HTTPException.""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + + # Mock Model Armor API response that triggers a block (SDP MATCH_FOUND) + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock( + return_value={ + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "inspectResult": { + "matchState": "MATCH_FOUND", + "findings": [ + { + "infoType": "PASSWORD", + "likelihood": "VERY_LIKELY", + } + ], + } + } + } + }, + } + } + ) + + guardrail._ensure_access_token_async = AsyncMock( + return_value=("test-token", "test-project") + ) + + with patch.object( + guardrail.async_handler, "post", AsyncMock(return_value=mock_response) + ): + + async def mock_stream(): + chunks = [ + litellm.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta( + content="My password is " + ) + ) + ] + ), + litellm.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="hunter2") + ) + ] + ), + ] + for chunk in chunks: + yield chunk + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "What's your password?"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + result_chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + result_chunks.append(chunk) + + # Should yield exactly one SSE error event (not raise HTTPException) + assert len(result_chunks) == 1 + error_data = json.loads(result_chunks[0].removeprefix("data: ")) + assert "error" in error_data + assert error_data["error"]["code"] == 400 + + +@pytest.mark.asyncio +async def test_model_armor_api_failure_returns_400(): + """Test that Model Armor API failures raise HTTP 400, not the upstream status code.""" + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + + # Mock a 500 response from the Model Armor GCP API + mock_response = AsyncMock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + + guardrail._ensure_access_token_async = AsyncMock( + return_value=("test-token", "test-project") + ) + + with patch.object( + guardrail.async_handler, "post", AsyncMock(return_value=mock_response) + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_model_armor_request( + content="test content", + source="user_prompt", + ) + + # Should be 400, NOT the upstream 500 + assert exc_info.value.status_code == 400 + assert "upstream 500" in str(exc_info.value.detail) + + def test_model_armor_ui_friendly_name(): """Test the UI-friendly name of the Model Armor guardrail""" from litellm.types.proxy.guardrails.guardrail_hooks.model_armor import ( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 45ec4b726a..781b559651 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4,7 +4,7 @@ from typing import AsyncGenerator from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import Request, status +from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, StreamingResponse import litellm @@ -899,6 +899,33 @@ class TestCommonRequestProcessingHelpers: assert content[0] == f"data: {json.dumps(expected_error_data)}\n\n" assert content[1] == "data: [DONE]\n\n" + async def test_create_streaming_response_generator_raises_http_exception( + self, + ): + """ + Test that when a generator raises HTTPException, the response preserves + the original status code instead of hardcoding 500. + """ + mock_gen = AsyncMock() + mock_gen.__anext__.side_effect = HTTPException( + status_code=400, detail="Content blocked by guardrail" + ) + + response = await create_response(mock_gen, "text/event-stream", {}) + assert response.status_code == 400 + content = await self.consume_stream(response) + import json + + expected_error_data = { + "error": { + "message": "Content blocked by guardrail", + "code": 400, + } + } + assert len(content) == 2 + assert content[0] == f"data: {json.dumps(expected_error_data)}\n\n" + assert content[1] == "data: [DONE]\n\n" + async def test_create_streaming_response_first_chunk_error_string_code(self): """ Test that when the first chunk contains a string error code, a JSON error response is returned