fix(guardrails): return HTTP 400 instead of 500 for Model Armor streaming blocks (#24693)

When Model Armor blocks a streaming response, it correctly raises
HTTPException(status_code=400) but create_response() catches it with a
bare except Exception and hardcodes a 500 response, discarding the
original status code.

Fix create_response() to preserve status_code from HTTPException instead
of hardcoding 500. Also update Model Armor's streaming hook to yield an
SSE error event instead of raising (matching the Prisma Airs pattern),
and fix make_model_armor_request() to return 400 for upstream API
failures instead of passing through the upstream status code.
This commit is contained in:
michelligabriele
2026-04-02 21:28:52 -07:00
committed by GitHub
parent 52a596d2a4
commit a6dfd02610
4 changed files with 179 additions and 9 deletions
@@ -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