[Cleanup][Full Duplex] Defer handoff recovery to PR 6170

Signed-off-by: Sy03 <1370724210@qq.com>

tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py

@@ -0,0 +1,89 @@
+"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
+
+``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
+handing control to upstream, because upstream ``Scheduler.schedule()`` raises
+``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
+admits in a finished state -- which kills the stage's engine core, not just
+the request.
+
+Regression for the sweep missing ``skipped_waiting``, the third queue
+upstream admits from. An aborted duplex session parked there was re-selected
+by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
+stage.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_POLICIES = [
+    pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
+    pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
+]
+
+
+class _StubRequest:
+    """Minimal ``Request`` stub with the surface the sweep exercises.
+
+    Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
+    relies on so the same stub works under both scheduling policies.
+    """
+
+    def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
+        self.request_id = request_id
+        self.status = status
+        self.priority = 0
+        self.arrival_time = arrival_time
+
+    def __lt__(self, other: _StubRequest) -> bool:
+        return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
+
+
+def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
+    """Bare scheduler with just the surface the sweep reads."""
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(policy)
+    for req in waiting:
+        scheduler.waiting.add_request(req)
+    scheduler.skipped_waiting = create_request_queue(policy)
+    for req in skipped_waiting:
+        scheduler.skipped_waiting.add_request(req)
+    scheduler.running = list(running)
+    return scheduler
+
+
+@pytest.mark.parametrize("policy", _POLICIES)
+def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
+    monkeypatch: pytest.MonkeyPatch,
+    policy: SchedulingPolicy,
+) -> None:
+    """The public schedule path must remove aborted skipped requests first."""
+    aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
+    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.num_waiting_for_streaming_input = 0
+    scheduler._process_pending_omni_inputs = lambda model_mode: None
+    scheduler._resync_streaming_input_counter = lambda: None
+    scheduler._should_defer_waiting_admission = lambda: False
+    scheduler._restore_omni_wait_queues = lambda: None
+    scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
+    scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
+    scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
+
+    sentinel = object()
+
+    def upstream_schedule(_self, _throttle_prefills: bool = False):
+        assert list(scheduler.skipped_waiting) == []
+        return sentinel
+
+    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
+
+    assert OmniARScheduler.schedule(scheduler) is sentinel

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -28,7 +28,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
     sched = OmniARScheduler.__new__(OmniARScheduler)
     sched._new_prompt_len_snapshot = {}
-    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
+    sched.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_id=stage_id, max_model_len=None),
+    )
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
@@ -210,6 +212,81 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.chunk_transfer_adapter.save_async.assert_not_called()
 
 
+def test_stale_async_frame_is_dropped_before_output_processing() -> None:
+    session = _make_request()
+    session.status = RequestStatus.RUNNING
+    session.num_in_flight_tokens = 2
+    session.num_computed_tokens = session.num_tokens + 1
+    session.num_output_placeholders = 1
+    session.async_tokens_to_discard = 1
+    session.sampling_params = SimpleNamespace(num_logprobs=1)
+    num_computed_tokens = session.num_computed_tokens
+    num_output_placeholders = session.num_output_placeholders
+
+    sched = MagicMock()
+    sched.requests = {session.request_id: session}
+    sched.perf_metrics = None
+    sched.structured_output_manager.should_advance.return_value = False
+
+    def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
+        request.async_tokens_to_discard = 0
+        return token_ids, False
+
+    sched._update_request_with_output.side_effect = discard_stale_output
+    sched._process_kv_transfer_trigger.return_value = False
+    sched.chunk_transfer_adapter = MagicMock()
+    sched.running = [session]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+
+    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.multimodal_outputs = None
+    model_runner_output.inter_stage_outputs = [{"hidden": object()}]
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {session.request_id: 0}
+    model_runner_output.routed_experts = None
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+    assert session.async_tokens_to_discard == 0
+    assert session.status == RequestStatus.RUNNING
+    assert session.num_computed_tokens == num_computed_tokens
+    assert session.num_output_placeholders == num_output_placeholders
+    sched.chunk_transfer_adapter.save_async.assert_not_called()
+
+    session.sampling_params = SamplingParams(max_tokens=8)
+    next_payload = {"hidden": object()}
+    sched._update_request_with_output.side_effect = None
+    sched._update_request_with_output.return_value = ([43], False)
+    model_runner_output.sampled_token_ids = [[43]]
+    model_runner_output.inter_stage_outputs = [next_payload]
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+    sched.chunk_transfer_adapter.save_async.assert_called_once_with(
+        next_payload,
+        session,
+        False,
+    )
+
+
 def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
     sched = _make_scheduler(stage_id=0)
     session = _make_request()
@@ -354,3 +431,34 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
     assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
     assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
     assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+
+
+def test_native_duplex_prompt_reset_prevents_stage1_context_overflow() -> None:
+    sched = _make_scheduler(stage_id=1)
+    sched.vllm_config.model_config.max_model_len = 4096
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        receives_chunks=False,
+        segment_finished_requests=set(),
+    )
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.resumable = True
+    session.prompt_token_ids = [0] * 3
+    session._all_token_ids.clear()
+    session._all_token_ids.extend([0] * 4095)
+    session._output_token_ids.extend([1] * 4080)
+    session.num_prompt_tokens = 3
+    session.num_computed_tokens = 4095
+
+    update = _make_update([0] * 16)
+    update.model_intermediate_buffer = {"native_duplex": True}
+
+    sched._update_request_as_session(session, update)
+
+    assert session.prompt_token_ids == [0] * 16
+    assert list(session._all_token_ids) == [0] * 16
+    assert session._output_token_ids == []
+    assert session.num_prompt_tokens == 16
+    assert session.num_computed_tokens == 0
+    assert session.model_intermediate_buffer == update.model_intermediate_buffer
+    assert session.status == RequestStatus.WAITING

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -76,6 +76,9 @@ def _make_scheduler(
         else:
             hf_config.voxcpm2_runtime_config = runtime_config
     sched.vllm_config = SimpleNamespace(model_config=model_config)
+    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
+    # abort sweep at the top of ``schedule()`` reads it.
+    sched.skipped_waiting = _MockQueue()
     return sched
 
 

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -0,0 +1,151 @@
+"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
+``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
+and in ``OmniARScheduler.schedule``.
+
+Upstream keeps that counter incrementally (``+1`` when a resumable request
+parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
+next update arrives or the request is finished) and
+``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
+queues. Omni rewrites ``request.status`` outside both hooks -- the
+chunk-transfer adapter's park/restore and
+``_realign_request_status_to_queues`` -- so a counted request can leave
+``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
+counter for the lifetime of the process.
+
+An inflated counter does not degrade throughput, it hangs the stage:
+``EngineCore.has_work()`` reads false while a live request sits in
+``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
+``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
+server answered the first conversation after a boot and was silent from the
+second on while still reporting healthy.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _StubRequest:
+    def __init__(self, request_id: str, status: RequestStatus) -> None:
+        self.request_id = request_id
+        self.status = status
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+
+_SCHEDULER_PARAMS = [
+    pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
+    pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
+]
+
+
+def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.requests = requests
+    scheduler.running = running
+    scheduler.waiting = waiting
+    scheduler.skipped_waiting = skipped if skipped is not None else []
+    scheduler.num_waiting_for_streaming_input = counter
+    return scheduler
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_clears_counter_left_by_a_stomped_status(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    """The abort case that wedged the talker.
+
+    The request was counted while parked for streaming input, then had its
+    status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
+    ``finish_requests`` took the running branch and never decremented. With
+    no parked requests left, the counter must land at 0.
+    """
+    aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={},  # upstream already dropped it
+        running=[],
+        waiting=[],
+        counter=1,
+    )
+
+    monkeypatch.setattr(
+        scheduler_mod.VLLMScheduler,
+        "finish_requests",
+        lambda self, request_ids, finished_status: [aborted],
+    )
+
+    scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
+
+    assert scheduler.num_waiting_for_streaming_input == 0
+
+
+def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
+    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
+    so the resync has to look at both or it under-counts and reintroduces the
+    spin the counter exists to prevent."""
+    parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
+    scheduler = _make_scheduler(
+        OmniARScheduler,
+        requests={parked.request_id: parked},
+        running=[],
+        waiting=[],
+        counter=0,
+        skipped=[parked],
+    )
+
+    scheduler._resync_streaming_input_counter()
+
+    assert scheduler.num_waiting_for_streaming_input == 1
+
+
+def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Pins placement in ``schedule``: the chunk-transfer adapter rewrites
+    status during ``process_pending_chunks``, so the resync has to run after
+    that and before upstream schedules, while ``self.waiting`` is still the
+    real queue."""
+    parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.requests = {parked.request_id: parked}
+    scheduler.running = [parked]
+    scheduler.waiting = []
+    scheduler.skipped_waiting = []
+    scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
+    scheduler.policy = "fcfs"
+
+    seen: dict[str, int] = {}
+
+    monkeypatch.setattr(
+        ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
+    )
+    monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
+    monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
+
+    def fake_schedule(self, throttle_prefills=False):
+        seen["counter"] = self.num_waiting_for_streaming_input
+        raise _StopScheduleError
+
+    class _StopScheduleError(Exception):
+        pass
+
+    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
+
+    with pytest.raises(_StopScheduleError):
+        OmniARScheduler.schedule(scheduler)
+
+    assert seen["counter"] == 0

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -546,7 +546,7 @@ def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
     assert second.codes is None
 
 
-def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
+def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
     request.resumable = True
@@ -565,6 +565,51 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt
     assert adapter.requests_num_chunks_sent["ext-stream"] == 0
 
 
+def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
+    request.resumable = True
+    request._omni_segment_generation = 0
+    adapter.requests_num_chunks_sent["ext-stream"] = 111
+
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+
+    late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream")
+    late.resumable = True
+    late.num_computed_tokens = 111
+    late._omni_segment_generation = 0
+    adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False)
+
+    next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream")
+    next_segment.resumable = True
+    next_segment.num_computed_tokens = 3
+    next_segment._omni_segment_generation = 1
+    adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False)
+
+    assert len(adapter._pending_save_reqs) == 2
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 3
+    assert adapter._pending_save_reqs[-1]["request"] is next_segment
+
+
+def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter):
+    adapter, connector = build_adapter(stage_id=1)
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct())
+    request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old")
+    request.resumable = True
+    request._omni_segment_generation = 0
+
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+
+    old_task = adapter._pending_save_reqs.popleft()
+    adapter._send_single_request(old_task)
+    connector.put.assert_called_once()
+
+    boundary_task = adapter._pending_save_reqs.popleft()
+    adapter._send_single_request(boundary_task)
+    assert connector.put.call_count == 2
+
+
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")

tests/engine/test_orchestrator.py

@@ -30,6 +30,7 @@ from vllm_omni.engine.messages import (
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
+    StreamingSegmentState,
     _build_terminal_empty_output,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -1951,7 +1952,7 @@ async def test_resumable_segment_boundary_builds_stage_metrics() -> None:
         final_stage_id=0,
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segment_finished = True
+    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
     req_state.stage_submit_ts[0] = time.time()
     orchestrator.request_states = {"req-stream": req_state}
     orchestrator.stage_pools = [pool]

tests/engine/test_orchestrator_segment_stage_keying.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Streaming segment-boundary state must be keyed by stage.
+
+``Orchestrator._orchestration_loop`` polls every stage into the same
+``OrchestratorRequestState``. While the segment-boundary fields lived in a single
+flat slot per request, whichever stage polled most recently overwrote the others,
+and the per-stage consumers then read back a different stage's boundary.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import (
+    Orchestrator,
+    OrchestratorRequestState,
+    StreamingSegmentState,
+)
+from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+THINKER_STAGE = 0
+TALKER_STAGE = 1
+
+
+@dataclass
+class _StageMetrics:
+    pipeline_timings: dict[str, float] = field(default_factory=dict)
+
+
+@dataclass
+class _StageOutput:
+    """Stand-in for a processed stage output, with the fields the router reads."""
+
+    request_id: str
+    finished: bool = False
+    error: object | None = None
+
+
+class _RecordingPool:
+    """Stage pool that records the requests it was asked to build metrics for."""
+
+    def __init__(self) -> None:
+        self.built_for: list[str] = []
+
+    def build_stage_metrics(self, outputs: list[_StageOutput], **_kwargs: object) -> _StageMetrics:
+        self.built_for.extend(output.request_id for output in outputs)
+        return _StageMetrics()
+
+
+def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState:
+    req_state = OrchestratorRequestState(
+        request_id=request_id,
+        sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)],
+        final_stage_id=TALKER_STAGE,
+    )
+    req_state.streaming.enabled = True
+    req_state.duplex_identity = DuplexRequestIdentity(
+        session_id="session-1",
+        fence=DuplexFence(session_id="session-1"),
+    )
+    return req_state
+
+
+def test_duplex_output_context_reads_only_its_own_stage_segment() -> None:
+    req_state = _duplex_req_state()
+    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(
+        finished=True,
+        token_ids=[11, 22],
+        output_metadata={"stage": "thinker"},
+    )
+    # The talker polls after the thinker and is still mid-segment.
+    req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False)
+
+    thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE)
+    talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE)
+
+    assert thinker is not None
+    assert talker is not None
+
+    # These are the assertions that fail on the flat field: the talker is written
+    # second, so the thinker loses its boundary to whichever stage polled last.
+    assert thinker.segment_finished is True
+    assert thinker.segment_token_ids == (11, 22)
+    assert thinker.segment_output_metadata == {"stage": "thinker"}
+
+    assert talker.segment_finished is False
+    assert talker.segment_token_ids == ()
+    assert talker.segment_output_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None:
+    req_state = _duplex_req_state()
+    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True)
+
+    talker_pool = _RecordingPool()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.request_states = {req_state.request_id: req_state}
+    orchestrator.stage_pools = [_RecordingPool(), talker_pool]
+
+    routed: list[_StageMetrics | None] = []
+
+    async def record_route(
+        _stage: int,
+        _replica: int,
+        _output: _StageOutput,
+        _state: OrchestratorRequestState,
+        metrics: _StageMetrics | None,
+    ) -> None:
+        routed.append(metrics)
+
+    orchestrator._route_output = record_route
+
+    # The talker's own output is neither finished nor at a segment boundary, so
+    # it must not be treated as one just because the thinker reached one.
+    await orchestrator._handle_processed_outputs(
+        TALKER_STAGE,
+        0,
+        [_StageOutput(request_id=req_state.request_id)],
+    )
+
+    assert talker_pool.built_for == []
+    assert routed == [None]

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
+    StreamingSegmentState,
     _OrchestratorDuplexStagePort,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -358,7 +359,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None:
         final_output_stage_ids={0},
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segment_finished = True
+    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
     output = SimpleNamespace(
         request_id=req_state.request_id,
         finished=True,

tests/entrypoints/openai/test_duplex_protocol.py

@@ -462,6 +462,16 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
     assert caps["stage_handoff_transport"] == "scheduler_data_plane"
 
 
+def test_duplex_overlap_policy_auto_remains_a_compatibility_value():
+    from vllm_omni.experimental.fullduplex.openai.protocol import (
+        DuplexOverlapPolicy,
+        DuplexSessionConfig,
+    )
+
+    assert DuplexOverlapPolicy.AUTO.value == "auto"
+    assert DuplexSessionConfig._normalize_overlap_policy("auto") == "auto"
+
+
 def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
     caps = DuplexCapabilities.minicpmo45_native(max_sessions=1).as_dict()
 

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -218,16 +218,21 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if stop_after_transfer and req_id in self.requests_needing_kv_transfer:
             self.pending_stop_after_extraction.add(req_id)
 
+    def _drop_aborted_queued_requests(self) -> None:
+        for queue in (self.waiting, self.skipped_waiting):
+            aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED]
+            if aborted:
+                queue.remove_requests(aborted)
+        self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED]
+
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         # Remove FINISHED_ABORTED requests before the upstream scheduler sees
         # them. Upstream vllm raises RuntimeError on this status; omni allows
         # async abort (e.g. client disconnect during TTS streaming) to leave
         # requests in the waiting/running queues temporarily.
-        for queue in (self.waiting, self.running):
+        self._drop_aborted_queued_requests()
-            for req in list(queue):
-                if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
-                    queue.remove(req)
         self._process_pending_omni_inputs(model_mode="ar")
+        self._resync_streaming_input_counter()
 
         original_waiting = None
         if self._should_defer_waiting_admission():
@@ -354,6 +359,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+
+            if generated_token_ids and getattr(request, "async_tokens_to_discard", 0) > 0:
+                self._update_request_with_output(request, generated_token_ids)
+                continue
+
             status_before_stop = request.status
             new_logprobs = None
             logprob_validation_failed = False
@@ -651,6 +661,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter and self.chunk_transfer_adapter.receives_chunks:
             self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None)
             if stage_id != 0:
+                session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
                 # Downstream async-chunk stages receive real payloads from the
                 # connector. This update only resumes polling for the next segment.
                 self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
@@ -682,6 +693,30 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             self._release_replaced_streaming_prompt_cache(session)
             self._replace_streaming_session(session, update)
             return
+        native_duplex_update = any(
+            isinstance(info, dict) and info.get("native_duplex") is True for info in update_infos
+        )
+        if native_duplex_update:
+            max_model_len = int(getattr(self.vllm_config.model_config, "max_model_len", 0) or 0)
+            current_prompt_len = len(getattr(session, "prompt_token_ids", None) or ())
+            computed_tokens = int(getattr(session, "num_computed_tokens", 0) or 0)
+            kept_output_len = max(0, computed_tokens - current_prompt_len)
+            next_prompt_len = len(update.prompt_token_ids or ())
+            candidate_len = current_prompt_len + kept_output_len + next_prompt_len
+            # Keep one context position available for the next sampled codec token.
+            if max_model_len > 0 and candidate_len >= max_model_len:
+                logger.info(
+                    "Resetting native duplex Stage-1 prompt before context overflow: "
+                    "req=%s current=%s kept_output=%s next=%s max_model_len=%s",
+                    session.request_id,
+                    current_prompt_len,
+                    kept_output_len,
+                    next_prompt_len,
+                    max_model_len,
+                )
+                self._replace_streaming_session(session, update)
+                return
+        session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
         super()._update_request_as_session(session, update)
         if hasattr(update, "model_intermediate_buffer"):
             session.model_intermediate_buffer = update.model_intermediate_buffer

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -146,6 +146,7 @@ class OmniSchedulerMixin:
 
     def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None:
         """Replace a downstream stage's placeholder with its next payload."""
+        session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
         adapter = getattr(self, "chunk_transfer_adapter", None)
         if adapter is not None:
             adapter.segment_finished_requests.discard(session.request_id)
@@ -645,6 +646,7 @@ class OmniSchedulerMixin:
         )
         finished = super().finish_requests(request_ids, finished_status)
         self._purge_finished_from_running(target_request_ids)
+        self._resync_streaming_input_counter()
 
         for request in finished:
             self._free_input_coordinator_request(request.request_id)
@@ -810,3 +812,18 @@ class OmniSchedulerMixin:
             return resumable_segment_stop and req.request_id not in target_request_ids
 
         self.running[:] = [req for req in self.running if keep_running(req)]
+
+    def _resync_streaming_input_counter(self) -> None:
+        counter = getattr(self, "num_waiting_for_streaming_input", None)
+        if counter is None:
+            return
+        parked = sum(
+            1
+            for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None))
+            if queue
+            for request in queue
+            if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+        )
+        if parked != counter:
+            logger.debug("[Omni] resynced streaming-input counter %s -> %s", counter, parked)
+            self.num_waiting_for_streaming_input = parked

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -95,6 +95,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         #   been read").
         self._held_non_active: deque[Any] = deque()
         self.requests_num_chunks_sent: dict[str, int] = defaultdict(int)
+        # Boundary tasks advance this before the background sender runs, so a
+        # late old-segment frame cannot restore the previous watermark.
+        self._segment_generation: dict[str, int] = defaultdict(int)
         self._pending_streaming_prefills: dict[str, dict] = {}
         # Monotonic timestamp of when each request last began waiting for a
         # chunk, refreshed every time one arrives.  Read by
@@ -188,11 +191,31 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             is_segment_finished: whether the segment of request is finished
         """
         is_finished = request.is_finished() and not request.resumable
+        if not hasattr(self, "_segment_generation"):
+            self._segment_generation = defaultdict(int)
+        external_req_id = request.external_req_id
+        raw_generation = getattr(request, "_omni_segment_generation", 0)
+        try:
+            generation = int(raw_generation)
+        except (TypeError, ValueError):
+            generation = 0
+        expected_generation = self._segment_generation.get(external_req_id, generation)
+        if generation < expected_generation:
+            logger.warning(
+                "Skip late save_async for request %s, segment_generation=%s, expected=%s",
+                external_req_id,
+                generation,
+                expected_generation,
+            )
+            return
+        self._segment_generation[external_req_id] = generation
 
         confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
 
         # If the request is preempted, skip the already saved chunks.
-        if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
+        if not is_segment_finished and confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(
+            request.external_req_id, 0
+        ):
             logger.warning(
                 f"Enqueue save_async for request {request.external_req_id}, "
                 f"request.num_computed_tokens={request.num_computed_tokens}, "
@@ -201,13 +224,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             )
             return
 
-        self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
         task = {
             "multimodal_output": multimodal_output,
             "request": request,
             "is_finished": is_finished,
             "is_segment_finished": is_segment_finished,
+            "segment_generation": generation,
         }
+        if is_segment_finished:
+            self._segment_generation[external_req_id] = generation + 1
+            self.requests_num_chunks_sent.pop(external_req_id, None)
+        else:
+            self.requests_num_chunks_sent[external_req_id] = confirmed_num_computed_tokens
         self._pending_save_reqs.append(task)
         with self._save_cond:
             self._save_cond.notify()
@@ -425,7 +453,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
-            self.requests_num_chunks_sent.pop(external_req_id, None)
             self.ramp_chunk_count.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
             if cached_ic is not None:
@@ -501,6 +528,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.request_payload.pop(external_req_id, None)
         self.code_prompt_token_ids.pop(external_req_id, None)
         self.requests_num_chunks_sent.pop(external_req_id, None)
+        self._segment_generation.pop(external_req_id, None)
         self.ramp_chunk_count.pop(external_req_id, None)
         self._pending_streaming_prefills.pop(external_req_id, None)
 

vllm_omni/engine/orchestrator.py

@@ -199,17 +199,37 @@ class OrchestratorRequestState:
 
 
 @dataclass
-class StreamingInputState:
+class StreamingSegmentState:
-    # Flag of streaming input request
+    """Streaming segment-boundary state for one stage of a request.
-    enabled: bool = False
+
+    Every stage of a multi-stage pipeline reaches its own segment boundaries
+    independently, so this is tracked per stage rather than per request.
+    """
+
     # Flag of segment of streaming input finished
-    segment_finished: bool = False
+    finished: bool = False
     # Tokens from the current raw segment boundary. The vLLM output processor
     # does not guarantee that EngineCoreOutput.new_token_ids survives on the
     # processed RequestOutput used by the routing layer.
-    segment_token_ids: list[int] = field(default_factory=list)
+    token_ids: list[int] = field(default_factory=list)
-    segment_output_metadata: dict[str, Any] = field(default_factory=dict)
+    output_metadata: dict[str, Any] = field(default_factory=dict)
-    # Streaming update prompt length
+
+
+@dataclass
+class StreamingInputState:
+    # Flag of streaming input request
+    enabled: bool = False
+    # Keyed by stage. ``_orchestration_loop`` records a boundary per stage poll
+    # into this one ``req_state``, while every consumer reads it for a specific
+    # ``stage_id``, so a flat slot is only correct as long as each read happens in
+    # the same iteration as its own stage's write. It is not for a stage that
+    # never records one (the ``poll_diffusion_output`` branch reaches
+    # ``_handle_processed_outputs`` without the raw-output loop) or for any
+    # consumer reading outside the writing iteration.
+    segments: dict[int, StreamingSegmentState] = field(default_factory=dict)
+    # Streaming update prompt length. NOT stage-keyed: its only consumer reads it
+    # through the streaming context on the stage-input-bridge path, not by stage
+    # id. Do not assume the rest of this struct is stage-safe.
     new_prompt_len_snapshot: int | None = None
     # Model/bridge-specific runtime states (e.g., thinker->talker)
     bridge_states: dict[str, Any] = field(default_factory=dict)
@@ -217,6 +237,14 @@ class StreamingInputState:
     # while the downstream input processor consumes upstream token output.
     source_token_decoder: Callable[..., str] | None = None
 
+    def segment(self, stage_id: int | None) -> StreamingSegmentState:
+        """Return ``stage_id``'s segment state, or a fresh empty one if unreported.
+
+        Fresh rather than shared: ``output_metadata`` is handed out by reference.
+        """
+        segment = self.segments.get(stage_id)
+        return segment if segment is not None else StreamingSegmentState()
+
 
 class _OrchestratorDuplexStagePort:
     """Adapts generic stage pools to the model-neutral duplex control plane."""
@@ -966,17 +994,18 @@ class Orchestrator:
                                 req_state = self.request_states.get(getattr(eco, "request_id", None))
                                 if req_state is None or not req_state.streaming.enabled:
                                     continue
-                                req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False))
+                                segment_finished = bool(getattr(eco, "is_segment_finished", False))
-                                req_state.streaming.segment_token_ids = (
-                                    self._coerce_int_list(getattr(eco, "new_token_ids", None))
-                                    if req_state.streaming.segment_finished
-                                    else []
-                                )
                                 raw_mm = self._completion_multimodal_output(eco, None)
-                                req_state.streaming.segment_output_metadata = (
+                                req_state.streaming.segments[stage_id] = StreamingSegmentState(
-                                    dict(raw_mm)
+                                    finished=segment_finished,
-                                    if req_state.streaming.segment_finished and isinstance(raw_mm, dict)
+                                    token_ids=(
-                                    else {}
+                                        self._coerce_int_list(getattr(eco, "new_token_ids", None))
+                                        if segment_finished
+                                        else []
+                                    ),
+                                    output_metadata=(
+                                        dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}
+                                    ),
                                 )
                                 req_state.streaming.new_prompt_len_snapshot = getattr(
                                     eco,
@@ -1044,7 +1073,7 @@ class Orchestrator:
                 continue
 
             stage_metrics = None
-            segment_finished = req_state.streaming.enabled and req_state.streaming.segment_finished
+            segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished
             if output.finished or segment_finished:
                 stage_metrics = pool.build_stage_metrics(
                     [output],
@@ -1378,6 +1407,7 @@ class Orchestrator:
         req_id = output.request_id
         finished = output.finished
         submit_ts = req_state.stage_submit_ts.get(stage_id)
+        segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished
         # CFG companion: stash output so parent can bundle [parent, *companions]
         # into source_outputs for the bridge (e.g. thinker2imagegen).
         if finished and self._cfg_tracker.is_companion(req_id):
@@ -1387,11 +1417,7 @@ class Orchestrator:
             return
 
         request_finished = False
-        if (
+        if finished and self.stage_pools[stage_id].final_output and not segment_finished:
-            finished
-            and self.stage_pools[stage_id].final_output
-            and not (req_state.streaming.enabled and req_state.streaming.segment_finished)
-        ):
             req_state.finished_final_output_stage_ids.add(stage_id)
             final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
             request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids)
@@ -1403,7 +1429,7 @@ class Orchestrator:
         # filter out again (the official implementation returns exactly one
         # result per audio chunk).
         is_duplex_stage0_segment = (
-            stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished
+            stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment(0).finished
         )
         if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment:
             await self.output_async_queue.put(
@@ -1415,7 +1441,10 @@ class Orchestrator:
                     metrics=stage_metrics,
                     finished=(
                         request_finished
-                        or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished)
+                        or (
+                            self._is_duplex_session_request(req_state)
+                            and req_state.streaming.segment(stage_id).finished
+                        )
                     ),
                     stage_submit_ts=submit_ts,
                 )
@@ -1450,7 +1479,7 @@ class Orchestrator:
             return
 
         if (
-            (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished))
+            (finished or segment_finished)
             and stage_id < req_state.final_stage_id
             and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1))
             and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled)
@@ -1589,7 +1618,9 @@ class Orchestrator:
     def _duplex_output_context(
         req_state: OrchestratorRequestState,
         *,
-        stage_id: int | None = None,
+        # Required: segment state is stage-keyed, so a caller that omitted this
+        # would silently read "no boundary" rather than the intended stage's.
+        stage_id: int | None,
     ) -> DuplexOutputContext | None:
         identity = req_state.duplex_identity
         if identity is None:
@@ -1600,15 +1631,16 @@ class Orchestrator:
         )
 
         fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence
+        segment = req_state.streaming.segment(stage_id)
         return DuplexOutputContext(
             identity=DuplexRequestIdentity(
                 session_id=identity.session_id,
                 fence=fence,
             ),
             final_stage_id=req_state.final_stage_id,
-            segment_finished=req_state.streaming.enabled and req_state.streaming.segment_finished,
+            segment_finished=req_state.streaming.enabled and segment.finished,
-            segment_token_ids=tuple(req_state.streaming.segment_token_ids),
+            segment_token_ids=tuple(segment.token_ids),
-            segment_output_metadata=req_state.streaming.segment_output_metadata,
+            segment_output_metadata=segment.output_metadata,
         )
 
     @staticmethod