[Bugfix][Frontend] Route synthetic aborts through final stage (#7006)

Signed-off-by: mershi <mershi@tencent.com>
Co-authored-by: mershi <mershi@tencent.com>

tests/entrypoints/test_omni_entrypoints.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import gc
 import queue
 from collections.abc import Callable
@@ -634,6 +635,43 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc
     assert "req-1" not in app.request_states
 
 
+@pytest.mark.asyncio
+async def test_async_omni_abort_yields_terminal_from_requested_final_stage(monkeypatch: pytest.MonkeyPatch):
+    stage_metadata = [
+        _stage_meta(stage_type="llm", final_output=False, final_output_type=None),
+        _stage_meta(stage_type="llm", final_output=True, final_output_type="audio"),
+    ]
+    engine = FakeAsyncOmniEngine(stage_metadata=stage_metadata)
+    _patch_engine(monkeypatch, engine)
+    app = AsyncOmni("dummy-model")
+
+    async def collect_outputs() -> list[OmniRequestOutput]:
+        return [
+            output
+            async for output in app.generate(
+                prompt="hello",
+                request_id="req-1",
+                output_modalities=["audio"],
+            )
+        ]
+
+    try:
+        generate_task = asyncio.create_task(collect_outputs())
+        while not engine.submitted:
+            await asyncio.sleep(0)
+        await app.abort("req-1")
+        outputs = await asyncio.wait_for(generate_task, timeout=1)
+    finally:
+        app.shutdown()
+
+    assert len(outputs) == 1
+    assert outputs[0].stage_id == 1
+    assert outputs[0].final_output_type == "audio"
+    assert outputs[0].finished is True
+    assert outputs[0].outputs[0].finish_reason == "abort"
+    assert "req-1" not in app.request_states
+
+
 @pytest.mark.asyncio
 async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
     engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)

vllm_omni/entrypoints/async_omni.py

@@ -236,6 +236,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient):
             req_state = ClientRequestState(
                 request_id=request_id,
                 external_request_id=external_request_id,
+                final_stage_id=final_stage_id_for_e2e,
             )
             req_state.metrics = metrics
             req_state.request_arrival_ts = wall_start_ts

vllm_omni/entrypoints/async_omni_base.py

@@ -479,7 +479,22 @@ class AsyncOmniBase(OmniBase):
             if state is not None and rid not in delivered:
                 queue = getattr(state, "queue", None)
                 if queue is not None:
-                    await state.queue.put(self._synthetic_abort_output_message(rid))
+                    final_stage_id = getattr(state, "final_stage_id", None)
+                    if final_stage_id is None:
+                        final_stage_id = 0
+                    stage_metadata = getattr(self, "_stage_meta_list", ())
+                    final_output_type = (
+                        getattr(stage_metadata[final_stage_id], "final_output_type", None)
+                        if final_stage_id < len(stage_metadata)
+                        else None
+                    ) or "text"
+                    await state.queue.put(
+                        self._synthetic_abort_output_message(
+                            rid,
+                            stage_id=final_stage_id,
+                            final_output_type=final_output_type,
+                        )
+                    )
                     delivered.add(rid)
         for rid in request_ids:
             self._record_request_failure_once(rid, reason="client_abort")
@@ -491,13 +506,18 @@ class AsyncOmniBase(OmniBase):
             logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
 
     @staticmethod
-    def _synthetic_abort_output_message(request_id: str) -> OutputMessage:
+    def _synthetic_abort_output_message(
+        request_id: str,
+        *,
+        stage_id: int,
+        final_output_type: str,
+    ) -> OutputMessage:
         """Terminal abort OutputMessage used when the engine returned none."""
         engine_output = OmniRequestOutput(
             request_id=request_id,
             finished=True,
-            stage_id=0,
-            final_output_type="text",
+            stage_id=stage_id,
+            final_output_type=final_output_type,
             outputs=[
                 CompletionOutput(
                     index=0,
@@ -512,7 +532,7 @@ class AsyncOmniBase(OmniBase):
         )
         return OutputMessage(
             request_id=request_id,
-            stage_id=0,
+            stage_id=stage_id,
             replica_id=None,
             engine_outputs=engine_output,
             metrics=None,

vllm_omni/entrypoints/client_request_state.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import asyncio
 
 from vllm_omni.metrics import OrchestratorAggregator
@@ -11,10 +14,12 @@ class ClientRequestState:
         request_id: str,
         external_request_id: str | None = None,
         queue: asyncio.Queue | None = None,
+        final_stage_id: int | None = None,
     ):
         self.request_id = request_id
         self.external_request_id = external_request_id
         self.stage_id: int | None = None
+        self.final_stage_id: int | None = final_stage_id
         self.queue = queue if queue is not None else asyncio.Queue()
         self.metrics: OrchestratorAggregator | None = None
         # Request-scoped idempotency guard for Prometheus failure counters.