[Bugfix][Realtime] Stop disconnected sessions from cycling through stages (#5388)

Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Sy03 <1370724210@qq.com>

tests/engine/test_orchestrator.py

@@ -1813,6 +1813,56 @@ async def test_stage_pool_submit_update_refreshes_output_processor_state() -> No
 
 
 @pytest.mark.asyncio
+async def test_handle_streaming_update_unknown_request_is_dropped() -> None:
+    """A streaming update for an untracked request must not resurrect it.
+
+    Updates always follow the first-chunk add_request through the same
+    ordered submission queue, so an unknown id can only mean the request
+    already finished or was aborted (e.g. the client disconnected).
+    Falling back to add_request created headless sessions that kept
+    cycling through the stages (issue #4271).
+    """
+
+    class RecordingPool:
+        def __init__(self) -> None:
+            self.calls: list[tuple[str, Any]] = []
+
+        async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int:
+            self.calls.append((request_id, prompt_text))
+            return 0
+
+    pool = RecordingPool()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.async_chunk = False
+    orchestrator.request_states = {}
+    orchestrator.stage_pools = [pool]
+    add_request_calls: list[str] = []
+
+    async def record_add_request(msg) -> None:
+        add_request_calls.append(msg.request_id)
+
+    orchestrator._handle_add_request = record_add_request
+
+    await orchestrator._handle_streaming_update(
+        StageSubmissionMessage(
+            type="streaming_update",
+            request_id="req-unknown",
+            prompt=SimpleNamespace(request_id="req-unknown", prompt_token_ids=[1], resumable=True),
+            original_prompt={"prompt": "segment-2"},
+            output_prompt_text="segment-2",
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+            preprocess_ms=0.0,
+            request_timestamp=time.time(),
+            enqueue_ts=time.perf_counter(),
+        )
+    )
+
+    assert pool.calls == []
+    assert add_request_calls == []
+    assert "req-unknown" not in orchestrator.request_states
+
+
 async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
     class RecordingPool:
         def __init__(self) -> None:

vllm_omni/engine/orchestrator.py

@@ -681,24 +681,16 @@ class Orchestrator:
         final_stage_id = msg.final_stage_id
         req_state = self.request_states.get(request_id)
         if req_state is None:
+            # Streaming updates always follow the first-chunk add_request
+            # through the same ordered submission queue, so an unknown id
+            # here can only mean the request already finished or was
+            # aborted (e.g. the client disconnected). Re-adding it would
+            # resurrect a headless session that keeps cycling through the
+            # stages with nobody consuming its outputs (issue #4271).
             logger.warning(
-                "[Orchestrator] streaming_update for unknown req=%s, falling back to add_request",
+                "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)",
                 request_id,
             )
-            fallback_msg = StageSubmissionMessage(
-                type="add_request",
-                request_id=msg.request_id,
-                prompt=msg.prompt,
-                original_prompt=msg.original_prompt,
-                output_prompt_text=msg.output_prompt_text,
-                sampling_params_list=msg.sampling_params_list,
-                final_stage_id=msg.final_stage_id,
-                final_output_stage_ids=msg.final_output_stage_ids,
-                preprocess_ms=msg.preprocess_ms,
-                request_timestamp=msg.request_timestamp,
-                enqueue_ts=msg.enqueue_ts,
-            )
-            await self._handle_add_request(fallback_msg)
             return
 
         if msg.sampling_params_list:

vllm_omni/entrypoints/async_omni.py

@@ -628,6 +628,13 @@ class AsyncOmni(EngineClient, OmniBase):
 
             logger.debug(f"[AsyncOmni] Request {request_id} completed")
 
+            # The input pump can outlive a normally-completed generation
+            # (e.g. a realtime client that keeps streaming after the final
+            # output); a live pump keeps issuing streaming updates for a
+            # request the orchestrator no longer tracks.
+            if input_stream_task is not None and not input_stream_task.done():
+                input_stream_task.cancel()
+
             self._log_summary_and_cleanup(request_id)
 
         except (asyncio.CancelledError, GeneratorExit):
@@ -1059,7 +1066,10 @@ class AsyncOmni(EngineClient, OmniBase):
         """Submit request IDs to be aborted to the engine."""
         await self.engine.abort_async(request_ids)
         for rid in request_ids:
-            self.request_states.pop(rid, None)
+            state = self.request_states.pop(rid, None)
+            input_stream_task = getattr(state, "input_stream_task", None)
+            if input_stream_task is not None and not input_stream_task.done():
+                input_stream_task.cancel()
         if self.log_stats:
             logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
 

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -138,6 +138,7 @@ class RealtimeConnection(VllmRealtimeConnection):
             is_streaming=True,
         )
 
+        result_gen = None
         try:
             result_gen = self.engine.generate(
                 prompt=streaming_input_gen,
@@ -196,6 +197,16 @@ class RealtimeConnection(VllmRealtimeConnection):
             logger.exception("Error in generation: %s", e)
             await self.send_error(str(e), "processing_error")
         finally:
+            # Close the generator explicitly so AsyncOmni.generate's cleanup
+            # (input-pump cancellation and engine-side abort) runs now rather
+            # than whenever the event loop garbage-collects the async
+            # generator; the delay window is where a disconnected session
+            # keeps cycling through the stages (issue #4271).
+            if result_gen is not None:
+                try:
+                    await result_gen.aclose()
+                except Exception:
+                    logger.exception("Failed to close realtime result generator")
             # Always send terminal event so clients don't hang forever.
             if self._is_connected and not audio_done_sent:
                 try:
@@ -206,4 +217,10 @@ class RealtimeConnection(VllmRealtimeConnection):
                 self.audio_queue.get_nowait()
 
     async def send_json(self, payload: dict):
-        await self.websocket.send_text(json.dumps(payload))
+        try:
+            await self.websocket.send_text(json.dumps(payload))
+        except Exception:
+            # A failed send means the client is gone; flag it so the
+            # generation loop stops instead of retrying into a dead socket.
+            self._is_connected = False
+            raise