[refactor] Refactor guardrail error handling - add 400 error code (#4297)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/openai_api/test_image_server.py

@@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
     parse_size,
 )
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.errors import GuardrailViolationError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -587,6 +588,24 @@ def test_generate_single_image(test_client):
     assert test_client.app.state.engine_client.captured_prompt["modalities"] == ["image"]
 
 
+def test_generate_images_guardrail_error_returns_400(test_client, mock_async_diffusion):
+    async def blocked_generate(**kwargs):
+        raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.")
+        yield MockGenerationResult([])  # pragma: no cover
+
+    mock_async_diffusion.generate = blocked_generate
+    response = test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "blocked prompt",
+            "n": 1,
+            "size": "1024x1024",
+        },
+    )
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails."
+
+
 def test_generate_images_async_omni_sampling_params(async_omni_test_client):
     """Test AsyncOmni path uses per-stage sampling params."""
     response = async_omni_test_client.post(
@@ -891,6 +910,31 @@ def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_
     assert payloads[2] == "[DONE]"
 
 
+def test_image_edits_streaming_guardrail_error_uses_400(streaming_image_edit_client):
+    async def generate_guardrail_error(prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
+        raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.")
+        yield MockStageResult(stage_id=1, final_output_type="image", images=[])  # pragma: no cover
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_guardrail_error
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "blocked prompt",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["object"] == "error"
+    assert payloads[0]["error"]["message"] == "Input was blocked by Cosmos3 guardrails."
+    assert payloads[0]["error"]["type"] == "BadRequestError"
+    assert payloads[0]["error"]["code"] == 400
+    assert payloads[1] == "[DONE]"
+
+
 def test_image_edits_streaming_rejects_single_stage(test_client):
     img_bytes = make_test_image_bytes((16, 16))
     response = test_client.post(