Merge main and preserve dispatch backlog across monitor windows

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

tests/engine/duplex/test_deadline_pacing.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Deterministic unit tests for deadline-aligned silence continuation.
+
+The engine-resident duplex scheduler in
+``vllm_omni/engine/duplex/session/runner.py`` aligns each silence
+continuation to ``submission_time_N + chunk_period`` and sleeps only the
+remaining budget. These tests cover the pure deadline arithmetic and the
+per-session reset semantics of the shared ``DuplexModelSessionState``
+contract.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.engine.duplex.session.runner import (
+    compute_silence_continuation_deadline,
+)
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import (
+    MiniCPMO45ServingSessionState,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("chunk_period_s", "now", "last_submit", "current_deadline", "delay_s", "next_silence_deadline"),
+    [
+        # First continuation anchors to the last real submission: unit N was
+        # submitted at t=1.0, audio N produced at 1.4 -> 0.6 s of sleep left.
+        (1.0, 1.4, 1.0, None, 0.6, 3.0),
+        # Only the remaining budget is slept (audio produced at 1.7).
+        (1.0, 1.7, 1.0, None, 0.3, 3.0),
+        # Overdue deadline: no sleep; the chain still advances by the period.
+        (1.0, 2.5, 1.0, None, 0.0, 3.0),
+        # No submission yet (None): the first continuation anchors to now.
+        (1.0, 0.4, None, None, 1.0, 2.4),
+        # Deadlines advance from the current deadline, not from now, so pipeline
+        # processing time does not accumulate as timer drift.
+        (1.0, 4.2, 3.0, 5.0, 0.8, 6.0),
+        # Two consecutive units that overrun their deadline by 0.2 s keep the
+        # 1 s cadence: no sleep, and the chain advances one period each.
+        (1.0, 2.2, 1.0, None, 0.0, 3.0),
+        (1.0, 3.2, 2.0, 3.0, 0.0, 4.0),
+        # A short overshoot (within one period of the current deadline) still
+        # chases the stale deadline: immediate submit, chain advances.
+        (1.0, 2.5, 1.0, 2.0, 0.0, 3.0),
+        # A long stall (more than one period past the current deadline) submits
+        # one continuation immediately and restarts from that submission.
+        (1.0, 5.0, 1.0, 2.0, 0.0, 6.0),
+        # The same recovery applies to the first continuation after a stall.
+        (1.0, 5.0, 1.0, None, 0.0, 6.0),
+        # Zero chunk period (guarded upstream): total, no sleep.
+        (0.0, 10.0, None, None, 0.0, 10.0),
+    ],
+)
+def test_compute_silence_continuation_deadline(
+    chunk_period_s: float,
+    now: float,
+    last_submit: float | None,
+    current_deadline: float | None,
+    delay_s: float,
+    next_silence_deadline: float,
+) -> None:
+    delay, next_dl = compute_silence_continuation_deadline(
+        chunk_period_s=chunk_period_s,
+        now=now,
+        last_submit=last_submit,
+        current_deadline=current_deadline,
+    )
+    assert delay == pytest.approx(delay_s)
+    assert next_dl == pytest.approx(next_silence_deadline)
+
+
+class TestSilenceDeadlineSessionState:
+    def test_clear_continuation_resets_the_deadline_chain(self) -> None:
+        # The native-duplex session state feeds the shared runner, so it must
+        # start with an unset chain and drop it at turn boundaries.
+        state = MiniCPMO45ServingSessionState()
+        assert state.last_native_submit_monotonic is None
+        assert state.silence_deadline_monotonic is None
+        state.last_native_submit_monotonic = 5.0
+        state.silence_deadline_monotonic = 6.0
+        state.clear_continuation()
+        assert state.last_native_submit_monotonic is None
+        assert state.silence_deadline_monotonic is None

tests/engine/duplex/test_deadline_pacing_scheduler.py

@@ -0,0 +1,370 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Focused scheduler tests for deadline-aligned silence continuation.
+
+These tests exercise the production ``_schedule_silence_continuation`` path
+(and the append acceptance callback it installs) rather than copying the
+scheduling arithmetic. The monotonic clock is a fake so the deadline math is
+deterministic, and ``asyncio.sleep`` is simulated by advancing that clock.
+Append completion is controlled by gating the stage port's ``submit``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from typing import Any
+
+import pytest
+
+import vllm_omni.engine.duplex.session.model_channel as model_channel_module
+import vllm_omni.engine.duplex.session.runner as runner_module
+from tests.engine.duplex.test_session_runner import (
+    Harness,
+    append_audio,
+    close_harness,
+    open_harness,
+    tts_output,
+)
+from vllm_omni.engine.duplex.contracts import DuplexStageSubmission
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+CHUNK_PERIOD_S = 1.0
+
+
+class FakeClock:
+    """A monotonic clock whose value the test controls directly."""
+
+    def __init__(self, start: float = 0.0) -> None:
+        self.value = start
+
+    def __call__(self) -> float:
+        return self.value
+
+
+async def _active_response_harness() -> Harness:
+    """Open a harness and drive a TTS segment so a response is active."""
+    h = await open_harness()
+    await h.run(append_audio())
+    request_id = h.stage0_request_id()
+    await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello"))
+    assert h.session.active_response_id is not None
+    return h
+
+
+def _continuation_kwargs(h: Harness) -> dict[str, object]:
+    """The scheduler arguments the model channel would pass for this session."""
+    return {
+        "request_id": h.session.active_request_id,
+        "owner_id": f"response:{h.session.active_response_id}",
+        "response_id": h.session.active_response_id,
+        "response_owned": True,
+        "expected_epoch": h.session.epoch,
+        "expected_model_turn_id": h.session.turn_id,
+    }
+
+
+def _install_fake_clock(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    clock: FakeClock,
+    real_sleep: Callable[..., Any],
+) -> list[float]:
+    """Point both modules at the fake clock and simulate ``asyncio.sleep``.
+
+    ``real_sleep`` is the unpatched ``asyncio.sleep`` captured before patching;
+    the simulated sleep advances the fake clock by the requested delay and then
+    yields control so the append task can run. The returned ``stall`` list is
+    added once to the next *positive* sleep so a test can model a wake that
+    happens late (after the planned deadline); zero-duration scheduling yields
+    do not consume it.
+    """
+    monkeypatch.setattr(runner_module.time, "monotonic", clock)
+    monkeypatch.setattr(model_channel_module.time, "monotonic", clock)
+    stall: list[float] = [0.0]
+
+    async def fake_sleep(delay: float) -> None:
+        clock.value += delay
+        if delay > 0:
+            clock.value += stall[0]
+            stall[0] = 0.0
+        await real_sleep(0)
+
+    monkeypatch.setattr(asyncio, "sleep", fake_sleep)
+    return stall
+
+
+@pytest.mark.asyncio
+async def test_normal_cadence_keeps_deadlines_aligned(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Processing time reduces the remaining sleep; consecutive deadlines align.
+
+    Each accepted unit becomes the new anchor (its submission time), and the
+    following continuation advances from the stored deadline rather than
+    resetting to now.
+    """
+    clock = FakeClock(start=100.0)
+    real_sleep = asyncio.sleep
+    _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep)
+    h = await _active_response_harness()
+    try:
+        # Real input accepted at t=100.0; the model processes for 0.4 s.
+        h.runner.model_state.last_native_submit_monotonic = 100.0
+        clock.value = 100.4
+
+        # First continuation: due at 101.0, so only 0.6 s of sleep remains.
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        first_task = h.runner.tasks.append_tail
+        assert first_task is not None
+        assert await first_task
+        # The accepted unit re-anchored the chain to its own submission time
+        # (101.0) and stored the following deadline (102.0).
+        assert h.runner.model_state.last_native_submit_monotonic == pytest.approx(101.0)
+        assert h.runner.model_state.silence_deadline_monotonic == pytest.approx(102.0)
+
+        # Model produces the continuation's audio; the next unit is planned
+        # from the stored deadline (no drift): the following deadline is 103.0.
+        clock.value = 101.3
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        second_task = h.runner.tasks.append_tail
+        assert second_task is not None
+        assert await second_task
+        assert h.runner.model_state.silence_deadline_monotonic == pytest.approx(103.0)
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_in_flight_real_append_skips_outdated_silence(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A real append accepted after the silence was planned re-anchors the chain.
+
+    The planned silence queues behind the real append; once the real append
+    completes and re-anchors, ``before_append`` detects the changed anchor
+    (numeric timestamp) and the outdated silence is skipped (no submission,
+    no callback overwrite).
+    """
+    clock = FakeClock(start=200.0)
+    real_sleep = asyncio.sleep
+    _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep)
+    h = await _active_response_harness()
+    state = h.runner.model_state
+    submissions_before = len(h.port.submissions)
+
+    original_submit = h.port.submit
+    entered = asyncio.Event()
+    release = asyncio.Event()
+
+    async def gated_submit(submission: DuplexStageSubmission) -> Any:
+        entered.set()
+        await release.wait()
+        return await original_submit(submission)
+
+    monkeypatch.setattr(h.port, "submit", gated_submit)
+    try:
+        # Existing anchor at t=200.0.
+        state.last_native_submit_monotonic = 200.0
+        state.silence_deadline_monotonic = None
+
+        # The real append starts at t+0.8 and blocks inside the gate.
+        clock.value = 200.8
+        real_append_task = await h.runner._start_append(
+            h.runner.model.silence_unit_payload(),
+            final=False,
+            silence_continuation=False,
+        )
+        await entered.wait()
+        assert not real_append_task.done()
+
+        # At t+1.0 the silence is planned using the existing anchor; its append
+        # task waits behind the real append.
+        clock.value = 201.0
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        silence_task = h.runner.tasks.append_tail
+        assert silence_task is not None
+        assert not silence_task.done()
+
+        # At t+1.1 the real append is released; both appends finish.
+        clock.value = 201.1
+        release.set()
+        assert await real_append_task
+        assert await silence_task
+
+        # Exactly one additional submission happened: the real append. The
+        # outdated silence was skipped, so its callback never overwrote the
+        # new anchor (the real append's submission time t+0.8), and the
+        # deadline the real append cleared stays cleared.
+        assert len(h.port.submissions) == submissions_before + 1
+        assert state.last_native_submit_monotonic == pytest.approx(200.8)
+        assert state.silence_deadline_monotonic is None
+
+        # A fresh continuation uses the new anchor. The clock is left at t+1.1:
+        # the scheduler sleeps the remaining 0.7 s (deadline t+1.8), which
+        # verifies re-anchoring controls the wait. Following deadline t+2.8.
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        fresh_task = h.runner.tasks.append_tail
+        assert fresh_task is not None
+        assert await fresh_task
+        assert state.last_native_submit_monotonic == pytest.approx(201.8)
+        assert state.silence_deadline_monotonic == pytest.approx(202.8)
+    finally:
+        release.set()
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_long_stall_saves_future_deadline_from_actual_submission(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A submission more than one period past the planned deadline restarts.
+
+    The schedule plans the following deadline at ``planned + period``, but the
+    actual submission happens much later. The accepted continuation must save
+    ``submit_time + period`` so the next unit is not immediately due.
+    """
+    clock = FakeClock(start=300.0)
+    real_sleep = asyncio.sleep
+    stall = _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep)
+    h = await _active_response_harness()
+    state = h.runner.model_state
+    try:
+        # Plan a continuation whose deadline is 1.0 s in the future. The wake
+        # is delayed 2 s past the planned submission (anchor + 3.0 during the
+        # scheduler's sleep), so the append accepts 2 s late.
+        anchor = clock.value
+        state.last_native_submit_monotonic = anchor
+        state.silence_deadline_monotonic = None
+        stall[0] = 2.0
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        append_task = h.runner.tasks.append_tail
+        assert append_task is not None
+        assert await append_task
+        # anchor + 3.0 > anchor + 2.0 (planned deadline + period) -> save
+        # submit_time + period = anchor + 4.0.
+        assert state.last_native_submit_monotonic == pytest.approx(anchor + 3.0)
+        assert state.silence_deadline_monotonic == pytest.approx(anchor + 4.0)
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_small_wakeup_delay_preserves_planned_deadline(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Ordinary wakeup jitter (within one period) keeps the planned cadence.
+
+    A submission 0.2 s past the planned deadline must not reset the chain: the
+    stored deadline stays at ``planned + period`` so drift does not accumulate.
+    """
+    clock = FakeClock(start=400.0)
+    real_sleep = asyncio.sleep
+    stall = _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep)
+    h = await _active_response_harness()
+    state = h.runner.model_state
+    try:
+        anchor = clock.value
+        state.last_native_submit_monotonic = anchor
+        state.silence_deadline_monotonic = None
+        stall[0] = 0.2
+
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        append_task = h.runner.tasks.append_tail
+        assert append_task is not None
+        assert await append_task
+
+        # Submitted at anchor + period + 0.2 (within one period): the planned
+        # following deadline is kept.
+        assert state.last_native_submit_monotonic == pytest.approx(anchor + CHUNK_PERIOD_S + 0.2)
+        assert state.silence_deadline_monotonic == pytest.approx(anchor + 2 * CHUNK_PERIOD_S)
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_exact_one_period_late_preserves_planned_deadline(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Exactly one period late is not "more than one period late".
+
+    The strict ``>`` comparison means a submission exactly at the planned
+    following deadline keeps the planned cadence (no extra catch-up unit).
+    """
+    clock = FakeClock(start=500.0)
+    real_sleep = asyncio.sleep
+    stall = _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep)
+    h = await _active_response_harness()
+    state = h.runner.model_state
+    try:
+        anchor = clock.value
+        state.last_native_submit_monotonic = anchor
+        state.silence_deadline_monotonic = None
+        stall[0] = 1.0
+
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        append_task = h.runner.tasks.append_tail
+        assert append_task is not None
+        assert await append_task
+
+        # Submitted exactly at the planned following deadline: kept.
+        assert state.last_native_submit_monotonic == pytest.approx(anchor + 2 * CHUNK_PERIOD_S)
+        assert state.silence_deadline_monotonic == pytest.approx(anchor + 2 * CHUNK_PERIOD_S)
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_failed_append_leaves_timing_unchanged(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A failed append never advances the timing chain."""
+    clock = FakeClock(start=600.0)
+    real_sleep = asyncio.sleep
+    _install_fake_clock(monkeypatch, clock=clock, real_sleep=real_sleep)
+    h = await _active_response_harness()
+    state = h.runner.model_state
+    try:
+        state.last_native_submit_monotonic = 600.0
+        state.silence_deadline_monotonic = 601.0
+        before_anchor = state.last_native_submit_monotonic
+        before_deadline = state.silence_deadline_monotonic
+        h.port.fail_submit = RuntimeError("boom")
+
+        scheduled = await h.runner._schedule_silence_continuation(
+            h.runner.model.silence_unit_payload(),
+            **_continuation_kwargs(h),
+        )
+        assert scheduled is True
+        append_task = h.runner.tasks.append_tail
+        assert append_task is not None
+        # The append fails: its task reports failure, and timing is untouched.
+        assert await append_task is False
+        assert state.last_native_submit_monotonic == before_anchor
+        assert state.silence_deadline_monotonic == before_deadline
+    finally:
+        await close_harness(h)

tests/engine/duplex/test_duplex_events.py

@@ -11,6 +11,7 @@ import inspect
 import pytest
 
 from vllm_omni.engine.duplex import events as events_module
+from vllm_omni.engine.duplex.commands import ClearOutputAudio, TruncateItem
 from vllm_omni.engine.duplex.events import (
     REALTIME_ERROR_TYPES_BY_CODE,
     AudioDelta,
@@ -40,6 +41,9 @@ from vllm_omni.engine.duplex.events import (
 from vllm_omni.engine.duplex.realtime_events import (
     RealtimeProjectionState,
     project_internal_event,
+    resolve_clear_output_audio,
+    resolve_truncate_item,
+    retrieve_item_events,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -373,6 +377,96 @@ def test_projection_of_output_audio_buffer_clear_emits_cleared_before_terminals(
     assert state.item_truncation_cursors["item_resp_1"] == (0, 250)
 
 
+@pytest.mark.parametrize("reason", ["turn_detected", "barge_in", "output_audio_buffer_clear"])
+@pytest.mark.parametrize("committed_ms", [1000, 10000])
+def test_completed_response_cancellation_only_clears_pending_playback(reason, committed_ms):
+    state = RealtimeProjectionState(session_id="duplex-playback")
+    project_internal_event(state, {"type": "response.created", "response_id": "resp_1"})
+    project_internal_event(state, {"type": "response.done", "response_id": "resp_1"})
+
+    cancelled = project_internal_event(
+        state,
+        {
+            "type": "audio.cancelled",
+            "response_id": "resp_1",
+            "reason": reason,
+            "committed_ms": committed_ms,
+            "playback": {"sent_ms": 10000, "committed_ms": committed_ms},
+        },
+    )
+
+    if reason == "output_audio_buffer_clear" or committed_ms < 10000:
+        assert _types(cancelled) == ["output_audio_buffer.cleared"]
+        assert cancelled[0].response_id == "resp_1"
+    else:
+        assert cancelled == []
+
+
+@pytest.mark.parametrize("explicit_response_id", [False, True])
+def test_clear_latest_completed_response_reaches_playback_owner(explicit_response_id):
+    state = RealtimeProjectionState(session_id="duplex-clear")
+    project_internal_event(state, {"type": "response.created", "response_id": "resp_1"})
+    project_internal_event(state, {"type": "response.done", "response_id": "resp_1"})
+
+    control = resolve_clear_output_audio(
+        state, ClearOutputAudio(response_id="resp_1" if explicit_response_id else None, event_id="clear_1")
+    )
+
+    # Generation completion does not prove the client has drained its audio.
+    # A cleared event must follow the engine clearing its playback cursor.
+    assert control.events == []
+    assert control.payloads == [
+        {
+            "type": "output_audio_buffer.clear",
+            "reason": "output_audio_buffer.clear",
+            "response_id": "resp_1",
+            "realtime_event_id": "clear_1",
+        }
+    ]
+
+
+@pytest.mark.parametrize("new_response_finished", [False, True])
+def test_late_clear_of_completed_response_does_not_touch_new_response(new_response_finished):
+    state = RealtimeProjectionState(session_id="duplex-clear")
+    project_internal_event(state, {"type": "response.created", "response_id": "resp_1"})
+    project_internal_event(state, {"type": "response.done", "response_id": "resp_1"})
+    project_internal_event(state, {"type": "response.created", "response_id": "resp_2"})
+    if new_response_finished:
+        project_internal_event(state, {"type": "response.done", "response_id": "resp_2"})
+
+    control = resolve_clear_output_audio(state, ClearOutputAudio(response_id="resp_1"))
+
+    assert control.payloads == []
+    assert _types(control.events) == ["output_audio_buffer.cleared"]
+    assert control.events[0].response_id == "resp_1"
+    assert state.active_response_id == (None if new_response_finished else "resp_2")
+    assert state.last_response_id == "resp_2"
+
+
+def test_one_truncate_command_truncates_retrieved_transcript_once():
+    state = RealtimeProjectionState(session_id="duplex-truncate")
+    state.conversation_items["item_resp_1"] = {
+        "id": "item_resp_1",
+        "type": "message",
+        "role": "assistant",
+        "status": "completed",
+        "content": [{"type": "output_audio", "transcript": "abcdefghij", "audio_duration_ms": 10_000}],
+    }
+
+    control = resolve_truncate_item(state, TruncateItem(item_id="item_resp_1", audio_end_ms=4_000))
+    # The runner executes this signal and projects its acknowledgement.
+    signal = control.payloads[0]
+    payload = signal["payload"]
+    assert isinstance(payload, dict)
+    assert signal["event"] == "conversation.item.truncate"
+    project_internal_event(state, {"type": "conversation.item.truncated", **payload})
+
+    retrieved = retrieve_item_events(state, {"item_id": "item_resp_1"})[0].to_realtime()
+    assert retrieved["item"]["content"][0]["transcript"] == "abcd"
+    # Keep the cursor: later output must not restore the untruncated item.
+    assert state.item_truncation_cursors["item_resp_1"] == (0, 4_000)
+
+
 def test_projection_leaves_error_to_typed_emit_sites():
     # Errors are constructed typed (``error_event``) at the emit site; the projector
     # has no ``error`` branch and falls back to the raw wrapper for such a dict.

tests/engine/duplex/test_engine_session.py

@@ -21,6 +21,7 @@ from vllm_omni.engine.duplex.config import (
 from vllm_omni.engine.duplex.contracts import DuplexFence
 from vllm_omni.engine.duplex.events import TurnEvent
 from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError
+from vllm_omni.engine.duplex.session.playback_ledger import apply_playback_ack
 from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import (
     minicpmo45_native_capabilities,
 )
@@ -241,6 +242,83 @@ def test_history_commit_with_ack_only_policy_defers_unacknowledged_text():
     assert f"item_{response_id}" in session.pending_history_item_ids
 
 
+@pytest.mark.parametrize("audio_complete", [False, True])
+def test_unaligned_response_keeps_empty_turn_before_later_user_input(audio_complete):
+    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
+    first = session.commit_audio_input(transcript="recite a poem")
+    response_id = session.begin_response()
+    session.reserve_history_item(f"item_{response_id}")
+    session.append_assistant_text("An answer the user has only partly heard")
+    session.mark_audio_sent(
+        duration_ms=10_000,
+        text_requires_complete_audio=True,
+        audio_complete=audio_complete,
+    )
+
+    assert session.end_response(commit_text=True) is None
+    session.barge_in()
+    session.clear_playback_cursor()
+    second = session.commit_audio_input(transcript="stop")
+    empty_answer = {"role": "assistant", "content": ""}
+    assert session.history == (first.message, empty_answer, second.message)
+
+    ack = apply_playback_ack(session, {"response_id": response_id, "played_ms": 6_981})
+    assert ack[0].to_realtime()["event"]["history_committed"] is False
+    assert session.history == (first.message, empty_answer, second.message)
+
+    if audio_complete:
+        ack = apply_playback_ack(session, {"response_id": response_id, "played_ms": 10_000})
+        assert ack[0].to_realtime()["event"]["history_committed"] is True
+        assert session.history == (
+            first.message,
+            {"role": "assistant", "content": "An answer the user has only partly heard"},
+            second.message,
+        )
+
+
+@pytest.mark.parametrize(("reserve_slot", "assistant_text"), [(True, ""), (False, "unplayed answer")])
+def test_unaligned_response_only_materializes_an_existing_answer_slot(reserve_slot, assistant_text):
+    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
+    first = session.commit_audio_input(transcript="first input")
+    response_id = session.begin_response()
+    if reserve_slot:
+        session.reserve_history_item(f"item_{response_id}")
+    session.append_assistant_text(assistant_text)
+    session.mark_audio_sent(duration_ms=10_000, text_requires_complete_audio=True, audio_complete=True)
+
+    assert session.end_response(commit_text=reserve_slot) is None
+    assert session.history == (first.message,)
+
+
+@pytest.mark.parametrize("operation", ["delete", "truncate"])
+def test_removing_one_empty_assistant_preserves_the_other_turn_boundary(operation):
+    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
+    user_inputs = []
+    response_ids = []
+    for transcript in ("first input", "second input"):
+        user_inputs.append(session.commit_audio_input(transcript=transcript).message)
+        response_id = session.begin_response()
+        response_ids.append(response_id)
+        session.reserve_history_item(f"item_{response_id}")
+        session.append_assistant_text("unplayed answer")
+        session.mark_audio_sent(duration_ms=10_000, text_requires_complete_audio=True, audio_complete=True)
+        session.end_response(commit_text=True)
+
+    item_id = f"item_{response_ids[1]}"
+    if operation == "delete":
+        assert session.delete_history_item(item_id) is True
+    else:
+        # A full ACK releases the pending snapshot. Truncation then edits the
+        # stored message into another empty dict before removing that item.
+        ack = apply_playback_ack(session, {"response_id": response_ids[1], "played_ms": 10_000})
+        assert ack[0].to_realtime()["event"]["history_committed"] is True
+        assert session.truncate_history_item(item_id, audio_end_ms=0, hard=True) is True
+    assert session.history == (user_inputs[0], {"role": "assistant", "content": ""}, user_inputs[1])
+    ack = apply_playback_ack(session, {"response_id": response_ids[1], "played_ms": 10_000})
+    assert ack[0].to_realtime()["error"]["code"] == "playback_item_not_found"
+    assert session.history == (user_inputs[0], {"role": "assistant", "content": ""}, user_inputs[1])
+
+
 def test_truncate_history_item_uses_response_alignment_marks():
     session = _session()
     response_id = session.begin_response()
@@ -567,10 +645,10 @@ def test_overlap_policy_defaults_and_invalid_values_to_listen_only():
     assert DuplexSessionConfig.from_event({"session": {"overlap_policy": "bogus"}}).overlap_policy == "listen_only"
 
 
-def test_capabilities_as_dict_reports_model_native_duplex_constants():
+def test_capabilities_as_dict_reports_declared_implementation():
     caps = DuplexCapabilities().as_dict()
 
-    assert caps["implementation_level"] == "model_native_duplex"
+    assert caps["implementation_level"] == "turn_based_duplex"
     assert caps["input_modes"] == ["append_audio_chunk"]
     assert caps["supports_kv_lease"] is False
     assert caps["supports_core_kv_lease"] is False

tests/engine/duplex/test_session_manager.py

@@ -120,6 +120,8 @@ class FakeModelSessionState(DuplexModelSessionState):
         self.continuation_units = 0
         self.pending_silence_task = None
         self.pending_silence_owner_id = None
+        self.last_native_submit_monotonic = None
+        self.silence_deadline_monotonic = None
 
     def retain_committed_audio(self, payload, *, operation_id, reserved_bytes=0) -> None:
         self.committed_audio_payload = payload
@@ -136,6 +138,8 @@ class FakeModelSessionState(DuplexModelSessionState):
     def clear_continuation(self) -> None:
         self.continuation_owner_id = None
         self.continuation_units = 0
+        self.last_native_submit_monotonic = None
+        self.silence_deadline_monotonic = None
 
 
 class FakeDataPlane(DuplexDataPlane):
@@ -192,7 +196,7 @@ class FakePlugin(DuplexModelPlugin):
 
     def capabilities(self, *, max_sessions: int) -> DuplexCapabilities:
         del max_sessions
-        return DuplexCapabilities(supports_input_append=True)
+        return DuplexCapabilities(supports_input_append=True, supports_core_resumable_request=True)
 
     def validate_client_extra_body(self, extra_body: object) -> None:
         pass
@@ -404,7 +408,9 @@ async def test_open_answers_with_capabilities_and_emits_session_created() -> Non
         assert result.control_id == "open-sid-open"
         assert result.session_id == "sid-open"
         assert result.lease_generation == 0
-        assert result.capabilities == DuplexCapabilities(supports_input_append=True)
+        assert result.capabilities == DuplexCapabilities(
+            supports_input_append=True, supports_core_resumable_request=True
+        )
         assert result.public_session is not None
         assert result.public_session["id"] == "sid-open"
         assert result.public_session["voice"] == "test"

tests/engine/duplex/test_vad_backend.py

@@ -107,8 +107,8 @@ def test_endpoint_decisions_match_the_serving_side_policy(
 # --------------------------------------------------------------------------- #
 
 
-def test_a_loud_frame_does_not_restart_the_silence_timer() -> None:
-    """Silero v6.2 hysteresis: once silence is running, a loud frame only delays it.
+def test_hysteresis_band_does_not_restart_the_silence_timer() -> None:
+    """Scores between the negative and activation thresholds preserve the timer.
 
     The frame at 0.45 sits above the negative threshold (0.35) but below the
     activation threshold. It must not cancel the pending endpoint, or a speaker
@@ -164,9 +164,14 @@ def _pcm16_b64(samples: np.ndarray) -> str:
 
 def test_pcm16_input_is_accepted_alongside_float32() -> None:
     config = SileroVADConfig(threshold=0.5)
-    tone = np.zeros(FRAME * 2, dtype=np.float32)
+    tone: np.ndarray = np.zeros(FRAME * 2, dtype=np.float32)
     seen: list[int] = []
-    vad = SileroStreamingVAD(config, frame_scorer=lambda frame: seen.append(frame.size) or 0.0)
+
+    def _score_frame(frame: np.ndarray) -> float:
+        seen.append(frame.size)
+        return 0.0
+
+    vad = SileroStreamingVAD(config, frame_scorer=_score_frame)
     vad.process_base64(_pcm16_b64(tone), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ)
     assert seen == [FRAME, FRAME]
 
@@ -272,7 +277,7 @@ def test_the_onnx_backend_holds_the_upstream_silero_v62_contract(monkeypatch, tm
         intra_op_num_threads = 0
 
     class FakeInferenceSession:
-        def __init__(self, path: str, *, providers: list[str], sess_options: object) -> None:
+        def __init__(self, path: str, *, providers: list[str], sess_options: FakeSessionOptions) -> None:
             self.providers = providers
             self.sess_options = sess_options
             self.calls: list[dict[str, np.ndarray]] = []
@@ -313,7 +318,7 @@ def test_the_onnx_backend_holds_the_upstream_silero_v62_contract(monkeypatch, tm
     assert session.calls[0]["sr"].item() == SAMPLE_RATE_HZ
 
     state = backend.new_state()
-    first_frame = np.arange(backend.frame_samples, dtype=np.float32)
+    first_frame: np.ndarray = np.arange(backend.frame_samples, dtype=np.float32)
     probability, state = backend.infer(first_frame, state)
     first_call = session.calls[1]
 
@@ -377,3 +382,17 @@ def test_the_documented_server_vad_fields_are_all_accepted() -> None:
         )
         is None
     )
+
+
+@pytest.mark.parametrize("speech_probability", [0.5, 0.9])
+def test_resumed_speech_cancels_pending_silence(speech_probability):
+    config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=96, min_speech_duration_ms=32)
+    # 32 ms pause, resumed speech, then a full 96 ms pause: one utterance.
+    scores = [0.9, 0.1, speech_probability, speech_probability, 0.1, 0.1, 0.1]
+    assert _drive(config, scores) == [("start", 0), ("stop", 224)]
+
+
+def test_short_pauses_do_not_split_continuous_speech_into_repeated_turns():
+    config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=500, min_speech_duration_ms=32)
+    scores = [0.9] + [0.1, 0.9, 0.9, 0.9] * 20
+    assert _drive(config, scores) == [("start", 0)]

tests/engine/test_arg_utils.py

@@ -539,7 +539,44 @@ def test_tensor_parallel_size_none_is_handled():
     engine_args = OmegaConf.create({"stage_id": 0, "engine_args": {"tensor_parallel_size": None}})
     args = build_engine_args_dict(
         engine_args,
-        model="snu-aidas/Dynin-Omni",
+        model="Qwen/Qwen2-VL-2B-Instruct",
     )
     assert isinstance(args, dict)
     assert "tensor_parallel_size" not in args
+
+
+# For https://github.com/vllm-project/vllm-omni/issues/7564
+def test_from_cli_args_preserves_text_encoder_tp_size():
+    """`--text-encoder-tp-size` must survive from_cli_args field filtering.
+
+    Library callers build engine args via ``OmniEngineArgs.from_cli_args``;
+    the dataclass field filter drops any namespace attribute the dataclass
+    does not declare, silently resetting the diffusion text-encoder TP to 1.
+    """
+    engine_args = OmniEngineArgs.from_cli_args(
+        SimpleNamespace(text_encoder_tp_size=2),
+    )
+    assert engine_args.text_encoder_tp_size == 2
+
+
+# For https://github.com/vllm-project/vllm-omni/issues/7564
+def test_text_encoder_tp_size_reaches_default_diffusion_parallel_config():
+    """The preserved CLI value must land in DiffusionParallelConfig.
+
+    Mirrors the library flow: from_cli_args -> kwargs -> the generic
+    diffusion fallback, which resolves the field through
+    ``DiffusionParallelConfig.from_stage_overrides``.
+    """
+    from dataclasses import asdict
+
+    from vllm_omni.config.config_factory import StageConfigFactory
+
+    engine_args = OmniEngineArgs.from_cli_args(
+        SimpleNamespace(text_encoder_tp_size=2),
+    )
+    stage_cfg = StageConfigFactory.create_default_diffusion(
+        {k: v for k, v in asdict(engine_args).items() if v is not None},
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config["text_encoder_tp_size"] == 2

tests/engine/test_async_omni_engine_abort_ack.py

@@ -66,9 +66,10 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture):
         assert msg.rpc_id == "abort-rpc-1"
         assert not task.done()
 
-        rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
+        partial_output = SimpleNamespace(request_id="req-a", finished=True)
+        rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True, abort_outputs=[partial_output]))
         result = await task
-        assert result == []
+        assert result == [partial_output]
 
     try:
         asyncio.run(_run())

tests/engine/test_async_omni_engine_stage_init.py

@@ -8,24 +8,66 @@ import json
 import os
 import time
 import types
+from dataclasses import dataclass, field
 
 import pytest
+from omegaconf import OmegaConf
+from vllm.v1.engine.utils import EngineZmqAddresses
 
+from vllm_omni.config.omni_config import OmniStageRuntimeConfig
 from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine import omni_engine_base as async_omni_engine_module
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.stage_engine_startup import StageReplicaResources
 from vllm_omni.engine.stage_init_utils import (
     LogicalStageInitPlan,
     ReplicaInitPlan,
     build_stage0_input_processor,
     compute_replica_layout,
     split_devices_for_replicas,
+    stage_runtime_env,
 )
 from vllm_omni.engine.stage_runtime import StageRuntime
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@dataclass
+class _FakeParallelConfig:
+    enable_fault_tolerance: bool = False
+    enable_elastic_ep: bool = False
+    data_parallel_size: int = 1
+    use_ray: bool = False
+
+
+@dataclass
+class _FakeVllmConfig:
+    parallel_config: _FakeParallelConfig = field(default_factory=_FakeParallelConfig)
+
+
+@dataclass
+class _FakeRuntimeConfig:
+    devices: str
+
+
+@dataclass
+class _FakeStageConfig:
+    stage_id: int
+    stage_type: str
+    engine_args: dict[str, object]
+    runtime: _FakeRuntimeConfig
+
+
+def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
+    env_key = "VLLM_OMNI_TEST_TYPED_STAGE_ENV"
+    monkeypatch.delenv(env_key, raising=False)
+
+    with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "typed-value"})):
+        assert os.environ[env_key] == "typed-value"
+
+    assert env_key not in os.environ
+
+
 def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch):
     engine = object.__new__(AsyncOmniEngine)
     engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True)
@@ -187,7 +229,7 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     )
 
 
-def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch):
+def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):
     import vllm_omni.engine.omni_engine_base as engine_mod
 
     engine = object.__new__(AsyncOmniEngine)
@@ -205,6 +247,7 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
     engine._omni_lb_policy = "random"
     engine.request_queue = types.SimpleNamespace()
     engine._log_stats = True
+    engine._client_config = engine_mod.OmniClientConfig(client_count=2, client_index=1, stage_addresses={})
     engine._parallel_stage_init = False
 
     captured: dict[str, object] = {}
@@ -220,6 +263,51 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
 
     assert captured["stage_init_timeout"] == 7
     assert captured["log_stats"] is True
+    assert captured["client_config"] is engine._client_config
+
+
+def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch):
+    import vllm_omni.engine.omni_engine_base as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.stage_configs = [types.SimpleNamespace()]
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.single_stage_mode = False
+    engine.async_chunk = False
+    engine.tokenizer = None
+    engine._single_stage_id_filter = None
+    engine._omni_master_address = None
+    engine._omni_master_port = None
+    engine._omni_dp_size_local = 1
+    engine._omni_heartbeat_timeout = 30.0
+    engine._omni_lb_policy = "random"
+    engine.request_queue = types.SimpleNamespace()
+    engine._log_stats = False
+    engine._parallel_stage_init = False
+
+    prompt_transform = object()
+    client = types.SimpleNamespace(
+        prompt_transform_func=prompt_transform,
+        prompt_expand_func=None,
+        default_sampling_params=types.SimpleNamespace(),
+        final_output=True,
+        final_output_type="text",
+        stage_type="llm",
+        model_stage="text_encoder",
+        is_comprehension=True,
+    )
+    pool = types.SimpleNamespace(
+        stage_client=client,
+        stage_vllm_config=None,
+        output_processor=None,
+    )
+    runtime = types.SimpleNamespace(stage_pools=[pool], initialize=lambda: None)
+    monkeypatch.setattr(engine_mod, "create_stage_runtime", lambda **_kwargs: runtime)
+
+    engine._initialize_stages(stage_init_timeout=7)
+
+    assert engine.prompt_transform_func is prompt_transform
 
 
 def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
@@ -776,6 +864,289 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
     assert captured["client_log_stats"] is True
 
 
+def test_stage_runtime_attaches_external_llm_client_without_launching_engine(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    client_addresses = {
+        "input_address": "ipc://stage0-input-1",
+        "output_address": "ipc://stage0-output-1",
+    }
+    client_config = {
+        "client_count": 2,
+        "client_index": 1,
+        "stage_addresses": {0: {0: client_addresses}},
+    }
+    runtime = StageRuntime(
+        stage_configs=[],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+        client_config=client_config,
+    )
+    parallel_config = _FakeParallelConfig()
+    vllm_config = _FakeVllmConfig(parallel_config=parallel_config)
+    plan = _make_llm_plan(0, stage_id=0, vllm_config=vllm_config).replicas[0]
+    captured: dict[str, object] = {}
+    stage_client = object()
+
+    def _capture_client(**kwargs):
+        captured.update(kwargs)
+        return stage_client
+
+    monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", _capture_client)
+    monkeypatch.setattr(
+        runtime_mod,
+        "launch_stage_replica",
+        lambda **_kwargs: pytest.fail("external client must not launch a stage engine"),
+    )
+
+    assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
+    assert captured["client_addresses"] == client_addresses
+    assert captured["client_count"] == 2
+    assert captured["client_index"] == 1
+    assert not hasattr(parallel_config, "_api_process_count")
+    assert not hasattr(parallel_config, "_api_process_rank")
+
+
+@pytest.mark.parametrize("stage_ids", [(0,), (0, 1)], ids=["single-stage", "multi-stage"])
+@pytest.mark.parametrize("client_count", [1, 2])
+def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeypatch, stage_ids, client_count):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = _make_stage_runtime()
+    stage_plans = []
+    for stage_id in stage_ids:
+        parallel_config = _FakeParallelConfig()
+        plan = _make_llm_plan(
+            stage_id,
+            stage_id=stage_id,
+            vllm_config=_FakeVllmConfig(parallel_config=parallel_config),
+        )
+        plan.replicas[0].engine_args_dict = {}
+        stage_plans.append(plan)
+    stage_plans[0].replicas[0].metadata.runtime_cfg = {"env": {"VLLM_OMNI_TEST_STAGE_RUNTIME_ENV": "enabled"}}
+
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda stage_id, _cfg: str(stage_id))
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
+
+    events: list[tuple[str, int]] = []
+    captured_launch_kwargs: list[dict[str, object]] = []
+    captured_launch_env: list[str | None] = []
+
+    @contextlib.contextmanager
+    def _fake_launch_stage_replica(**kwargs):
+        stage_id = kwargs["stage_id"]
+        events.append(("enter", stage_id))
+        captured_launch_kwargs.append(kwargs)
+        captured_launch_env.append(os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"))
+        yield StageReplicaResources(
+            addresses=EngineZmqAddresses(
+                inputs=[f"ipc://stage{stage_id}-input-{idx}" for idx in range(client_count)],
+                outputs=[f"ipc://stage{stage_id}-output-{idx}" for idx in range(client_count)],
+            ),
+        )
+        events.append(("exit", stage_id))
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
+
+    with runtime.launch_stage_engines(client_count) as launch:
+        expected = (
+            [("enter", stage_id) for stage_id in stage_ids]
+            if client_count > 1
+            else [(event, stage_id) for stage_id in stage_ids for event in ("enter", "exit")]
+        )
+        assert events == expected
+        for client_index, config in enumerate(launch.client_configs):
+            assert config["client_count"] == client_count
+            assert config["client_index"] == client_index
+            for stage_id in stage_ids:
+                assert config["stage_addresses"][stage_id][0] == {
+                    "input_address": f"ipc://stage{stage_id}-input-{client_index}",
+                    "output_address": f"ipc://stage{stage_id}-output-{client_index}",
+                }
+
+    assert events == (
+        [(event, stage_id) for event in ("enter", "exit") for stage_id in stage_ids] if client_count > 1 else expected
+    )
+    assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_count")
+    assert not hasattr(stage_plans[0].replicas[0].stage_vllm_config.parallel_config, "_api_process_rank")
+    assert all(
+        kwargs["watched_frontend_processes"] is (launch.watched_frontend_processes if client_count > 1 else None)
+        for kwargs in captured_launch_kwargs
+    )
+    assert captured_launch_env == ["enabled" if stage_id == 0 else None for stage_id in stage_ids]
+    assert os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV") is None
+
+
+@pytest.mark.parametrize("client_count", [1, 2])
+def test_stage_runtime_overlapping_devices_acquire_real_locks_once(monkeypatch, tmp_path, client_count):
+    import fcntl
+
+    import vllm_omni.engine.stage_init_utils as init_utils
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = _make_stage_runtime()
+    plans = [_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) for stage_id in (0, 1)]
+    for stage_id, plan in enumerate(plans):
+        plan.replicas[0].engine_args_dict = {"tensor_parallel_size": stage_id + 1}
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: plans)
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda sid, _cfg: "0" if sid == 0 else "0,1")
+    monkeypatch.setattr(
+        init_utils, "device_init_lock_path", lambda device_id: str(tmp_path / f"device-{device_id}.lock")
+    )
+    monkeypatch.setattr(
+        init_utils.time, "sleep", lambda _: pytest.fail("Overlapping stages must not wait on their own lock")
+    )
+
+    @contextlib.contextmanager
+    def launch(**kwargs):
+        # Real flock on a separate descriptor proves the runtime holds each lock.
+        for device_id in range(kwargs["stage_id"] + 1):
+            with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
+                with pytest.raises(BlockingIOError):
+                    fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        yield StageReplicaResources(
+            addresses=EngineZmqAddresses(
+                inputs=[f"ipc://input-{i}" for i in range(client_count)],
+                outputs=[f"ipc://output-{i}" for i in range(client_count)],
+            )
+        )
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
+    with runtime.launch_stage_engines(client_count):
+        pass
+    for device_id in (0, 1):
+        with open(tmp_path / f"device-{device_id}.lock", "a") as probe:
+            fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            fcntl.flock(probe, fcntl.LOCK_UN)
+
+
+@pytest.mark.parametrize("diffusion_stage_id", [0, 1], ids=["diffusion-only", "enginecore-to-diffusion"])
+def test_stage_runtime_multi_api_rejects_diffusion_before_launch(monkeypatch, diffusion_stage_id):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = _make_stage_runtime()
+    stage_plans = [
+        _make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig())
+        for stage_id in range(diffusion_stage_id)
+    ]
+    stage_plans.append(_make_diffusion_plan(diffusion_stage_id, stage_id=diffusion_stage_id))
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
+    monkeypatch.setattr(
+        runtime_mod,
+        "acquire_device_locks",
+        lambda *_args: pytest.fail("unsupported topology must fail before acquiring devices"),
+    )
+    monkeypatch.setattr(
+        runtime_mod,
+        "launch_stage_replica",
+        lambda **_kwargs: pytest.fail("unsupported topology must not partially launch EngineCore stages"),
+    )
+
+    with pytest.raises(ValueError, match="diffusion stage\\(s\\) are not supported"):
+        with runtime.launch_stage_engines(2):
+            pytest.fail("unsupported topology must not yield a launch handle")
+
+
+def test_stage_runtime_multi_api_maps_stage_devices_from_launcher_visibility(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = _make_stage_runtime()
+    stage_plan = _make_llm_plan(
+        0,
+        stage_id=0,
+        vllm_config=_FakeVllmConfig(),
+    )
+    replica = stage_plan.replicas[0]
+    replica.engine_args_dict = {}
+    replica.metadata.runtime_cfg = {"devices": "0"}
+
+    device_env = runtime_mod.current_omni_platform.device_control_env_var
+    monkeypatch.setenv(device_env, "5")
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
+    captured_devices: list[str | None] = []
+
+    @contextlib.contextmanager
+    def _fake_launch_stage_replica(**kwargs):
+        captured_devices.append(kwargs["stage_visible_devices"])
+        yield StageReplicaResources(
+            addresses=EngineZmqAddresses(
+                inputs=["ipc://input-0", "ipc://input-1"],
+                outputs=["ipc://output-0", "ipc://output-1"],
+            ),
+        )
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
+
+    with runtime.launch_stage_engines(2):
+        pass
+
+    assert captured_devices == ["5"]
+    assert os.environ[device_env] == "5"
+
+
+def test_stage_runtime_multi_api_rejects_elastic_ep(monkeypatch):
+    runtime = _make_stage_runtime()
+    stage_plan = _make_llm_plan(
+        0,
+        stage_id=0,
+        vllm_config=_FakeVllmConfig(parallel_config=_FakeParallelConfig(enable_elastic_ep=True)),
+    )
+    stage_plan.replicas[0].engine_args_dict = {}
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
+
+    with pytest.raises(ValueError, match="enable-elastic-ep"):
+        with runtime.launch_stage_engines(2):
+            pass
+
+
+def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_exit(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = _make_stage_runtime()
+    stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
+    stage_plan.replicas[0].engine_args_dict = {}
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [stage_plan])
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_args: None)
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args: [])
+
+    events: list[str] = []
+
+    class _Manager:
+        def shutdown(self) -> None:
+            events.append("shutdown")
+
+    @contextlib.contextmanager
+    def _fake_launch_stage_replica(**_kwargs):
+        try:
+            yield StageReplicaResources(
+                manager=_Manager(),
+                addresses=EngineZmqAddresses(
+                    inputs=["tcp://127.0.0.1:0", "tcp://127.0.0.1:1"],
+                    outputs=["tcp://127.0.0.1:2", "tcp://127.0.0.1:3"],
+                ),
+            )
+        except RuntimeError as exc:
+            events.append(f"exceptional-exit:{exc}")
+            raise
+        else:
+            events.append("normal-exit")
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
+
+    with pytest.raises(RuntimeError, match="deferred TCP addresses"):
+        with runtime.launch_stage_engines(2):
+            pass
+
+    assert events == [
+        "shutdown",
+        "exceptional-exit:Stage 0 returned deferred TCP addresses; multi-API launch requires fixed ports or IPC addresses",
+    ]
+
+
 def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 
@@ -1022,7 +1393,7 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
         num_replicas=1,
         launch_mode="local",
         stage_cfg=types.SimpleNamespace(engine_args={}, runtime=types.SimpleNamespace(devices="0")),
-        metadata=types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}),
+        metadata=types.SimpleNamespace(stage_id=0, stage_type="llm", runtime_cfg={"devices": "0"}),
         stage_connector_spec={},
         omni_kv_connector=(None, None, None),
         stage_vllm_config=fake_vllm_config,
@@ -1714,6 +2085,51 @@ def test_inject_kv_stage_info_infers_receiver_tp_topology():
     assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
 
 
+def test_inject_kv_stage_info_updates_typed_connector_config():
+    from vllm_omni.config.omni_config import (
+        OmniStageConnectorConfig,
+        OmniStageDiffusionParallelConfig,
+        VllmOmniDiffusionStageConfig,
+    )
+    from vllm_omni.config.stage_config import StageExecutionType, StagePipelineConfig
+    from vllm_omni.engine.stage_init_utils import inject_kv_stage_info
+    from vllm_omni.entrypoints.utils import inject_omni_kv_config
+
+    stage0 = VllmOmniDiffusionStageConfig(
+        stage_pipeline_config=StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+        ),
+        connector_config=OmniStageConnectorConfig(
+            omni_kv_config={
+                "need_send_cache": True,
+                "omni_from_stage": "0",
+                "omni_to_stage": "1",
+            }
+        ),
+        parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=4),
+    )
+    stage1 = VllmOmniDiffusionStageConfig(
+        stage_pipeline_config=StagePipelineConfig(
+            stage_id=1,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(0,),
+        ),
+        connector_config=OmniStageConnectorConfig(omni_kv_config={"need_recv_cache": True}),
+        parallel_config=OmniStageDiffusionParallelConfig(tensor_parallel_size=2),
+    )
+
+    inject_omni_kv_config(stage0, {"kv_connector": "P2pNcclConnector"}, "0", "1")
+    inject_kv_stage_info(stage0, 0, [stage0, stage1])
+
+    assert stage0.connector_config.omni_kv_config["stage_id"] == 0
+    assert stage0.connector_config.omni_kv_config["connector_config"] == {"kv_connector": "P2pNcclConnector"}
+    assert stage0.connector_config.omni_kv_config["engine_input_source"] == []
+    assert stage0.connector_config.omni_kv_config["rank_mapping"] == {"from_tp": 4, "to_tp": 2}
+
+
 def test_extract_legacy_stage_metadata_rocm_does_not_inject_diffusion_attention(monkeypatch):
     """ROCm default attention logic only applies to LLM stages, not diffusion."""
     from vllm_omni.engine.stage_init_utils import extract_legacy_stage_metadata
@@ -1829,3 +2245,308 @@ def test_port_from_zmq_address_parsing():
     assert _port_from_zmq_address(None) is None
     assert _port_from_zmq_address("ipc:///tmp/sock") is None
     assert _port_from_zmq_address("tcp://host:not-a-port") is None
+
+
+@pytest.mark.parametrize("parallel", [False, True])
+@pytest.mark.parametrize("failure", ["ready", "attach", None])
+def test_single_api_common_launch_ownership_and_rollback(monkeypatch, parallel, failure):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = _make_stage_runtime()
+    runtime._parallel_stage_init = parallel
+    plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig()).replicas[0]
+    plan.engine_args_dict = {}
+    events = []
+
+    class Manager:
+        def shutdown(self):
+            events.append("shutdown")
+
+    manager = Manager()
+    addresses = EngineZmqAddresses(inputs=["ipc://input"], outputs=["ipc://output"])
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
+
+    def acquire(*_):
+        events.append("lock")
+        return [42]
+
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", acquire)
+    monkeypatch.setattr(runtime_mod, "release_device_locks", lambda _: events.append("unlock"))
+
+    @contextlib.contextmanager
+    def launch(**kwargs):
+        assert kwargs["num_api_servers"] == 1
+        assert kwargs["omni_parallel_stage_init"] is parallel
+        events.append("spawn")
+        yield StageReplicaResources(manager=manager, addresses=addresses)
+        assert runtime._replica_launch_lock.locked() is False
+        events.append("ready")
+        if failure == "ready":
+            raise RuntimeError("ready")
+
+    def attach(**kwargs):
+        assert "ready" in events
+        assert kwargs["engine_manager"] is manager
+        events.append("attach")
+        if failure == "attach":
+            raise RuntimeError("attach")
+        return manager
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
+    monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", attach)
+    if failure:
+        with pytest.raises(RuntimeError, match=failure):
+            runtime._initialize_local_llm_replica(plan, 1)
+        assert events.count("shutdown") == 1
+    else:
+        assert runtime._initialize_local_llm_replica(plan, 1) is manager
+        assert "shutdown" not in events
+    assert events.count("lock") == events.count("unlock") == (0 if parallel else 1)
+
+
+@pytest.mark.parametrize("client_count", [1, 2])
+def test_common_launch_parallel_admission_before_spawn(monkeypatch, client_count):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = _make_stage_runtime()
+    runtime._parallel_stage_init = True
+    plan = _make_llm_plan(0, stage_id=0, vllm_config=_FakeVllmConfig())
+    plan.replicas[0].engine_args_dict = {}
+    events = []
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: [plan])
+    monkeypatch.setattr(runtime, "_reject_unguardable_executors", lambda _: events.append("guard"))
+    monkeypatch.setattr(runtime, "_run_stage_admission", lambda _: events.append("admit"))
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
+
+    def forbidden_lock(*_):
+        pytest.fail("parallel initialization must use child phase locks")
+
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", forbidden_lock)
+
+    @contextlib.contextmanager
+    def launch(**kwargs):
+        assert events == ["guard", "admit"]
+        assert kwargs["omni_parallel_stage_init"] is True
+        events.append("spawn")
+        yield StageReplicaResources(
+            addresses=EngineZmqAddresses(
+                inputs=[f"ipc://in{i}" for i in range(client_count)],
+                outputs=[f"ipc://out{i}" for i in range(client_count)],
+            )
+        )
+        assert not runtime._replica_launch_lock.locked()
+        events.append("ready")
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
+    with runtime.launch_stage_engines(client_count):
+        pass
+    assert events == ["guard", "admit", "spawn", "ready"]
+
+
+@pytest.mark.parametrize(
+    "stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)]
+)
+def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes, expected):
+    engine = object.__new__(AsyncOmniEngine)
+    stages = [OmegaConf.create({"engine_args": {"async_chunk": mode}}) for mode in stage_modes]
+    monkeypatch.setattr(async_omni_engine_module.StageConfigFactory, "get_pipeline_config", lambda *a, **k: None)
+    monkeypatch.setattr(engine, "_resolve_stage_configs", lambda *a, **k: (None, stages))
+    monkeypatch.setattr(engine, "_set_pipeline_runtime_config", lambda *a: None)
+
+    class ConfigResolvedError(Exception):
+        pass
+
+    def stop_before_queues(*args, **kwargs):
+        raise ConfigResolvedError
+
+    monkeypatch.setattr(async_omni_engine_module.janus, "Queue", stop_before_queues)
+    with pytest.raises(ConfigResolvedError):
+        engine.__init__("test-model")
+    assert engine.async_chunk is expected
+
+
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+def test_dist_stage_runtime_applies_local_dp_to_stage_config(typed):
+    from vllm_omni.config.omni_config import VllmOmniARStageConfig
+    from vllm_omni.config.stage_config import StagePipelineConfig
+    from vllm_omni.engine.stage_runtime import DistStageRuntime
+
+    if typed:
+        stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar"))
+        runtime_cfg = stage.runtime_config
+    else:
+        runtime_cfg = types.SimpleNamespace(num_replicas=1)
+        stage = types.SimpleNamespace(stage_id=0, runtime=runtime_cfg)
+    runtime = DistStageRuntime(
+        stage_configs=[stage],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+        single_stage_id_filter=0,
+        omni_master_address="127.0.0.1",
+        omni_master_port=12345,
+        omni_dp_size_local=2,
+    )
+
+    runtime._validate_single_stage_mode_replica_constraints()
+
+    assert runtime_cfg.num_replicas == 2
+
+
+@pytest.mark.parametrize("num_replicas", [1, 2, 3])
+@pytest.mark.parametrize("kv_owner", [None, "connector_config", "diffusion_config"])
+def test_typed_diffusion_replicas_share_one_config_between_planning_and_launch(mocker, num_replicas, kv_owner):
+    from vllm_omni.config.config_factory import StageConfigFactory
+    from vllm_omni.engine import stage_runtime as runtime_module
+
+    stage = StageConfigFactory.create_typed_default_diffusion(
+        "generic-diffusion", {"model_class_name": "QwenImagePipeline"}
+    ).stage_configs[0]
+    if kv_owner is not None:
+        from vllm.config import KVTransferConfig
+
+        getattr(stage, kv_owner).kv_transfer_config = KVTransferConfig(
+            kv_connector="MooncakeConnector", kv_role="kv_consumer", engine_id="dit"
+        )
+    devices = ",".join(str(i) for i in range(num_replicas))
+    stage.runtime_config.devices = devices
+    stage.runtime_config.num_replicas = num_replicas
+    runtime = StageRuntime(
+        stage_configs=[stage],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+    )
+    mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
+    mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
+    mocker.patch.object(runtime, "_stage_device_scope", side_effect=lambda *_: contextlib.nullcontext())
+    client = mocker.Mock()
+    launch = mocker.patch.object(runtime_module, "launch_diffusion_stage_replica", return_value=(client, None))
+
+    # Replanning must not inherit the first replica's narrowed device slice.
+    for _ in range(2):
+        counts, device_map = compute_replica_layout([stage])
+        plans = runtime._build_logical_stage_init_plans(
+            omni_transfer_config=None,
+            replicas_per_stage=counts,
+            replica_devices_map=device_map,
+        )
+        replicas = plans[0].replicas
+        assert stage.runtime_config.devices == devices
+        assert len({id(plan.stage_cfg) for plan in replicas}) == num_replicas
+        for i, plan in enumerate(replicas):
+            assert plan.metadata.runtime_cfg is plan.stage_cfg.runtime_config
+            assert plan.stage_cfg.runtime_config.devices == str(i)
+            if num_replicas > 1 or kv_owner is not None:
+                assert plan.stage_cfg is not stage
+            if kv_owner is not None:
+                from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+                assert getattr(stage, kv_owner).kv_transfer_config.engine_id == "dit"
+                engine_args = build_engine_args_dict_from_omni_stage_config(plan.stage_cfg, "dummy-model")
+                assert engine_args["kv_transfer_config"].engine_id == f"dit-s0-r{i}"
+            assert runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) is client
+            assert launch.call_args.kwargs["stage_config"] is plan.stage_cfg
+            assert launch.call_args.kwargs["metadata"] is plan.metadata
+            assert launch.call_args.kwargs["use_inline"] is (num_replicas == 1)
+
+    assert launch.call_count == 2 * num_replicas
+
+
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize("num_replicas", [1, 2])
+def test_native_kv_producer_replica_identity_and_bootstrap_are_isolated(mocker, typed, num_replicas):
+    from vllm.config import KVTransferConfig
+
+    from vllm_omni.config.omni_config import VllmOmniARStageConfig
+    from vllm_omni.config.stage_config import StagePipelineConfig
+    from vllm_omni.engine import stage_runtime as runtime_module
+
+    kv_config = KVTransferConfig(
+        kv_connector="MooncakeConnector",
+        kv_role="kv_producer",
+        engine_id="ar",
+        kv_ip="127.0.0.1",
+        kv_connector_extra_config={"bootstrap_port": 9100},
+    )
+    if typed:
+        stage = VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="ar"))
+        stage.connector_config.kv_transfer_config = kv_config
+        runtime_cfg = stage.runtime_config
+    else:
+        runtime_cfg = OmegaConf.create({"devices": "0,1", "env": {"KEEP": "value"}})
+        stage = types.SimpleNamespace(stage_id=0, engine_args={"kv_transfer_config": kv_config}, runtime=runtime_cfg)
+        metadata = _make_llm_metadata(0)
+        mocker.patch.object(
+            runtime_module,
+            "extract_legacy_stage_metadata",
+            side_effect=lambda cfg: types.SimpleNamespace(**{**metadata.__dict__, "runtime_cfg": cfg.runtime}),
+        )
+    runtime_cfg.devices = "0,1"
+    runtime_cfg.env = {"KEEP": "value"}
+    runtime = StageRuntime(
+        stage_configs=[stage],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+    )
+    mocker.patch.object(runtime_module, "get_stage_connector_spec", return_value={})
+    mocker.patch.object(runtime_module, "resolve_omni_kv_config_for_stage", return_value=(None, None, None))
+    for builder in ("build_engine_args_dict", "build_engine_args_dict_from_omni_stage_config"):
+        mocker.patch.object(runtime_module, builder, return_value={})
+    mocker.patch.object(
+        runtime_module,
+        "build_vllm_config",
+        return_value=(types.SimpleNamespace(kv_transfer_config=kv_config), object),
+    )
+
+    for _ in range(2):
+        plans = runtime._build_logical_stage_init_plans(None, [num_replicas], {})
+        for i, replica in enumerate(plans[0].replicas):
+            config = replica.stage_vllm_config.kv_transfer_config
+            assert config.engine_id == f"ar-s0-r{i}"
+            assert config.kv_connector_extra_config["bootstrap_addr"] == f"http://127.0.0.1:{9100 + i}"
+            assert dict(replica.metadata.runtime_cfg.env) == {
+                "KEEP": "value",
+                "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(9100 + i),
+            }
+            if typed:
+                assert replica.metadata.runtime_cfg is replica.stage_cfg.runtime_config
+        assert kv_config.engine_id == "ar"
+        assert kv_config.kv_connector_extra_config == {"bootstrap_port": 9100}
+        assert dict(runtime_cfg.env) == {"KEEP": "value"}
+
+
+@pytest.mark.parametrize(
+    "roles, sources, async_chunk, valid",
+    [
+        (["kv_producer", "kv_consumer"], [0], False, True),
+        (["kv_producer", "kv_consumer"], [0], True, False),
+        ([None, "kv_consumer"], [0], False, False),
+        (["kv_producer", None], [0], False, False),
+        ([None, "kv_producer", "kv_consumer"], [1], False, False),
+        (["kv_producer", "kv_consumer"], [1], False, False),
+        ([None, None], [0], False, True),
+    ],
+)
+def test_native_kv_topology_rejects_silent_legacy_fallback(roles, sources, async_chunk, valid):
+    plans = []
+    for stage_id, role in enumerate(roles):
+        diffusion = stage_id == len(roles) - 1
+        config = types.SimpleNamespace(kv_role=role) if role else None
+        replica = types.SimpleNamespace(
+            metadata=types.SimpleNamespace(stage_type="diffusion" if diffusion else "llm", engine_input_source=sources),
+            stage_vllm_config=None if diffusion else types.SimpleNamespace(kv_transfer_config=config),
+            stage_cfg=types.SimpleNamespace(engine_args={"kv_transfer_config": config}),
+        )
+        plans.append(types.SimpleNamespace(stage_id=stage_id, replicas=[replica]))
+    runtime = object.__new__(StageRuntime)
+    runtime._async_chunk = async_chunk
+    if valid:
+        runtime._validate_native_kv_topology(plans)
+    else:
+        with pytest.raises(ValueError, match="two-stage"):
+            runtime._validate_native_kv_topology(plans)

tests/engine/test_duplex_orchestrator.py

@@ -515,3 +515,48 @@ async def test_a_dead_replica_closes_the_sessions_it_was_serving() -> None:
     messages = [output_q.get_nowait() for _ in range(output_q.qsize())]
     types = [getattr(getattr(m, "event", None), "type", type(m).__name__) for m in messages]
     assert types[-1] in {"session.expired", "session.closed"}, f"the client needs a terminal event, got {types}"
+
+
+@pytest.mark.asyncio
+async def test_turn_plugin_processes_multimodal_prompt_before_stage_submission(monkeypatch):
+    from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens
+    from vllm_omni.model_executor.models.qwen3_omni.duplex.plugin import Qwen3OmniDuplexPlugin
+
+    orchestrator, clients, rpc_q, _ = _build()
+    plugin = Qwen3OmniDuplexPlugin(_encode_audio)
+    plugin.processor = SimpleNamespace(apply_chat_template=lambda messages, **kwargs: "audio prompt")
+    orchestrator.plugin = orchestrator.session_manager.plugin = plugin
+    seen = []
+
+    def process_inputs(**kwargs):
+        seen.append(kwargs)
+        return build_engine_core_request_from_tokens(
+            request_id=kwargs["request_id"],
+            prompt={"prompt_token_ids": [1, 2]},
+            params=kwargs["params"],
+            model_config=orchestrator.stage_pools[0].stage_vllm_config.model_config,
+            resumable=kwargs["resumable"],
+        )
+
+    monkeypatch.setattr(
+        orchestrator, "_get_stage_input_processor", lambda stage_id: SimpleNamespace(process_inputs=process_inputs)
+    )
+    await orchestrator._dispatch_message(
+        OpenDuplexSessionMessage(
+            control_id="open-qwen",
+            session_id=SESSION_ID,
+            session_config=DuplexSessionConfig(model="qwen", modalities=["text", "audio"]),
+        )
+    )
+    assert (await rpc_q.get()).ok
+    try:
+        await _submit(orchestrator, _append_audio())
+        assert not clients[0].add_request_calls
+        await _submit(orchestrator, commands.Commit(final=True, create_response=True))
+        assert len(seen) == 1
+        assert seen[0]["prompt"]["multi_modal_data"]["audio"][0][0].shape == (16000,)
+        submitted = clients[0].add_request_calls[0][0]
+        assert submitted.resumable is False
+        assert not orchestrator.request_states[submitted.request_id].streaming.enabled
+    finally:
+        await _close(orchestrator, rpc_q)

tests/engine/test_orchestrator.py

@@ -22,6 +22,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.engine import OmniEngineCoreOutput
+from vllm_omni.engine.errors import NativeKVHandoffError
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
     AbortResultMessage,
@@ -310,7 +311,7 @@ def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
     return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None)
 
 
-def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs:
+def _terminal_engine_core_outputs(request_id: str, timestamp: float = 1.0) -> EngineCoreOutputs:
     return EngineCoreOutputs(
         outputs=[
             EngineCoreOutput(
@@ -319,7 +320,7 @@ def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs:
                 finish_reason=FinishReason.STOP,
             )
         ],
-        timestamp=1.0,
+        timestamp=timestamp,
         finished_requests={request_id},
     )
 
@@ -519,6 +520,7 @@ async def _enqueue_add_request(
     original_prompt,
     sampling_params_list,
     final_stage_id: int,
+    final_output_stage_ids: list[int] | None = None,
 ) -> None:
     orchestrator_fixture.request_sync_q.put_nowait(
         StageSubmissionMessage(
@@ -529,6 +531,7 @@ async def _enqueue_add_request(
             output_prompt_text=None,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
+            final_output_stage_ids=final_output_stage_ids,
             preprocess_ms=0.0,
             request_timestamp=time.time(),
             enqueue_ts=time.perf_counter(),
@@ -788,6 +791,10 @@ async def test_run_async_chunk(orchestrator_factory) -> None:
 
         stage1.push_engine_core_outputs(_engine_core_outputs("stage1-final", 3.0))
 
+        await _wait_for(
+            lambda: orchestrator_fixture.orchestrator.request_states["req-async"].pending_final_output is not None
+        )
+        stage0.push_engine_core_outputs(_engine_core_outputs("stage0-final", 3.1))
         output_msg = await _get_output_message(orchestrator_fixture)
 
         assert output_msg.request_id == "req-async"
@@ -1414,6 +1421,105 @@ async def test_handle_streaming_update_unknown_request_is_dropped() -> None:
     assert "req-unknown" not in orchestrator.request_states
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("native_kv", [False, True])
+async def test_add_request_attaches_native_kv_ticket_before_dispatch(mocker, native_kv) -> None:
+    sampling = SamplingParams(extra_args={"keep": "sampling"})
+    prompt = SimpleNamespace(sampling_params=SamplingParams(extra_args={"keep": "prompt"}))
+    source = StagePool(
+        0,
+        FakeStageClient(),
+        stage_vllm_config=SimpleNamespace(
+            kv_transfer_config=SimpleNamespace(kv_role="kv_producer") if native_kv else None
+        ),
+    )
+    target = StagePool(1, FakeStageClient(stage_type="diffusion", final_output=True))
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=[source, target],
+    )
+
+    async def check_dispatch(request_id, req_state, request, **kwargs):
+        assert req_state.native_kv_transfer_id == ("xfer-req" if native_kv else None)
+        for params, preserved in ((sampling, "sampling"), (request.sampling_params, "prompt")):
+            assert params.extra_args["keep"] == preserved
+            if native_kv:
+                assert params.extra_args["kv_transfer_params"] == {
+                    "transfer_id": "xfer-req",
+                    "do_remote_decode": True,
+                    "do_remote_prefill": False,
+                }
+            else:
+                assert "kv_transfer_params" not in params.extra_args
+        return 0
+
+    dispatch = mocker.patch.object(source, "submit_initial", side_effect=check_dispatch)
+    await orchestrator._handle_add_request(
+        StageSubmissionMessage(
+            type="add_request",
+            request_id="req",
+            prompt=prompt,
+            original_prompt={},
+            output_prompt_text=None,
+            sampling_params_list=[sampling, OmniDiffusionSamplingParams()],
+            final_stage_id=1,
+            preprocess_ms=0,
+            request_timestamp=0,
+            enqueue_ts=0,
+        )
+    )
+    dispatch.assert_awaited_once()
+
+
+def test_native_handoff_uses_bound_replica_and_reports_missing_binding(mocker) -> None:
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.stage_pools = [mocker.Mock()]
+    orchestrator.stage_pools[0].get_bound_client.return_value = None
+    req_state = SimpleNamespace(native_kv_transfer_id="xfer-req")
+    output = SimpleNamespace(kv_transfer_params={"num_transfer_tokens": 4})
+
+    with pytest.raises(NativeKVHandoffError, match="bound AR replica"):
+        orchestrator._diffusion_submit_kwargs("req", 0, SimpleNamespace(engine_input_source=[0]), req_state, output)
+    orchestrator.stage_pools[0].get_bound_client.assert_called_once_with("req")
+    # The pool default points at replica 0; this request actually used replica 1.
+    orchestrator.stage_pools[0].stage_vllm_config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(engine_id="ar-0")
+    )
+    config = SimpleNamespace(engine_id="ar-1", kv_connector_extra_config={"bootstrap_addr": "http://host:8999"})
+    orchestrator.stage_pools[0].get_bound_client.return_value = SimpleNamespace(
+        vllm_config=SimpleNamespace(kv_transfer_config=config)
+    )
+    params = orchestrator._diffusion_submit_kwargs(
+        "req", 0, SimpleNamespace(engine_input_source=[0]), req_state, output
+    )
+    assert params["kv_transfer_params"]["remote_engine_id"] == "ar-1"
+    assert params["kv_transfer_params"]["remote_bootstrap_addr"] == "http://host:8999"
+    assert "remote_engine_id" not in output.kv_transfer_params
+
+
+@pytest.mark.asyncio
+async def test_native_handoff_failure_is_request_scoped(mocker):
+    orchestrator = object.__new__(Orchestrator)
+    fail = mocker.patch.object(orchestrator, "_fail_request_client_error", new_callable=mocker.AsyncMock)
+
+    def lost_binding():
+        raise NativeKVHandoffError("bound AR replica is unavailable")
+
+    assert not await orchestrator._dispatch_or_fail_request(
+        lost_binding, req_id="req", stage_id=1, operation="inter-stage forward"
+    )
+    fail.assert_awaited_once_with(
+        "req",
+        1,
+        "bound AR replica is unavailable",
+        status_code=502,
+        error_type="NativeKVHandoffError",
+        release_owners=True,
+    )
+
+
 async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
     class RecordingPool:
         def __init__(self) -> None:

tests/engine/test_orchestrator_kv_sender_info.py

@@ -414,3 +414,67 @@ def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas():
 
     assert all(isinstance(request, OmniEngineCoreRequest) for request in submitted.values())
     assert {key: request.payload_sender_info for key, request in submitted.items()} == endpoints
+
+
+@pytest.mark.asyncio
+async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker):
+    from vllm import PoolingParams, RequestOutput
+    from vllm.config import VllmConfig
+
+    from vllm_omni.config.model import OmniModelConfig
+    from vllm_omni.engine import OmniEngineCoreRequest
+
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.async_chunk = True
+    orchestrator._pd_pair = None
+    orchestrator._cfg_tracker = CfgCompanionTracker()
+    orchestrator.output_async_queue = asyncio.Queue()
+    orchestrator._duplex_output_decision = lambda *args: None
+    orchestrator._is_duplex_session_request = lambda state: False
+    orchestrator._forward_to_next_stage = mocker.AsyncMock()
+    orchestrator._cleanup_request_ids = mocker.AsyncMock()
+    # A sender stage stands in for Code2Wav, whose audio is also a final output.
+    audio_pool = _build_sender_pool(0, {})
+    audio_pool.stage_client.final_output = True
+    aligner = _DummyDiffusionStage()
+    aligner_pool = StagePool(
+        1,
+        aligner,
+        stage_vllm_config=mocker.Mock(
+            spec=VllmConfig, model_config=mocker.Mock(spec=OmniModelConfig, async_chunk=False)
+        ),
+    )
+    orchestrator.stage_pools = [audio_pool, aligner_pool]
+    state = OrchestratorRequestState(
+        request_id="aligned",
+        prompt={"additional_information": {"text": ["Hello world"]}},
+        sampling_params_list=[SamplingParams(), PoolingParams(task="token_classify")],
+        final_stage_id=1,
+        final_output_stage_ids={0, 1},
+    )
+
+    await orchestrator._prewarm_async_chunk_stages(
+        "aligned", mocker.Mock(spec=OmniEngineCoreRequest, prompt_token_ids=[1, 2]), state
+    )
+    assert aligner.calls == []
+    assert 1 not in state.stage_submit_ts
+
+    output = RequestOutput("aligned", None, [1, 2], None, [], False)
+    await orchestrator._route_output(0, 0, output, state, None)
+    orchestrator._forward_to_next_stage.assert_not_awaited()
+    output.finished = True
+    await orchestrator._route_output(0, 0, output, state, None)
+    orchestrator._forward_to_next_stage.assert_awaited_once_with(
+        "aligned",
+        0,
+        output,
+        state,
+        src_replica_id=0,
+        is_streaming_session=False,
+        is_final_update=False,
+    )
+    orchestrator._cleanup_request_ids.assert_not_awaited()
+    await orchestrator._route_output(1, 0, output, state, None)
+    orchestrator._cleanup_request_ids.assert_awaited_once()
+    messages = [orchestrator.output_async_queue.get_nowait() for _ in range(3)]
+    assert [message.finished for message in messages] == [False, False, True]

tests/engine/test_orchestrator_monitor.py

@@ -65,6 +65,31 @@ def test_note_loop_and_flush_json(tmp_path, monkeypatch):
     assert payload["replicas"]["stage=0,replica=0"]["inflight"] == [10]
 
 
+@pytest.mark.parametrize("drain", [False, True])
+def test_dispatch_backlog_carries_across_windows(tmp_path, monkeypatch, drain):
+    out_path = tmp_path / "orch_monitor.json"
+    monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
+    now = [0.0]
+    monkeypatch.setattr(time, "monotonic", lambda: now[0])
+    monitor = OrchestratorMonitor(replica_sampler=lambda: {})
+    monitor.set_dispatch_queue_size(3)
+    monitor.set_dispatch_queue_size(1)
+
+    now[0] = monitor_mod._WINDOW_S
+    monitor.note_loop(idle=False)
+    if drain:
+        monitor.set_dispatch_queue_size(0)
+    now[0] += monitor_mod._WINDOW_S
+    monitor.note_loop(idle=False)
+    now[0] += monitor_mod._WINDOW_S
+    monitor.flush()
+
+    windows = json.loads(out_path.read_text())["windows"]
+    remaining = 0 if drain else 1
+    assert windows["dispatch_queue_size"] == [1, remaining, remaining]
+    assert windows["dispatch_queue_high_water"] == [3, 1, remaining]
+
+
 def test_register_replica_backfills_prior_windows(tmp_path, monkeypatch):
     out_path = tmp_path / "orch_monitor.json"
     monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))

tests/engine/test_output_processor.py

@@ -259,6 +259,17 @@ def test_native_text_tpot_only_accepts_positive_finished_metric(
     assert processor.pop_native_text_metrics("r")["vllm_tpot_ms"] == expected_tpot_ms
 
 
+def test_suppressed_native_tokens_use_authoritative_count_even_without_stats(monkeypatch):
+    monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
+    processor = object.__new__(MultimodalOutputProcessor)
+    processor._native_text_metrics_by_request = {}
+    state = _make_state(RequestOutputKind.FINAL_ONLY)
+    processor._update_stats_from_output(
+        state, SimpleNamespace(new_token_ids=[2150], num_generation_tokens=2048), None, None
+    )
+    assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 2048
+
+
 def test_delta_drains_output_modality_per_step():
     """DELTA drains the mm_type key (output modality) but preserves hidden-state keys."""
     s = _make_state(RequestOutputKind.DELTA)
@@ -616,6 +627,7 @@ def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypa
         finish_reason=None,
         stop_reason=None,
         kv_transfer_params=None,
+        ec_transfer_params=None,
         routed_experts=None,
         num_cached_tokens=0,
     )
@@ -631,6 +643,7 @@ def _make_mm_only_output_processor(monkeypatch):
     processor = object.__new__(MultimodalOutputProcessor)
     processor.output_modality = OutputModality.AUDIO
     processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)}
+    processor.tracing_enabled = False
     monkeypatch.setattr(
         VLLMOutputProcessor,
         "process_outputs",
@@ -660,6 +673,7 @@ def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool):
         finish_reason=FinishReason.STOP,
         stop_reason=None,
         kv_transfer_params=None,
+        ec_transfer_params=None,
         routed_experts=None,
         num_cached_tokens=0,
         is_segment_finished=is_segment_finished,
@@ -888,3 +902,24 @@ def test_abort_snapshot_leaves_state_until_commit():
     assert "0_parent" not in processor.request_states
     assert "parent" not in processor.parent_requests
     assert "parent" not in processor.external_req_ids
+
+
+def test_cumulative_audio_terminal_output_supplies_full_aligner_waveform():
+    from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner
+
+    state = _make_state(RequestOutputKind.CUMULATIVE)
+    for index, count in enumerate((800, 1200, 400)):
+        state.add_multimodal_tensor(torch.full((1, count), float(index)), mm_type=AUDIO)
+        output = state.make_request_output([index], None, FinishReason.STOP if index == 2 else None, None)
+        assert output is not None
+        if index < 2:
+            assert code2wav2aligner([output], {}) == []
+    prompt = {"additional_information": {"text": ["Hello world"], "language": ["English"]}}
+    aligned = code2wav2aligner([output], prompt)[0]
+    waveform, sr = aligned["multi_modal_data"]["audio"]
+    assert sr == 24000
+    assert waveform.shape == (2400,)
+    assert (waveform[:800] == 0).all()
+    assert (waveform[800:2000] == 1).all()
+    assert (waveform[2000:] == 2).all()
+    assert aligned["additional_information"]["aligner_audio_duration_ms"] == [100.0]

tests/engine/test_output_processor_mrv2_text.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Regression tests for MRv2 Omni AR text detach and generation payload outputs."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.engine import OmniEngineCoreOutput
+from vllm_omni.engine.output_modality import OutputModality
+from vllm_omni.engine.output_processor import (
+    MultimodalOutputProcessor,
+    OmniRequestState,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Detokenizer:
+    def __init__(self):
+        self.output_token_ids = []
+        self.output_text = ""
+        self.update_calls = []
+
+    def update(self, token_ids, stop):
+        self.update_calls.append((list(token_ids), stop))
+        self.output_token_ids.extend(token_ids)
+        self.output_text += "X" * len(token_ids)
+        return None
+
+    def get_next_output_text(self, finished, delta):
+        return self.output_text
+
+    def num_output_tokens(self):
+        return len(self.output_token_ids)
+
+
+def _make_processor(output_modality, detokenizer=None):
+    state = OmniRequestState(
+        request_id="r",
+        external_req_id="r",
+        parent_req=None,
+        request_index=0,
+        lora_request=None,
+        output_kind=RequestOutputKind.FINAL_ONLY,
+        prompt="prompt" if detokenizer else None,
+        prompt_token_ids=[1] if detokenizer else [],
+        prompt_embeds=None,
+        logprobs_processor=MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None)
+        if detokenizer
+        else None,
+        detokenizer=detokenizer,
+        max_tokens_param=None,
+        arrival_time=0.0,
+        queue=None,
+        log_stats=False,
+        stream_interval=1,
+    )
+    processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False, output_modality=output_modality)
+    processor.request_states["r"] = state
+    processor.external_req_ids["r"].append("r")
+    return processor, state
+
+
+def _rehydrated_output(*, pooling_output=None, finish_reason=FinishReason.STOP):
+    """Build the post-StagePool shape without violating vLLM's wire schema."""
+    return SimpleNamespace(
+        request_id="r",
+        new_token_ids=[42],
+        pooling_output=pooling_output,
+        finish_reason=finish_reason,
+        stop_reason=None,
+        kv_transfer_params=None,
+        ec_transfer_params=None,
+        routed_experts=None,
+        trace_headers=None,
+        prefill_stats=None,
+        spec_decode_metrics=None,
+        new_sampling_mask=None,
+        num_nans_in_logits=0,
+        multimodal_output=None,
+        is_segment_finished=False,
+        is_non_final_audio_chunk=False,
+        output_type=None,
+        num_generation_tokens=None,
+        finished=True,
+    )
+
+
+def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload():
+    detokenizer = _Detokenizer()
+    processor, state = _make_processor(OutputModality.LATENT, detokenizer)
+
+    # StagePool has already rehydrated the MRv2 carrier.  Use a structural
+    # output here because upstream EngineCoreOutput strictly validates
+    # pooling_output as torch.Tensor | None on construction.
+    output = _rehydrated_output(pooling_output={"hidden": torch.ones(1, 4)})
+
+    processed = processor.process_outputs([output])
+
+    assert detokenizer.update_calls == [([42], True)]
+    completion = processed.request_outputs[0].outputs[0]
+    assert list(completion.token_ids) == [42]
+    assert completion.text == "X"
+    assert not state.mm_accumulated.is_empty
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_generation_stage_mm_delivered_only_with_terminal(streaming):
+    # Non-terminal chunks accumulate; the terminal output carries the audio.
+    # When the first output is already terminal it is delivered directly.
+    processor, _ = _make_processor(OutputModality.AUDIO)
+    audio = torch.ones(1, 320)
+    sr = torch.tensor(24000, dtype=torch.int32)
+    mm_output = OmniEngineCoreOutput(
+        request_id="r",
+        new_token_ids=[],
+        multimodal_output={"model_outputs": audio, "sr": sr},
+        finish_reason=None if streaming else FinishReason.STOP,
+    )
+
+    processed = processor.process_outputs([mm_output])
+    if streaming:
+        assert processed.request_outputs == []
+        processed = processor.process_outputs(
+            [OmniEngineCoreOutput(request_id="r", new_token_ids=[], finish_reason=FinishReason.STOP)]
+        )
+
+    assert len(processed.request_outputs) == 1
+    completion = processed.request_outputs[0].outputs[0]
+    assert completion.text == ""
+    assert torch.equal(completion.multimodal_output["audio"], audio)
+    assert completion.multimodal_output["sr"].item() == 24000

tests/engine/test_parallel_stage_init.py

@@ -306,6 +306,23 @@ def test_check_admission_admits_diffusion_with_explicit_util():
     assert ledgers[0].kv_budget_bytes == int(cap * 0.3)
 
 
+def test_check_admission_reads_typed_diffusion_cache_config():
+    cap = 80 * 1024**3
+    diff = _local_diffusion_replica(0, 0, "0", util=None)
+    diff.stage_cfg.cache_config = types.SimpleNamespace(gpu_memory_utilization=0.35)
+    plan = LogicalStageInitPlan(stage_idx=0, stage_id=0, replicas=[diff])
+
+    ledgers = check_admission(
+        [plan],
+        resolve_physical_devices=lambda _r: [0],
+        device_total_memory=lambda _d: cap,
+        external_reserve_bytes=0,
+        safety_margin_bytes=0,
+    )
+
+    assert ledgers[0].kv_budget_bytes == int(cap * 0.35)
+
+
 def test_check_admission_unresolved_local_raises():
     """A local replica invisible to the ledger must fail admission, not skip it:
     it would otherwise initialize concurrently with unbounded demand."""

tests/engine/test_request_snapshot.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import copy
+
+import pytest
+
+from vllm_omni.engine.request_snapshot import copy_request_snapshot
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_full_snapshot_preserves_fields_aliases_and_isolation():
+    samples = [0.25, 0.5]
+    source = {"prompt_token_ids": [1, 2], "additional_information": {"samples": samples}, "alias": samples}
+    result = copy_request_snapshot(source)
+    assert result == copy.deepcopy(source)
+    assert result["alias"] is result["additional_information"]["samples"]
+    result["alias"].append(0.75)
+    result["prompt_token_ids"][0] = 9
+    assert samples == [0.25, 0.5]
+    assert source["prompt_token_ids"] == [1, 2]

tests/engine/test_single_stage_mode.py

@@ -16,7 +16,13 @@ from pytest_mock import MockerFixture
 from vllm.v1.engine.utils import EngineZmqAddresses
 
 from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, PipelineConfig, StagePipelineConfig
+from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig
+from vllm_omni.config.stage_config import (
+    DuplexSessionRuntimeConfig,
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
 from vllm_omni.engine.stage_engine_startup import (
@@ -57,6 +63,10 @@ def _make_llm_plan(
     vllm_config: Any | None = None,
 ) -> LogicalStageInitPlan:
     stage_cfg = _make_stage_cfg(stage_id)
+    if launch_mode == "remote":
+        stage_cfg = VllmOmniARStageConfig(
+            stage_pipeline_config=StagePipelineConfig(stage_id=stage_id, model_stage="thinker"),
+        )
     metadata = SimpleNamespace(
         stage_id=stage_id,
         stage_type="llm",
@@ -98,6 +108,16 @@ def _make_diffusion_plan(
     launch_mode: str,
 ) -> LogicalStageInitPlan:
     stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion")
+    if launch_mode == "remote":
+        stage_cfg = VllmOmniDiffusionStageConfig(
+            stage_pipeline_config=StagePipelineConfig(
+                stage_id=stage_id,
+                model_stage="diffusion",
+                execution_type=StageExecutionType.DIFFUSION,
+                final_output=True,
+                final_output_type="image",
+            ),
+        )
     metadata = SimpleNamespace(
         stage_id=stage_id,
         stage_type="diffusion",
@@ -983,6 +1003,7 @@ class TestSingleStageReplicaInitialization:
         assert mock_connect.call_args.kwargs["stage_id"] == 1
         assert mock_connect.call_args.kwargs["replica_id"] == 0
         assert client_kwargs["log_stats"] is True
+        assert client_kwargs["metadata"].model_stage == "thinker"
         assert events == ["enter", "exit", "attach"]
 
     def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture):
@@ -1148,11 +1169,9 @@ class TestSingleStageReplicaInitialization:
 
             return _ctx()
 
-        remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata
         plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0]
         sentinel_client = SimpleNamespace()
 
-        mocker.patch.object(runtime_mod, "extract_legacy_stage_metadata", return_value=remote_metadata)
         mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect)
         mock_from_addresses = mocker.patch(
             "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses",
@@ -1169,6 +1188,7 @@ class TestSingleStageReplicaInitialization:
             replica_id=0,
         )
         mock_from_addresses.assert_called_once()
+        assert mock_from_addresses.call_args.args[0].final_output_type == "image"
 
     def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture):
         import vllm_omni.engine.stage_runtime as runtime_mod
@@ -1401,7 +1421,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             yield mocker.Mock()
 
         mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx())
-        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with connect_remote_engine_cores(
             vllm_config=vllm_config,
             omni_master_server=omni_master_server,
@@ -1436,7 +1456,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             yield mocker.Mock()
 
         mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx())
-        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with connect_remote_engine_cores(
             vllm_config=vllm_config,
             omni_master_server=omni_master_server,
@@ -1462,7 +1482,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             yield mocker.Mock()
 
         mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx())
-        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
 
         with connect_remote_diffusion_proc(
             omni_master_server=omni_master_server,
@@ -1474,7 +1494,8 @@ class TestConnectRemoteEngineCoresCoordinator:
         omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2)
         omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2)
         mock_wait.assert_called_once()
-        _, core_engines, parallel_config, *_ = mock_wait.call_args.args
+        _, core_engines, parallel_config, _, _, launch = mock_wait.call_args.args
+        assert launch.addresses is omni_master_server.get_zmq_addresses.return_value
         assert core_engines[0].local is False
         assert parallel_config.data_parallel_size_local == 0
 
@@ -1515,7 +1536,7 @@ class TestLaunchOmniCoreEngines:
             "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager",
             return_value=local_engine_manager,
         )
-        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with _launch_omni_core_engines(
             vllm_config=vllm_config,
             executor_class=mocker.Mock(),
@@ -1592,7 +1613,7 @@ class TestLaunchOmniCoreEngines:
             "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager",
             return_value=mocker.Mock(),
         )
-        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with _launch_omni_core_engines(
             vllm_config=vllm_config,
             executor_class=mocker.Mock(),

tests/engine/test_stage_device_layout.py

@@ -18,6 +18,7 @@ import types
 from unittest import mock
 
 import pytest
+from pytest_mock import MockerFixture
 
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.stage_init_utils import (
@@ -156,6 +157,43 @@ def test_replica_guard_and_splitter_share_local_world_size(devices, expected):
     assert replica_devices_map == {0: expected}
 
 
+@pytest.mark.parametrize("typed", [False, True], ids=["legacy", "typed"])
+@pytest.mark.parametrize("local_dp", [None, 0, 1, 2])
+@pytest.mark.parametrize("pool_devices", [False, True], ids=["template", "pool"])
+def test_replica_layout_preserves_local_dp_semantics(typed, local_dp, pool_devices):
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig
+    from vllm_omni.engine.stage_engine_startup import get_headless_replica_devices
+
+    engine_args = {
+        "tensor_parallel_size": 2,
+        "data_parallel_size": 4,
+        "data_parallel_size_local": local_dp,
+        "pipeline_parallel_size": 2,
+    }
+    width = 2 * max(1, local_dp if local_dp is not None else 4) * 2
+    devices = ",".join(str(i) for i in range(width * (2 if pool_devices else 1)))
+    if typed:
+        pipeline = PipelineConfig(
+            model_type="test",
+            model_arch="TestModel",
+            stages=(StagePipelineConfig(stage_id=0, model_stage="ar", final_output=True),),
+        )
+        stage = VllmOmniConfig.from_pipeline_config(
+            pipeline,
+            cli_overrides={
+                f"stage_0_{key}": value for key, value in {**engine_args, "devices": devices, "num_replicas": 2}.items()
+            },
+        ).stage_configs[0]
+    else:
+        stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args)
+    expected = [",".join(str(i) for i in range(r * width, (r + 1) * width)) for r in range(2)]
+
+    assert get_stage_devices_per_replica(stage) == width
+    assert compute_replica_layout([stage]) == ([2], {0: expected})
+    assert get_headless_replica_devices(stage, stage_id=0, omni_dp_size_local=2) == expected
+
+
 def test_non_tp_layout_error_uses_generic_guidance():
     """A PP mismatch must not be reported as a top-level TP broadcast."""
     stage = _stage(0, devices="0")
@@ -203,7 +241,7 @@ def test_build_vllm_config_fails_before_engine_config_on_mismatch():
     get_class.assert_not_called()
 
 
-def test_build_vllm_config_proceeds_on_consistent_layout():
+def test_build_vllm_config_proceeds_on_consistent_layout(mocker: MockerFixture) -> None:
     """The guard must not false-positive: a consistent single-GPU stage flows past
     it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual."""
     stage = _stage(1, devices="1")
@@ -212,18 +250,30 @@ def test_build_vllm_config_proceeds_on_consistent_layout():
         model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()),
     )
     sentinel_executor = object()
-    with (
-        mock.patch.object(
-            stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config
-        ) as create_engine_config,
-        mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor),
-        mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant),
-    ):
-        vllm_config, executor_class = build_vllm_config(
-            stage,
-            model="dummy-model",
-            engine_args_dict={"tensor_parallel_size": 1},
-        )
+    captured_api_process_config: dict[str, int] = {}
+
+    def capture_create_engine_config(engine_args, *args, **kwargs):
+        captured_api_process_config["count"] = engine_args._api_process_count
+        captured_api_process_config["rank"] = engine_args._api_process_rank
+        return fake_config
+
+    create_engine_config = mocker.patch.object(
+        stage_init_utils.OmniEngineArgs,
+        "create_engine_config",
+        autospec=True,
+        side_effect=capture_create_engine_config,
+    )
+    mocker.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor)
+    mocker.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant)
+
+    vllm_config, executor_class = build_vllm_config(
+        stage,
+        model="dummy-model",
+        engine_args_dict={"tensor_parallel_size": 1},
+        api_process_count=3,
+        api_process_rank=2,
+    )
     create_engine_config.assert_called_once()
     assert vllm_config is fake_config
     assert executor_class is sentinel_executor
+    assert captured_api_process_config == {"count": 3, "rank": 2}

tests/engine/test_stage_engine_args.py

@@ -8,6 +8,7 @@ import sys
 import types
 from dataclasses import fields, replace
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 from pydantic.fields import FieldInfo
@@ -29,6 +30,7 @@ from vllm_omni.config.omni_config import (
     OmniStageSchedulerConfig,
     VllmOmniARStageConfig,
     VllmOmniConfig,
+    extract_diffusion_stage_config_kwargs,
 )
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
 from vllm_omni.config.stage_config import (
@@ -60,9 +62,14 @@ _DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
 def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, object]:
     backend_fields = fields(config_cls)
     backend_field_names = {backend_field.name for backend_field in backend_fields}
-    backend_config = config_cls(
-        **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names}
-    )
+    # This is a schema/normalization comparison, not a Hub download test.
+    with (
+        patch("vllm.engine.arg_utils.get_model_path", side_effect=lambda model, *_: model),
+        patch("vllm_omni.diffusion.data.get_model_path", side_effect=lambda model, *_: model),
+    ):
+        backend_config = config_cls(
+            **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names}
+        )
     effective_values: dict[str, object] = {}
     for backend_field in backend_fields:
         value = getattr(backend_config, backend_field.name)
@@ -77,6 +84,9 @@ _DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusi
 _TOPOLOGY_ONLY_ENGINE_ARGS = frozenset({"inline_diffusion"})
 _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
     {
+        "final_output",
+        "use_v2_model_runner",
+        "supports_native_mrv2_data_plane",
         "active_stream_window",
         "codec_frame_rate_hz",
         "custom_voice_dir",
@@ -95,6 +105,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
         "num_weight_load_threads",
         "omni_kv_config",
         "parallel_config",
+        "pooling_output_decoder",
         "silence_ban_frames",
         "subtalker_sampling_params",
         "task_type",
@@ -366,6 +377,21 @@ def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
     assert inherited_defaults.isdisjoint(engine_args)
 
 
+def test_typed_llm_projection_omits_diffusion_only_and_process_only_defaults():
+    stage_config = VllmOmniARStageConfig(
+        stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+    )
+
+    engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+    assert {
+        "enable_multithread_weight_load",
+        "num_weight_load_threads",
+        "disable_autocast",
+        "log_level",
+    }.isdisjoint(engine_args)
+
+
 def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
     monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
 
@@ -551,7 +577,9 @@ def test_engine_args_consume_stage_diffusion_attention_shorthand(tmp_path):
         assert engine_args.get("diffusion_attention_backend") is None
         assert isinstance(engine_args["diffusion_attention_config"], AttentionConfig)
         assert engine_args["diffusion_attention_config"].default.backend == "TORCH_SDPA"
-        od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+        od_config = OmniDiffusionConfig.from_kwargs(
+            **extract_diffusion_stage_config_kwargs(engine_args, stage_id=2, include_engine_adapter_metadata=True)
+        )
         assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
 
 
@@ -584,7 +612,9 @@ def test_engine_args_apply_cli_attention_shorthand_over_yaml_config(tmp_path, ya
         attention_config = engine_args["diffusion_attention_config"]
         assert attention_config.default.backend == "TORCH_SDPA"
         assert attention_config.per_role["cross"].backend == "SAGE_ATTN"
-        od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
+        od_config = OmniDiffusionConfig.from_kwargs(
+            **extract_diffusion_stage_config_kwargs(engine_args, stage_id=2, include_engine_adapter_metadata=True)
+        )
         assert od_config.diffusion_attention_config.default.backend == "TORCH_SDPA"
 
 

tests/engine/test_terminal_stage_metrics_order.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Terminal audio must not retire a request before upstream terminal metrics."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def build():
+    o = Orchestrator.__new__(Orchestrator)
+    o.stage_pools = [
+        SimpleNamespace(
+            final_output=False, build_stage_metrics=lambda *a, **k: SimpleNamespace(stage_id=0, num_tokens_out=42)
+        ),
+        SimpleNamespace(
+            final_output=True, build_stage_metrics=lambda *a, **k: SimpleNamespace(stage_id=1, num_tokens_out=0)
+        ),
+    ]
+    state = OrchestratorRequestState(
+        request_id="r",
+        final_stage_id=1,
+        final_output_stage_ids={1},
+        stage_submit_ts={0: 0, 1: 0},
+        sampling_params_list=[None, None],
+    )
+    o.request_states = {"r": state}
+    o.output_async_queue = asyncio.Queue()
+    o.async_chunk = True
+    o._pd_pair = None
+    o._cfg_tracker = MagicMock()
+    o._cfg_tracker.is_companion.return_value = False
+    o._cfg_tracker.cleanup_parent.return_value = []
+    o._is_duplex_session_request = lambda s: False
+    o._duplex_output_decision = lambda *a: None
+    o._stage_receives_async_chunks = lambda s: True
+    cleaned = []
+
+    async def cleanup(ids):
+        cleaned.extend(ids)
+        for r in ids:
+            o.request_states.pop(r, None)
+
+    o._cleanup_request_ids = cleanup
+    return o, state, cleaned
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("order", [(0, 1), (1, 0)])
+async def test_terminal_metrics_precede_final_audio_for_either_arrival_order(order):
+    o, state, cleaned = build()
+    for stage in order:
+        await o._handle_processed_outputs(stage, 0, [SimpleNamespace(request_id="r", finished=True, error=None)])
+        if stage == 1 and not state.finished_stage_ids.issuperset({0, 1}):
+            assert o.output_async_queue.empty()
+            assert "r" in o.request_states
+            assert cleaned == []
+    metrics = o.output_async_queue.get_nowait()
+    final = o.output_async_queue.get_nowait()
+    assert type(metrics).__name__ == "StageMetricsMessage"
+    assert metrics.stage_id == 0 and metrics.metrics.num_tokens_out == 42
+    assert type(final).__name__ == "OutputMessage" and final.finished
+    assert o.output_async_queue.empty()
+    assert cleaned == ["r"]
+    assert state.pending_final_output is None
+
+
+@pytest.mark.asyncio
+async def test_cancelled_pending_output_is_not_published_by_late_upstream():
+    o, state, cleaned = build()
+    await o._handle_processed_outputs(1, 0, [SimpleNamespace(request_id="r", finished=True, error=None)])
+    assert state.pending_final_output is not None
+    await o._cleanup_request_ids(["r"])
+    await o._handle_processed_outputs(0, 0, [SimpleNamespace(request_id="r", finished=True, error=None)])
+    assert o.output_async_queue.empty()
+    assert cleaned == ["r"]