From c4220f0979cb07df4443e2688392171fd53f150f Mon Sep 17 00:00:00 2001
From: lif <1835304752@qq.com>
Date: Wed, 28 Jan 2026 13:11:15 +0800
Subject: [PATCH] [BugFix] Add /health and /v1/models endpoints for diffusion
 mode, fix streaming compatibility (#454)

Signed-off-by: majiayu000 <1835304752@qq.com>
---
 .../openai_api/test_image_server.py           |  59 +++++-
 tests/test_outputs.py                         | 171 ++++++++++++++++++
 vllm_omni/entrypoints/openai/api_server.py    |  88 ++++++++-
 vllm_omni/outputs.py                          |  52 ++++++
 4 files changed, 358 insertions(+), 12 deletions(-)
 create mode 100644 tests/test_outputs.py

diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 45f4b33ad..86331e085 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -124,6 +124,8 @@ class FakeAsyncOmni:
 def mock_async_diffusion():
     """Mock AsyncOmniDiffusion instance that returns fake images"""
     mock = Mock()
+    mock.is_running = True  # For health endpoint
+    mock.check_health = AsyncMock()  # For LLM mode health check
 
     async def generate(**kwargs):
         # Return n PIL images wrapped in result object
@@ -147,8 +149,9 @@ def test_client(mock_async_diffusion):
 
     # Set up app state with diffusion engine
     app.state.engine_client = mock_async_diffusion
+    app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
     app.state.stage_configs = [{"stage_type": "diffusion"}]
-    app.state.served_model_names = "Qwen/Qwen-Image"
+    app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
 
     return TestClient(app)
 
@@ -169,10 +172,58 @@ def async_omni_test_client():
     return TestClient(app)
 
 
-@pytest.mark.skip(reason="Async API server uses different health check mechanism")
 def test_health_endpoint(test_client):
-    """Test health check endpoint - skipped for async server"""
-    pass
+    """Test health check endpoint for diffusion mode"""
+    response = test_client.get("/health")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["status"] == "healthy"
+
+
+def test_health_endpoint_no_engine():
+    """Test health check endpoint when no engine is initialized"""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+    # Don't set any engine
+
+    client = TestClient(app)
+    response = client.get("/health")
+    assert response.status_code == 503
+    data = response.json()
+    assert data["status"] == "unhealthy"
+
+
+def test_models_endpoint(test_client):
+    """Test /v1/models endpoint for diffusion mode"""
+    response = test_client.get("/v1/models")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["object"] == "list"
+    assert len(data["data"]) == 1
+    assert data["data"][0]["id"] == "Qwen/Qwen-Image"
+    assert data["data"][0]["object"] == "model"
+
+
+def test_models_endpoint_no_engine():
+    """Test /v1/models endpoint when no engine is initialized"""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+    # Don't set any engine
+
+    client = TestClient(app)
+    response = client.get("/v1/models")
+    assert response.status_code == 200
+    data = response.json()
+    assert data["object"] == "list"
+    assert len(data["data"]) == 0
 
 
 def test_generate_single_image(test_client):
diff --git a/tests/test_outputs.py b/tests/test_outputs.py
new file mode 100644
index 000000000..c0c6201f6
--- /dev/null
+++ b/tests/test_outputs.py
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for OmniRequestOutput class."""
+
+from unittest.mock import Mock
+
+from PIL import Image
+
+from vllm_omni.outputs import OmniRequestOutput
+
+
+class TestOmniRequestOutput:
+    """Tests for OmniRequestOutput class."""
+
+    def test_from_diffusion(self):
+        """Test creating output from diffusion model."""
+        images = [Image.new("RGB", (64, 64), color="red")]
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=images,
+            prompt="a cat",
+            metrics={"steps": 50},
+        )
+        assert output.request_id == "test-123"
+        assert output.images == images
+        assert output.prompt == "a cat"
+        assert output.metrics == {"steps": 50}
+        assert output.is_diffusion_output
+        assert output.num_images == 1
+
+    def test_from_pipeline(self):
+        """Test creating output from pipeline stage."""
+        mock_request_output = Mock()
+        mock_request_output.request_id = "pipeline-123"
+        mock_request_output.prompt_token_ids = [1, 2, 3]
+        mock_request_output.outputs = [Mock()]
+        mock_request_output.encoder_prompt_token_ids = None
+        mock_request_output.prompt_logprobs = None
+        mock_request_output.num_cached_tokens = 10
+        mock_request_output.kv_transfer_params = None
+        mock_request_output.multimodal_output = {"image": Mock()}
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.request_id == "pipeline-123"
+        assert output.stage_id == 0
+        assert output.final_output_type == "text"
+        assert output.is_pipeline_output
+
+    def test_prompt_token_ids_property(self):
+        """Test prompt_token_ids property for streaming compatibility."""
+        mock_request_output = Mock()
+        mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.prompt_token_ids == [1, 2, 3, 4, 5]
+
+    def test_prompt_token_ids_none_when_no_request_output(self):
+        """Test prompt_token_ids returns None when no request_output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[],
+            prompt="a cat",
+        )
+        assert output.prompt_token_ids is None
+
+    def test_outputs_property(self):
+        """Test outputs property for chat completion compatibility."""
+        mock_output = Mock()
+        mock_request_output = Mock()
+        mock_request_output.outputs = [mock_output]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.outputs == [mock_output]
+
+    def test_outputs_empty_when_no_request_output(self):
+        """Test outputs returns empty list when no request_output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[],
+            prompt="a cat",
+        )
+        assert output.outputs == []
+
+    def test_encoder_prompt_token_ids_property(self):
+        """Test encoder_prompt_token_ids property."""
+        mock_request_output = Mock()
+        mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.encoder_prompt_token_ids == [10, 20, 30]
+
+    def test_num_cached_tokens_property(self):
+        """Test num_cached_tokens property."""
+        mock_request_output = Mock()
+        mock_request_output.num_cached_tokens = 42
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.num_cached_tokens == 42
+
+    def test_multimodal_output_property(self):
+        """Test multimodal_output property."""
+        mock_request_output = Mock()
+        mock_audio = Mock()
+        expected_output = {"audio": mock_audio}
+        mock_request_output.multimodal_output = expected_output
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="audio",
+            request_output=mock_request_output,
+        )
+
+        assert output.multimodal_output is expected_output
+
+    def test_to_dict_diffusion(self):
+        """Test to_dict for diffusion output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[Image.new("RGB", (64, 64), color="red")],
+            prompt="a cat",
+            metrics={"steps": 50},
+        )
+        result = output.to_dict()
+
+        assert result["request_id"] == "test-123"
+        assert result["finished"] is True
+        assert result["final_output_type"] == "image"
+        assert result["num_images"] == 1
+        assert result["prompt"] == "a cat"
+
+    def test_to_dict_pipeline(self):
+        """Test to_dict for pipeline output."""
+        mock_request_output = Mock()
+        mock_request_output.request_id = "pipeline-123"
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+        result = output.to_dict()
+
+        assert result["request_id"] == "pipeline-123"
+        assert result["finished"] is True
+        assert result["final_output_type"] == "text"
+        assert result["stage_id"] == 0
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 3c44da8b3..a47a864cd 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -666,14 +666,6 @@ def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
     return request.app.state.openai_serving_speech
 
 
-# Remove the original /v1/chat/completions route before registering our own
-# This prevents duplicate route registration warnings in FastAPI logs.
-for route in router.routes[:]:
-    if hasattr(route, "path") and route.path == "/v1/chat/completions":
-        router.routes.remove(route)
-        break
-
-
 @router.post(
     "/v1/chat/completions",
     dependencies=[Depends(validate_json_request)],
@@ -763,6 +755,86 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
 
 
+# Health and Model endpoints for diffusion mode
+
+
+# Remove existing health endpoint if present (from vllm imports)
+# to ensure our handler takes precedence
+_remove_route_from_router(router, "/health")
+
+
+@router.get("/health")
+async def health(raw_request: Request) -> JSONResponse:
+    """Health check endpoint that works for both LLM and diffusion modes.
+
+    Returns 200 OK if the server is healthy.
+    For LLM mode: delegates to engine_client health check
+    For diffusion mode: checks if diffusion_engine is running
+    """
+    # Check if we're in diffusion mode
+    diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None)
+    if diffusion_engine is not None:
+        # Diffusion mode health check
+        if hasattr(diffusion_engine, "is_running") and diffusion_engine.is_running:
+            return JSONResponse(content={"status": "healthy"})
+        return JSONResponse(
+            content={"status": "unhealthy", "reason": "Diffusion engine is not running"},
+            status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+        )
+
+    # LLM mode - delegate to engine_client
+    engine_client = getattr(raw_request.app.state, "engine_client", None)
+    if engine_client is not None:
+        await engine_client.check_health()
+        return JSONResponse(content={"status": "healthy"})
+
+    return JSONResponse(
+        content={"status": "unhealthy", "reason": "No engine initialized"},
+        status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+    )
+
+
+# Remove existing models endpoint if present (from vllm imports)
+# to ensure our handler takes precedence
+_remove_route_from_router(router, "/v1/models")
+
+
+@router.get("/v1/models")
+async def show_available_models(raw_request: Request) -> JSONResponse:
+    """Show available models endpoint that works for both LLM and diffusion modes.
+
+    Returns model information in OpenAI-compatible format.
+    """
+    # Check if we're in diffusion mode
+    diffusion_model_name = getattr(raw_request.app.state, "diffusion_model_name", None)
+    if diffusion_model_name is not None:
+        # Diffusion mode - return the loaded model
+        return JSONResponse(
+            content={
+                "object": "list",
+                "data": [
+                    {
+                        "id": diffusion_model_name,
+                        "object": "model",
+                        "created": 0,
+                        "owned_by": "vllm-omni",
+                        "permission": [],
+                    }
+                ],
+            }
+        )
+
+    # LLM mode - delegate to openai_serving_models
+    openai_serving_models = getattr(raw_request.app.state, "openai_serving_models", None)
+    if openai_serving_models is not None:
+        models = await openai_serving_models.show_available_models()
+        return JSONResponse(content=models.model_dump())
+
+    return JSONResponse(
+        content={"object": "list", "data": []},
+    )
+
+
 # Image generation API endpoints
 
 
diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py
index 24b168287..958cebd3e 100644
--- a/vllm_omni/outputs.py
+++ b/vllm_omni/outputs.py
@@ -123,6 +123,58 @@ class OmniRequestOutput:
         """Return the number of generated images."""
         return len(self.images)
 
+    # Pass-through properties keep vLLM serving codepaths compatible with
+    # OmniRequestOutput for pipeline outputs (Issue #345).
+    @property
+    def prompt_token_ids(self) -> list[int] | None:
+        """Return prompt token IDs from the underlying request output.
+
+        This property is required for compatibility with vLLM's streaming
+        chat completion generator which checks res.prompt_token_ids.
+        """
+        if self.request_output is not None:
+            return getattr(self.request_output, "prompt_token_ids", None)
+        return None
+
+    @property
+    def outputs(self) -> list[Any]:
+        """Return outputs from the underlying request output.
+
+        This property is required for compatibility with vLLM's streaming
+        and non-streaming chat completion generators.
+        """
+        if self.request_output is not None:
+            return getattr(self.request_output, "outputs", [])
+        return []
+
+    @property
+    def encoder_prompt_token_ids(self) -> list[int] | None:
+        """Return encoder prompt token IDs from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "encoder_prompt_token_ids", None)
+        return None
+
+    @property
+    def prompt_logprobs(self) -> Any:
+        """Return prompt logprobs from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "prompt_logprobs", None)
+        return None
+
+    @property
+    def num_cached_tokens(self) -> int | None:
+        """Return number of cached tokens from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "num_cached_tokens", None)
+        return None
+
+    @property
+    def kv_transfer_params(self) -> Any:
+        """Return KV transfer params from the underlying request output."""
+        if self.request_output is not None:
+            return getattr(self.request_output, "kv_transfer_params", None)
+        return None
+
     @property
     def is_diffusion_output(self) -> bool:
         """Check if this is a diffusion model output."""