[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/engine/test_orchestrator_error_handling.py

@@ -161,3 +161,113 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -
     finally:
         orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
         orchestrator_fixture.thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_diffusion_client_error_output_propagates_status_code(orchestrator_factory) -> None:
+    """A client-error OmniRequestOutput (e.g. a guardrail 4xx) from a diffusion
+    stage must surface as an ErrorMessage that carries status_code/error_type,
+    so the frontend can map it to the correct HTTP response instead of a 500.
+    """
+    stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    orchestrator_fixture = orchestrator_factory([stage0])
+    params = OmniDiffusionSamplingParams()
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-blocked",
+            prompt={"prompt": "blocked"},
+            original_prompt={"prompt": "blocked"},
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        # Push a client-error output from the diffusion stage.
+        stage0.push_diffusion_output(
+            OmniRequestOutput.from_error(
+                "req-blocked",
+                "Input was blocked by Cosmos3 guardrails.",
+                status_code=400,
+                error_type="BadRequestError",
+            )
+        )
+
+        msg = await _get_any_output_message(orchestrator_fixture)
+
+        assert isinstance(msg, ErrorMessage)
+        assert msg.request_id == "req-blocked"
+        assert msg.stage_id == 0
+        assert msg.fatal is False
+        assert msg.status_code == 400
+        assert msg.error_type == "BadRequestError"
+        assert msg.error == "Input was blocked by Cosmos3 guardrails."
+
+        # Request state should be cleaned up.
+        await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states)
+    finally:
+        orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        orchestrator_fixture.thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("status_code", "error_type"),
+    [
+        pytest.param(429, "RateLimitError", id="429-too-many-requests"),
+        pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"),
+        pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"),
+        pytest.param(403, "PermissionDeniedError", id="403-forbidden"),
+    ],
+)
+async def test_diffusion_client_error_output_propagates_non_400_status(
+    orchestrator_factory,
+    status_code: int,
+    error_type: str,
+) -> None:
+    """A diffusion-stage client error with a non-400 4xx status must be serialized
+    into an ErrorMessage that preserves that exact status_code/error_type.
+    """
+    stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    orchestrator_fixture = orchestrator_factory([stage0])
+    params = OmniDiffusionSamplingParams()
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-blocked",
+            prompt={"prompt": "blocked"},
+            original_prompt={"prompt": "blocked"},
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        # Push a non-400 client-error output from the diffusion stage.
+        stage0.push_diffusion_output(
+            OmniRequestOutput.from_error(
+                "req-blocked",
+                "client error from diffusion stage",
+                status_code=status_code,
+                error_type=error_type,
+            )
+        )
+
+        msg = await _get_any_output_message(orchestrator_fixture)
+
+        assert isinstance(msg, ErrorMessage)
+        assert msg.request_id == "req-blocked"
+        assert msg.stage_id == 0
+        assert msg.fatal is False
+        assert msg.status_code == status_code
+        assert msg.error_type == error_type
+        assert msg.error == "client error from diffusion stage"
+
+        # Request state should be cleaned up.
+        await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states)
+    finally:
+        orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        orchestrator_fixture.thread.join(timeout=5)

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(

tests/entrypoints/openai_api/test_video_server.py

@@ -31,6 +31,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
 from vllm_omni.entrypoints.openai.storage import LocalStorageManager
 from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
+from vllm_omni.errors import GuardrailViolationError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -1019,6 +1020,25 @@ def test_invalid_lora_returns_400(test_client):
     assert "lora object" in failed["error"]["message"].lower()
 
 
+def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture):
+    mocker.patch.object(
+        OmniOpenAIServingVideo,
+        "generate_video_bytes",
+        side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."),
+    )
+    response = test_client.post("/v1/videos", data={"prompt": "blocked prompt"})
+    assert response.status_code == 200
+
+    video_id = response.json()["id"]
+    failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
+    assert failed["error"]["code"] == 400
+    assert failed["error"]["message"] == "Input was blocked by Cosmos3 guardrails."
+
+    retrieve = test_client.get(f"/v1/videos/{video_id}")
+    assert retrieve.status_code == 400
+    assert retrieve.json()["error"]["code"] == 400
+
+
 def test_unsupported_image_reference_file_id_returns_400(test_client):
     response = test_client.post(
         "/v1/videos",
@@ -1521,6 +1541,20 @@ def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):
     assert "GPU exploded" in response.json()["detail"]
 
 
+def test_sync_guardrail_error_returns_400(test_client, mocker: MockerFixture):
+    mocker.patch.object(
+        OmniOpenAIServingVideo,
+        "generate_video_bytes",
+        side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."),
+    )
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "blocked prompt"},
+    )
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails."
+
+
 def test_sync_does_not_create_store_entry(test_client, mocker: MockerFixture):
     """The sync endpoint should NOT leave any record in VIDEO_STORE — it is
     stateless by design."""

tests/entrypoints/test_omni_entrypoints.py

@@ -18,6 +18,12 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
+from vllm_omni.errors import (
+    OmniClientError,
+    client_error_from_metadata,
+    client_error_metadata,
+    is_client_error_status,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -604,6 +610,71 @@ async def test_async_omni_propagates_fatal_error_context(monkeypatch: pytest.Mon
     assert getattr(exc_info.value, "error_stage_id") == 2
 
 
+def _enqueue_client_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
+    engine.output_q.put_nowait(
+        ErrorMessage(
+            request_id=msg["request_id"],
+            stage_id=2,
+            error="Input was blocked by Cosmos3 guardrails.",
+            status_code=400,
+            error_type="BadRequestError",
+        )
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_omni_propagates_client_error_status(monkeypatch: pytest.MonkeyPatch):
+    """A non-fatal client error from the orchestrator must be routed to the
+    requesting generate() call (not raised in the shared dispatcher) and
+    surface as an OmniClientError carrying status_code/error_type."""
+    engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_client_error_message)
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        with pytest.raises(OmniClientError) as exc_info:
+            async for _ in app.generate(prompt="blocked", request_id="req-1"):
+                pass
+    finally:
+        app.shutdown()
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.error_type == "BadRequestError"
+    assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails."
+
+
+def _enqueue_server_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
+    engine.output_q.put_nowait(
+        ErrorMessage(
+            request_id=msg["request_id"],
+            stage_id=2,
+            error="GPU exploded",
+        )
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_omni_propagates_server_error_as_runtime(monkeypatch: pytest.MonkeyPatch):
+    """A non-fatal error WITHOUT a 4xx status_code is a server fault: it must
+    surface as a RuntimeError (-> HTTP 500), not an OmniClientError (-> 400).
+    Guards against the fix over-broadly mapping every error to a client error."""
+    engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_server_error_message)
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        # OmniClientError subclasses ValueError, not RuntimeError, so matching
+        # RuntimeError here also proves it was NOT raised as a client error.
+        with pytest.raises(RuntimeError) as exc_info:
+            async for _ in app.generate(prompt="boom", request_id="req-1"):
+                pass
+    finally:
+        app.shutdown()
+
+    assert not isinstance(exc_info.value, OmniClientError)
+    assert str(exc_info.value) == "GPU exploded"
+
+
 def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch):
     sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
     engine = FakeAsyncOmniEngine(
@@ -761,6 +832,140 @@ def test_non_fatal_error_raises_runtime():
         base._handle_output_message(msg)
 
 
+def test_non_fatal_client_error_raises_omni_client_error():
+    """A non-fatal ErrorMessage carrying a 4xx status_code (e.g. a guardrail
+    block) must surface as an OmniClientError with the metadata intact, not a
+    bare RuntimeError. This covers the offline/sync Omni consumer path."""
+    base = _make_base()
+    msg = ErrorMessage(
+        error="Input was blocked by Cosmos3 guardrails.",
+        status_code=400,
+        error_type="BadRequestError",
+    )
+
+    with pytest.raises(OmniClientError) as exc_info:
+        base._handle_output_message(msg)
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.error_type == "BadRequestError"
+    assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails."
+
+
+_NON_400_CLIENT_ERRORS = [
+    pytest.param(429, "RateLimitError", id="429-too-many-requests"),
+    pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"),
+    pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"),
+    pytest.param(403, "PermissionDeniedError", id="403-forbidden"),
+]
+
+
+@pytest.mark.parametrize(
+    ("status_code", "error_type"), [pytest.param(400, "BadRequestError", id="400")] + _NON_400_CLIENT_ERRORS
+)
+def test_client_error_metadata_round_trip_preserves_4xx(status_code: int, error_type: str):
+    original = OmniClientError("blocked", status_code=status_code, error_type=error_type)
+
+    # Outbound: what the broad `except Exception` handlers serialize.
+    carried_status, carried_type = client_error_metadata(original)
+    assert carried_status == status_code
+    assert carried_type == error_type
+    assert is_client_error_status(carried_status)
+
+    # Inbound: reconstruction at the consuming end.
+    rebuilt = client_error_from_metadata("blocked", status_code=carried_status, error_type=carried_type)
+    assert isinstance(rebuilt, OmniClientError)
+    assert rebuilt.status_code == status_code
+    assert rebuilt.error_type == error_type
+    assert str(rebuilt) == "blocked"
+
+
+@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS)
+def test_non_fatal_client_error_preserves_non_400_status(status_code: int, error_type: str):
+    base = _make_base()
+    msg = ErrorMessage(
+        error="client side failure",
+        status_code=status_code,
+        error_type=error_type,
+    )
+
+    with pytest.raises(OmniClientError) as exc_info:
+        base._handle_output_message(msg)
+
+    assert exc_info.value.status_code == status_code
+    assert exc_info.value.error_type == error_type
+    assert str(exc_info.value) == "client side failure"
+
+
+@pytest.mark.parametrize(
+    "status_code",
+    [
+        pytest.param(None, id="no-status"),
+        pytest.param(399, id="399-below-4xx"),
+        pytest.param(500, id="500-server-error"),
+        pytest.param(503, id="503-service-unavailable"),
+    ],
+)
+def test_non_fatal_non_4xx_status_raises_runtime(status_code: int | None):
+    base = _make_base()
+    msg = ErrorMessage(
+        error="server side failure",
+        status_code=status_code,
+        error_type="ShouldBeIgnoredForNon4xx",
+    )
+
+    with pytest.raises(RuntimeError) as exc_info:
+        base._handle_output_message(msg)
+
+    assert not isinstance(exc_info.value, OmniClientError)
+    assert str(exc_info.value) == "server side failure"
+
+
+def _make_enqueue_client_error(
+    status_code: int,
+    error_type: str,
+    error_text: str,
+) -> Callable[[FakeAsyncOmniEngine, dict[str, Any]], None]:
+    def _enqueue(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
+        engine.output_q.put_nowait(
+            ErrorMessage(
+                request_id=msg["request_id"],
+                stage_id=2,
+                error=error_text,
+                status_code=status_code,
+                error_type=error_type,
+            )
+        )
+
+    return _enqueue
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS)
+async def test_async_omni_propagates_non_400_client_error_status(
+    monkeypatch: pytest.MonkeyPatch,
+    status_code: int,
+    error_type: str,
+):
+    error_text = f"blocked with {status_code}"
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=_make_enqueue_client_error(status_code, error_type, error_text),
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        with pytest.raises(OmniClientError) as exc_info:
+            async for _ in app.generate(prompt="blocked", request_id="req-1"):
+                pass
+    finally:
+        app.shutdown()
+
+    assert exc_info.value.status_code == status_code
+    assert exc_info.value.error_type == error_type
+    assert str(exc_info.value) == error_text
+
+
 def test_async_omni_errored_property_alive():
     omni = object.__new__(AsyncOmni)
     omni.engine = SimpleNamespace(

vllm_omni/diffusion/data.py

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
+from vllm_omni.errors import client_error_metadata
 from vllm_omni.quantization import build_quant_config
 
 if TYPE_CHECKING:
@@ -1062,6 +1063,8 @@ class DiffusionOutput:
     trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None
     trajectory_decoded: list[Image.Image] | None = None
     error: str | None = None
+    error_status_code: int | None = None
+    error_type: str | None = None
     aborted: bool = False
     abort_message: str | None = None
 
@@ -1102,6 +1105,15 @@ class DiffusionOutput:
         if self.custom_output:
             self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()}
 
+    @classmethod
+    def from_exception(cls, exc: BaseException) -> "DiffusionOutput":
+        status_code, error_type = client_error_metadata(exc)
+        return cls(
+            error=str(exc),
+            error_status_code=status_code,
+            error_type=error_type,
+        )
+
 
 class DiffusionRequestAbortedError(RuntimeError):
     """Raised when a diffusion request ends via user-visible abort."""

vllm_omni/diffusion/diffusion_engine.py

@@ -34,6 +34,7 @@ from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusio
 from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
 from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
 from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
+from vllm_omni.errors import client_error_from_metadata, is_client_error_status
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -231,6 +232,12 @@ class DiffusionEngine:
         if output.aborted:
             raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
         if output.error:
+            if is_client_error_status(output.error_status_code):
+                raise client_error_from_metadata(
+                    output.error,
+                    status_code=output.error_status_code,
+                    error_type=output.error_type,
+                )
             raise RuntimeError(output.error)
         logger.debug("Generation completed successfully.")
 
@@ -507,7 +514,7 @@ class DiffusionEngine:
                             request_id=request_id,
                             step_index=None,
                             finished=True,
-                            result=DiffusionOutput(error=str(exc)),
+                            result=DiffusionOutput.from_exception(exc),
                         )
                         for request_id in sched_output.scheduled_request_ids
                     ]
@@ -664,7 +671,7 @@ class DiffusionEngine:
                         request_id=request_id,
                         step_index=None,
                         finished=True,
-                        result=DiffusionOutput(error=str(exc)),
+                        result=DiffusionOutput.from_exception(exc),
                     )
 
                 self._process_aborts_queue()

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -22,6 +22,7 @@ from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.errors import client_error_metadata
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -143,11 +144,13 @@ class InlineStageDiffusionClient(StageClientBase):
             logger.info("request_id: %s aborted: %s", request_id, str(e))
         except Exception as e:
             logger.exception("Diffusion request %s failed: %s", request_id, e)
-            error_output = OmniRequestOutput.from_diffusion(
+            status_code, error_type = client_error_metadata(e)
+            error_output = OmniRequestOutput.from_error(
                 request_id=request_id,
-                images=[],
+                error_message=str(e),
+                status_code=status_code,
+                error_type=error_type,
             )
-            error_output.error = str(e)
             self._output_queue.put_nowait(error_output)
         finally:
             self._tasks.pop(request_id, None)
@@ -248,11 +251,13 @@ class InlineStageDiffusionClient(StageClientBase):
             logger.info("request_id: %s aborted: %s", request_id, str(e))
         except Exception as e:
             logger.exception("Batch diffusion request %s failed: %s", request_id, e)
-            error_output = OmniRequestOutput.from_diffusion(
+            status_code, error_type = client_error_metadata(e)
+            error_output = OmniRequestOutput.from_error(
                 request_id=request_id,
-                images=[],
+                error_message=str(e),
+                status_code=status_code,
+                error_type=error_type,
             )
-            error_output.error = str(e)
             self._output_queue.put_nowait(error_output)
         finally:
             self._tasks.pop(request_id, None)

vllm_omni/diffusion/models/cosmos3/guardrails.py

@@ -20,6 +20,7 @@ import torch
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.progress_bar import _is_rank_zero
+from vllm_omni.errors import GuardrailViolationError
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -67,7 +68,7 @@ def _build_text_guardrail(checker: Any) -> TextGuardrailFn:
     def text_guardrail(prompt: str) -> None:
         if not checker.check_text_safety(prompt):
             # CosmosSafetyChecker logs the specific reason at CRITICAL.
-            raise ValueError("Guardrail blocked prompt")
+            raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.")
 
     return text_guardrail
 

vllm_omni/diffusion/stage_diffusion_client.py

@@ -228,7 +228,9 @@ class StageDiffusionClient(StageClientBase):
             elif msg_type == "error":
                 req_id = msg.get("request_id")
                 rpc_id = msg.get("rpc_id")
-                error_msg = msg.get("error")
+                error_msg = msg.get("error") or "Unknown diffusion subprocess error."
+                status_code = msg.get("status_code")
+                error_type = msg.get("error_type")
                 logger.error(
                     "[StageDiffusionClient] stage-%s [rep-%s] subprocess error for %s: %s",
                     self.stage_id,
@@ -245,7 +247,14 @@ class StageDiffusionClient(StageClientBase):
                 # Route request errors as error outputs so the Orchestrator
                 # sees the request complete (instead of hanging forever).
                 if req_id is not None:
-                    self._output_queue.put_nowait(OmniRequestOutput.from_error(req_id, error_msg))
+                    self._output_queue.put_nowait(
+                        OmniRequestOutput.from_error(
+                            req_id,
+                            error_msg,
+                            status_code=status_code,
+                            error_type=error_type,
+                        )
+                    )
 
     # Fields that are subprocess-local and cannot be serialized across
     # process boundaries.  They are recreated in the subprocess with

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -33,6 +33,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
 )
 from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
 from vllm_omni.engine.stage_init_utils import set_death_signal
+from vllm_omni.errors import client_error_metadata
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -400,12 +401,15 @@ class StageDiffusionProc:
                 )
             except Exception as e:
                 logger.exception("Diffusion request %s failed: %s", request_id, e)
+                status_code, error_type = client_error_metadata(e)
                 await response_socket.send(
                     encoder.encode(
                         {
                             "type": "error",
                             "request_id": request_id,
                             "error": str(e),
+                            "status_code": status_code,
+                            "error_type": error_type,
                         }
                     )
                 )
@@ -483,12 +487,15 @@ class StageDiffusionProc:
                             )
                         except Exception as e:
                             logger.exception("Batch diffusion request %s failed: %s", rid, e)
+                            status_code, error_type = client_error_metadata(e)
                             await response_socket.send(
                                 encoder.encode(
                                     {
                                         "type": "error",
                                         "request_id": rid,
                                         "error": str(e),
+                                        "status_code": status_code,
+                                        "error_type": error_type,
                                     }
                                 )
                             )

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -801,7 +801,7 @@ class WorkerProc:
                 except Exception as e:
                     logger.error(f"Error processing RPC: {e}", exc_info=True)
                     if self.result_mq is not None:
-                        self.return_result(DiffusionOutput(error=str(e)))
+                        self.return_result(DiffusionOutput.from_exception(e))
 
             elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                 logger.info("Worker %s: Received shutdown message", self.gpu_id)
@@ -817,7 +817,7 @@ class WorkerProc:
                         f"Error executing forward in event loop: {e}",
                         exc_info=True,
                     )
-                    output = DiffusionOutput(error=str(e))
+                    output = DiffusionOutput.from_exception(e)
 
                 try:
                     self.return_result(output)

vllm_omni/engine/messages.py

@@ -73,6 +73,8 @@ class UnregisterRemoteReplicaMessage(EngineQueueMessage, kw_only=True):
 class ErrorMessage(EngineQueueMessage, kw_only=True):
     type: Literal["error"] = "error"
     error: str
+    status_code: int | None = None
+    error_type: str | None = None
     fatal: bool = False
     request_id: str | None = None
     stage_id: int | None = None

vllm_omni/engine/orchestrator.py

@@ -809,6 +809,8 @@ class Orchestrator:
                 request_id=parent_id,
                 stage_id=stage_id,
                 error=output.error,
+                status_code=getattr(output, "error_status_code", None),
+                error_type=getattr(output, "error_type", None),
             )
         )
         await self._cleanup_request_ids(

vllm_omni/entrypoints/async_omni.py

@@ -33,6 +33,7 @@ from vllm_omni.entrypoints.omni_base import (
     OmniBase,
     OmniEngineDeadError,
 )
+from vllm_omni.errors import client_error_metadata
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
 from vllm_omni.outputs import OmniRequestOutput
@@ -483,10 +484,13 @@ class AsyncOmni(EngineClient, OmniBase):
             except (asyncio.CancelledError, GeneratorExit):
                 cancelled = True
             except Exception as error:
+                status_code, error_type = client_error_metadata(error)
                 await req_state.queue.put(
                     ErrorMessage(
                         request_id=request_id,
                         error=str(error),
+                        status_code=status_code,
+                        error_type=error_type,
                     )
                 )
             finally:
@@ -589,7 +593,7 @@ class AsyncOmni(EngineClient, OmniBase):
                         result.error,
                         error_stage_id=result.stage_id,
                     )
-                raise RuntimeError(result.error)
+                self._raise_nonfatal_error_message(result)
 
             if not isinstance(result, OutputMessage):
                 logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result)
@@ -662,6 +666,18 @@ class AsyncOmni(EngineClient, OmniBase):
                         await self.event_resolver.resolve(msg)
                         continue
 
+                    if isinstance(msg, ErrorMessage) and not msg.fatal:
+                        req_state = self.request_states.get(msg.request_id)
+                        if req_state is not None:
+                            await req_state.queue.put(msg)
+                        else:
+                            logger.warning(
+                                "[%s] dropping non-fatal error for unknown req %s",
+                                self._name,
+                                msg.request_id,
+                            )
+                        continue
+
                     should_continue, _, stage_id, req_state = self._handle_output_message(msg)
                     if should_continue:
                         continue

vllm_omni/entrypoints/omni_base.py

@@ -20,6 +20,7 @@ from vllm_omni.engine.messages import (
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e
+from vllm_omni.errors import raise_client_error_or
 from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_finalize
 from vllm_omni.metrics.prometheus import OmniPrometheusMetrics
 from vllm_omni.metrics.stats import OrchestratorAggregator
@@ -372,7 +373,7 @@ class OmniBase(PDDisaggregationMixin):
                     msg.error,
                     error_stage_id=msg.stage_id,
                 )
-            raise RuntimeError(msg.error)
+            self._raise_nonfatal_error_message(msg)
 
         if not isinstance(msg, OutputMessage):
             logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg.type)
@@ -408,6 +409,15 @@ class OmniBase(PDDisaggregationMixin):
 
         return False, req_id, stage_id, req_state
 
+    def _raise_nonfatal_error_message(self, msg: ErrorMessage) -> None:
+        """Raise the exception for a non-fatal, request-scoped error message."""
+        raise_client_error_or(
+            msg.error,
+            status_code=msg.status_code,
+            error_type=msg.error_type,
+            fallback=RuntimeError,
+        )
+
     def _check_engine_output_error(
         self,
         result: OutputMessage,
@@ -424,6 +434,8 @@ class OmniBase(PDDisaggregationMixin):
         error_text = getattr(engine_outputs, "error", None)
         if error_text is None:
             return
+        status_code = engine_outputs.error_status_code
+        error_type = engine_outputs.error_type
         logger.error(
             "[%s] Stage error for req=%s stage-%s: %s",
             self.__class__.__name__,
@@ -437,7 +449,12 @@ class OmniBase(PDDisaggregationMixin):
                 error_text,
                 error_stage_id=stage_id,
             )
-        raise EngineGenerateError(error_text)
+        raise_client_error_or(
+            error_text,
+            status_code=status_code,
+            error_type=error_type,
+            fallback=EngineGenerateError,
+        )
 
     def _process_single_result(
         self,

vllm_omni/entrypoints/openai/api_server.py

@@ -142,6 +142,7 @@ from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url, decode_input_reference
 from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI
+from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
@@ -1812,6 +1813,9 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         return _create_engine_error_json_response(raw_request, exc)
     except HTTPException:
         raise
+    except OmniClientError as e:
+        logger.info("Client error during image generation: %s", e)
+        raise HTTPException(status_code=e.status_code, detail=e.message) from e
     except ValueError as e:
         logger.error(f"Validation error: {e}")
         raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
@@ -2167,6 +2171,9 @@ async def edit_images(
         return _create_engine_error_json_response(raw_request, exc)
     except HTTPException:
         raise
+    except OmniClientError as e:
+        logger.info("Client error during image edit: %s", e)
+        raise HTTPException(status_code=e.status_code, detail=e.message) from e
     except ValueError as e:
         logger.error(f"Validation error: {e}")
         raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e))
@@ -2667,6 +2674,9 @@ def _video_error_from_exception(exc: Exception) -> VideoError:
         message = str(exc.detail) if exc.detail else str(exc)
         return VideoError(code=exc.status_code, message=message)
 
+    if isinstance(exc, OmniClientError):
+        return VideoError(code=exc.status_code, message=exc.message)
+
     if isinstance(exc, (EngineGenerateError, EngineDeadError)):
         err = create_error_response(exc)
         return VideoError(code=err.error.code, message=err.error.message)
@@ -3021,6 +3031,9 @@ async def create_video_sync(
         return _create_engine_error_json_response(raw_request, exc)
     except HTTPException:
         raise
+    except OmniClientError as exc:
+        logger.info("Client error during sync video generation: %s", exc)
+        raise HTTPException(status_code=exc.status_code, detail=exc.message) from exc
     except Exception as exc:
         logger.exception("Sync video generation failed for request_id=%s", request_id)
         raise HTTPException(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -113,6 +113,7 @@ from vllm_omni.entrypoints.openai.utils import (
     resolve_diffusion_od_config,
     validate_requested_speaker,
 )
+from vllm_omni.errors import OmniClientError
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate
@@ -2869,6 +2870,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     "[OmniOpenAIServingChat] Engine dead during streaming image edit, "
                     "but cannot terminate server (no raw_request context)."
                 )
+        except OmniClientError as exc:
+            logger.info("Client error during streaming image edit: %s", exc)
+            chunk = ImageEditStreamError(
+                created=created,
+                model=model,
+                error={
+                    "message": exc.message,
+                    "type": exc.error_type,
+                    "code": exc.status_code,
+                },
+            )
+            yield f"data: {chunk.model_dump_json()}\n\n"
         except Exception as exc:
             logger.exception("Streaming image edit failed: %s", exc)
             chunk = ImageEditStreamError(
@@ -3265,6 +3278,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             return response
 
+        except OmniClientError as e:
+            logger.info("Client error during diffusion chat completion: %s", e)
+            return self._create_error_response(
+                e.message,
+                err_type=e.error_type,
+                status_code=e.status_code,
+            )
         except Exception as e:
             logger.exception("Diffusion chat completion failed: %s", e)
             return self._create_error_response(

vllm_omni/errors.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Request-scoped client error types shared across vLLM-Omni entrypoints."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from http import HTTPStatus
+from typing import NoReturn
+
+DEFAULT_CLIENT_ERROR_TYPE = "BadRequestError"
+
+
+class OmniClientError(ValueError):
+    """
+    Request-scoped error that should be surfaced as a 4xx response.
+    One example of using OmniClientError is GuardrailViolationError, which is captured and resurfaced
+    as HTTP 400 error code, instead of a generic 500. OmniClientError should be used for any exceptions
+    which need to be resurfaced as 4xx, as opposed to EngineDeadError/EngineGenerateError which are resurfaced as 500.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status_code: int = HTTPStatus.BAD_REQUEST.value,
+        error_type: str = DEFAULT_CLIENT_ERROR_TYPE,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.status_code = int(status_code)
+        self.error_type = error_type
+
+
+class GuardrailViolationError(OmniClientError):
+    """Raised when a model guardrail rejects request content."""
+
+
+def client_error_metadata(exc: BaseException) -> tuple[int | None, str | None]:
+    if isinstance(exc, OmniClientError):
+        return exc.status_code, exc.error_type
+    return None, None
+
+
+def client_error_from_metadata(
+    message: str,
+    *,
+    status_code: int | None,
+    error_type: str | None,
+) -> OmniClientError:
+    return OmniClientError(
+        message,
+        status_code=status_code or HTTPStatus.BAD_REQUEST.value,
+        error_type=error_type or DEFAULT_CLIENT_ERROR_TYPE,
+    )
+
+
+def is_client_error_status(status_code: int | None) -> bool:
+    return status_code is not None and 400 <= int(status_code) < 500
+
+
+def raise_client_error_or(
+    message: str,
+    *,
+    status_code: int | None,
+    error_type: str | None,
+    fallback: Callable[[str], BaseException],
+) -> NoReturn:
+    """Raise a client error for 4xx statuses, otherwise raise ``fallback(message)``.
+
+    Centralizes the "client-error-or-fallback" decision shared by the engine
+    error paths so the status mapping lives in one place.
+    """
+    if is_client_error_status(status_code):
+        raise client_error_from_metadata(
+            message,
+            status_code=status_code,
+            error_type=error_type,
+        )
+    raise fallback(message)

vllm_omni/outputs.py

@@ -102,12 +102,17 @@ class OmniRequestOutput:
 
     # error handling
     error: str | None = None
+    error_status_code: int | None = None
+    error_type: str | None = None
 
     @classmethod
     def from_error(
         cls,
         request_id: str,
         error_message: str,
+        *,
+        status_code: int | None = None,
+        error_type: str | None = None,
     ) -> "OmniRequestOutput":
         """Create a terminal error output.
 
@@ -122,6 +127,8 @@ class OmniRequestOutput:
             request_id=request_id,
             finished=True,
             error=error_message,
+            error_status_code=status_code,
+            error_type=error_type,
         )
 
     @classmethod