Merge main and preserve event-driven defaults in shared engine bases

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

tests/benchmarks/duplex/test_omni_duplex_eval.py

@@ -64,6 +64,177 @@ def test_hugging_face_subset_names_are_splits(monkeypatch):
     assert sample.task_type == "correction"
 
 
+def test_local_hf_dataset_directory_is_routed_to_hf_loader(tmp_path, monkeypatch):
+    # An existing local Hugging Face dataset layout (directory with a
+    # ``data/*.parquet`` file) must be loaded through ``datasets.load_dataset``
+    # instead of being parsed as a JSON/JSONL manifest.
+    hf_dir = tmp_path / "Omni-DuplexEval"
+    data_dir = hf_dir / "data"
+    data_dir.mkdir(parents=True)
+    (data_dir / "train-00000-of-00001.parquet").write_bytes(b"PAR1")
+    calls = []
+
+    def fake_load_dataset(name, **kwargs):
+        calls.append((name, kwargs))
+        return [{"id": "pr", "split": "PR_correction", "question_text": "Correct this."}]
+
+    monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+    samples = load_samples(hf_dir)
+    assert calls == [(str(hf_dir), {})]
+    assert len(samples) == 1
+    assert samples[0].task_type == "correction"
+
+
+def test_local_parquet_file_is_routed_to_hf_loader(tmp_path, monkeypatch):
+    parquet = tmp_path / "samples.parquet"
+    parquet.write_bytes(b"PAR1")
+    calls = []
+
+    def fake_load_dataset(name, **kwargs):
+        calls.append((name, kwargs))
+        return [{"id": "pr", "split": "PR_correction", "question_text": "Correct this."}]
+
+    monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+    samples = load_samples(str(parquet))
+    assert calls == [("parquet", {"data_files": str(parquet)})]
+    assert len(samples) == 1
+    assert samples[0].family == "pr"
+
+
+def test_hf_collapse_to_train_keeps_row_split_identity(tmp_path, monkeypatch):
+    # datasets.load_dataset collapses a config-less local mirror (a bare data/
+    # directory) into a single generic ``train`` split. Rows that already carry
+    # their own split/subset identity must keep it so family inference is not
+    # poisoned by the generic name (regression for the direct-read path).
+    hf_dir = tmp_path / "Omni-DuplexEval"
+    data_dir = hf_dir / "data"
+    data_dir.mkdir(parents=True)
+    (data_dir / "train-00000-of-00021.parquet").write_bytes(b"PAR1")
+
+    def fake_load_dataset(name, **kwargs):
+        return {
+            "train": [
+                {"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
+                {"id": "pr-subset", "subset": "PR_event_reminder", "question_text": "Remind me."},
+                {"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
+            ]
+        }
+
+    monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+    samples = load_samples(hf_dir)
+    by_id = {sample.id: sample for sample in samples}
+    assert by_id["pr"].split == "PR_correction"
+    assert by_id["pr"].family == "pr"
+    assert by_id["pr"].task_type == "correction"
+    assert by_id["pr-subset"].split == "PR_event_reminder"
+    assert by_id["pr-subset"].family == "pr"
+    assert by_id["pr-subset"].task_type == "proactive_reminder"
+    assert by_id["rtd"].split == "RTD_OCR"
+    assert by_id["rtd"].family == "rtd"
+
+
+def test_hf_config_layout_stamps_split_name_when_row_has_none(tmp_path, monkeypatch):
+    # With a proper config-per-split mirror the loader still stamps the Hugging
+    # Face split name onto rows that do not identify themselves (unchanged).
+    hf_dir = tmp_path / "Omni-DuplexEval"
+    data_dir = hf_dir / "data"
+    data_dir.mkdir(parents=True)
+    (data_dir / "RTD_OCR-00000-of-00001.parquet").write_bytes(b"PAR1")
+    (data_dir / "PR_correction-00000-of-00001.parquet").write_bytes(b"PAR1")
+
+    def fake_load_dataset(name, **kwargs):
+        return {
+            "RTD_OCR": [{"id": "rtd", "question_text": "Read this."}],
+            "PR_correction": [{"id": "pr", "question_text": "Correct this."}],
+        }
+
+    monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+    samples = load_samples(hf_dir)
+    by_id = {sample.id: sample for sample in samples}
+    assert by_id["rtd"].split == "RTD_OCR"
+    assert by_id["rtd"].family == "rtd"
+    assert by_id["pr"].split == "PR_correction"
+    assert by_id["pr"].family == "pr"
+    assert by_id["pr"].task_type == "correction"
+
+
+def test_hf_collapse_to_train_without_row_identity_raises_actionable_error(tmp_path, monkeypatch):
+    # A row that carries no split/family/task_type and lands in a generic
+    # collapsed ``train`` split cannot be routed; the loader must raise a clear,
+    # actionable ValueError instead of the bare family_for_split message.
+    hf_dir = tmp_path / "Omni-DuplexEval"
+    data_dir = hf_dir / "data"
+    data_dir.mkdir(parents=True)
+    (data_dir / "train-00000-of-00021.parquet").write_bytes(b"PAR1")
+
+    def fake_load_dataset(name, **kwargs):
+        return {"train": [{"id": "x", "question_text": "What is this?"}]}
+
+    monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+    with pytest.raises(ValueError, match="Hugging Face splits"):
+        load_samples(hf_dir)
+
+
+def test_non_manifest_file_raises_clear_error(tmp_path):
+    # A plain non-JSON file (e.g. a README) must surface a clear ValueError
+    # listing the accepted inputs instead of a raw JSONDecodeError.
+    bad = tmp_path / "README.md"
+    bad.write_text("# readme\nnot a manifest", encoding="utf-8")
+    with pytest.raises(ValueError, match="JSON/JSONL manifest"):
+        load_samples(str(bad))
+
+
+def _write_parquet(path, rows):
+    pyarrow = pytest.importorskip("pyarrow")
+    pyarrow.parquet.write_table(pyarrow.Table.from_pylist(rows), str(path))
+
+
+def test_local_parquet_split_filters_rows_by_identity(tmp_path):
+    # Regression: a single ``.parquet`` file is exposed by ``load_dataset`` as a
+    # generic ``train`` physical split, so ``--split RTD_OCR`` used to raise
+    # ``Unknown split "RTD_OCR"``. The requested split must instead filter the
+    # rows by their preserved identity.
+    pytest.importorskip("datasets")
+    parquet = tmp_path / "samples.parquet"
+    _write_parquet(
+        parquet,
+        [
+            {"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
+            {"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
+        ],
+    )
+    assert {sample.id for sample in load_samples(str(parquet))} == {"rtd", "pr"}
+    rtd = load_samples(str(parquet), split="RTD_OCR")
+    assert [sample.id for sample in rtd] == ["rtd"]
+    assert rtd[0].family == "rtd"
+    pr = load_samples(str(parquet), split="PR_correction")
+    assert [sample.id for sample in pr] == ["pr"]
+    assert pr[0].family == "pr"
+    assert pr[0].task_type == "correction"
+
+
+def test_local_directory_named_splits_resolve_without_row_identity(tmp_path):
+    # Regression: a single-configuration mirror whose data files are named after
+    # the benchmark splits must resolve each split for real (no fake loader) and
+    # without any per-row identity.
+    pytest.importorskip("datasets")
+    hf_dir = tmp_path / "Omni-DuplexEval"
+    data_dir = hf_dir / "data"
+    data_dir.mkdir(parents=True)
+    _write_parquet(data_dir / "RTD_OCR-00000-of-00001.parquet", [{"id": "rtd", "question_text": "Read this."}])
+    _write_parquet(
+        data_dir / "PR_correction-00000-of-00001.parquet",
+        [{"id": "pr", "question_text": "Correct this."}],
+    )
+    by_id = {sample.id: sample for sample in load_samples(hf_dir)}
+    assert by_id["rtd"].split == "RTD_OCR"
+    assert by_id["rtd"].family == "rtd"
+    assert by_id["pr"].split == "PR_correction"
+    assert by_id["pr"].family == "pr"
+    assert by_id["pr"].task_type == "correction"
+    assert [sample.id for sample in load_samples(hf_dir, split="RTD_OCR")] == ["rtd"]
+
+
 def test_response_aliases_and_clock_guard():
     assert split_text("One. Two!") == ["One.", "Two!"]
     assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
@@ -141,3 +312,113 @@ def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
         judge_video_mode="frame-sample",
     )
     assert score["content"]["frame_count"] == 2
+
+
+def test_local_directory_collapsed_train_with_row_identity_and_split_filter(tmp_path):
+    # Reviewer-pointed gap: a collapsed local mirror (data/train-*.parquet)
+    # whose rows carry their own split identity must correctly filter by
+    # --split RTD_OCR.
+    pytest.importorskip("datasets")
+    hf_dir = tmp_path / "Omni-DuplexEval"
+    data_dir = hf_dir / "data"
+    data_dir.mkdir(parents=True)
+    _write_parquet(
+        data_dir / "train-00000-of-00001.parquet",
+        [
+            {"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
+            {"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
+        ],
+    )
+    all_samples = load_samples(hf_dir)
+    assert {s.id for s in all_samples} == {"rtd", "pr"}
+    rtd = load_samples(hf_dir, split="RTD_OCR")
+    assert [s.id for s in rtd] == ["rtd"]
+    assert rtd[0].family == "rtd"
+
+
+def test_local_directory_mistyped_split_raises_clear_error(tmp_path):
+    # A named-shard mirror with --split RTD_TYPO must raise ValueError
+    # listing the observed splits instead of silently returning nothing.
+    pytest.importorskip("datasets")
+    hf_dir = tmp_path / "Omni-DuplexEval"
+    data_dir = hf_dir / "data"
+    data_dir.mkdir(parents=True)
+    _write_parquet(
+        data_dir / "RTD_OCR-00000-of-00001.parquet",
+        [{"id": "rtd", "question_text": "Read this."}],
+    )
+    _write_parquet(
+        data_dir / "PR_correction-00000-of-00001.parquet",
+        [{"id": "pr", "question_text": "Correct this."}],
+    )
+    with pytest.raises(ValueError) as exc_info:
+        load_samples(hf_dir, split="RTD_TYPO")
+    message = str(exc_info.value)
+    assert "matched zero" in message
+    assert "RTD_OCR" in message or "PR_correction" in message
+
+
+def test_manifest_without_split_identity_split_override_still_works(tmp_path):
+    # Manifest rows without split/subset/config identity must keep the
+    # base override behaviour: --split RTD_OCR stamps the split so rows
+    # are kept instead of rejected (regression test for P2).
+    manifest = tmp_path / "samples.json"
+    manifest.write_text(
+        json.dumps([{"id": "a", "video": "clip.mp4"}, {"id": "b", "video": "clip2.mp4"}]),
+        encoding="utf-8",
+    )
+    samples = load_samples(manifest, split="RTD_OCR")
+    assert len(samples) == 2
+    assert all(s.split == "RTD_OCR" for s in samples)
+    assert all(s.family == "rtd" for s in samples)
+
+
+def test_iterable_without_split_identity_split_override_still_works():
+    # Iterable rows without split/subset/config identity must keep the
+    # base override behaviour (regression test for F5/P2).
+    data = [{"id": "a", "video": "clip.mp4"}, {"id": "b", "video": "clip2.mp4"}]
+    samples = load_samples(data, split="RTD_OCR")
+    assert len(samples) == 2
+    assert all(s.split == "RTD_OCR" for s in samples)
+    assert all(s.family == "rtd" for s in samples)
+
+
+def test_audio_feature_decode_false_preserves_bytes_key(tmp_path):
+    # When a parquet file carries an Audio feature column, the loader
+    # must recast it with decode=False so rows contain the ``bytes`` key
+    # instead of ``array``/``sampling_rate``.
+    pytest.importorskip("datasets")
+    pytest.importorskip("pyarrow")
+    import datasets
+
+    audio_bytes = b"\x00\x01\x02"
+    ds = datasets.Dataset.from_dict(
+        {
+            "id": ["a"],
+            "question_text": ["Hello"],
+            "task_type": ["pr_correction"],
+            "audio": [{"bytes": audio_bytes, "path": None}],
+        },
+        features=datasets.Features(
+            {
+                "id": datasets.Value("string"),
+                "question_text": datasets.Value("string"),
+                "task_type": datasets.Value("string"),
+                "audio": datasets.Audio(sampling_rate=16000),
+            }
+        ),
+    )
+    parquet = tmp_path / "audio_sample.parquet"
+    ds.to_parquet(str(parquet))
+
+    samples = load_samples(str(parquet))
+    assert len(samples) == 1
+    raw = samples[0].raw
+    assert raw is not None
+    audio_val = raw.get("audio")
+    assert audio_val is not None
+    # The Audio feature was decoded with decode=False, so the struct
+    # keeps the ``bytes`` key that materialize_media needs.
+    assert isinstance(audio_val, dict) and "bytes" in audio_val
+    assert audio_val["bytes"] == audio_bytes
+    assert "path" in audio_val

tests/benchmarks/duplex/test_omni_duplex_eval_media.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_media import materialize_media
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dict_with_bytes_materializes_a_real_file(tmp_path):
+    # A Hugging Face Audio/Video feature value {"bytes": ..., "path": ...} used
+    # to short-circuit on ``path`` (the artifact name inside the dataset, which
+    # does not exist on disk), so downstream ffmpeg saw a missing file. Inline
+    # bytes must win and be written to a real file whose content equals bytes.
+    payload = b"\x00\xff\x01RIFFfake-content"
+    resolved = materialize_media({"bytes": payload, "path": "question_audio.wav"}, tmp_path, "s1_question", ".wav")
+    assert resolved.is_file()
+    assert resolved.read_bytes() == payload
+    assert resolved.parent == tmp_path
+    assert resolved.name == "s1_question.wav"
+
+
+def test_dict_extension_comes_from_path_safely(tmp_path):
+    payload = b"flac-ish-bytes"
+    resolved = materialize_media({"bytes": payload, "path": "audio/clip.FLAC"}, tmp_path, "s1", ".wav")
+    assert resolved.name == "s1.flac"
+    assert resolved.read_bytes() == payload
+    # A path with no usable extension falls back to the caller's suffix.
+    resolved = materialize_media({"bytes": payload, "path": "audio/clip"}, tmp_path, "s2", ".wav")
+    assert resolved.name == "s2.wav"
+    assert resolved.read_bytes() == payload
+
+
+def test_materialization_is_idempotent_and_updated(tmp_path):
+    payload = b"abc"
+    first = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
+    second = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
+    assert first == second
+    assert first.read_bytes() == payload
+    replacement = b"xyz"
+    third = materialize_media({"bytes": replacement, "path": "q.wav"}, tmp_path, "s", ".wav")
+    assert third == first
+    assert third.read_bytes() == replacement
+
+
+def test_str_path_and_path_only_mapping_returned_as_is(tmp_path):
+    existing = tmp_path / "clip.mp4"
+    existing.write_bytes(b"video")
+    assert materialize_media(str(existing), tmp_path, "s", ".mp4") == existing
+    assert materialize_media({"path": str(existing)}, tmp_path, "s", ".mp4") == existing
+
+
+def test_raw_bytes_value_is_written(tmp_path):
+    payload = b"raw-media"
+    resolved = materialize_media(payload, tmp_path, "r", ".bin")
+    assert resolved.name == "r.bin"
+    assert resolved.read_bytes() == payload
+
+
+def test_unresolvable_value_raises_clear_error(tmp_path):
+    with pytest.raises(ValueError, match="cannot materialize"):
+        materialize_media(None, tmp_path, "s", ".wav")
+    with pytest.raises(ValueError, match="cannot materialize"):
+        materialize_media({"path": None}, tmp_path, "s", ".wav")

tests/benchmarks/patch/test_patch.py

@@ -135,6 +135,7 @@ class MockResponse:
 async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch):
     class FakeRealtimeClient:
         last_instance = None
+        instances: list = []
 
         def __init__(self, url):
             assert url == "ws://localhost:8000/v1/realtime?duplex=1"
@@ -143,6 +144,9 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
             self.sent = []
             self.response_count = 0
             self.ack_count = 0
+            self.silence_seconds = []
+            self.closed = 0
+            type(self).instances.append(self)
             type(self).last_instance = self
 
         async def __aenter__(self):
@@ -155,10 +159,22 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
             assert model == "openbmb/MiniCPM-o-4_5"
             self.configure_kwargs = kwargs
 
+        async def stream_silence(self, *, seconds, chunk_ms=200, until=None):
+            # A model-native session speaks off its seeded context once audio
+            # units arrive; the silence itself carries no content. The probe
+            # stops the silence as soon as the turn settles.
+            self.silence_seconds.append(seconds)
+            self._emit_response()
+            assert until is not None and until()
+            return 1.0
+
         async def send(self, event):
             self.sent.append(event)
             if event["type"] != "response.create":
                 return
+            self._emit_response()
+
+        def _emit_response(self):
             self.response_count += 1
             response_id = f"resp-{self.response_count}"
             now = time.monotonic()
@@ -210,7 +226,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
             self.ack_count += 1
 
         async def close_session(self, **_kwargs):
-            return None
+            self.closed += 1
 
     monkeypatch.setattr(
         "vllm_omni.benchmarks.patch.patch._RealtimeTTSProbe",
@@ -239,30 +255,27 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
         session=None,
     )
 
-    client = FakeRealtimeClient.last_instance
-    assert client.configure_kwargs["native_duplex"] is False
-    assert client.configure_kwargs["extra_body"] == {
-        "ref_audio": "data:audio/wav;base64,AAAA",
-        "return_stage_metrics": True,
-    }
-    assert client.sent == [
-        event
-        for index in range(4)
-        for event in (
-            {
-                "type": "conversation.item.create",
-                "item": {
-                    "type": "message",
-                    "role": "user",
-                    "content": [{"type": "input_text", "text": f"text {index}"}],
-                },
-            },
-            {"type": "response.create"},
-        )
-    ]
-    assert client.ack_count == 4
+    # One session per utterance: a model-native duplex session takes its text
+    # once, in the session context, so it cannot be re-seeded for a second one.
+    sessions = FakeRealtimeClient.instances
+    assert len(sessions) == 4
+    for index, client in enumerate(sessions):
+        assert "native_duplex" not in client.configure_kwargs
+        assert client.configure_kwargs["extra_body"] == {
+            "ref_audio": "data:audio/wav;base64,AAAA",
+            "return_stage_metrics": True,
+            "duplex_initial_user_text": f"text {index}",
+            "force_listen_count": 0,
+        }
+        # The target text rides the session context, never a conversation item:
+        # a text-only response.create is rejected by a model-native session.
+        assert client.sent == []
+        assert client.silence_seconds and all(seconds > 0 for seconds in client.silence_seconds)
+        assert client.ack_count == 1
+        assert client.closed == 1
+
     assert output.success is True
-    assert output.generated_text == "turn 1 turn 2 turn 3 turn 4"
+    assert output.generated_text == "turn 1 turn 1 turn 1 turn 1"
     assert output.audio_duration == pytest.approx(0.4)
     assert output.ttft > 0
     assert output.audio_ttfp > output.ttft
@@ -281,11 +294,12 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
             "session_id": session_id,
             "request_index": index,
             "utterance_id": f"utt-{index}",
-            "response_id": f"resp-{index + 1}",
+            # Response ids are per session, and each utterance now gets its own.
+            "response_id": "resp-1",
             "source": "client_monotonic_receive",
             "measurement_origin": {
-                "ttft": "conversation.item.create client send to first non-empty text delta",
-                "ttfp": "conversation.item.create client send to first audio packet",
+                "ttft": "first silence append client send to first non-empty text delta",
+                "ttfp": "first silence append client send to first audio packet",
                 "rtf": "request-start-to-last-audio receive time divided by emitted audio duration",
                 "tpot": "Stage-0 engine mean time per output token",
             },

tests/benchmarks/test_omniinteract.py

@@ -467,6 +467,27 @@ def test_tolerated_playback_ack_rejection_is_a_warning_not_a_failure():
         oi._raise_if_session_terminated(collector, 0, warnings=warnings)
 
 
+def test_our_own_close_is_expected_even_though_the_server_stamps_a_reason():
+    """``session.close`` is answered with ``session.closed`` carrying ``client_close``.
+
+    The guard exists to catch a session that ended for a reason we did not ask
+    for. A close we requested is not that, whether or not the server names it.
+    """
+    collector = _collector(({"type": "session.closed", "reason": "client_close"}, 1.0))
+    oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
+
+    collector = _collector(({"type": "session.closed", "event": {"reason": "client_close"}}, 1.0))
+    oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
+
+    collector = _collector(({"type": "session.closed", "reason": "disconnect"}, 1.0))
+    with pytest.raises(RuntimeError, match="Unexpected session.closed: disconnect"):
+        oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
+
+    collector = _collector(({"type": "session.expired", "reason": "timeout"}, 1.0))
+    with pytest.raises(RuntimeError, match="session.expired: timeout"):
+        oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
+
+
 @pytest.mark.parametrize(
     ("event", "match"),
     [
@@ -561,8 +582,8 @@ class _RealtimeClient(oi._RealtimeSession):
 
     instances: list[_RealtimeClient] = []
 
-    def __init__(self, config: oi.OmniInteractBenchmarkConfig, session_id: str, reference_audio: str):
-        super().__init__(config, session_id, reference_audio)
+    def __init__(self, config: oi.OmniInteractBenchmarkConfig, reference_audio: str):
+        super().__init__(config, reference_audio)
         self.acks: list[tuple[str, int]] = []
         self.instances.append(self)
 
@@ -621,7 +642,7 @@ async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monke
     assert "autostart=0" in _RealtimeClient.instances[-1].url
     session_config = _RealtimeClient.instances[-1].session_config
     assert session_config.extra_body["custom"] == "value"
-    assert session_config.extra_body["native_duplex"] is True
+    assert "native_duplex" not in session_config.extra_body
     assert session_config.ref_audio == "data:audio/wav;base64,ref"
     acks = _RealtimeClient.instances[-1].acks
     # Cumulative acks for the one response: an optional 0 ms checkpoint the

tests/benchmarks/test_realtime_tts_probe.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""The Seed-TTS Realtime probe queues for a duplex session slot instead of failing."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.benchmarks.patch import patch as bench_patch
+from vllm_omni.clients import duplex as duplex_client
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _AdmissionGate:
+    """Stands in for ``DuplexClient``: refuses admission ``refusals`` times, then admits."""
+
+    instances: list[_AdmissionGate] = []
+    refusals = 0
+    refusal_code = "resource_exhausted"
+
+    def __init__(self, url: str, **kwargs: object) -> None:
+        del url, kwargs
+        self.entered = False
+        type(self).instances.append(self)
+
+    async def __aenter__(self) -> _AdmissionGate:
+        cls = type(self)
+        if cls.refusals > 0:
+            cls.refusals -= 1
+            raise duplex_client.DuplexProtocolError("duplex_session_capacity_exhausted: limit=4", code=cls.refusal_code)
+        self.entered = True
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        del exc_type, exc, tb
+
+    async def events(self):  # the collector's consume loop ends at once
+        return
+        yield  # pragma: no cover
+
+
+@pytest.fixture
+def gate(monkeypatch: pytest.MonkeyPatch) -> type[_AdmissionGate]:
+    _AdmissionGate.instances = []
+    _AdmissionGate.refusals = 0
+    _AdmissionGate.refusal_code = "resource_exhausted"
+    monkeypatch.setattr(duplex_client, "DuplexClient", _AdmissionGate)
+    return _AdmissionGate
+
+
+async def _configure(probe: bench_patch._RealtimeTTSProbe) -> None:
+    try:
+        await probe.configure("test-model", auto_response=True)
+    finally:
+        await probe.__aexit__(None, None, None)
+
+
+@pytest.mark.asyncio
+async def test_configure_waits_for_a_free_session_slot(gate) -> None:
+    """Two refusals, then a slot: the request is served, not counted as failed."""
+    gate.refusals = 2
+    probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
+    await _configure(probe)
+    assert len(gate.instances) == 3
+    assert gate.instances[-1].entered is True
+    assert probe._client is gate.instances[-1]
+
+
+@pytest.mark.asyncio
+async def test_configure_raises_other_rejections_at_once(gate) -> None:
+    gate.refusals = 1
+    gate.refusal_code = "unsupported_turn_detection"
+    probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
+    with pytest.raises(duplex_client.DuplexProtocolError):
+        await _configure(probe)
+    assert len(gate.instances) == 1
+
+
+@pytest.mark.asyncio
+async def test_configure_gives_up_when_the_slot_wait_budget_is_spent(gate, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(bench_patch._RealtimeTTSProbe, "_SESSION_SLOT_WAIT_S", 0.0)
+    gate.refusals = 5
+    probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
+    with pytest.raises(duplex_client.DuplexProtocolError) as excinfo:
+        await _configure(probe)
+    assert excinfo.value.code == "resource_exhausted"
+    assert len(gate.instances) == 1
+
+
+class _SilenceSink:
+    """Stands in for a configured ``DuplexClient``: records every appended chunk."""
+
+    def __init__(self) -> None:
+        self.config = duplex_client.SessionConfig()
+        self.chunks: list[bytes] = []
+
+    async def append_audio(self, pcm: bytes, *, is_speech: bool | None = None, video_frames=None) -> None:
+        del video_frames
+        assert is_speech is False
+        self.chunks.append(pcm)
+
+
+@pytest.mark.asyncio
+async def test_stream_silence_stops_at_the_turn_end(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Silence is fed until the turn settles, not for the whole cap."""
+    probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
+    sink = _SilenceSink()
+    probe._client = sink
+
+    async def no_wait(_: float) -> None:
+        return
+
+    monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait)
+    streamed = await probe.stream_silence(seconds=12.0, chunk_ms=200, until=lambda: len(sink.chunks) >= 3)
+    assert len(sink.chunks) == 3
+    assert streamed == pytest.approx(0.6)
+    assert all(len(chunk) == sink.config.input_audio.byte_count(200) for chunk in sink.chunks)
+
+
+@pytest.mark.asyncio
+async def test_stream_silence_runs_to_the_cap_when_the_turn_never_settles(monkeypatch: pytest.MonkeyPatch) -> None:
+    probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
+    sink = _SilenceSink()
+    probe._client = sink
+
+    async def no_wait(_: float) -> None:
+        return
+
+    monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait)
+    streamed = await probe.stream_silence(seconds=1.0, chunk_ms=200, until=lambda: False)
+    assert len(sink.chunks) == 5
+    assert streamed == pytest.approx(1.0)
+
+
+class _TurnEvents:
+    def __init__(self, response_ids: list[str], audio: dict[str, bytes], text: dict[str, str] | None = None) -> None:
+        self.response_ids = response_ids
+        self._audio = audio
+        self._text = text or {}
+
+    def audio_bytes(self, response_id: str) -> bytes:
+        return self._audio.get(response_id, b"")
+
+    def response_text(self, response_id: str) -> str:
+        return self._text.get(response_id, "")
+
+
+def test_turn_response_is_the_first_one_with_audio_even_if_the_model_spoke_again(caplog) -> None:
+    events = _TurnEvents(["r0", "r1", "r2"], {"r1": b"\x01\x00", "r2": b"\x02\x00"})
+    with caplog.at_level("WARNING", logger=bench_patch.logger.name):
+        assert bench_patch._seed_tts_turn_response_id(events, 0, 0) == "r1"
+    assert "spoke again" in caplog.text
+
+
+def test_turn_response_ignores_responses_from_before_the_turn() -> None:
+    events = _TurnEvents(["r0", "r1"], {"r0": b"\x01\x00", "r1": b"\x02\x00"})
+    assert bench_patch._seed_tts_turn_response_id(events, 1, 0) == "r1"
+
+
+def test_turn_without_audio_is_an_error() -> None:
+    events = _TurnEvents(["r0"], {})
+    with pytest.raises(RuntimeError, match="no audio response"):
+        bench_patch._seed_tts_turn_response_id(events, 0, 0)
+
+
+def test_stall_report_says_the_model_never_answered() -> None:
+    events = _TurnEvents(["r0"], {"r0": b"\x01\x00"})
+    report = bench_patch._seed_tts_turn_stall_report(events, 1, 3, 30.0)
+    assert "turn 3 never started a response" in report
+    assert "30.0s of silence" in report
+
+
+def test_stall_report_lists_a_response_that_never_finished() -> None:
+    events = _TurnEvents(["r1"], {"r1": b"\x01\x00\x02\x00"}, {"r1": "hello"})
+    report = bench_patch._seed_tts_turn_stall_report(events, 0, 0, 4.2)
+    assert "started 1 response(s) after 4.2s of silence" in report
+    assert "r1 (4 audio bytes, text 'hello')" in report

tests/clients/test_duplex_client.py

@@ -1,6 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Unit tests for the public duplex client (fake transport, no server)."""
+"""Unit tests for the public duplex clients (fake transports, no server).
+
+``DuplexClient`` runs over a fake websocket; ``InlineDuplexClient`` over a
+fake ``DuplexOmni`` whose handle records typed commands and replays typed
+events. The demux / handshake tests that only need "feed a server event" are
+parameterized over both clients.
+"""
 
 from __future__ import annotations
 
@@ -8,6 +14,8 @@ import asyncio
 import base64
 import json
 import wave
+from collections.abc import Callable
+from typing import Any
 from urllib.parse import parse_qs, urlsplit
 
 import pytest
@@ -16,8 +24,11 @@ from vllm_omni.clients.duplex import (
     AudioFormat,
     ConnectionResumed,
     DuplexClient,
+    DuplexClientBase,
+    DuplexConnectionError,
     DuplexProtocolError,
     DuplexSessionClosedError,
+    ErrorEvent,
     EventCollector,
     ReconnectPolicy,
     SessionConfig,
@@ -31,17 +42,26 @@ from vllm_omni.clients.duplex import (
     summarize_session_request_metrics,
     write_pcm16_wav,
 )
+from vllm_omni.clients.inline_duplex import InlineDuplexClient
+from vllm_omni.engine.duplex import commands as duplex_commands
+from vllm_omni.engine.duplex import events as duplex_events
+from vllm_omni.engine.duplex.messages import DuplexSessionError
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+SESSION_ID = "duplex-0123456789abcdef"
+# session.created is not journaled by the server, so it carries no server_event_seq.
 SESSION_CREATED = {
     "type": "session.created",
-    "session": {"session_id": "sess-1"},
-    "incarnation": 0,
+    "session": {"id": SESSION_ID, "model": "test-model"},
+    "attachment_generation": 1,
     "resume_token": "tok-1",
-    "server_event_seq": 1,
 }
-SESSION_CLOSED = {"type": "session.closed", "session_id": "sess-1", "server_event_seq": 99}
+SESSION_CLOSED = {"type": "session.closed", "session_id": SESSION_ID, "reason": "client_close", "server_event_seq": 99}
+
+
+# ---------------------------------------------------------------------------
+# Fake websocket transport
 
 
 class FakeSocket:
@@ -88,14 +108,139 @@ def make_client(*sockets: FakeSocket, **kwargs) -> tuple[DuplexClient, list[str]
     return client, calls
 
 
+# ---------------------------------------------------------------------------
+# Fake DuplexOmni for the inline client
+
+
+class _WireEvent:
+    """A typed-event stand-in that renders a fixed wire payload."""
+
+    def __init__(self, raw: dict[str, object]) -> None:
+        self.raw = raw
+
+    @property
+    def is_terminal(self) -> bool:
+        return self.raw.get("type") in {"session.closed", "session.expired"}
+
+    def to_realtime(self) -> dict[str, object]:
+        return dict(self.raw)
+
+
+class FakeInlineHandle:
+    def __init__(self, session_id: str) -> None:
+        self.session_id = session_id
+        self.closed = False
+        self.close_calls = 0
+        self.commands: list[duplex_commands.DuplexCommand] = []
+        self.consumer_active = False
+        self._outbox: asyncio.Queue = asyncio.Queue()
+
+    def deliver(self, event: Any) -> None:
+        self._outbox.put_nowait(event)
+
+    def feed(self, raw: dict[str, object]) -> None:
+        self.deliver(_WireEvent(raw))
+
+    async def submit(self, command: duplex_commands.DuplexCommand) -> None:
+        if self.closed:
+            raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id)
+        self.commands.append(command)
+
+    async def events(self):
+        self.consumer_active = True
+        try:
+            while True:
+                event = await self._outbox.get()
+                yield event
+                if event.is_terminal:
+                    return
+        finally:
+            self.consumer_active = False
+
+    async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None:
+        self.close_calls += 1
+        if self.closed:
+            return
+        self.closed = True
+        self.deliver(duplex_events.SessionClosed(session_id=self.session_id, reason=reason))
+
+
+class FakeInlineOmni:
+    def __init__(self) -> None:
+        self.opened: list[dict[str, object]] = []
+        self.handles: list[FakeInlineHandle] = []
+
+    async def open_session(self, config: dict[str, object]) -> FakeInlineHandle:
+        self.opened.append(dict(config))
+        handle = FakeInlineHandle(SESSION_ID)
+        self.handles.append(handle)
+        handle.deliver(
+            duplex_events.SessionCreated(
+                session_id=SESSION_ID,
+                session={"id": SESSION_ID, "model": config.get("model")},
+            )
+        )
+        return handle
+
+
+def make_inline_client(**kwargs) -> tuple[InlineDuplexClient, FakeInlineOmni]:
+    omni = FakeInlineOmni()
+    kwargs.setdefault("handshake_timeout_s", 5.0)
+    return InlineDuplexClient(omni, model="test-model", **kwargs), omni
+
+
+# ---------------------------------------------------------------------------
+# One rig over both clients: ``feed`` a server event, drive ``client``
+
+
+class Rig:
+    def __init__(self, client: DuplexClientBase, feed: Callable[[dict[str, object]], None]) -> None:
+        self.client = client
+        self.feed = feed
+
+
+def _websocket_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig:
+    sock = FakeSocket()
+    if handshake is not None:
+        sock.feed(handshake)
+    client, _ = make_client(sock)
+    return Rig(client, sock.feed)
+
+
+def _inline_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig:
+    client, omni = make_inline_client()
+    original_open = omni.open_session
+
+    async def open_session(config: dict[str, object]) -> FakeInlineHandle:
+        handle = await original_open(config)
+        if handshake is not SESSION_CREATED:
+            # Replace the default session.created with the requested first event.
+            handle._outbox = asyncio.Queue()
+            if handshake is not None:
+                handle.feed(handshake)
+        return handle
+
+    omni.open_session = open_session  # type: ignore[method-assign]
+
+    def feed(raw: dict[str, object]) -> None:
+        omni.handles[-1].feed(raw)
+
+    return Rig(client, feed)
+
+
+@pytest.fixture(params=["websocket", "inline"])
+def make_rig(request: pytest.FixtureRequest) -> Callable[..., Rig]:
+    return _websocket_rig if request.param == "websocket" else _inline_rig
+
+
 # ---------------------------------------------------------------------------
 # SessionConfig
 
 
 def test_session_config_payload_defaults():
-    payload = SessionConfig().to_session_payload(model="m", session_id="s")
+    payload = SessionConfig().to_session_payload(model="m")
     assert payload["model"] == "m"
-    assert payload["session_id"] == "s"
+    assert "session_id" not in payload
     assert payload["modalities"] == ["audio", "text"]
     assert payload["input_audio_format"] == "pcm16"
     assert payload["output_audio_format"] == "pcm16"
@@ -137,19 +282,59 @@ def test_audio_format_math():
 
 
 # ---------------------------------------------------------------------------
-# Handshake and lifecycle
+# Handshake and lifecycle (both clients)
 
 
-async def test_handshake_adopts_session_and_acks():
+@pytest.mark.asyncio
+async def test_handshake_adopts_server_allocated_session_id(make_rig):
+    rig = make_rig()
+    async with rig.client as client:
+        assert client.session_id == SESSION_ID
+        assert client.session_info == {"id": SESSION_ID, "model": "test-model"}
+        rig.feed(SESSION_CLOSED)
+
+
+@pytest.mark.asyncio
+async def test_handshake_error_raises_protocol_error(make_rig):
+    rig = make_rig(handshake={"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}})
+    with pytest.raises(DuplexProtocolError) as excinfo:
+        async with rig.client:
+            pass
+    assert excinfo.value.code == "unsupported_audio_format"
+
+
+@pytest.mark.asyncio
+async def test_session_expired_raises_from_event_stream(make_rig):
+    rig = make_rig()
+    async with rig.client as client:
+        rig.feed(
+            {"type": "session.expired", "session_id": SESSION_ID, "reason": "lease_expired", "server_event_seq": 2}
+        )
+        with pytest.raises(DuplexSessionClosedError) as excinfo:
+            async for _ in client.events():
+                pass
+        assert "lease_expired" in excinfo.value.reason
+
+
+# ---------------------------------------------------------------------------
+# Websocket transport specifics
+
+
+@pytest.mark.asyncio
+async def test_websocket_handshake_sends_session_update_and_acks_sequenced_events():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
     client, calls = make_client(sock)
     async with client:
-        assert client.session_id == "sess-1"
+        assert client.session_id == SESSION_ID
         assert client.resume_token == "tok-1"
         assert calls == ["ws://test-host:8099/v1/realtime?duplex=1&model=test-model&autostart=0"]
         assert sock.sent[0]["type"] == "session.update"
         assert sock.sent[0]["session"]["model"] == "test-model"
+        assert "session_id" not in sock.sent[0]["session"]
+        # session.created is unjournaled (no seq): nothing to ack yet.
+        assert _acks(sock) == []
+        sock.feed({"type": "response.listen", "server_event_seq": 1})
         await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 1} in _acks(sock))
         close_task = asyncio.create_task(client.close())
         await _drain(lambda: "session.close" in sock.sent_types())
@@ -158,28 +343,7 @@ async def test_handshake_adopts_session_and_acks():
     assert "session.close" in sock.sent_types()
 
 
-async def test_handshake_error_raises_protocol_error():
-    sock = FakeSocket()
-    sock.feed({"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}})
-    client, _ = make_client(sock)
-    with pytest.raises(DuplexProtocolError) as excinfo:
-        async with client:
-            pass
-    assert excinfo.value.code == "unsupported_audio_format"
-
-
-async def test_session_expired_raises_from_event_stream():
-    sock = FakeSocket()
-    sock.feed(SESSION_CREATED)
-    client, _ = make_client(sock)
-    async with client:
-        sock.feed({"type": "session.expired", "reason": "lease_expired", "server_event_seq": 2})
-        with pytest.raises(DuplexSessionClosedError) as excinfo:
-            async for _ in client.events():
-                pass
-        assert "lease_expired" in excinfo.value.reason
-
-
+@pytest.mark.asyncio
 async def test_resync_required_drops_resume_credential():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -187,7 +351,7 @@ async def test_resync_required_drops_resume_credential():
     async with client:
         assert client.resume_token == "tok-1"
         stream = client.events()
-        sock.feed({"type": "session.resync_required", "session_id": "sess-1", "server_event_seq": 2})
+        sock.feed({"type": "session.resync_required", "session_id": SESSION_ID, "server_event_seq": 2})
         async for event in stream:
             if event.type == "session.resync_required":
                 break
@@ -202,9 +366,10 @@ async def test_resync_required_drops_resume_credential():
 
 
 # ---------------------------------------------------------------------------
-# Input events
+# Input events (websocket wire shape)
 
 
+@pytest.mark.asyncio
 async def test_append_audio_tracks_cumulative_end_ms():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -224,6 +389,7 @@ async def test_append_audio_tracks_cumulative_end_ms():
         sock.feed(SESSION_CLOSED)
 
 
+@pytest.mark.asyncio
 async def test_append_audio_strips_video_frame_data_url_prefix():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -238,6 +404,7 @@ async def test_append_audio_strips_video_frame_data_url_prefix():
         sock.feed(SESSION_CLOSED)
 
 
+@pytest.mark.asyncio
 async def test_stream_pcm_chunking():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -249,6 +416,7 @@ async def test_stream_pcm_chunking():
         sock.feed(SESSION_CLOSED)
 
 
+@pytest.mark.asyncio
 async def test_interruption_primitives_send_documented_events():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -262,22 +430,21 @@ async def test_interruption_primitives_send_documented_events():
 
 
 # ---------------------------------------------------------------------------
-# Response demultiplexing
+# Response demultiplexing (both clients)
 
 
 def _b64(data: bytes) -> str:
     return base64.b64encode(data).decode("ascii")
 
 
-async def test_response_handle_flow():
-    sock = FakeSocket()
-    sock.feed(SESSION_CREATED)
+@pytest.mark.asyncio
+async def test_response_handle_flow(make_rig):
+    rig = make_rig()
     chunk = b"\x00\x01" * 2400  # 100 ms of pcm16 @ 24 kHz
-    client, _ = make_client(sock)
-    async with client:
-        sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
-        sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
-        sock.feed(
+    async with rig.client as client:
+        rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
+        rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
+        rig.feed(
             {
                 "type": "response.output_audio.delta",
                 "response_id": "resp-1",
@@ -286,7 +453,7 @@ async def test_response_handle_flow():
                 "server_event_seq": 4,
             }
         )
-        sock.feed(
+        rig.feed(
             {
                 "type": "response.output_audio_transcript.delta",
                 "response_id": "resp-1",
@@ -294,7 +461,7 @@ async def test_response_handle_flow():
                 "server_event_seq": 5,
             }
         )
-        sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6})
+        rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6})
 
         async for response in client.responses():
             chunks = [piece async for piece in response.audio()]
@@ -304,49 +471,46 @@ async def test_response_handle_flow():
             assert response.played_ms == pytest.approx(100.0)
             assert response.finished
             break
-        sock.feed(SESSION_CLOSED)
+        rig.feed(SESSION_CLOSED)
 
 
-async def test_listen_decision_yields_finished_silent_handle():
-    sock = FakeSocket()
-    sock.feed(SESSION_CREATED)
-    client, _ = make_client(sock)
-    async with client:
-        sock.feed({"type": "response.listen", "server_event_seq": 2})
+@pytest.mark.asyncio
+async def test_listen_decision_yields_finished_silent_handle(make_rig):
+    rig = make_rig()
+    async with rig.client as client:
+        rig.feed({"type": "response.listen", "server_event_seq": 2})
         async for response in client.responses():
             assert response.decision == "listen"
             assert response.finished
             assert [piece async for piece in response.audio()] == []
             break
-        sock.feed(SESSION_CLOSED)
+        rig.feed(SESSION_CLOSED)
 
 
-async def test_listen_terminates_active_response():
-    sock = FakeSocket()
-    sock.feed(SESSION_CREATED)
-    client, _ = make_client(sock)
-    async with client:
-        sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
-        sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3})
+@pytest.mark.asyncio
+async def test_listen_terminates_active_response(make_rig):
+    rig = make_rig()
+    async with rig.client as client:
+        rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
+        rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3})
         async for response in client.responses():
             await response.wait(timeout_s=5.0)
             assert response.decision == "listen"
             break
-        sock.feed(SESSION_CLOSED)
+        rig.feed(SESSION_CLOSED)
 
 
-async def test_id_less_listen_never_closes_an_in_flight_response():
+@pytest.mark.asyncio
+async def test_id_less_listen_never_closes_an_in_flight_response(make_rig):
     # An id-less listen is a standalone decision beat (e.g. a silence-skip
     # while a response streams); it must surface as its own finished handle
     # and leave the in-flight response to its real terminal.
-    sock = FakeSocket()
-    sock.feed(SESSION_CREATED)
-    client, _ = make_client(sock)
-    async with client:
-        sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
-        sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
-        sock.feed({"type": "response.listen", "server_event_seq": 4})
-        sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5})
+    rig = make_rig()
+    async with rig.client as client:
+        rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
+        rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
+        rig.feed({"type": "response.listen", "server_event_seq": 4})
+        rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5})
         seen: list[tuple[str | None, str | None, bool]] = []
         async for response in client.responses():
             await response.wait(timeout_s=5.0)
@@ -355,20 +519,19 @@ async def test_id_less_listen_never_closes_an_in_flight_response():
                 break
         assert ("resp-1", "speak", True) in seen
         assert (None, "listen", True) in seen
-        sock.feed(SESSION_CLOSED)
+        rig.feed(SESSION_CLOSED)
 
 
-async def test_terminal_listen_keeps_spoken_decision():
+@pytest.mark.asyncio
+async def test_terminal_listen_keeps_spoken_decision(make_rig):
     # A terminal listen after the response spoke is the model yielding the
     # turn; the handle must stay decision="speak" with its audio intact.
-    sock = FakeSocket()
-    sock.feed(SESSION_CREATED)
+    rig = make_rig()
     chunk = b"\x00\x01" * 2400
-    client, _ = make_client(sock)
-    async with client:
-        sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
-        sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
-        sock.feed(
+    async with rig.client as client:
+        rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
+        rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
+        rig.feed(
             {
                 "type": "response.output_audio.delta",
                 "response_id": "resp-1",
@@ -377,24 +540,23 @@ async def test_terminal_listen_keeps_spoken_decision():
                 "server_event_seq": 4,
             }
         )
-        sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5})
+        rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5})
         async for response in client.responses():
             chunks = [piece async for piece in response.audio()]
             assert chunks == [chunk]
             assert response.decision == "speak"
             assert response.finished
             break
-        sock.feed(SESSION_CLOSED)
+        rig.feed(SESSION_CLOSED)
 
 
-async def test_error_event_surfaces_on_response_path():
+@pytest.mark.asyncio
+async def test_error_event_surfaces_on_response_path(make_rig):
     # A rejected send produces no response; a consumer waiting in responses()
     # must see the rejection instead of waiting forever.
-    sock = FakeSocket()
-    sock.feed(SESSION_CREATED)
-    client, _ = make_client(sock)
-    async with client:
-        sock.feed(
+    rig = make_rig()
+    async with rig.client as client:
+        rig.feed(
             {
                 "type": "error",
                 "error": {"code": "input_audio_buffer_empty", "message": "empty commit"},
@@ -405,9 +567,37 @@ async def test_error_event_surfaces_on_response_path():
             async for _ in client.responses():
                 pass
         assert excinfo.value.code == "input_audio_buffer_empty"
+        rig.feed(SESSION_CLOSED)
+
+
+@pytest.mark.asyncio
+async def test_close_after_error_reports_the_server_reason_to_a_sender():
+    # A client streaming input sits in send(), not in responses(), so the
+    # error text has to survive into the close marker or the session looks
+    # like it ended for no reason.
+    sock = FakeSocket()
+    sock.feed(SESSION_CREATED)
+    client, _ = make_client(sock)
+    async with client:
+        sock.feed(
+            {
+                "type": "error",
+                "error": {
+                    "code": "runtime_data_plane_stream_failed",
+                    "message": "context_length_exceeded: streaming session prompt would grow to 41207 tokens",
+                },
+                "server_event_seq": 2,
+            }
+        )
         sock.feed(SESSION_CLOSED)
+        await _drain(lambda: client._closed.is_set())
+        with pytest.raises(DuplexSessionClosedError) as excinfo:
+            await client.send({"type": "input_audio_buffer.append", "audio": ""})
+    assert "runtime_data_plane_stream_failed" in str(excinfo.value)
+    assert "context_length_exceeded" in str(excinfo.value)
 
 
+@pytest.mark.asyncio
 async def test_slow_audio_consumer_drops_oldest_instead_of_stalling():
     from vllm_omni.clients.duplex import AudioDelta, ResponseHandle
 
@@ -425,24 +615,24 @@ async def test_slow_audio_consumer_drops_oldest_instead_of_stalling():
 
 
 # ---------------------------------------------------------------------------
-# Resume
+# Resume (websocket only)
 
 
+@pytest.mark.asyncio
 async def test_resume_after_transport_drop():
     first = FakeSocket()
     first.feed(SESSION_CREATED)
     second = FakeSocket()
     # Mirror the real resume activation payload
     # (entrypoints/duplex/serving.py, activation_payload_factory): it carries
-    # no nested "session" object, only the identity fields.
+    # the identity fields plus the (empty) public session object.
     second.feed(
         {
             "type": "session.resumed",
-            "session_id": "sess-1",
-            "incarnation": 0,
-            "attachment_generation": 1,
+            "session_id": SESSION_ID,
+            "session": {},
+            "attachment_generation": 2,
             "resume_token": "tok-2",
-            "server_event_seq": 5,
         }
     )
     client, calls = make_client(
@@ -453,34 +643,39 @@ async def test_resume_after_transport_drop():
     async with client:
         stream = client.events()
 
-        async def take_two():
+        async def collect_until_resumed():
             seen = []
             async for event in stream:
                 seen.append(event)
-                if len(seen) == 2:
+                if isinstance(event, SessionResumed):
                     return seen
 
-        consumer = asyncio.create_task(take_two())
+        consumer = asyncio.create_task(collect_until_resumed())
         await asyncio.sleep(0)  # let the consumer subscribe before the drop
         await asyncio.sleep(0)
+        first.feed({"type": "response.listen", "server_event_seq": 7})
+        await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 7} in _acks(first))
         first.feed(RuntimeError("transport dropped"))
         seen = await asyncio.wait_for(consumer, timeout=5.0)
-        assert isinstance(seen[0], ConnectionResumed)
-        assert isinstance(seen[1], SessionResumed)
+        assert [type(event) for event in seen[-2:]] == [ConnectionResumed, SessionResumed]
         assert client.resume_token == "tok-2"
+        # The superseded transport is closed, not leaked.
+        assert first.closed
         # The activation payload has no session object; the info captured at
         # the original handshake must survive the resume.
-        assert client.session_info == {"session_id": "sess-1"}
+        assert client.session_info == {"id": SESSION_ID, "model": "test-model"}
         assert len(calls) == 2
         resume = second.sent[0]
-        assert resume["type"] == "session.resume"
-        assert resume["session_id"] == "sess-1"
-        assert resume["incarnation"] == 0
-        assert resume["resume_token"] == "tok-1"
-        assert resume["last_received_server_event_seq"] == 1
+        assert resume == {
+            "type": "session.resume",
+            "session_id": SESSION_ID,
+            "resume_token": "tok-1",
+            "last_received_server_event_seq": 7,
+        }
         second.feed(SESSION_CLOSED)
 
 
+@pytest.mark.asyncio
 async def test_no_reconnect_policy_surfaces_closed():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -493,6 +688,7 @@ async def test_no_reconnect_policy_surfaces_closed():
                 pass
 
 
+@pytest.mark.asyncio
 async def test_heartbeat_survives_transient_send_failure():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -514,6 +710,7 @@ async def test_heartbeat_survives_transient_send_failure():
         sock.feed(SESSION_CLOSED)
 
 
+@pytest.mark.asyncio
 async def test_aexit_sends_session_close_on_error_path():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -534,6 +731,7 @@ def test_target_url_forces_autostart_off():
     assert query["autostart"] == ["0"]
 
 
+@pytest.mark.asyncio
 async def test_fatal_resume_error_gives_up():
     first = FakeSocket()
     first.feed(SESSION_CREATED)
@@ -553,6 +751,121 @@ async def test_fatal_resume_error_gives_up():
         assert len(calls) == 2  # no retry after a fatal resume error
 
 
+# ---------------------------------------------------------------------------
+# Inline client over a fake DuplexOmni
+
+
+@pytest.mark.asyncio
+async def test_inline_handshake_opens_session_with_payload_and_adopts_id():
+    client, omni = make_inline_client(config=SessionConfig(ref_audio="data:audio/wav;base64,AAA="))
+    async with client:
+        assert omni.opened == [client.config.to_session_payload(model="test-model")]
+        assert omni.opened[0]["ref_audio"] == "data:audio/wav;base64,AAA="
+        assert client.session_id == SESSION_ID
+        handle = omni.handles[0]
+        # The session was opened with the full session object already; the
+        # base handshake's session.update must not be re-submitted as a patch
+        # (a ref_audio patch would be rejected by the runner).
+        assert handle.commands == []
+        await client.send({"type": "session.update", "session": {"instructions": "later"}})
+        assert [type(command) for command in handle.commands] == [duplex_commands.UpdateSession]
+        assert handle.commands[0].patch == {"instructions": "later"}
+
+
+@pytest.mark.asyncio
+async def test_inline_client_submits_typed_commands_in_order():
+    client, omni = make_inline_client()
+    async with client:
+        handle = omni.handles[0]
+        pcm = b"\x01\x00" * 1600
+        await client.append_audio(pcm, is_speech=True)
+        await client.commit(create_response=True)
+        await client.cancel_response("resp-1")
+        await client.ack_playback(120.9, response_id="resp-1")
+        assert [type(command) for command in handle.commands] == [
+            duplex_commands.AppendAudio,
+            duplex_commands.Commit,
+            duplex_commands.CancelResponse,
+            duplex_commands.AckPlayback,
+        ]
+        append, commit, cancel, ack = handle.commands
+        assert append.is_speech is True
+        assert append.audio_end_ms == 100
+        assert append.event_id
+        assert commit.final is True and commit.create_response is True
+        assert cancel.response_id == "resp-1"
+        assert ack.played_ms == 120 and ack.response_id == "resp-1"
+
+
+@pytest.mark.asyncio
+async def test_inline_session_close_maps_to_handle_close():
+    client, omni = make_inline_client()
+    async with client:
+        pass
+    handle = omni.handles[0]
+    assert handle.closed
+    assert handle.close_calls >= 1
+    assert not any(isinstance(command, duplex_commands.CloseSession) for command in handle.commands)
+    # The typed session.closed reached the client as the expected end.
+    with pytest.raises(DuplexSessionClosedError):
+        await client.send({"type": "session.heartbeat"})
+
+
+@pytest.mark.asyncio
+async def test_inline_client_forwards_heartbeats_and_drops_event_acks():
+    client, omni = make_inline_client()
+    async with client:
+        handle = omni.handles[0]
+        await client.send({"type": "session.heartbeat"})
+        await client.send({"type": "session.event_ack", "server_event_seq": 3})
+        # A heartbeat is a lease touch the engine must see; the ack is a journal concern.
+        assert [type(command).__name__ for command in handle.commands] == ["Heartbeat"]
+
+
+@pytest.mark.asyncio
+async def test_inline_malformed_payload_surfaces_as_error_event():
+    client, omni = make_inline_client()
+    async with client:
+        handle = omni.handles[0]
+        # Subscribe before sending: the translation error is dispatched inline.
+        first_event = asyncio.create_task(client.wait_for("error", timeout_s=5.0))
+        await asyncio.sleep(0)
+        event_id = await client.send({"type": "playback.ack"})  # played_ms missing
+        event = await first_event
+        assert isinstance(event, ErrorEvent)
+        assert event.code == "bad_event"
+        assert event.related_event_id == event_id
+        assert handle.commands == []
+        with pytest.raises(DuplexProtocolError) as excinfo:
+            async for _ in client.responses():
+                pass
+        assert excinfo.value.code == "bad_event"
+
+
+@pytest.mark.asyncio
+async def test_inline_teardown_cancels_pump_and_closes_handle():
+    client, omni = make_inline_client()
+    async with client:
+        handle = omni.handles[0]
+        await _drain(lambda: handle.consumer_active)
+        pump = client._pump_task
+        assert pump is not None and not pump.done()
+    assert pump.done()
+    assert not handle.consumer_active
+    assert handle.closed
+
+
+@pytest.mark.asyncio
+async def test_inline_send_after_handle_closed_raises_connection_error():
+    client, omni = make_inline_client()
+    async with client:
+        handle = omni.handles[0]
+        handle.closed = True  # the engine ended the session without an event yet
+        with pytest.raises(DuplexConnectionError, match="not open"):
+            await client.send({"type": "input_audio_buffer.clear"})
+        handle.closed = False
+
+
 # ---------------------------------------------------------------------------
 # Collector
 
@@ -600,8 +913,7 @@ def test_build_realtime_url_with_model_extra_query():
     url = build_realtime_url(
         "ws://localhost:8099/v1/realtime?custom=1",
         "openbmb/MiniCPM-o-4_5",
-        session_id="session-a",
-        extra_query={"native_duplex": "1"},
+        extra_query={"probe": "1"},
     )
 
     query = parse_qs(urlsplit(url).query)
@@ -609,8 +921,7 @@ def test_build_realtime_url_with_model_extra_query():
         "custom": ["1"],
         "duplex": ["1"],
         "model": ["openbmb/MiniCPM-o-4_5"],
-        "native_duplex": ["1"],
-        "session_id": ["session-a"],
+        "probe": ["1"],
     }
 
 
@@ -619,12 +930,11 @@ def test_build_realtime_url_resume_only_when_autostart_disabled():
         "ws://localhost:8099/v1/realtime?duplex=1",
         "openbmb/MiniCPM-o-4_5",
         autostart=False,
-        extra_query={"native_duplex": "1"},
     )
 
     query = parse_qs(urlsplit(url).query)
     assert query["autostart"] == ["0"]
-    assert query["native_duplex"] == ["1"]
+    assert "session_id" not in query
 
 
 def test_event_collector_partitions_audio_by_response():
@@ -989,6 +1299,7 @@ def test_duplex_unit_boundary_and_residual_math():
     assert chunk_period_ms([{"type": "session.created", "session": {}}]) == 1000
 
 
+@pytest.mark.asyncio
 async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
     sock = FakeSocket()
     sock.feed(SESSION_CREATED)
@@ -1011,16 +1322,16 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
         sock.feed(SESSION_CLOSED)
 
 
-def test_build_realtime_url_native_duplex_flag_and_http_scheme():
-    url = build_realtime_url("http://localhost:8099/v1/realtime", None, native_duplex=None)
+def test_build_realtime_url_rewrites_http_scheme_and_rejects_others():
+    url = build_realtime_url("http://localhost:8099/v1/realtime", None)
     parts = urlsplit(url)
     assert parts.scheme == "ws"
     assert parse_qs(parts.query) == {"duplex": ["1"]}
 
-    url = build_realtime_url("https://host/v1/realtime", "m", native_duplex=False)
+    url = build_realtime_url("https://host/v1/realtime", "m", autostart=True)
     parts = urlsplit(url)
     assert parts.scheme == "wss"
-    assert parse_qs(parts.query)["native_duplex"] == ["0"]
+    assert parse_qs(parts.query) == {"duplex": ["1"], "model": ["m"], "autostart": ["1"]}
 
     with pytest.raises(ValueError):
         build_realtime_url("ftp://host/v1/realtime", "m")

tests/clients/test_model_session_configs.py

@@ -18,13 +18,13 @@ from vllm_omni.clients.personaplex import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def test_minicpmo45_session_config_matches_native_duplex_deployment():
+def test_minicpmo45_session_config_matches_duplex_deployment():
     config = create_minicpmo45_session_config(ref_audio="data:audio/wav;base64,AAA=", temperature=0.0)
     payload = config.to_session_payload(model="openbmb/MiniCPM-o-4_5")
     assert payload["ref_audio"] == "data:audio/wav;base64,AAA="
     assert payload["overlap_policy"] == "listen_only"
     assert payload["playback_commit_policy"] == "ack_only"
-    assert payload["extra_body"]["native_duplex"] is True
+    assert "native_duplex" not in payload["extra_body"]
     assert payload["extra_body"]["auto_response"] is True
     assert payload["temperature"] == 0.0
 
@@ -41,8 +41,9 @@ def test_personaplex_session_config_matches_deployment():
 @pytest.mark.parametrize(
     "duplex_package",
     [
+        # PersonaPlex and Nemotron VoiceChat rejoin this list with the follow-up
+        # PRs that port them to the duplex plugin framework.
         "vllm_omni.model_executor.models.minicpmo_4_5.duplex",
-        "vllm_omni.model_executor.models.personaplex.duplex",
     ],
 )
 def test_model_duplex_packages_stay_clear_of_client_library(duplex_package: str):

tests/config/test_config_factory.py

@@ -1328,7 +1328,11 @@ class TestDeployConfigLoading:
         stages = merge_pipeline_deploy(pipeline, deploy)
 
         assert deploy.session_mode == "duplex"
-        assert deploy.active_stream_window == max_sessions
+        # ``0`` disables the limiter: the configured session/stage capacity is
+        # what bounds the pipeline. Any positive window serializes concurrent
+        # audio first-packet generation (measured: TTFP -23% at 0 vs 4 on
+        # L20X, 128 requests / concurrency 8).
+        assert deploy.active_stream_window == 0
         assert deploy.duplex_session.max_sessions == max_sessions
         assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
         assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
@@ -1558,7 +1562,10 @@ stages:
         monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
 
         deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
-        assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
+        assert deploy.stages[1].engine_extras["model_loaded"] == {
+            "text_encoder": False,
+            "vae_encoder": False,
+        }
         stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
         resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
 
@@ -1583,12 +1590,20 @@ stages:
         deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
         stages = merge_pipeline_deploy(pipeline, deploy)
 
-        assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3TextEncoder"
+        assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder"
         assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline"
         assert stages[0].yaml_runtime["num_replicas"] == 1
         assert stages[1].yaml_runtime["num_replicas"] == 1
-        assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False}
+        assert stages[1].yaml_engine_args["model_loaded"] == {
+            "text_encoder": False,
+            "vae_encoder": False,
+        }
         assert stages[0].yaml_engine_args["max_num_seqs"] == 1
+        assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == {
+            "text_encoder": {"parallel_mode": "tp"},
+            "video_vae": {"parallel_mode": "patch"},
+            "audio_vae": {"parallel_mode": "leader"},
+        }
         assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
         assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
         parallel = stages[1].yaml_engine_args["parallel_config"]
@@ -1600,6 +1615,9 @@ stages:
 
         turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
         turbo_stages = merge_pipeline_deploy(pipeline, turbo)
+        assert turbo_stages[0].yaml_engine_args["model_arch"] == stages[0].yaml_engine_args["model_arch"]
+        assert turbo_stages[1].yaml_engine_args["model_loaded"] == stages[1].yaml_engine_args["model_loaded"]
+        assert turbo_stages[0].yaml_engine_args["hf_overrides"] == stages[0].yaml_engine_args["hf_overrides"]
         turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
         assert turbo_sampling["num_inference_steps"] == 5
         assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
@@ -2587,6 +2605,26 @@ class TestPlatformOverrides:
         assert rocm.stages[0].enforce_eager is None
         assert rocm.stages[1].enforce_eager is True
 
+    @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"])
+    @pytest.mark.parametrize("platform", ["cuda", "npu", "rocm"])
+    def test_qwen3_tts_default_code2wav_dtype_is_bf16(self, deploy_name, platform):
+        deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
+        deploy = _apply_platform_overrides(deploy, platform=platform)
+        stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy)
+
+        assert stages[1].yaml_engine_args["dtype"] == "bfloat16"
+
+    @pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"])
+    @pytest.mark.parametrize("platform", ["cuda", "npu"])
+    @pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"])
+    def test_qwen3_tts_propagates_explicit_code2wav_dtype(self, deploy_name, platform, dtype):
+        deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
+        deploy.stages[1].engine_extras["dtype"] = dtype
+        deploy = _apply_platform_overrides(deploy, platform=platform)
+        stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy)
+
+        assert stages[1].yaml_engine_args["dtype"] == dtype
+
     def test_higgs_audio_v3_rocm_uses_triton_attention(self):
         deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
 

tests/config/test_environment_variables.py

@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
     """Make an inventory expansion an explicit review decision."""
     category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
     assert category_counts == {
-        EnvironmentVariableCategory.PUBLIC_OMNI: 27,
+        EnvironmentVariableCategory.PUBLIC_OMNI: 28,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 64,

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -149,6 +149,9 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
     for name in _MIXIN_UPDATE_HELPERS:
         setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler))
     scheduler._cleanup_kv_tracking = MethodType(OmniARScheduler._cleanup_kv_tracking, scheduler)
+    scheduler._emit_streaming_context_overflow_outputs = MethodType(
+        OmniARScheduler._emit_streaming_context_overflow_outputs, scheduler
+    )
     scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
     scheduler.make_stats = lambda *args, **kwargs: None
     return scheduler

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from collections import defaultdict
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -17,6 +18,8 @@ import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.engine import FinishReason
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
@@ -958,3 +961,303 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
     assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
     sched.skipped_waiting.remove_requests.assert_not_called()
     sched._enqueue_waiting_request.assert_not_called()
+
+
+def _make_live_session_scheduler(*, max_model_len: int) -> OmniARScheduler:
+    """A scheduler whose finish path is real: the queues, ``finish_requests``,
+    ``_free_request`` and the finished-request bookkeeping that wakes an idle
+    engine all run. Only block freeing and the encoder cache are mocked."""
+    sched = _make_admission_scheduler(max_model_len=max_model_len)
+    del sched._free_request  # the admission helper mocks it; run the real one
+    sched._omits_kv_transfer_cache = {}
+    sched.connector = None
+    sched.perf_metrics = None
+    sched.recompute_kv_load_failures = False
+    sched.finished_req_ids = set()
+    sched.finished_req_ids_dict = defaultdict(set)
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.waiting_for_transfer_free = set()
+    sched.pending_stop_after_extraction = set()
+    sched.requests_needing_kv_transfer = {}
+    sched._kv_wait_start_ts = {}
+    sched.kv_cache_manager = SimpleNamespace(take_events=lambda: None, estimate_cached_tokens=lambda _request: 0)
+    return sched
+
+
+def _run_idle_step(sched: OmniARScheduler):
+    """One ``update_from_output`` with nothing scheduled: the step an idle
+    engine runs once ``has_finished_requests()`` reports a freed session."""
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={},
+        routed_experts=None,
+    )
+    return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+
+def _park_session(sched: OmniARScheduler, session: Request) -> None:
+    """Put a session where a 1 fps video stream leaves it between appends:
+    prompt partly computed, waiting for streaming input, in admission."""
+    sched.requests[session.request_id] = session
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 0
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    sched.skipped_waiting.add_request(session)
+    sched.num_waiting_for_streaming_input = 1
+
+
+def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_session() -> None:
+    """A native duplex session grows its prompt on every append; once the
+    extension cannot fit the model the worker crashes copying the prompt.
+    The scheduler must drop the update and fail only this session, at once,
+    through the real finish path: a parked session does not make the engine
+    schedule, so what wakes the engine is the freed request showing up in
+    ``has_finished_requests()``, and the step that follows must carry one
+    explicit ERROR rather than the synthesized ABORT."""
+    sched = _make_live_session_scheduler(max_model_len=8)
+    session = _make_request()
+    _park_session(sched, session)
+
+    sched._update_request_as_session(session, _make_update([10, 20, 30]))
+
+    assert session.prompt_token_ids == [1, 2, 3]
+    assert session.num_prompt_tokens == 3
+    assert session.status == RequestStatus.FINISHED_ERROR
+    assert session.request_id not in sched.requests
+    assert len(sched.waiting) == 0
+    assert len(sched.skipped_waiting) == 0
+    assert sched.num_waiting_for_streaming_input == 0
+    sched._free_request_blocks.assert_called_once_with(session)
+    assert sched.finished_req_ids == {session.request_id}
+    assert sched.has_finished_requests()
+    client_index, reason = sched._streaming_context_overflow[session.request_id]
+    assert client_index == session.client_index
+    assert reason.startswith("context_length_exceeded: ")
+    assert "9 tokens" in reason and "max_model_len 8" in reason
+
+    engine_core_outputs = _run_idle_step(sched)
+
+    (output,) = engine_core_outputs[session.client_index].outputs
+    assert output.request_id == session.request_id
+    assert output.finish_reason == FinishReason.ERROR
+    assert output.stop_reason == reason
+    assert output.new_token_ids == []
+    assert engine_core_outputs[session.client_index].finished_requests == {session.request_id}
+    assert sched._streaming_context_overflow == {}
+    # Nothing is left to report once the error went out.
+    assert _run_idle_step(sched) == {}
+
+
+def test_stage0_streaming_update_that_fills_max_model_len_exactly_finishes_the_session() -> None:
+    """An exactly full prompt leaves no room for the token every duplex step
+    samples: upstream's running budget drops to -1, which ``schedule()`` does
+    not catch, so the exact fill is an overflow too."""
+    sched = _make_scheduler(stage_id=0)
+    sched.max_model_len = 8
+    sched.finish_requests = MagicMock()
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 0
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.prompt_token_ids == [1, 2, 3]
+    assert session.num_prompt_tokens == 3
+    sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR)
+    _, reason = sched._streaming_context_overflow[session.request_id]
+    assert "8 tokens" in reason and "max_model_len 8" in reason
+
+
+def test_stage0_streaming_update_that_leaves_room_to_sample_is_applied() -> None:
+    sched = _make_scheduler(stage_id=0)
+    sched.max_model_len = 9
+    sched.finish_requests = MagicMock()
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 0
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
+    assert session.num_prompt_tokens == 8
+    assert session.status == RequestStatus.WAITING
+    sched.finish_requests.assert_not_called()
+    assert not getattr(sched, "_streaming_context_overflow", {})
+
+
+def test_stage0_streaming_update_keeps_room_for_every_token_sampled_per_step() -> None:
+    """With speculative decoding a step samples more than one token, and the
+    prompt must leave room for all of them."""
+    sched = _make_scheduler(stage_id=0)
+    sched.max_model_len = 9
+    sched.num_sampled_tokens_per_step = 2
+    sched.finish_requests = MagicMock()
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 0
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.prompt_token_ids == [1, 2, 3]
+    sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR)
+    assert "8 tokens" in sched._streaming_context_overflow[session.request_id][1]
+
+
+def test_stage0_streaming_update_overflow_counts_an_uncomputed_prompt() -> None:
+    """A preempted session keeps its prompt even though num_computed_tokens
+    fell behind it; the projection must use the longer of the two."""
+    sched = _make_scheduler(stage_id=0)
+    sched.max_model_len = 4
+    sched.finish_requests = MagicMock()
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.num_computed_tokens = 0
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.prompt_token_ids == [1, 2, 3]
+    sched.finish_requests.assert_called_once()
+    assert "5 tokens" in sched._streaming_context_overflow[session.request_id][1]
+
+
+def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
+    """A scheduler with real admission queues, so nothing about the terminal
+    handling below is mocked away except block freeing."""
+    sched = _make_scheduler(stage_id=0)
+    sched.max_model_len = max_model_len
+    sched.policy = SchedulingPolicy.FCFS
+    sched.waiting = create_request_queue(sched.policy)
+    sched.skipped_waiting = create_request_queue(sched.policy)
+    sched.running = []
+    sched.requests = {}
+    sched._free_request = MagicMock()  # needs a KV manager; not what this covers
+    return sched
+
+
+def _make_queued_stop(session: Request, update: StreamingUpdate) -> None:
+    """Put the request in the state upstream hands to _handle_stopped_request:
+    an append that arrived during generation is queued, and the stop status is
+    still on the request."""
+    from collections import deque
+
+    session.resumable = True
+    session.streaming_queue = deque([update])
+    session.status = RequestStatus.FINISHED_STOPPED
+
+
+def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -> None:
+    sched = _make_admission_scheduler(max_model_len=4)
+    session = _make_request()
+    sched.requests[session.request_id] = session
+    session.num_computed_tokens = 3
+    _make_queued_stop(session, _make_update([10, 20]))
+
+    finished = sched._handle_stopped_request(session)
+
+    assert finished is True
+    assert len(sched.waiting) == 0
+    assert len(sched.skipped_waiting) == 0
+    assert session.status == RequestStatus.FINISHED_ERROR
+    assert session.resumable is False
+    assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
+    # update_from_output frees every request reported as finished; freeing here
+    # as well deleted it from self.requests twice (KeyError in _free_blocks).
+    sched._free_request.assert_not_called()
+
+
+def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
+    sched = _make_admission_scheduler(max_model_len=64)
+    session = _make_request()
+    sched.requests[session.request_id] = session
+    session.num_computed_tokens = 3
+    _make_queued_stop(session, _make_update([10, 20]))
+
+    finished = sched._handle_stopped_request(session)
+
+    assert finished is False
+    assert len(sched.waiting) == 1
+    assert session.status == RequestStatus.WAITING
+    assert not getattr(sched, "_streaming_context_overflow", {})
+    sched._free_request.assert_not_called()
+
+
+def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:
+    """The overflow may already be on record when the stop is handled (an
+    earlier update of the same session). Upstream still re-enqueues the
+    request, so the override must take it back out either way instead of
+    leaving a FINISHED_ERROR request for admission to trip over."""
+    sched = _make_admission_scheduler(max_model_len=64)
+    session = _make_request()
+    sched.requests[session.request_id] = session
+    session.num_computed_tokens = 3
+    sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")}
+    _make_queued_stop(session, _make_update([10, 20]))  # this update itself fits
+
+    finished = sched._handle_stopped_request(session)
+
+    assert finished is True
+    assert len(sched.waiting) == 0
+    assert len(sched.skipped_waiting) == 0
+    assert session.status == RequestStatus.FINISHED_ERROR
+    assert session.resumable is False
+    sched._free_request.assert_not_called()
+
+
+def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanced() -> None:
+    """Same as above with nothing queued: upstream parks the request as
+    waiting for streaming input and counts it; the dequeue must uncount it."""
+    from collections import deque
+
+    sched = _make_admission_scheduler(max_model_len=64)
+    session = _make_request()
+    sched.requests[session.request_id] = session
+    session.num_computed_tokens = 3
+    sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")}
+    session.resumable = True
+    session.streaming_queue = deque()
+    session.status = RequestStatus.FINISHED_STOPPED
+
+    finished = sched._handle_stopped_request(session)
+
+    assert finished is True
+    assert len(sched.skipped_waiting) == 0
+    assert sched.num_waiting_for_streaming_input == 0
+    assert session.status == RequestStatus.FINISHED_ERROR
+    sched._free_request.assert_not_called()
+
+
+def test_context_overflow_emits_an_error_output_with_the_reason() -> None:
+    sched = _make_scheduler(stage_id=0)
+    sched._streaming_context_overflow = {"req-a": (2, "context_length_exceeded: too long")}
+    outputs: dict[int, list] = {}
+
+    sched._emit_streaming_context_overflow_outputs(outputs)
+
+    (output,) = outputs[2]
+    assert output.request_id == "req-a"
+    assert output.finish_reason == FinishReason.ERROR
+    assert output.stop_reason == "context_length_exceeded: too long"
+    assert output.new_token_ids == []
+    assert sched._streaming_context_overflow == {}
+
+    sched._emit_streaming_context_overflow_outputs(outputs)
+    assert len(outputs[2]) == 1

tests/dfx/perf/tests/test_minicpmo_4_5.json

@@ -26,8 +26,8 @@
       {
         "dataset_name": "seed-tts",
         "dataset_path": "zhaochenyang20/seed-tts-eval",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
+        "backend": "openai-realtime-tts",
+        "endpoint": "/v1/realtime",
         "num_prompts": [
           32,
           64,
@@ -42,73 +42,7 @@
         "disable_shuffle": true,
         "trust_remote_code": true,
         "seed_tts_locale": "en",
-        "extra_body": {
-          "modalities": [
-            "text",
-            "audio"
-          ],
-          "chat_template_kwargs": {
-            "enable_thinking": false,
-            "use_tts_template": true
-          }
-        },
-        "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
-        "baseline": {
-          "H100": {
-            "request_throughput": [
-              0.9731,
-              1.1986,
-              1.686
-            ],
-            "mean_ttft_ms": [
-              152.5665,
-              270.2559,
-              301.455
-            ],
-            "mean_e2el_ms": [
-              1027.2953,
-              3306.3735,
-              4686.5366
-            ],
-            "mean_audio_ttfp_ms": [
-              458.1553,
-              2030.2335,
-              2068.6233
-            ],
-            "mean_audio_rtf": [
-              0.2357,
-              0.791,
-              1.0324
-            ]
-          },
-          "A3": {
-            "request_throughput": [
-              0.5383,
-              0.6042,
-              0.7547
-            ],
-            "mean_ttft_ms": [
-              333.2633,
-              533.8656,
-              547.3388
-            ],
-            "mean_e2el_ms": [
-              1857.2154,
-              6600.5101,
-              10450.6811
-            ],
-            "mean_audio_ttfp_ms": [
-              986.4666,
-              3411.0754,
-              3352.7515
-            ],
-            "mean_audio_rtf": [
-              0.4423,
-              1.5734,
-              2.3024
-            ]
-          }
-        }
+        "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
       }
     ]
   },
@@ -138,8 +72,8 @@
       {
         "dataset_name": "seed-tts",
         "dataset_path": "zhaochenyang20/seed-tts-eval",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
+        "backend": "openai-realtime-tts",
+        "endpoint": "/v1/realtime",
         "num_prompts": [
           32,
           64,
@@ -154,16 +88,6 @@
         "disable_shuffle": true,
         "trust_remote_code": true,
         "seed_tts_locale": "en",
-        "extra_body": {
-          "modalities": [
-            "text",
-            "audio"
-          ],
-          "chat_template_kwargs": {
-            "enable_thinking": false,
-            "use_tts_template": true
-          }
-        },
         "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
       }
     ]

tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json

@@ -43,43 +43,7 @@
         "extra_body": {
           "save_duplex_request_metrics": true
         },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration",
-        "baseline": {
-          "H100": {
-            "request_throughput": [
-              0.3038
-            ],
-            "mean_ttft_ms": [
-              139.997
-            ],
-            "mean_e2el_ms": [
-              3269.8807
-            ],
-            "mean_audio_ttfp_ms": [
-              456.073
-            ],
-            "mean_audio_rtf": [
-              0.2254
-            ]
-          },
-          "A3": {
-            "request_throughput": [
-              0.1594
-            ],
-            "mean_ttft_ms": [
-              274.164
-            ],
-            "mean_e2el_ms": [
-              6250.7696
-            ],
-            "mean_audio_ttfp_ms": [
-              994.422
-            ],
-            "mean_audio_rtf": [
-              0.46163
-            ]
-          }
-        }
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
       }
     ]
   }

tests/dfx/perf/tests/test_minicpmo_4_5_ready.json

@@ -26,8 +26,8 @@
       {
         "dataset_name": "seed-tts",
         "dataset_path": "zhaochenyang20/seed-tts-eval",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
+        "backend": "openai-realtime-tts",
+        "endpoint": "/v1/realtime",
         "num_prompts": [
           10,
           40
@@ -40,41 +40,7 @@
         "disable_shuffle": true,
         "trust_remote_code": true,
         "seed_tts_locale": "en",
-        "extra_body": {
-          "modalities": [
-            "text",
-            "audio"
-          ],
-          "chat_template_kwargs": {
-            "enable_thinking": false,
-            "use_tts_template": true
-          }
-        },
-        "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
-        "baseline": {
-          "H100": {
-            "request_throughput": [
-              0.9731,
-              1.1986
-            ],
-            "mean_ttft_ms": [
-              152.5665,
-              270.2559
-            ],
-            "mean_e2el_ms": [
-              1027.2953,
-              3306.3735
-            ],
-            "mean_audio_ttfp_ms": [
-              458.1553,
-              2030.2335
-            ],
-            "mean_audio_rtf": [
-              0.2357,
-              0.791
-            ]
-          }
-        }
+        "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
       }
     ]
   }

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -157,14 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
                 "ref_audio": None,
                 "ref_text": None,
             },
-            # NOTE: This is delegating down the TTS adapter, so the message is both specific
-            # to the Qwen3TTS adapter and dependent on whether or not the model has any speakers
-            # (as opposed to having valida speakers, but getting an invalid speaker). it would
-            # be more ideal to validate that this is landing on `.validate()` on the adapter
-            # and ensuring that the returned message matches what we would get from calling the
-            # adapter validate directly.
-            ("no speakers configured",),
-            id="customvoice_invalid_voice",
+            ("Base checkpoint does not support task_type='CustomVoice'",),
+            id="base_checkpoint_customvoice_task_mismatch",
         ),
         pytest.param(
             {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"

tests/diffusion/attention/test_block_sparse_ops.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import sys
+import types
+
+import pytest
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from vllm_omni.diffusion.attention.ops.block_sparse import (
+    build_prefix_dense_block_map,
+    fastvideo_block_sparse_attn_bshd,
+    mean_pool_tiles,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("prefix_blocks,topk", [(0, 1), (2, 0), (2, 1), (2, 3), (2, 5)])
+def test_block_map_keeps_prefix_dense_and_selects_non_prefix_keys(prefix_blocks, topk):
+    sparse_blocks = 3
+    blocks = prefix_blocks + sparse_blocks
+    scores = torch.arange(2 * blocks * blocks, dtype=torch.float32).reshape(1, 2, blocks, blocks)
+    block_map = build_prefix_dense_block_map(scores, prefix_blocks, sparse_blocks, topk)
+    expected = torch.zeros_like(scores, dtype=torch.bool)
+    # Scores increase with key index, so top-k always selects the last keys.
+    for query_block in range(blocks):
+        for key_block in range(blocks):
+            if (
+                query_block < prefix_blocks
+                or key_block < prefix_blocks
+                or key_block >= blocks - min(topk, sparse_blocks)
+            ):
+                expected[..., query_block, key_block] = True
+    assert torch.equal(block_map, expected)
+
+
+@pytest.mark.parametrize("block_size", [4, 32, 64])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_pooling_uses_valid_tile_sizes_and_fp32_accumulation(block_size, dtype):
+    sizes = torch.tensor([block_size, block_size - 1, 0], dtype=torch.int32)
+    tiles = torch.arange(2 * 3 * block_size * 2 * 3).reshape(2, 3, block_size, 2, 3).to(dtype)
+    tiles[:, 1, -1] = 0
+    tiles[:, 2] = 0
+    expected = torch.stack(
+        [tiles[:, 0].float().mean(dim=1), tiles[:, 1, :-1].float().mean(dim=1), torch.zeros(2, 2, 3)], dim=2
+    )
+    result = mean_pool_tiles(tiles.flatten(1, 2), sizes, block_size)
+    assert result.dtype == torch.float32
+    torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+def test_provider_wrapper_restores_transport_rows_and_bshd_layout(monkeypatch):
+    monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False)
+    calls = []
+
+    def block_sparse_attn(q, k, v, block_map, sizes):
+        calls.append((q.shape, block_map.shape, sizes.tolist()))
+        assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
+        assert sizes.dtype == torch.int32
+        return q + k + v, None
+
+    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    setattr(module, "block_sparse_attn", block_sparse_attn)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
+    query = torch.randn(1, 4 * 64, 2, 8)
+    block_map = torch.ones(1, 2, 4, 4, dtype=torch.bool)
+    sizes = torch.tensor([64, 17, 64, 0])
+    output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
+    expected = torch.zeros_like(query)
+    expected[:, : 3 * 64] = query[:, : 3 * 64] * 3
+    torch.testing.assert_close(output, expected)
+    assert calls == [(torch.Size([1, 2, 3 * 64, 8]), torch.Size([1, 2, 3, 3]), [64, 17, 64])]
+
+
+def test_custom_op_fake_preserves_metadata_without_loading_provider(monkeypatch):
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", None)
+    with FakeTensorMode():
+        query = torch.empty(1, 4 * 64, 2, 8, dtype=torch.bfloat16)
+        block_map = torch.empty(1, 2, 4, 4, dtype=torch.bool)
+        sizes = torch.empty(4, dtype=torch.int32)
+        output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
+        assert output.shape == query.shape
+        assert output.dtype == query.dtype
+        assert output.device == query.device
+
+
+@pytest.mark.parametrize("error_type", [ImportError, RuntimeError, torch.AcceleratorError])
+def test_native_provider_failure_only_recovers_without_accelerator_fault(monkeypatch, mocker, error_type):
+    monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1")
+    error = error_type("native provider failure")
+    native = mocker.Mock()
+    native.is_supported.return_value = True
+    native.block_sparse_attn_sm100a.side_effect = error
+    provider = types.ModuleType("fastvideo_kernel")
+    setattr(provider, "block_sparse_attn_sm100a", native)
+    index = types.ModuleType("fastvideo_kernel.triton_kernels.index")
+    setattr(
+        index,
+        "map_to_index",
+        mocker.Mock(return_value=(torch.zeros(1, dtype=torch.int32), torch.ones(1, dtype=torch.int32))),
+    )
+    triton = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    fallback = mocker.Mock(side_effect=lambda q, *args: (q.clone(), None))
+    setattr(triton, "block_sparse_attn", fallback)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", provider)
+    monkeypatch.setitem(
+        sys.modules, "fastvideo_kernel.triton_kernels", types.ModuleType("fastvideo_kernel.triton_kernels")
+    )
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.triton_kernels.index", index)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", triton)
+    query = torch.arange(128 * 2 * 8, dtype=torch.float32).reshape(1, 128, 2, 8)
+    block_map = torch.ones(1, 2, 2, 2, dtype=torch.bool)
+    sizes = torch.tensor([64, 64], dtype=torch.int32)
+
+    if error_type is torch.AcceleratorError:
+        with pytest.raises(torch.AcceleratorError, match="native provider failure") as raised:
+            fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
+        assert raised.value is error
+        fallback.assert_not_called()
+    else:
+        output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
+        torch.testing.assert_close(output, query, rtol=0, atol=0)
+        fallback.assert_called_once()

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-import math
 import sys
 import types
 from typing import Any
@@ -12,14 +11,10 @@ import torch
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
     PackedPaddingMetadata,
-    VideoTokenLayout,
-    VideoTokenSpan,
 )
 from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
     FastVideoVSABackend,
     FastVideoVSAImpl,
-    _build_h3_block_map,
-    _get_h3_tile_metadata,
 )
 from vllm_omni.diffusion.attention.backends.registry import (
     DiffusionAttentionBackendEnum,
@@ -218,144 +213,8 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
     assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
 
 
-def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
-    extra: dict[str, Any] = {
-        "vsa_h3_prefix_segments": prefix_segments,
-    }
-    if gate is not None:
-        extra["gate_compress"] = gate
-    return AttentionMetadata(
-        packed_padding=packed_padding,
-        video_layout=VideoTokenLayout(
-            used_len=sum(prefix_segments) + math.prod(video_shape),
-            video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
-        ),
-        extra=extra,
-    )
-
-
-def _fake_block_sparse_kernel(monkeypatch, calls):
-    """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
-
-    def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
-        calls["block_map"] = block_map.clone()
-        calls["variable_block_sizes"] = variable_block_sizes.tolist()
-        calls["q_shape"] = tuple(q.shape)
-        blocks = block_map.shape[-1]
-        assert blocks * 64 == q.shape[2]
-        # Expand the per-block map to a per-token mask so the stand-in honours
-        # the sparsity contract the real kernel implements, and drop the rows
-        # each partial block pads with, which the real kernel skips through
-        # variable_block_sizes.
-        token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
-        within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
-        token_mask = token_mask & within.reshape(1, 1, 1, -1)
-        out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
-        return out.to(q.dtype), None
-
-    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
-    setattr(module, "block_sparse_attn", block_sparse_attn)
-    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
-    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
-
-
-def _h3_impl(**backend_kwargs):
-    kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
-    kwargs.update(backend_kwargs)
-    return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
-
-
-def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
-    calls: dict[str, Any] = {}
-    _fake_block_sparse_kernel(monkeypatch, calls)
-    impl = _h3_impl(topk=64)
-    prefix_segments, video_shape = (5, 70), (2, 4, 4)
-    seq_len = sum(prefix_segments) + 2 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
-
-    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-
-    assert output.shape == query.shape
-    # topk covers every video block, so the block map is dense and the tiled
-    # kernel must reproduce plain attention over the same rows.
-    reference = torch.nn.functional.scaled_dot_product_attention(
-        *(x.float().transpose(1, 2) for x in (query, query, query))
-    ).transpose(1, 2)
-    torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
-    # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
-    assert calls["variable_block_sizes"] == [5, 64, 6, 32]
-
-
-def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
-    calls: dict[str, Any] = {}
-    _fake_block_sparse_kernel(monkeypatch, calls)
-    impl = _h3_impl(topk=1)
-    prefix_segments, video_shape = (5,), (8, 4, 4)
-    seq_len = 5 + 8 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
-
-    impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-
-    # 1 prefix block + 2 video blocks. The odd count is padded to an even one
-    # for the paired-CTA kernel contract, and the Triton route slices that
-    # transport-only partner off again before launching.
-    block_map = calls["block_map"]
-    assert block_map.shape[-1] == 3
-    assert calls["variable_block_sizes"] == [5, 64, 64]
-    assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
-    assert block_map[..., :1].all(), "prefix keys are exempt from selection"
-    assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
-
-    # The same odd geometry must still untile back onto the original rows.
-    dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-    reference = torch.nn.functional.scaled_dot_product_attention(
-        *(x.float().transpose(1, 2) for x in (query, query, query))
-    ).transpose(1, 2)
-    torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
-
-
-def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
-    _fake_block_sparse_kernel(monkeypatch, {})
-    impl = _h3_impl(topk=64)
-    prefix_segments, video_shape = (5,), (2, 4, 4)
-    seq_len = 5 + 2 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
-
-    without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
-    zero_gate = torch.zeros_like(query)
-    with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
-    live_gate = torch.full_like(query, 0.5)
-    with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
-
-    torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
-    assert not torch.allclose(with_live_gate.float(), without_gate.float())
-
-
-def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
-    _fake_block_sparse_kernel(monkeypatch, {})
-    impl = _h3_impl(topk=64)
-    prefix_segments, video_shape = (5,), (2, 4, 4)
-    valid = 5 + 2 * 4 * 4
-    torch.manual_seed(0)
-    query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
-    padding = PackedPaddingMetadata(
-        q_length=valid,
-        kv_length=valid,
-        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
-        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
-    )
-
-    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
-
-    assert output.shape == query.shape
-    assert torch.count_nonzero(output[:, valid:]) == 0
-
-
 def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
-    # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
+    # The backend advertises supports_packed_mask_free, so a model may skip
     # building the padding mask. Every fallback out of the VSA path must honour
     # that contract itself; SDPA reads attn_mask and nothing else.
     impl = FastVideoVSAImpl(
@@ -374,6 +233,8 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
         cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
     )
 
+    monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
+
     baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
     perturbed_input = query.clone()
     perturbed_input[:, valid:] += 100.0
@@ -384,22 +245,3 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
     assert baseline.shape == query.shape
     torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
     assert torch.count_nonzero(baseline[:, valid:]) == 0
-
-
-def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
-    partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
-        (5, 70, 9), (5, 6, 6), torch.device("cpu")
-    )
-    assert sizes[:4].tolist() == [5, 64, 6, 9]
-    assert prefix_blocks == 4
-    assert video_blocks == 8
-    assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
-    assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
-
-
-def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
-    scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
-    block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
-    assert block_map[:, :, :2].all()
-    assert block_map[..., :2].all()
-    assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()

tests/diffusion/attention/test_impl_overrides.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.diffusion.attention import layer
+from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSABackend, FastVideoVSAImpl
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import AttentionSpec
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _SpecializedImpl(FastVideoVSAImpl):
+    pass
+
+
+@pytest.fixture
+def local_attention(monkeypatch):
+    monkeypatch.setattr(layer, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention())
+
+
+@pytest.mark.parametrize("explicit", [False, True])
+@pytest.mark.parametrize("override", [False, True])
+def test_specialization_preserves_selected_backend_and_options(mocker, local_attention, explicit, override):
+    spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=7) if explicit else None
+    select = mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, spec))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        role="video.self",
+        role_category="self",
+        qkv_layout="BSND",
+        impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl} if override else None,
+    )
+
+    assert attention.attn_backend is FastVideoVSABackend
+    assert type(attention.attention) is (_SpecializedImpl if override else FastVideoVSAImpl)
+    assert attention.attention.topk == (7 if explicit else 64)
+    assert attention.attention.qkv_layout == "BSND"
+    assert attention.backend_pref == "FASTVIDEO_VSA"
+    assert attention.backend_explicit is explicit
+    assert attention.attn_spec is spec
+    assert isinstance(attention.parallel_strategy, NoParallelAttention)
+    select.assert_called_once_with(
+        role="video.self", head_size=8, attention_config=None, role_category="self", allow_trtllm_default=False
+    )
+
+
+def test_unselected_specialization_does_not_replace_dense_backend(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(SDPABackend, None))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+    )
+    assert attention.attn_backend is SDPABackend
+    assert type(attention.attention) is SDPABackend.get_impl_cls()
+
+
+def test_specialization_does_not_bypass_selection_errors(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", side_effect=ImportError("optional kernel unavailable"))
+    with pytest.raises(ImportError, match="optional kernel unavailable"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+        )
+
+
+def test_specialization_must_extend_selected_implementation(mocker, local_attention):
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, None))
+    with pytest.raises(TypeError, match="must subclass the selected implementation"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": SDPABackend.get_impl_cls()},
+        )
+
+
+def test_platform_backend_capabilities_survive_model_specialization(mocker, local_attention):
+    from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import MiniMaxH3VSAImpl
+
+    class PlatformBackend(FastVideoVSABackend):
+        @classmethod
+        def supports_packed_mask_free(cls):
+            return False
+
+        @staticmethod
+        def get_supported_head_sizes():
+            return [128]
+
+    # A registry/platform may refine capabilities without changing the
+    # implementation. The model must not replace those capability decisions.
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
+    attention = layer.Attention(
+        num_heads=2,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        causal=False,
+        impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl},
+    )
+    assert attention.attn_backend is PlatformBackend
+    assert not attention.attn_backend.supports_packed_mask_free()
+    assert attention.attn_backend.get_supported_head_sizes() == [128]
+    assert type(attention.attention) is MiniMaxH3VSAImpl
+
+
+def test_incompatible_platform_implementation_is_not_discarded(mocker, local_attention):
+    class PlatformImpl(FastVideoVSAImpl):
+        pass
+
+    class PlatformBackend(FastVideoVSABackend):
+        @staticmethod
+        def get_impl_cls():
+            return PlatformImpl
+
+    mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
+    with pytest.raises(TypeError, match="selected implementation .*PlatformImpl"):
+        layer.Attention(
+            num_heads=2,
+            head_size=8,
+            softmax_scale=8**-0.5,
+            causal=False,
+            impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
+        )

tests/diffusion/cache/test_seacache.py

@@ -0,0 +1,635 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import datetime
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.cache.seacache import (
+    SeaCacheBackend,
+    SeaCacheConfig,
+    SeaCacheRootHook,
+    apply_sea_cache_hook,
+)
+from vllm_omni.diffusion.cache.seacache.sea_filter import (
+    apply_sea_filter,
+    extrapolate_residual,
+    indicator_distance,
+)
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext, get_extractor
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TinyCosmos3Transformer(torch.nn.Module):
+    """Small model that exposes a cache-neutral execution boundary."""
+
+    def _run_gen_layers(self, hidden_gen: torch.Tensor) -> torch.Tensor:
+        return hidden_gen * 2
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_ids: torch.Tensor | None = None,
+        text_mask: torch.Tensor | None = None,
+        video_shape: tuple[int, int, int] | None = None,
+        noisy_frame_mask: torch.Tensor | None = None,
+        control_latents: list[torch.Tensor] | torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del timestep, text_ids, text_mask, video_shape, noisy_frame_mask
+        if kwargs:
+            raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
+        controls = (
+            []
+            if control_latents is None
+            else [control_latents]
+            if isinstance(control_latents, torch.Tensor)
+            else list(control_latents)
+        )
+        inputs = [*controls, hidden_states]
+        gen_input = torch.cat(
+            [value.movedim(1, -1).flatten(1, 3) for value in inputs],
+            dim=1,
+        )
+        return self._run_gen_layers(gen_input)
+
+
+def _extract_tiny_cosmos3_context(
+    module: TinyCosmos3Transformer,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    text_ids: torch.Tensor | None = None,
+    text_mask: torch.Tensor | None = None,
+    video_shape: tuple[int, int, int] | None = None,
+    noisy_frame_mask: torch.Tensor | None = None,
+    control_latents: list[torch.Tensor] | torch.Tensor | None = None,
+    **kwargs,
+) -> CacheContext:
+    del timestep, text_ids, text_mask, video_shape
+    if kwargs:
+        raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
+    controls = (
+        []
+        if control_latents is None
+        else [control_latents]
+        if isinstance(control_latents, torch.Tensor)
+        else list(control_latents)
+    )
+    vision_items = [*controls, hidden_states]
+    gen_input = torch.cat(
+        [value.movedim(1, -1).flatten(1, 3) for value in vision_items],
+        dim=1,
+    )
+    return CacheContext(
+        modulated_input=gen_input,
+        hidden_states=gen_input,
+        encoder_hidden_states=None,
+        temb=torch.zeros_like(gen_input[:, 0]),
+        run_transformer_blocks=lambda: (module._run_gen_layers(gen_input),),
+        postprocess=lambda output: output,
+        extra_states={
+            "sea_cache_latents": vision_items,
+            "sea_cache_noisy_frame_mask": noisy_frame_mask,
+        },
+    )
+
+
+class Cosmos3OmniDiffusersPipeline:
+    def __init__(self) -> None:
+        self.transformer = TinyCosmos3Transformer()
+        self._current_step_index: int | None = None
+        self._current_sigma: float | None = None
+        self._num_timesteps: int | None = None
+
+    @property
+    def current_step_index(self) -> int | None:
+        return self._current_step_index
+
+    @property
+    def current_sigma(self) -> float | None:
+        return self._current_sigma
+
+    @property
+    def num_timesteps(self) -> int | None:
+        return self._num_timesteps
+
+
+def _latent(value: float) -> torch.Tensor:
+    return torch.full((1, 2, 2, 2, 2), value)
+
+
+def _run_step(
+    transformer: TinyCosmos3Transformer,
+    timestep: int,
+    value: float,
+    *,
+    hook: SeaCacheRootHook | None = None,
+    context: str = "cond",
+    control: float | None = None,
+    noisy_frame_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    controls = None if control is None else [_latent(control)]
+    cache_context = hook.cache_context(context) if hook is not None else nullcontext()
+    with torch.inference_mode(), cache_context:
+        return transformer(
+            hidden_states=_latent(value),
+            timestep=torch.tensor([timestep]),
+            noisy_frame_mask=noisy_frame_mask,
+            control_latents=controls,
+        )
+
+
+def _apply_test_hook(
+    transformer: TinyCosmos3Transformer,
+    metadata: SimpleNamespace,
+    config: SeaCacheConfig | None = None,
+) -> SeaCacheRootHook:
+    return apply_sea_cache_hook(
+        transformer,
+        config or SeaCacheConfig(threshold=100.0),
+        current_step_callback=lambda: metadata.step,
+        current_sigma_callback=lambda: metadata.sigma,
+        num_inference_steps_callback=lambda: metadata.num_steps,
+        extractor_fn=_extract_tiny_cosmos3_context,
+    )
+
+
+def test_config_validation() -> None:
+    assert SeaCacheConfig().threshold == 0.25
+    with pytest.raises(ValueError, match="residual_order"):
+        SeaCacheConfig(residual_order=-1)
+    with pytest.raises(ValueError, match="max_consecutive_cached"):
+        SeaCacheConfig(max_consecutive_cached=-1)
+
+
+def test_sea_filter_matches_reference_equation() -> None:
+    hidden = torch.randn(3, 4, 5, 2, dtype=torch.float32)
+    sigma = 0.4
+    power_exp = 3.0
+
+    spectrum = torch.fft.fftn(hidden, dim=(0, 1, 2))
+    gain = None
+    for axis in (0, 1, 2):
+        frequencies = torch.fft.fftfreq(hidden.shape[axis], dtype=torch.float32)
+        clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16)
+        axis_gain = (1.0 - sigma) * clean_power / ((1.0 - sigma) ** 2 * clean_power + sigma**2 + 1e-16)
+        shape = [1] * hidden.ndim
+        shape[axis] = hidden.shape[axis]
+        gain = axis_gain.reshape(shape) if gain is None else gain * axis_gain.reshape(shape)
+    assert gain is not None
+    gain = gain / gain.mean()
+    expected = torch.fft.ifftn(spectrum * gain, dim=(0, 1, 2)).real
+
+    torch.testing.assert_close(
+        apply_sea_filter(hidden, sigma=sigma, power_exp=power_exp),
+        expected,
+    )
+
+
+def test_indicator_distance_and_linear_extrapolation() -> None:
+    previous = [torch.ones(2, 2)]
+    current = [torch.full((2, 2), 1.5)]
+    assert indicator_distance(current, previous) == pytest.approx(0.5)
+    assert indicator_distance([torch.ones(3)], previous) == float("inf")
+
+    history = [
+        (0, torch.full((2, 2), 2.0)),
+        (2, torch.full((2, 2), 6.0)),
+    ]
+    torch.testing.assert_close(
+        extrapolate_residual(history, step=3, order=1),
+        torch.full((2, 2), 8.0),
+    )
+    torch.testing.assert_close(
+        extrapolate_residual(history, step=3, order=0),
+        torch.full((2, 2), 6.0),
+    )
+    quadratic_history = [
+        (0, torch.zeros(2, 2)),
+        (1, torch.ones(2, 2)),
+        (2, torch.full((2, 2), 4.0)),
+    ]
+    torch.testing.assert_close(
+        extrapolate_residual(quadratic_history, step=3, order=2),
+        torch.full((2, 2), 9.0),
+    )
+
+
+def test_hook_skips_middle_steps_and_forces_endpoints() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=4)
+    hook = _apply_test_hook(
+        transformer,
+        metadata,
+        SeaCacheConfig(threshold=100.0, max_consecutive_cached=2),
+    )
+    hook.refresh(transformer)
+
+    for step, timestep in enumerate((1000, 750, 500, 250)):
+        metadata.step = step
+        metadata.sigma = timestep / 1000
+        _run_step(transformer, timestep, 1.0 - step * 0.01, hook=hook)
+
+    assert hook.full_count == 2
+    assert hook.skip_count == 2
+    assert [step for step, _ in hook.state_manager._states["cond"].history] == [0, 3]
+
+
+def test_parameter_sharded_hook_skips_when_all_shard_ranks_agree(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook._parameter_sharded = True
+    fs_group = object()
+    reduced_groups: list[object] = []
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_fs_group",
+        lambda: SimpleNamespace(world_size=2, device_group=fs_group),
+    )
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
+
+    def all_reduce(decision, *, op, group):
+        assert op == torch.distributed.ReduceOp.MAX
+        reduced_groups.append(group)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
+
+    _run_step(transformer, 1000, 1.0, hook=hook)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.99, hook=hook)
+
+    assert hook.full_count == 1
+    assert hook.skip_count == 1
+    assert reduced_groups == [fs_group, fs_group]
+
+
+def test_parameter_sharded_peer_forces_full_compute(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook._parameter_sharded = True
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_fs_group",
+        lambda: SimpleNamespace(world_size=2, device_group=object()),
+    )
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
+    reduce_count = 0
+
+    def all_reduce(decision, *, op, group):
+        nonlocal reduce_count
+        del op, group
+        reduce_count += 1
+        if reduce_count == 2:
+            decision.fill_(1)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
+
+    _run_step(transformer, 1000, 1.0, hook=hook)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.99, hook=hook)
+
+    assert hook.full_count == 2
+    assert hook.skip_count == 0
+    assert hook.state_manager._states["cond"].accumulated_distance == 0.0
+
+
+def _uneven_cfg_sharded_worker(
+    rank: int,
+    init_method: str,
+    result_queue: torch.multiprocessing.Queue,
+) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    torch.distributed.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank,
+        world_size=4,
+        timeout=datetime.timedelta(seconds=10),
+    )
+    try:
+        fs_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        sp_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        cfg_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 2], [1, 3])
+        ]
+        fs_group = next(group for ranks, group in fs_groups if rank in ranks)
+        sp_group = next(group for ranks, group in sp_groups if rank in ranks)
+        cfg_group = next(group for ranks, group in cfg_groups if rank in ranks)
+        parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
+        parallel_state._SP = SimpleNamespace(world_size=2, device_group=sp_group)
+
+        hook = SeaCacheRootHook(SeaCacheConfig())
+        hook._parameter_sharded = True
+        hook._synchronize_compute(False, torch.device("cpu"))
+        if rank in (0, 1):
+            hook._synchronize_compute(False, torch.device("cpu"))
+
+        gathered = [torch.zeros(1) for _ in range(2)]
+        torch.distributed.all_gather(
+            gathered,
+            torch.tensor([rank], dtype=torch.float32),
+            group=cfg_group,
+        )
+        result_queue.put((rank, [int(value.item()) for value in gathered]))
+    finally:
+        parallel_state._FS = None
+        parallel_state._SP = None
+        torch.distributed.destroy_process_group()
+
+
+def test_parameter_sharded_hook_handles_uneven_cfg_branch_dispatch() -> None:
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    result_queue = manager.Queue()
+    torch.multiprocessing.spawn(
+        _uneven_cfg_sharded_worker,
+        args=(get_distributed_init_method("seacache_uneven_cfg_"), result_queue),
+        nprocs=4,
+    )
+
+    results = sorted(result_queue.get(timeout=2) for _ in range(4))
+    assert results == [
+        (0, [0, 2]),
+        (1, [1, 3]),
+        (2, [0, 2]),
+        (3, [1, 3]),
+    ]
+
+
+def _hybrid_sp_worker(
+    rank: int,
+    init_method: str,
+    result_queue: torch.multiprocessing.Queue,
+) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    torch.distributed.init_process_group(
+        "gloo",
+        init_method=init_method,
+        rank=rank,
+        world_size=4,
+        timeout=datetime.timedelta(seconds=10),
+    )
+    try:
+        fs_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        ulysses_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 1], [2, 3])
+        ]
+        ring_groups = [
+            (ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
+            for ranks in ([0, 2], [1, 3])
+        ]
+        sp_group = torch.distributed.new_group(
+            [0, 1, 2, 3],
+            timeout=datetime.timedelta(seconds=10),
+        )
+        fs_group = next(group for ranks, group in fs_groups if rank in ranks)
+        ulysses_group = next(group for ranks, group in ulysses_groups if rank in ranks)
+        ring_group = next(group for ranks, group in ring_groups if rank in ranks)
+        parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
+        parallel_state._SP = SimpleNamespace(
+            world_size=4,
+            device_group=sp_group,
+            ulysses_world_size=2,
+            ulysses_group=ulysses_group,
+            ring_world_size=2,
+            ring_group=ring_group,
+        )
+
+        compute = rank == 3
+        non_sharded_hook = SeaCacheRootHook(SeaCacheConfig())
+        non_sharded_result = non_sharded_hook._synchronize_compute(compute, torch.device("cpu"))
+
+        sharded_hook = SeaCacheRootHook(SeaCacheConfig())
+        sharded_hook._parameter_sharded = True
+        sharded_result = sharded_hook._synchronize_compute(compute, torch.device("cpu"))
+        result_queue.put((rank, non_sharded_result, sharded_result))
+    finally:
+        parallel_state._FS = None
+        parallel_state._SP = None
+        torch.distributed.destroy_process_group()
+
+
+def test_hook_synchronizes_full_hybrid_sequence_parallel_group() -> None:
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    result_queue = manager.Queue()
+    torch.multiprocessing.spawn(
+        _hybrid_sp_worker,
+        args=(get_distributed_init_method("seacache_hybrid_sp_"), result_queue),
+        nprocs=4,
+    )
+
+    results = sorted(result_queue.get(timeout=2) for _ in range(4))
+    assert results == [
+        (0, True, True),
+        (1, True, True),
+        (2, True, True),
+        (3, True, True),
+    ]
+
+
+def test_parameter_sharded_hook_fails_open_without_distributed_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    hook = SeaCacheRootHook(SeaCacheConfig())
+    hook._parameter_sharded = True
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+
+    assert hook._synchronize_compute(False, torch.device("cpu")) is True
+
+
+def test_hook_keeps_three_transfer_branches_separate() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+
+    contexts = ("cond", "cond_no_control", "uncond")
+    for step, (timestep, value) in enumerate(((1000, 1.0), (500, 0.9))):
+        metadata.step = step
+        metadata.sigma = timestep / 1000
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[0], control=0.5)
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[1])
+        _run_step(transformer, timestep, value, hook=hook, context=contexts[2], control=0.5)
+
+    assert set(hook.state_manager._states) == set(contexts)
+    assert hook.full_count == 3
+    assert hook.skip_count == 3
+    assert len(hook.state_manager._states["cond"].previous_indicator) == 2
+    assert len(hook.state_manager._states["cond_no_control"].previous_indicator) == 1
+
+
+def test_hook_fails_open_without_noisy_vision() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    all_clean = torch.zeros(1, 1, 2, 1, 1)
+
+    _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
+    metadata.step = 1
+    metadata.sigma = 0.5
+    _run_step(transformer, 500, 0.9, hook=hook, noisy_frame_mask=all_clean)
+
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+
+
+def test_conditioning_only_model_failure_runs_forward_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    forward_calls = 0
+
+    def fail_forward(hidden_gen: torch.Tensor) -> torch.Tensor:
+        nonlocal forward_calls
+        forward_calls += 1
+        raise RuntimeError("model forward failed")
+
+    monkeypatch.setattr(transformer, "_run_gen_layers", fail_forward)
+    all_clean = torch.zeros(1, 1, 2, 1, 1)
+
+    with pytest.raises(RuntimeError, match="model forward failed"):
+        _run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
+
+    assert forward_calls == 1
+
+
+def test_hook_fails_open_without_explicit_context() -> None:
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=2)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+
+    _run_step(transformer, 1000, 1.0)
+
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+    assert hook.state_manager._states == {}
+
+
+def test_hook_uses_exact_sigma_callback(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.cache.seacache import hook as hook_module
+
+    transformer = TinyCosmos3Transformer()
+    metadata = SimpleNamespace(step=0, sigma=0.37, num_steps=2)
+    observed_sigmas: list[float] = []
+    original_filter = hook_module.apply_sea_filter
+
+    def recording_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float) -> torch.Tensor:
+        observed_sigmas.append(sigma)
+        return original_filter(hidden_states, sigma, power_exp)
+
+    monkeypatch.setattr(hook_module, "apply_sea_filter", recording_filter)
+    hook = _apply_test_hook(transformer, metadata)
+    hook.refresh(transformer)
+    _run_step(transformer, timestep=999, value=1.0, hook=hook)
+
+    assert observed_sigmas
+    assert all(sigma == pytest.approx(0.37) for sigma in observed_sigmas)
+
+
+def test_backend_selector_and_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
+
+    class FSDPTinyCosmos3Transformer(TinyCosmos3Transformer):
+        pass
+
+    assert get_extractor(FSDPTinyCosmos3Transformer) is _extract_tiny_cosmos3_context
+    backend = get_cache_backend(
+        "sea_cache",
+        {
+            "sea_threshold": 0.4,
+            "sea_residual_order": 0,
+        },
+    )
+    assert isinstance(backend, SeaCacheBackend)
+    assert backend.config.sea_threshold == 0.4
+
+    pipeline = Cosmos3OmniDiffusersPipeline()
+    backend.enable(pipeline)
+    hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
+    assert isinstance(hook, SeaCacheRootHook)
+    assert callable(getattr(pipeline, "_cache_context_factory", None))
+    for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
+        assert not hasattr(pipeline.transformer, name)
+    pipeline._current_step_index = 0
+    pipeline._current_sigma = 1.0
+    pipeline._num_timesteps = 7
+    _run_step(pipeline.transformer, 1000, 1.0, hook=hook)
+    assert hook.full_count == 1
+
+    backend.refresh(pipeline, num_inference_steps=7)
+    assert hook.full_count == 0
+    assert hook.skip_count == 0
+    assert hook.state_manager._states == {}
+
+
+def test_backend_uses_resolved_pipeline_metadata_not_refresh_argument(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
+    pipeline = Cosmos3OmniDiffusersPipeline()
+    backend = SeaCacheBackend(DiffusionCacheConfig())
+    backend.enable(pipeline)
+    hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
+    assert isinstance(hook, SeaCacheRootHook)
+    backend.refresh(pipeline, num_inference_steps=35)
+    pipeline._num_timesteps = 50
+    pipeline._current_step_index = 0
+    pipeline._current_sigma = 1.0
+    _run_step(pipeline.transformer, 1000, 1.0, hook=hook)
+    assert hook.full_count == 1
+    assert hook.num_inference_steps_callback is not None
+    assert hook.num_inference_steps_callback() == 50
+
+    pipeline._current_step_index = 1
+    pipeline._current_sigma = 0.98
+    _run_step(pipeline.transformer, 980, 0.99, hook=hook)
+    assert hook.full_count == 1
+    assert hook.skip_count == 1
+
+
+def test_shared_config_defaults() -> None:
+    config = DiffusionCacheConfig()
+    assert config.sea_threshold == 0.25
+    assert config.sea_residual_order == 1
+    assert config.sea_max_consecutive_cached == 2
+    assert config.sea_power_exp == 3.0

tests/diffusion/cache/test_teacache_extractors.py

@@ -68,12 +68,12 @@ class BaseExtractorTest(ABC):
         pass
 
     @abstractmethod
-    def get_module(self):
+    def get_module(self, module, /):
         """Return model module instance."""
         pass
 
     @abstractmethod
-    def get_sample_inputs(self):
+    def get_sample_inputs(self, sample_inputs, /):
         """Return sample inputs for model."""
         pass
 
@@ -737,10 +737,10 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
         monkeypatch,
     ):
         """Strict SP must not mix local TeaCache state with gathered block rows."""
-        from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
         from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
         from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
         from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.layers import indexed_modulation
 
         seq_len = sample_inputs["x"].shape[1]
         local_len = seq_len // 2
@@ -844,13 +844,13 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
         )
         monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
 
-        original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
+        original_indexed_scale_shift = indexed_modulation.indexed_scale_shift_
 
         def capture_indexed_scale_shift(hidden, shift, scale, indices):
             captured_modulation_indices.append(indices.clone())
             return original_indexed_scale_shift(hidden, shift, scale, indices)
 
-        monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
+        monkeypatch.setattr(indexed_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
 
         local_img_pos = torch.tensor([0, 1])
         local_audio_pos = torch.tensor([2, 3])

tests/diffusion/layers/test_qwen_select01_modulation.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.layers import qwen_select01_modulation
+from vllm_omni.diffusion.layers.qwen_select01_modulation import (
+    can_use_qwen_select01_triton,
+    fused_layernorm_select01,
+    fused_residual_layernorm_select01,
+    select01_modulation_native,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("platform_is_cuda", "tensor_is_cuda", "has_triton", "expected"),
+    [
+        (True, True, True, True),
+        (False, True, True, False),
+        (True, False, True, False),
+        (True, True, False, False),
+    ],
+)
+def test_qwen_select01_triton_requires_cuda_platform(
+    monkeypatch,
+    platform_is_cuda: bool,
+    tensor_is_cuda: bool,
+    has_triton: bool,
+    expected: bool,
+):
+    class FakePlatform:
+        @staticmethod
+        def is_cuda() -> bool:
+            return platform_is_cuda
+
+    class FakeTensor:
+        is_cuda = tensor_is_cuda
+
+    monkeypatch.setattr(qwen_select01_modulation, "current_platform", FakePlatform())
+    monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", has_triton)
+
+    assert can_use_qwen_select01_triton(FakeTensor()) is expected
+
+
+def _make_inputs(
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, ...]:
+    generator = torch.Generator(device=device)
+    generator.manual_seed(202817)
+    x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
+    residual = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
+    residual_gate = torch.randn(batch_size, 1, hidden_size, dtype=dtype, device=device, generator=generator)
+    mod_params = torch.randn(batch_size * 2, hidden_size * 3, dtype=dtype, device=device, generator=generator)
+    index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int64, device=device, generator=generator)
+    return x, residual, residual_gate, mod_params, index
+
+
+def _reference_layernorm_select01(
+    x: torch.Tensor,
+    mod_params: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    scale, shift, gate = select01_modulation_native(mod_params, index)
+    out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype)
+    return out * (1 + scale) + shift, gate
+
+
+def _reference_residual_layernorm_select01(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    residual_gate: torch.Tensor,
+    mod_params: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    scale, shift, gate = select01_modulation_native(mod_params, index)
+    residual_out = residual + residual_gate * x
+    out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype)
+    return out * (1 + scale) + shift, residual_out, gate
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("broadcast_index", [False, True])
+def test_qwen_select01_native_fallback_matches_reference(dtype: torch.dtype, broadcast_index: bool):
+    device = torch.device("cpu")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 64, dtype, device)
+    if broadcast_index:
+        index = index[:1]
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    ("dtype", "atol", "rtol"),
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 5e-2, 5e-2),
+    ],
+)
+@pytest.mark.parametrize("seq_len", [1, 7, 128, 1024])
+def test_qwen_select01_triton_matches_reference(dtype: torch.dtype, atol: float, rtol: float, seq_len: int):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, seq_len, 256, dtype, device)
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    ("dtype", "atol", "rtol"),
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 5e-2, 5e-2),
+    ],
+)
+def test_qwen_select01_triton_matches_reference_for_structured_multi_batch_index(
+    dtype: torch.dtype,
+    atol: float,
+    rtol: float,
+):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, _ = _make_inputs(3, 16, 256, dtype, device)
+    index = torch.tensor(
+        [
+            [0] * 16,
+            [1] * 16,
+            [0, 1] * 8,
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    ("dtype", "atol", "rtol"),
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 5e-2, 5e-2),
+    ],
+)
+def test_qwen_select01_triton_broadcasts_single_index_batch(dtype: torch.dtype, atol: float, rtol: float):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 256, dtype, device)
+    index = index[:1]
+    eps = 1e-6
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate.expand_as(residual),
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_qwen_select01_cuda_falls_back_without_triton(monkeypatch):
+    device = torch.device("cuda:0")
+    x, residual, residual_gate, mod_params, index = _make_inputs(2, 7, 256, torch.bfloat16, device)
+    residual_gate = residual_gate.expand_as(residual)
+    eps = 1e-6
+
+    def fail_if_launched(*args, **kwargs):
+        pytest.fail("Triton launcher must not run when HAS_TRITON is false")
+
+    monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", False)
+    monkeypatch.setattr(qwen_select01_modulation, "_launch_layernorm_select01", fail_if_launched)
+    monkeypatch.setattr(qwen_select01_modulation, "_launch_residual_layernorm_select01", fail_if_launched)
+
+    actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
+    ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
+    torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
+
+    actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
+        x,
+        residual,
+        residual_gate,
+        mod_params,
+        index,
+        eps,
+    )
+    torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
+    torch.testing.assert_close(actual_residual, ref_residual, atol=0, rtol=0)
+    torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -440,29 +440,86 @@ def test_transformer_rejects_prompt_tuning():
         BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
 
 
-def _native_to_checkpoint_name(name: str) -> str:
-    """Inverse of ``load_weights`` remapping: native param name -> diffusers name.
+_QKV_FANOUT = {
+    "to_qkv": ("to_q", "to_k", "to_v"),
+    "img_to_qkv": ("img_to_q", "img_to_k", "img_to_v"),
+    "instruct_to_qkv": ("instruct_to_q", "instruct_to_k", "instruct_to_v"),
+}
+
+
+def _native_to_checkpoint_weights(name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]:
+    """Inverse of ``load_weights`` remapping: native param -> diffusers weights.
+
+    The diffusers checkpoint stores fused projections as separate matrices, so a
+    single merged native param (QKV or FFN gate/up) fans out into several
+    checkpoint weights. Returns ``(diffusers_name, value)`` pairs whose values,
+    when fed through ``load_weights``, reassemble into the original ``param``.
 
     - ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
     - promoted joint-attention projections move back under ``.processor.``.
     """
+    q_size = NUM_HEADS * HEAD_DIM
+    kv_size = NUM_KV_HEADS * HEAD_DIM
+
+    # Fused QKV projections split back into per-matrix q/k/v weights.
+    for token, (q_name, k_name, v_name) in _QKV_FANOUT.items():
+        marker = f".{token}."
+        if marker in name:
+            q, k, v = param[:q_size], param[q_size : q_size + kv_size], param[q_size + kv_size :]
+            results = []
+            for sub_name, value in ((q_name, q), (k_name, k), (v_name, v)):
+                ckpt_name = name.replace(marker, f".{sub_name}.")
+                if token != "to_qkv":
+                    # Joint-attention projections live under `.processor.` upstream.
+                    ckpt_name = ckpt_name.replace(".img_instruct_attn.", ".img_instruct_attn.processor.")
+                results.append((ckpt_name, value))
+            return results
+
+    # Fused FFN gate/up splits into linear_1 (gate) / linear_3 (input).
+    if ".gate_up_proj." in name:
+        inner = param.shape[0] // 2
+        return [
+            (name.replace(".gate_up_proj.", ".linear_1."), param[:inner]),
+            (name.replace(".gate_up_proj.", ".linear_3."), param[inner:]),
+        ]
+
+    # Default: 1:1 with the existing diffusers name promotions.
     if ".to_out." in name:
         name = name.replace(".to_out.", ".to_out.0.")
-    for proj in (
-        "img_to_q",
-        "img_to_k",
-        "img_to_v",
-        "instruct_to_q",
-        "instruct_to_k",
-        "instruct_to_v",
-        "instruct_out",
-        "img_out",
-    ):
+    for proj in ("instruct_out", "img_out"):
         token = f".img_instruct_attn.{proj}."
         if token in name:
             name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
             break
-    return name
+    return [(name, param)]
+
+
+def test_transformer_exposes_stacked_params_mapping():
+    """The packed -> sub-layer mapping must be discoverable from the module tree.
+
+    ``diffusion/lora/loader.py`` and the quantized weight loaders read
+    ``stacked_params_mapping`` off the model, so it has to exist before (and
+    independently of) ``load_weights``.
+    """
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _BOOGU_STACKED_PARAMS_MAPPING,
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+
+    assert tuple(model.stacked_params_mapping) == _BOOGU_STACKED_PARAMS_MAPPING
+    # A per-instance copy, so consumers cannot mutate the module constant.
+    assert model.stacked_params_mapping is not _BOOGU_STACKED_PARAMS_MAPPING
+
+    # Every entry is a (param, shard, shard_id) triple.
+    for param_name, shard_name, shard_id in model.stacked_params_mapping:
+        assert param_name.startswith(".") and shard_name.startswith(".")
+        assert shard_id in {"q", "k", "v", 0, 1}
+
+    # The mapping targets exactly the fused projections this port creates.
+    mapped_leaves = {param.strip(".").split(".")[-1] for param, _, _ in model.stacked_params_mapping}
+    assert mapped_leaves == {"to_qkv", "img_to_qkv", "instruct_to_qkv", "gate_up_proj"}
 
 
 def test_transformer_load_weights_round_trip():
@@ -473,24 +530,59 @@ def test_transformer_load_weights_round_trip():
     model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
     native_params = dict(model.named_parameters())
 
-    # Build synthetic diffusers-named weights (one per native parameter).
-    checkpoint_weights = {}
+    # Build synthetic diffusers-named weights. Fused native projections fan out
+    # into the separate matrices the diffusers checkpoint stores.
+    expected: dict[str, torch.Tensor] = {}
+    checkpoint_weights: dict[str, torch.Tensor] = {}
     for native_name, param in native_params.items():
-        checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
+        full = torch.randn_like(param)
+        expected[native_name] = full
+        for ckpt_name, value in _native_to_checkpoint_weights(native_name, full):
+            checkpoint_weights[ckpt_name] = value
 
-    # The remapping must be a bijection over the parameter set.
-    assert len(checkpoint_weights) == len(native_params)
+    # Every checkpoint name is unique; merged params fan out to >=1 weight.
+    assert len(checkpoint_weights) >= len(native_params)
 
     loaded = model.load_weights(list(checkpoint_weights.items()))
 
     # No missing / unexpected parameters.
     assert loaded == set(native_params.keys())
 
-    # Values landed on the right parameters (TP=1: weight_loader copies verbatim).
+    # Values landed on the right parameters (TP=1: weight_loader copies verbatim;
+    # fused shards are placed at their q/k/v / gate/up offsets).
     reloaded = dict(model.named_parameters())
     for native_name in native_params:
-        expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
-        assert torch.allclose(reloaded[native_name], expected)
+        assert torch.allclose(reloaded[native_name], expected[native_name])
+
+
+def test_transformer_load_weights_rejects_partial_fused_params():
+    """A fused parameter must not count as loaded until all its shards arrive.
+
+    The loader compares parameter *names*, so reporting ``to_qkv`` complete
+    after a single ``to_q`` would let a checkpoint carrying only ``to_q`` start
+    up with the ``k``/``v`` slices left uninitialized.
+    """
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    native_params = dict(model.named_parameters())
+
+    complete_name = "noise_refiner.0.attn.to_qkv.weight"
+    partial_name = "noise_refiner.0.feed_forward.gate_up_proj.weight"
+    assert complete_name in native_params
+    assert partial_name in native_params
+
+    weights = list(_native_to_checkpoint_weights(complete_name, native_params[complete_name]))
+    # Only the gate half of the FFN, so its fused parameter stays incomplete.
+    gate_ckpt, gate_value = _native_to_checkpoint_weights(partial_name, native_params[partial_name])[0]
+    weights.append((gate_ckpt, gate_value))
+
+    loaded = model.load_weights(weights)
+
+    assert complete_name in loaded
+    assert partial_name not in loaded
 
 
 def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
@@ -513,12 +605,15 @@ def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
     try:
         model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
         native_params = dict(model.named_parameters())
-        loaded_name = next(iter(native_params))
-        checkpoint_name = _native_to_checkpoint_name(loaded_name)
+        # A non-fused parameter, so one checkpoint entry loads it completely.
+        loaded_name = "x_embedder.weight"
+        assert loaded_name in native_params
+        (checkpoint_name, checkpoint_value), *_ = _native_to_checkpoint_weights(loaded_name, native_params[loaded_name])
+        checkpoint_value = torch.randn_like(checkpoint_value)
 
         loaded = model.load_weights(
             [
-                (checkpoint_name, torch.randn_like(native_params[loaded_name])),
+                (checkpoint_name, checkpoint_value),
                 ("unexpected.weight", torch.ones(1)),
             ]
         )

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
 import json
 import sys
 import types
+from contextlib import contextmanager
 from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
@@ -85,6 +86,7 @@ class StubScheduler:
         flow_shift: float = 1.0,
     ) -> None:
         self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
+        self.sigmas = self.timesteps.float() / 10
         self.config = SimpleNamespace(
             num_train_timesteps=1000,
             flow_shift=flow_shift,
@@ -113,6 +115,7 @@ class StubScheduler:
         else:
             assert num_inference_steps is not None
             self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device)
+        self.sigmas = self.timesteps.float()
 
     def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
         del kwargs
@@ -250,7 +253,7 @@ def passthrough_progress_bar(iterable):
 
 @pytest.fixture(autouse=True)
 def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
-    module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
+    module: Any = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
     module.is_guardrails_enabled = lambda od_config, sampling_params=None: False
     module.ensure_initialized = lambda od_config: None
     module.check_text_safety = lambda text: None
@@ -283,6 +286,8 @@ def make_cosmos3_pipeline():
         pipeline.is_edge_model = False
         pipeline._guidance_scale = None
         pipeline._num_timesteps = None
+        pipeline._current_step_index = None
+        pipeline._current_sigma = None
         pipeline._cosmos3_branch_caches = None
         pipeline._cache_dit_requires_paired_cfg = False
         pipeline._sound_tokenizer = None
@@ -300,7 +305,7 @@ def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
 
 
 def make_sampling_params(**overrides: Any) -> SimpleNamespace:
-    values = {
+    values: dict[str, Any] = {
         "height": None,
         "width": None,
         "num_frames": None,
@@ -814,7 +819,7 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
     t_list = [1.0, 0.75, 0.5, 0.25]
     scheduler_dir = tmp_path / "scheduler"
     scheduler_dir.mkdir()
-    scheduler_config = {"_class_name": scheduler_class_name}
+    scheduler_config: dict[str, Any] = {"_class_name": scheduler_class_name}
     if expected_distilled:
         scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list}
     (scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config))
@@ -1531,7 +1536,7 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke
         fused_add_rms_norm: list[str]
         custom_op: list[str]
 
-    fake_kernel = types.ModuleType("vllm.config.kernel")
+    fake_kernel: Any = types.ModuleType("vllm.config.kernel")
     fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig
     monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel)
 
@@ -2154,6 +2159,47 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
     torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
 
 
+def test_diffuse_publishes_exact_seacache_metadata_and_cfg_contexts(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.scheduler.sigmas = torch.tensor([0.91, 0.17])
+    observations: list[tuple[str, int | None, float | None, int | None]] = []
+
+    class RecordingHook:
+        @contextmanager
+        def cache_context(self, name: str):
+            sigma = pipeline.current_sigma
+            observations.append(
+                (
+                    name,
+                    pipeline.current_step_index,
+                    None if sigma is None else float(sigma),
+                    pipeline.num_timesteps,
+                )
+            )
+            yield
+
+    pipeline._cache_context_factory = RecordingHook().cache_context
+    pipeline.diffuse(
+        latents=torch.zeros(1, 2, 1, 1, 1),
+        timesteps=torch.tensor([900, 100]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+    )
+
+    assert observations == [
+        ("cond", 0, pytest.approx(0.91), 2),
+        ("uncond", 0, pytest.approx(0.91), 2),
+        ("cond", 1, pytest.approx(0.17), 2),
+        ("uncond", 1, pytest.approx(0.17), 2),
+    ]
+    assert pipeline.current_step_index is None
+    assert pipeline.current_sigma is None
+
+
 def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     pipeline._use_session_state = True
@@ -2213,6 +2259,53 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
     torch.testing.assert_close(result, torch.full_like(latents, 254.0))
 
 
+def test_diffuse_transfer_uses_named_seacache_contexts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.scheduler.sigmas = torch.tensor([0.42])
+    contexts: list[tuple[str, int | None, float | None, int | None]] = []
+
+    class RecordingHook:
+        @contextmanager
+        def cache_context(self, name: str):
+            sigma = pipeline.current_sigma
+            contexts.append(
+                (
+                    name,
+                    pipeline.current_step_index,
+                    None if sigma is None else float(sigma),
+                    pipeline.num_timesteps,
+                )
+            )
+            yield
+
+    pipeline._cache_context_factory = RecordingHook().cache_context
+    latents = torch.zeros(1, 2, 1, 1, 1)
+    velocity_mask = torch.ones(1, 1, 1, 1, 1)
+    pipeline.diffuse_transfer(
+        latents=latents,
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        control_guidance=1.5,
+        control_guidance_interval=None,
+        control_latents=[torch.zeros_like(latents)],
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
+        velocity_mask=velocity_mask,
+        condition_latents=torch.zeros_like(latents),
+    )
+
+    assert contexts == [
+        ("cond", 0, pytest.approx(0.42), 1),
+        ("cond_no_control", 0, pytest.approx(0.42), 1),
+        ("uncond", 0, pytest.approx(0.42), 1),
+    ]
+    assert pipeline.current_step_index is None
+    assert pipeline.current_sigma is None
+
+
 def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     pipeline._use_session_state = True
@@ -2583,7 +2676,7 @@ def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -
 
 class TestForwardRouting:
     def _install_forward_stubs(self, pipeline):
-        captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []}
+        captured: dict[str, Any] = {"diffuse_calls": [], "prepare_calls": []}
 
         def fake_format(
             prompt,

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -513,6 +513,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.cache.teacache.extractors import extract_cosmos3_context
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    class TrackingNorm(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.calls = 0
+
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            self.calls += 1
+            return hidden_states + 5.0
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    norm = TrackingNorm()
+    model.norm_moe_gen = norm
+    captured: dict[str, torch.Tensor] = {}
+
+    def run_gen_layers(hidden_gen: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        captured["input"] = hidden_gen.detach().clone()
+        return hidden_gen + 2.0
+
+    monkeypatch.setattr(model, "_run_gen_layers", run_gen_layers)
+    forward_kwargs = {
+        "hidden_states": torch.zeros(1, 2, 1, 2, 2),
+        "timestep": torch.tensor([1.0]),
+        "text_ids": torch.tensor([[1, 2]], dtype=torch.long),
+        "text_mask": torch.ones(1, 2, dtype=torch.long),
+        "video_shape": (1, 2, 2),
+        "fps": 24.0,
+    }
+
+    full_output = model(**forward_kwargs)
+
+    assert norm.calls == 1
+    norm.calls = 0
+
+    ctx = extract_cosmos3_context(model, **forward_kwargs)
+    execution_input = ctx.hidden_states.detach().clone()
+    execution_output = ctx.run_transformer_blocks()[0]
+    residual = execution_output - execution_input
+    extracted_output = ctx.postprocess(execution_output)
+
+    assert norm.calls == 1
+    expected_residual = ((captured["input"] + 2.0) + 5.0) - captured["input"]
+    assert torch.equal(residual, expected_residual)
+    assert torch.equal(extracted_output, full_output)
+
+    norm.calls = 0
+
+    def fail_if_gen_layers_run(*args, **kwargs):
+        del args, kwargs
+        raise AssertionError("SeaCache hit unexpectedly executed GEN layers")
+
+    monkeypatch.setattr(model, "_run_gen_layers", fail_if_gen_layers_run)
+
+    cached_ctx = extract_cosmos3_context(model, **forward_kwargs)
+    cached_output = cached_ctx.postprocess(cached_ctx.hidden_states + residual)
+
+    assert norm.calls == 0
+    assert torch.equal(cached_output, full_output)
+    for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
+        assert not hasattr(model, name)
+
+
+def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    class CountingNorm(nn.Module):
+        def __init__(self, original: nn.Module) -> None:
+            super().__init__()
+            self.original = original
+            self.calls = 0
+
+        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+            self.calls += 1
+            return self.original(hidden_states)
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    norm = CountingNorm(model.norm_moe_gen)
+    model.norm_moe_gen = norm
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert norm.calls == 1
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+
+
 def test_qwen_gen_mlp_remains_gated() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
 

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -2711,3 +2711,125 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
     assert "from .pipeline import" in lingbot_init
     assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init
     assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init
+
+
+# ---------------------------------------------------------------------------
+# a prompt is encoded once, whichever path asks for it
+# ---------------------------------------------------------------------------
+def _counting_text_encoder(pipeline) -> list[str]:
+    """Give ``pipeline`` its real encode_prompt over stubs that record each encode."""
+    del pipeline.encode_prompt  # _pipeline() replaces it; these tests need the real one
+    encoded: list[str] = []
+
+    def tokenizer(texts, **kwargs):
+        del kwargs
+        encoded.extend(texts)
+        return SimpleNamespace(
+            input_ids=torch.arange(512).view(1, 512),
+            attention_mask=torch.ones(1, 512, dtype=torch.long),
+        )
+
+    def text_encoder(input_ids, attention_mask):
+        del input_ids, attention_mask
+        return SimpleNamespace(last_hidden_state=torch.full((1, 512, 8), float(len(encoded))))
+
+    pipeline.tokenizer = tokenizer
+    pipeline.text_encoder = text_encoder
+    return encoded
+
+
+def test_the_text_encoder_runs_once_per_distinct_prompt() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    encoded = _counting_text_encoder(pipeline)
+
+    def encode(prompt, *, length=512, dtype=torch.float32):
+        return pipeline.encode_prompt(prompt, max_sequence_length=length, dtype=dtype)
+
+    first = encode("move through the room")
+    assert encode("move through the room") is first
+    # The key is the text the tokenizer actually sees, so spacing alone is not new.
+    assert encode("  move   through the room ") is first
+    assert encoded == ["move through the room"]
+
+    encode("enter the snowy valley")
+    # dtype and sequence length change the tensor, so each is its own entry.
+    encode("move through the room", dtype=torch.float16)
+    encode("move through the room", length=256)
+    assert encoded == [
+        "move through the room",
+        "enter the snowy valley",
+        "move through the room",
+        "move through the room",
+    ]
+
+
+def test_the_prompt_encode_cache_is_bounded() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    encoded = _counting_text_encoder(pipeline)
+    size = module._PROMPT_EMBEDS_CACHE_SIZE
+
+    for index in range(size + 1):
+        pipeline.encode_prompt(f"prompt {index}", max_sequence_length=512, dtype=torch.float32)
+    assert len(pipeline._prompt_embeds_cache) == size
+
+    # The oldest entry was the one evicted, so it is the one encoded again.
+    pipeline.encode_prompt("prompt 0", max_sequence_length=512, dtype=torch.float32)
+    pipeline.encode_prompt(f"prompt {size}", max_sequence_length=512, dtype=torch.float32)
+    assert encoded == [f"prompt {index}" for index in range(size + 1)] + ["prompt 0"]
+
+
+def test_realtime_ticks_reuse_the_encode_and_still_invalidate_on_a_prompt_change(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    pipeline._ar_diffusion_kv_state = object()
+    encoded = _counting_text_encoder(pipeline)
+    invalidations: list[bool] = []
+
+    def ar_text_caches(prompt_embeds, *, invalidate):
+        del prompt_embeds
+        invalidations.append(invalidate)
+        return [SimpleNamespace()]
+
+    monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches)
+    monkeypatch.setattr(
+        pipeline,
+        "_generate_block",
+        lambda **kwargs: torch.randn((1, 16, 3, 2, 2), generator=kwargs["generator"]),
+    )
+
+    for chunk_index in range(3):
+        pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=chunk_index))))
+    switched = _prompt()
+    switched["prompt"] = "enter the snowy valley"
+    pipeline(
+        _request(
+            sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=3, prompt=switched["prompt"])),
+            prompt=switched,
+        )
+    )
+
+    assert encoded == ["move through the room", "enter the snowy valley"]
+    # A new prompt must still drop the cross-attention K/V built from the old one.
+    assert invalidations == [False, False, False, True]
+
+
+def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    encoded = _counting_text_encoder(pipeline)
+
+    states = [_stepwise_state(request_id=request_id, num_frames=21) for request_id in ("req-a", "req-b")]
+    for state in states:
+        with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+            pipeline.prepare_encode(state)
+
+    assert encoded == ["move through the room"]
+    assert states[0].prompt_embeds is states[1].prompt_embeds

tests/diffusion/models/magi2/test_native_compile.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from tests.diffusion.models.magi2.test_native_packing import (
+    _longer_text_tensors,
+    _sampler_tensors,
+    _tiny_model,
+    _tiny_sampler,
+)
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def test_layer_regions_bracket_the_eager_kernels() -> None:
+    moe_layer, dense_layer = _tiny_model().block.layers
+
+    assert moe_layer.region_methods == ("_attention_input", "_moe_input", "_moe_output")
+    assert dense_layer.region_methods == ("_attention_input", "_dense_output")
+
+
+def test_compile_regions_fullgraph_matches_eager_bitwise() -> None:
+    model = _tiny_model()
+    sampler = _tiny_sampler(model)
+    steps = [
+        sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
+        for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
+    ]
+    expected = [sampler.forward(step) for step in steps]
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    first = sampler.forward(steps[0])
+    graphs = torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        second = sampler.forward(steps[1])
+
+    for (video, audio), (video_ref, audio_ref) in zip((first, second), expected, strict=True):
+        assert torch.equal(video, video_ref)
+        assert torch.equal(audio, audio_ref)
+    # pre adapter, post adapter, three MoE-layer regions, two dense-layer regions
+    assert graphs == 7
+
+
+def test_pipeline_setup_compile_forwards_config_dynamic() -> None:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
+
+    calls: list[dict[str, object]] = []
+
+    class _Transformer(torch.nn.Module):
+        def compile_regions(self, **compile_kwargs: object) -> None:
+            calls.append(compile_kwargs)
+
+    pipeline = object.__new__(Magi2Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = OmniDiffusionConfig(
+        model="sand-ai/MAGI-2-preview",
+        model_class_name="Magi2Pipeline",
+        diffusion_compile_dynamic=False,
+    )
+    pipeline.transformer = _Transformer()
+
+    pipeline.setup_compile()
+
+    assert calls == [{"fullgraph": True, "dynamic": False, "options": {"emulate_precision_casts": True}}]
+
+
+def test_pipeline_setup_compile_warns_when_granularity_is_not_regional(caplog: pytest.LogCaptureFixture) -> None:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
+
+    calls: list[dict[str, object]] = []
+
+    class _Transformer(torch.nn.Module):
+        def compile_regions(self, **compile_kwargs: object) -> None:
+            calls.append(compile_kwargs)
+
+    pipeline = object.__new__(Magi2Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = OmniDiffusionConfig(
+        model="sand-ai/MAGI-2-preview",
+        model_class_name="Magi2Pipeline",
+        diffusion_compile_granularity="full",
+    )
+    pipeline.transformer = _Transformer()
+
+    pipeline.setup_compile()
+
+    assert calls == [{"fullgraph": True, "dynamic": True, "options": {"emulate_precision_casts": True}}]
+    assert "diffusion_compile_granularity='full'" in caplog.text
+
+
+def test_compiled_graphs_are_shared_across_layers_of_one_kind() -> None:
+    model = _tiny_model(moe_layers=2)
+    sampler = _tiny_sampler(model)
+    step = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([900.0]), cfg_config=CFGConfig())
+    expected = sampler.forward(step)
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    actual = sampler.forward(step)
+
+    assert torch.equal(actual[0], expected[0])
+    # Two MoE layers and two dense layers still compile one graph per region
+    # kind: pre adapter, post adapter, three MoE regions, two dense regions.
+    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
+
+
+def test_compiled_regions_survive_new_requests_without_recompiling() -> None:
+    model = _tiny_model()
+    sampler = _tiny_sampler(model)
+    t = torch.tensor([900.0])
+    first_request = sampler.prepare_model_input(**_sampler_tensors(0), t=t, cfg_config=CFGConfig())
+    second_request = sampler.prepare_model_input(**_sampler_tensors(1), t=t, cfg_config=CFGConfig())
+    longer_request = sampler.prepare_model_input(**_longer_text_tensors(2), t=t, cfg_config=CFGConfig())
+    expected = [sampler.forward(request) for request in (first_request, second_request, longer_request)]
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    actual = [sampler.forward(first_request)]
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        actual.append(sampler.forward(second_request))
+        actual.append(sampler.forward(longer_request))
+
+    for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
+        assert torch.equal(video, video_ref)
+        assert torch.equal(audio, audio_ref)
+    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7

tests/diffusion/models/magi2/test_native_compile_cuda.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from tests.diffusion.models.magi2.test_native_packing import (
+    _longer_text_tensors,
+    _sampler_tensors,
+    _tiny_model,
+    _tiny_sampler,
+)
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.core_model,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
+]
+
+
+def _cuda_steps(sampler, tensors=_sampler_tensors):
+    timesteps = (torch.tensor([900.0]), torch.tensor([600.0]), torch.tensor([300.0]))
+    return [
+        sampler.prepare_model_input(
+            **{name: value.cuda() for name, value in tensors(seed).items()},
+            t=t,
+            cfg_config=CFGConfig(),
+        )
+        for seed, t in enumerate(timesteps)
+    ]
+
+
+def test_inductor_regions_match_eager_in_deterministic_mode(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("MAGI2_DETERMINISTIC", "1")
+    model = _tiny_model(params_dtype=torch.bfloat16).cuda()
+    sampler = _tiny_sampler(model)
+    steps = _cuda_steps(sampler)
+    expected = [sampler.forward(step) for step in steps]
+    new_request_step = _cuda_steps(sampler)[0]
+    longer_text_step = _cuda_steps(sampler, _longer_text_tensors)[0]
+    expected_new_request = sampler.forward(new_request_step)
+    expected_longer_text = sampler.forward(longer_text_step)
+
+    torch._dynamo.reset()
+    graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+    # The pipeline compiles with inductor emulating eager's intermediate
+    # low-precision casts, so fused bf16 chains keep eager's rounding.
+    model.compile_regions(fullgraph=True, dynamic=True, options={"emulate_precision_casts": True})
+    first = sampler.forward(steps[0])
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        compiled = [first, sampler.forward(steps[1]), sampler.forward(steps[2])]
+        repeated = [sampler.forward(step) for step in steps]
+        # A second request repacks the same shapes; a third changes the
+        # text length and therefore every token count.
+        compiled_new_request = sampler.forward(new_request_step)
+        compiled_longer_text = sampler.forward(longer_text_step)
+
+    # Emulated precision casts restore eager's rounding boundaries; what is
+    # left is one-ulp transcendental noise that some Triton backends carry.
+    for (video, audio), (video_ref, audio_ref) in zip(compiled, expected, strict=True):
+        torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
+        torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
+    for (video, audio), (video_again, audio_again) in zip(compiled, repeated, strict=True):
+        assert torch.equal(video, video_again)
+        assert torch.equal(audio, audio_again)
+    for (video, audio), (video_ref, audio_ref) in (
+        (compiled_new_request, expected_new_request),
+        (compiled_longer_text, expected_longer_text),
+    ):
+        torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
+        torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
+    assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7

tests/diffusion/models/magi2/test_native_compile_distributed.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import os
+import tempfile
+from contextlib import ExitStack, contextmanager
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+from torch.distributed import init_device_mesh
+from torch.distributed.fsdp import MixedPrecisionPolicy
+
+import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dlo_module
+from tests.diffusion.models.magi2.test_native_distributed_parity import (
+    _current_group,
+    _new_groups,
+    _patched_groups,
+)
+from tests.diffusion.models.magi2.test_native_packing import _tiny_config, _tiny_model, _tiny_sampler
+from tests.diffusion.offloader.helpers import patch_offload_runtime
+from vllm_omni.diffusion.distributed.hsdp import shard_model
+from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
+from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
+from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
+from vllm_omni.diffusion.offloader.base import OffloadConfig
+from vllm_omni.diffusion.offloader.config import OffloadStrategy
+
+_WORLD_SIZE = 4
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+@contextmanager
+def _cpu_offload_runtime():
+    with pytest.MonkeyPatch.context() as monkeypatch:
+        patch_offload_runtime(monkeypatch, dlo_module.current_omni_platform, synchronize=True)
+        yield
+
+
+def _enable_dlo(
+    model: Magi2PreviewTransformer,
+    *,
+    dp_group: dist.ProcessGroup | None,
+    dp_size: int,
+    dp_rank: int,
+) -> None:
+    pipeline = Magi2Pipeline.__new__(Magi2Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.transformer = model
+    backend = dlo_module.DistributedLayerwiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            pin_cpu_memory=False,
+            dp_size=dp_size,
+            dlo_use_allgather=dp_size > 1,
+        ),
+        torch.device("cpu"),
+    )
+    backend.dp_group = dp_group
+    backend.rank = dp_rank
+    backend.enable(pipeline)
+    assert backend.enabled
+
+
+def _enable_hsdp(model: Magi2PreviewTransformer) -> None:
+    mesh = init_device_mesh("cpu", mesh_shape=(1, _WORLD_SIZE), mesh_dim_names=("replicate", "shard"))
+    ignored = {
+        parameter for name in model._hsdp_ignored_modules for parameter in model.get_submodule(name).parameters()
+    }
+    shard_model(
+        model,
+        mesh=mesh,
+        mp_policy=MixedPrecisionPolicy(cast_forward_inputs=False),
+        hsdp_shard_conditions=model._hsdp_shard_conditions,
+        ignored_params=ignored,
+    )
+    model.requires_grad_(False)
+
+
+def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
+    # 14 tokens per CFG branch, so the packed sequence splits evenly across
+    # two Ulysses ranks; gloo cannot gather uneven shards.
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    return {
+        "latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
+        "audio_latent": torch.randn(1, 5, 4, generator=generator),
+        "txt_feat": torch.randn(1, 3, 4, generator=generator),
+        "null_txt_feat": torch.randn(1, 3, 4, generator=generator),
+    }
+
+
+def _steps(sampler):
+    return [
+        sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
+        for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
+    ]
+
+
+def _run(model: Magi2PreviewTransformer, *, compile_regions: bool) -> list[tuple[torch.Tensor, torch.Tensor]]:
+    sampler = _tiny_sampler(model)
+    steps = _steps(sampler)
+    if not compile_regions:
+        return [sampler.forward(step) for step in steps]
+    torch._dynamo.reset()
+    model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
+    outputs = [sampler.forward(steps[0])]
+    with torch._dynamo.config.patch(error_on_recompile=True):
+        outputs.append(sampler.forward(steps[1]))
+    return outputs
+
+
+def _assert_same(actual, expected, *, exact: bool, label: str) -> None:
+    for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
+        if exact:
+            assert torch.equal(video, video_ref), label
+            assert torch.equal(audio, audio_ref), label
+        else:
+            torch.testing.assert_close(video, video_ref, atol=2e-5, rtol=2e-4, msg=label)
+            torch.testing.assert_close(audio, audio_ref, atol=2e-5, rtol=2e-4, msg=label)
+
+
+def _distributed_worker(rank: int, rendezvous: str) -> None:
+    torch.set_num_threads(1)
+    dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=_WORLD_SIZE)
+    try:
+        singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
+        with _patched_groups(singleton, singleton):
+            oracle = _tiny_model()
+            checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()]
+            expected = _run(oracle, compile_regions=False)
+
+        sp2_groups = _new_groups(((0, 1), (2, 3)))
+        dp2_groups = _new_groups(((0, 2), (1, 3)))
+        sp_group = _current_group(rank, sp2_groups)
+        dp_ranks, dp_process_group = next((ranks, group) for ranks, group in dp2_groups if rank in ranks)
+
+        def dlo_rank_local(model: Magi2PreviewTransformer) -> None:
+            _enable_dlo(model, dp_group=None, dp_size=1, dp_rank=0)
+
+        def dlo_allgather(model: Magi2PreviewTransformer) -> None:
+            _enable_dlo(model, dp_group=dp_process_group, dp_size=len(dp_ranks), dp_rank=dp_ranks.index(rank))
+
+        variants = (
+            ("dlo_rank_local_sp2", sp_group, dlo_rank_local),
+            ("dlo_allgather_dp2sp2", sp_group, dlo_allgather),
+            ("hsdp4", singleton, _enable_hsdp),
+        )
+        for name, ulysses_group, enable in variants:
+            with ExitStack() as stack:
+                stack.enter_context(_patched_groups(singleton, ulysses_group))
+                stack.enter_context(_cpu_offload_runtime())
+                outputs = {}
+                for compile_regions in (False, True):
+                    model = Magi2PreviewTransformer(_tiny_config())
+                    assert model.load_weights(checkpoint) == set(model.state_dict())
+                    enable(model)
+                    outputs[compile_regions] = _run(model, compile_regions=compile_regions)
+            _assert_same(outputs[False], expected, exact=False, label=f"{name} eager vs single-rank oracle")
+            _assert_same(outputs[True], outputs[False], exact=True, label=f"{name} compiled vs eager")
+            dist.barrier()
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(
+    not dist.is_available() or not dist.is_gloo_available(),
+    reason="requires torch.distributed gloo",
+)
+def test_compiled_regions_match_eager_under_dlo_and_hsdp() -> None:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
+        with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
+            mp.spawn(_distributed_worker, args=(rendezvous,), nprocs=_WORLD_SIZE, join=True)

tests/diffusion/models/magi2/test_native_packing.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
+    Magi2MHCConfig,
+    Magi2MoEConfig,
+    Magi2PreviewConfig,
+)
+from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm
+from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE
+from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
+from vllm_omni.diffusion.models.magi2.preview_data_proxy import (
+    Magi2DataProxy,
+    Magi2PreviewDataProxyConfig,
+)
+from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _tiny_config(params_dtype: torch.dtype = torch.float32, *, moe_layers: int = 1) -> Magi2PreviewConfig:
+    # The leading layers are multimodal with MoE, the rest single-modality
+    # dense, so every layer belongs to one of two kinds.
+    moe_indices = tuple(range(moe_layers))
+    return Magi2PreviewConfig(
+        num_layers=2 * moe_layers,
+        hidden_size=16,
+        head_dim=8,
+        num_query_groups=2,
+        video_in_channels=4,
+        audio_in_channels=4,
+        text_in_channels=4,
+        intermediate_factor=2,
+        multimodal_layers=moe_indices,
+        params_dtype=params_dtype,
+        mhc=Magi2MHCConfig(num_streams=2),
+        moe=Magi2MoEConfig(
+            num_heads=2,
+            num_experts=4,
+            top_k=2,
+            expert_intermediate_size=8,
+            shared_expert_intermediate_size=8,
+            modality_shared_expert_intermediate_size=8,
+            layers=moe_indices,
+        ),
+    )
+
+
+def _tiny_model(
+    seed: int = 11,
+    params_dtype: torch.dtype = torch.float32,
+    *,
+    moe_layers: int = 1,
+) -> Magi2PreviewTransformer:
+    model = Magi2PreviewTransformer(_tiny_config(params_dtype, moe_layers=moe_layers))
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    with torch.no_grad():
+        for parameter in model.parameters():
+            parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
+        for module in model.modules():
+            if isinstance(module, MultiModalityRMSNorm):
+                module.weight.zero_()
+            elif isinstance(module, Magi2MultiHeadMoE):
+                module.router.expert_bias.zero_()
+                module.router.expert_bias_ema.zero_()
+    return model
+
+
+def _tiny_sampler(model: torch.nn.Module) -> Magi2PreviewSampler:
+    return Magi2PreviewSampler(model, Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=8)))
+
+
+def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    return {
+        "latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
+        "audio_latent": torch.randn(1, 5, 4, generator=generator),
+        "txt_feat": torch.randn(1, 3, 4, generator=generator),
+        "null_txt_feat": torch.randn(1, 2, 4, generator=generator),
+    }
+
+
+def _longer_text_tensors(seed: int) -> dict[str, torch.Tensor]:
+    tensors = _sampler_tensors(seed)
+    generator = torch.Generator(device="cpu").manual_seed(seed + 100)
+    tensors["txt_feat"] = torch.randn(1, 4, 4, generator=generator)
+    return tensors
+
+
+def test_prepare_model_input_keeps_lengths_on_host() -> None:
+    sampler = Magi2PreviewSampler(torch.nn.Identity())
+    model_input = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([500.0]), cfg_config=CFGConfig())
+
+    assert model_input.audio_feat_len == [5, 5]
+    assert model_input.txt_feat_len == [3, 2]
+    assert model_input.ref_audio_feat_len == [0, 0]
+    assert model_input.ref_video_feat_len == [0, 0]
+
+    positive, negative = Magi2PreviewSampler._split_cfg_model_input(model_input)
+    assert positive.txt_feat_len == [3]
+    assert negative.txt_feat_len == [2]
+
+
+def test_pre_adapter_embeds_directly_in_checkpoint_dtype() -> None:
+    model = _tiny_model(params_dtype=torch.bfloat16)
+    packed = torch.randn(6, 4)
+    indices = torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 5])
+
+    with torch.no_grad():
+        hidden = model.pre_adapter(packed, *indices)
+
+    assert hidden.dtype == torch.bfloat16
+    assert hidden.shape == (6, model.config.virtual_width)

tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+import time
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+def _collective_worker(rank, rendezvous, backend):
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+    from vllm_omni.errors import OmniClientError
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
+    from vllm_omni.platforms import current_omni_platform
+
+    torch.set_num_threads(1)
+    device = torch.device("cuda", rank) if backend == "nccl" else torch.device("cpu")
+    if backend == "nccl":
+        current_omni_platform.set_device(device)
+    dist.init_process_group(
+        backend, init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=45)
+    )
+    try:
+        module._dit_rank_world = lambda: (dist.group.WORLD, rank, 2)
+
+        class VisualCodec:
+            fail = False
+
+            def is_distributed_enabled(self):
+                return True
+
+            def collective(self):
+                value = torch.ones(1, device=device)
+                dist.all_reduce(value)
+                assert value.item() == 2
+
+            def encode_image(self, image):
+                self.collective()
+                assert image.size == (32, 32)
+                return torch.ones(1, 96)
+
+            def encode_video(self, frames):
+                self.collective()
+                assert frames.shape == (4, 32, 32, 3)
+                if self.fail and rank == 1:
+                    raise ValueError("codec failed on rank 1")
+                return torch.full((1, 96), 2.0), (1, 2, 2)
+
+        class AudioCodec:
+            def encode_waveform(self, waveform, sample_rate):
+                assert rank == 0
+                assert sample_rate == 32000
+                assert waveform.shape == (64000,)
+                return torch.full((160, 32), 3.0), 80
+
+        model = object.__new__(module.MiniMaxH3Pipeline)
+        torch.nn.Module.__init__(model)
+        model.device = device
+        model.od_config = SimpleNamespace(enable_layerwise_offload=False)
+        model.video_vae = VisualCodec()
+        model.audio_vae = AudioCodec()
+        model.partition = "combined"
+        model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+        model.text_encoder_tp_size = 2
+        model.text_encoder_group = SimpleNamespace(
+            world_size=2, ranks=[0, 1], rank_in_group=rank, device_group=dist.group.WORLD
+        )
+        model.processor = SimpleNamespace(image_processor=SimpleNamespace(merge_size=2))
+        model.tokenizer = lambda text, **kwargs: {"input_ids": [1, 2, 3]}
+
+        def encode_text(ids, kwargs):
+            assert ids.tolist() == [1, 2, 3]
+            assert not kwargs
+            hidden = torch.ones(3, 5120, device=device, dtype=torch.bfloat16)
+            dist.all_reduce(hidden)
+            return hidden
+
+        model._encode_text_hidden = encode_text
+        sampling = OmniDiffusionSamplingParams(
+            height=32, width=32, num_frames=96, extra_args={"task": "t2va", "aspect_ratio": "1:1"}
+        )
+        text_conditioning = model._prepare_local_conditioning({"prompt": "test"}, sampling)
+        torch.testing.assert_close(text_conditioning.hidden_states, torch.full_like(text_conditioning.hidden_states, 2))
+        assert text_conditioning.visual_condition is None
+
+        media = (
+            MiniMaxH3EncoderMediaInput(
+                task="ref2va",
+                height=32,
+                width=32,
+                num_frames=107,
+                latent_t=32,
+                audio_t=178,
+                images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
+                videos=(torch.zeros(4, 32, 32, 3, dtype=torch.uint8),),
+                video_audios=(None,),
+                audios=((torch.zeros(64000), 32000),),
+            )
+            if rank == 0
+            else None
+        )
+        conditioning = model._encode_local_media(media)
+        assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video", "audio"]
+        torch.testing.assert_close(
+            conditioning.visual_condition.cpu(), torch.cat([torch.ones(1, 96), torch.full((1, 96), 2.0)])
+        )
+        torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0))
+        assert conditioning.audio_condition_lengths == (80,)
+
+        model.video_vae.fail = True
+        with pytest.raises(OmniClientError, match="codec failed on rank 1"):
+            model._encode_local_media(media)
+        with pytest.raises(OmniClientError, match="non-empty prompt"):
+            model._prepare_local_conditioning({"prompt": ""}, sampling)
+        model.video_vae.fail = False
+        assert model._encode_local_media(media).visual_condition_shapes == ((1, 2, 2), (1, 2, 2))
+    finally:
+        dist.destroy_process_group()
+
+
+def _run_collectives(tmp_path, backend):
+    context = mp.spawn(_collective_worker, args=(str(tmp_path / "rendezvous"), backend), nprocs=2, join=False)
+    deadline = time.monotonic() + 150
+    try:
+        while not context.join(timeout=1):
+            if time.monotonic() > deadline:
+                pytest.fail("MiniMax H3 collectives did not finish within 150 seconds")
+    finally:
+        for process in context.processes:
+            if process.is_alive():
+                process.terminate()
+            process.join(timeout=5)
+
+
+@pytest.mark.cpu
+def test_two_rank_conditioning_and_errors_gloo(tmp_path):
+    _run_collectives(tmp_path, "gloo")
+
+
+@pytest.mark.cuda
+@pytest.mark.parallel
+@pytest.mark.skipif(
+    os.environ.get("VLLM_TEST_MINIMAX_H3_COLLECTIVES") != "1",
+    reason="set VLLM_TEST_MINIMAX_H3_COLLECTIVES=1 with two reserved GPUs",
+)
+def test_two_rank_conditioning_and_errors_nccl(tmp_path):
+    if torch.accelerator.device_count() < 2:
+        pytest.skip("requires two GPUs")
+    _run_collectives(tmp_path, "nccl")

tests/diffusion/models/minimax_h3/test_minimax_h3_execution_modes.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import Mock
+
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture
+def pipeline():
+    model = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(model)
+    model.load_text_encoder = True
+    model.load_vae_encoder = True
+    model.device = torch.device("cpu")
+    model.partition = "combined"
+    model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    model.default_video_shift = 12.0
+    model.default_audio_shift = 3.0
+    model.od_config = SimpleNamespace(enable_layerwise_offload=False)
+    model._quality_policy = SimpleNamespace(resolve=Mock(return_value=SimpleNamespace(cache_dit=None)))
+    model._cache_dit_runtime = SimpleNamespace(prepare=Mock())
+    model.encode_prompt = Mock(
+        return_value=(torch.ones(2, 5120, dtype=torch.bfloat16), torch.ones(2, dtype=torch.long))
+    )
+    model.video_vae = Mock()
+    model.video_vae.is_distributed_enabled.return_value = False
+    model.video_vae.encode_image.side_effect = lambda image: torch.ones((image.height // 32) * (image.width // 32), 96)
+    model.audio_vae = Mock()
+    model.transformer = torch.nn.Identity()
+    model.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    model.decode = Mock(return_value=(torch.zeros(1, 3, 1, 32, 32), torch.zeros(1)))
+    return model
+
+
+def _request(task):
+    prompt: dict[str, Any] = {"prompt": "A person waves."}
+    if task != "t2va":
+        prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))}
+    sampling = OmniDiffusionSamplingParams(
+        height=32, width=32, num_frames=96, num_inference_steps=2, extra_args={"task": task, "aspect_ratio": "1:1"}
+    )
+    return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
+
+
+@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"])
+def test_single_stage_encodes_without_stage_zero(pipeline, task):
+    output = pipeline.forward(_request(task))
+    pipeline.encode_prompt.assert_called_once()
+    assert pipeline.video_vae.encode_image.call_count == int(task != "t2va")
+    pipeline.diffuse.assert_called_once()
+    assert pipeline.diffuse.call_args.kwargs["task"] == task
+    video, audio = output.output
+    torch.testing.assert_close(video, torch.zeros(1, 1, 32, 32, 3, dtype=torch.uint8))
+    torch.testing.assert_close(audio, pipeline.decode.return_value[1])
+
+
+def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch):
+    import numpy as np
+
+    from vllm_omni.data_entry_keys import flatten_payload
+    from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import (
+        MiniMaxH3EncoderConditioning,
+        MiniMaxH3TextConditioning,
+    )
+    from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
+    from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion, prepare_encoder_prompt
+
+    frames = np.zeros((48, 32, 32, 3), dtype=np.uint8)
+    decode_frames = Mock(return_value=frames)
+    sampler = Mock(wraps=processing.sample_reference_video_frames)
+    monkeypatch.setattr(processing, "load_video_frames", decode_frames)
+    monkeypatch.setattr(processing, "sample_reference_video_frames", sampler)
+    monkeypatch.setattr(
+        processing,
+        "prepare_reference_videos",
+        Mock(
+            return_value=[{"prepared_path": "prepared.mp4", "original_path": "original.mp4", "input_has_audio": True}]
+        ),
+    )
+    monkeypatch.setattr(processing, "load_video_audio", Mock(return_value=(torch.full((64000,), 3.0), 32000)))
+    pipeline.video_vae.encode_video.return_value = (torch.full((1, 96), 2.0), (1, 2, 2))
+
+    def encode_audio(waveform, sample_rate):
+        length = round(waveform.shape[-1] / sample_rate * 40)
+        return torch.full((length * 2, 32), float(waveform[0])), length
+
+    pipeline.audio_vae.encode_waveform.side_effect = encode_audio
+    request = _request("ref2va")
+    request.prompts[0]["multi_modal_data"].update(video="original.mp4", audio=(torch.full((192000,), 4.0), 32000))
+    pipeline.forward(request)
+    local_kwargs = pipeline.diffuse.call_args.kwargs
+
+    transformed = prepare_encoder_prompt(request.prompts[0], [request.sampling_params])
+    stage_zero = object.__new__(MiniMaxH3Encoder)
+    torch.nn.Module.__init__(stage_zero)
+    stage_zero._component_leader = True
+    stage_zero.video_vae = pipeline.video_vae
+    stage_zero.audio_vae = pipeline.audio_vae
+    media = stage_zero._encode_media(stage_zero._media_input(transformed["additional_information"]))
+    text = MiniMaxH3TextConditioning(*pipeline.encode_prompt.return_value)
+    payload = MiniMaxH3EncoderConditioning.from_components(text, media).to_omni_payload()
+    source = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=flatten_payload(payload))])
+    bridged = encoder2diffusion([source], request.prompts[0])
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
+    pipeline.encode_prompt.reset_mock()
+    pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params))
+    external_kwargs = pipeline.diffuse.call_args.kwargs
+    for key, value in local_kwargs.items():
+        if isinstance(value, torch.Tensor):
+            torch.testing.assert_close(value, external_kwargs[key])
+        else:
+            assert value == external_kwargs[key]
+    pipeline.encode_prompt.assert_not_called()
+    assert decode_frames.call_count == sampler.call_count == 2
+    assert all(call.kwargs["decoded_frames"] is frames for call in sampler.call_args_list)
+    assert [block["kind"] for block in local_kwargs["ref_blocks"]] == ["image", "video_audio", "audio"]
+    assert local_kwargs["audio_condition_lengths"] == [80, round(107 / 24 * 40)]
+    assert all(
+        call.args[0].shape[-1] == round(107 / 24 * 32000)
+        for call in pipeline.audio_vae.encode_waveform.call_args_list[1::2]
+    )

tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py

@@ -6,7 +6,7 @@ import torch
 from vllm.triton_utils import HAS_TRITON
 
 from tests.helpers.mark import hardware_marks
-from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
+from vllm_omni.diffusion.layers.indexed_modulation import (
     _MAX_1D_GRID_SIZE,
     _iter_row_chunks,
     _launch_row_chunks,

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -9,72 +9,25 @@ from unittest.mock import Mock
 
 import pytest
 import torch
-from PIL import Image
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders():
+def test_h3_full_pipeline_profiler_includes_local_encoding():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     targets = MiniMaxH3Pipeline._PROFILER_TARGETS
-    assert "_encode_visual_conditions" in targets
-    assert "_encode_reference_audio_conditions" in targets
-    assert "video_vae.decode_latent" in targets
-    assert "audio_vae.decode_latent" in targets
-    assert "_encode_video_conditions" not in targets
-    assert "_encode_video_audio_conditions" not in targets
-    assert "_encode_audio_conditions" not in targets
-
-
-@pytest.mark.parametrize("distributed", [False, True])
-def test_manual_component_offload_wins_over_requested_model_offload(distributed):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.od_config = SimpleNamespace(
-        enable_cpu_offload=True,
-        enable_layerwise_offload=not distributed,
-        enable_distributed_layerwise_offload=distributed,
-    )
-    pipeline._model_cpu_offload_modules = []
-    pipeline.text_encoder = Mock()
-    expected = torch.ones(2, 3)
-    pipeline.text_encoder.encode_ids.return_value = expected
-    component = Mock()
-
-    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
-    with pipeline._component_on_device(component):
-        component.load_to_device.assert_called_once_with()
-        component.offload_to_cpu.assert_not_called()
-
-    assert actual is expected
-    pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
-    component.offload_to_cpu.assert_called_once_with()
-
-
-def test_requested_model_offload_without_backend_uses_resident_components():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.od_config = SimpleNamespace(
-        enable_cpu_offload=True,
-        enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=False,
-    )
-    pipeline._model_cpu_offload_modules = []
-    pipeline.text_encoder = Mock()
-    expected = torch.ones(2, 3)
-    pipeline.text_encoder.encode_ids.return_value = expected
-
-    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
-
-    assert actual is expected
-    pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.encode_ids.assert_called_once()
+    assert targets == [
+        "encode_prompt",
+        "_encode_local_media",
+        "diffuse",
+        "decode",
+        "video_vae.decode_latent",
+        "audio_vae.decode_latent",
+        "prepare_encode",
+        "denoise_step",
+        "post_decode",
+    ]
 
 
 def test_decoded_output_is_unchanged_without_active_model_offload():
@@ -194,7 +147,7 @@ def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypa
         _extract_prompt=lambda raw: ("a prompt", {}),
         _extract_text_conditioning=lambda raw: None,
         _extract_prepared_reference_videos=lambda raw: None,
-        _prepare_request_inputs=lambda **kwargs: context,
+        _prepare_request_inputs=lambda *args, **kwargs: context,
         _denoise_kwargs=lambda ctx: {},
         diffuse=diffuse,
         decode=decode,
@@ -243,7 +196,7 @@ def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypat
         _extract_prompt=lambda raw: ("a prompt", {}),
         _extract_text_conditioning=lambda raw: None,
         _extract_prepared_reference_videos=lambda raw: None,
-        _prepare_request_inputs=lambda **kwargs: context,
+        _prepare_request_inputs=lambda *args, **kwargs: context,
         _denoise_kwargs=lambda ctx: {},
         diffuse=diffuse,
         decode=decode,
@@ -318,7 +271,8 @@ def test_output_owner_rank_matches_executor_contract(monkeypatch, available, ini
     assert MiniMaxH3Pipeline._is_output_owner_rank() is expected
 
 
-def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
+@pytest.mark.parametrize("load_text_encoder", [True, False])
+def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text_encoder):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
 
@@ -326,9 +280,11 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     torch.nn.Module.__init__(pipeline)
     pipeline.transformer = torch.nn.Linear(2, 2)
     pipeline.transformers_ref = torch.nn.Linear(2, 2)
-    pipeline.text_encoder = torch.nn.Linear(2, 2)
     pipeline.video_vae = torch.nn.Linear(2, 2)
     pipeline.audio_vae = torch.nn.Linear(2, 2)
+    pipeline._encoder_modules = ["text_encoder"] if load_text_encoder else []
+    if load_text_encoder:
+        pipeline.text_encoder = torch.nn.Linear(2, 2)
     apply_offload = Mock()
     remove_offload = Mock()
     monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
@@ -341,7 +297,7 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     )
 
     dits = [pipeline.transformer, pipeline.transformers_ref]
-    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+    stages = [*([pipeline.text_encoder] if load_text_encoder else []), pipeline.video_vae, pipeline.audio_vae]
     apply_offload.assert_called_once_with(
         dit_modules=dits,
         encoder_modules=stages,
@@ -356,6 +312,46 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
     remove_offload.assert_called_once_with([*dits, *stages])
 
 
+@pytest.mark.parametrize("decode_fails", [False, True])
+def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    component = torch.nn.Linear(2, 2)
+    events: list[tuple[Any, ...]] = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    pipeline._model_cpu_offload_modules = [component]
+
+    def decode():
+        with pipeline._component_on_device(component):
+            events.append(("decode", component))
+            if decode_fails:
+                raise RuntimeError("decode failed")
+
+    if decode_fails:
+        with pytest.raises(RuntimeError, match="decode failed"):
+            decode()
+    else:
+        decode()
+
+    assert events == [
+        ("activate", component),
+        ("decode", component),
+        ("offload", component),
+    ]
+
+
 def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
@@ -416,44 +412,20 @@ def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
     apply_offload.assert_not_called()
 
 
-@pytest.mark.parametrize("decode_fails", [False, True])
-def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+def _encode_media(pipeline, **kwargs):
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
+    from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
 
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    component = torch.nn.Linear(2, 2)
-    events: list[tuple[Any, ...]] = []
-
-    @contextmanager
-    def record_component(value):
-        events.append(("activate", value))
-        try:
-            yield
-        finally:
-            events.append(("offload", value))
-
-    monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    pipeline._model_cpu_offload_modules = [component]
-
-    def decode():
-        with pipeline._component_on_device(component):
-            events.append(("decode", component))
-            if decode_fails:
-                raise RuntimeError("decode failed")
-
-    if decode_fails:
-        with pytest.raises(RuntimeError, match="decode failed"):
-            decode()
-    else:
-        decode()
-
-    assert events == [
-        ("activate", component),
-        ("decode", component),
-        ("offload", component),
-    ]
+    media = MiniMaxH3EncoderMediaInput(
+        task="ref2va", height=32, width=32, num_frames=124, latent_t=31, audio_t=208, **kwargs
+    )
+    return encode_media(
+        media,
+        video_vae=getattr(pipeline, "video_vae", None),
+        audio_vae=getattr(pipeline, "audio_vae", None),
+        emit_conditioning=True,
+        component_scope=pipeline._component_on_device,
+    )
 
 
 @pytest.mark.parametrize("encode_fails", [False, True])
@@ -463,7 +435,6 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
     pipeline.video_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.video_vae]
     events = []
@@ -484,17 +455,15 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
 
     pipeline.video_vae.encode_image.side_effect = encode_image
     monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
-    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
-    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+    images = (torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8))
 
     if encode_fails:
         with pytest.raises(RuntimeError, match="encode failed"):
-            pipeline._encode_visual_conditions(images, None, video_count=0)
+            _encode_media(pipeline, images=images)
     else:
-        rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0)
-        assert rows.shape == (3, 4)
-        assert shapes == [(1, 1, 1), (1, 1, 2)]
+        result = _encode_media(pipeline, images=images)
+        assert result.visual_condition.shape == (3, 4)
+        assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2))
 
     assert events == [
         ("activate", pipeline.video_vae),
@@ -510,9 +479,7 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
     pipeline.video_vae = Mock()
-    pipeline.video_vae.is_distributed_enabled.return_value = False
     pipeline._model_cpu_offload_modules = [pipeline.video_vae]
     events = []
 
@@ -529,31 +496,26 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
         return torch.ones(1, 4)
 
     def encode_video(frames):
-        events.append(("video", frames))
+        events.append(("video", frames.shape))
         return torch.ones(3, 4), (3, 2, 2)
 
     pipeline.video_vae.encode_image.side_effect = encode_image
     pipeline.video_vae.encode_video.side_effect = encode_video
     monkeypatch.setattr(module, "sequential_offload_component", record_component)
-    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
-    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
-    monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}")
-    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
-    prepared_videos = [{"prepared_path": "reference.mp4"}]
-
-    rows, shapes = pipeline._encode_visual_conditions(
-        images,
-        prepared_videos,
-        video_count=1,
+    result = _encode_media(
+        pipeline,
+        images=(torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8)),
+        videos=(torch.zeros(12, 32, 32, 3, dtype=torch.uint8),),
+        video_audios=(None,),
     )
 
-    assert rows.shape == (5, 4)
-    assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)]
+    assert result.visual_condition.shape == (5, 4)
+    assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2), (3, 2, 2))
     assert events == [
         ("activate", pipeline.video_vae),
         ("image", (16, 16)),
         ("image", (32, 16)),
-        ("video", "frames:reference.mp4"),
+        ("video", (12, 32, 32, 3)),
         ("offload", pipeline.video_vae),
     ]
 
@@ -567,7 +529,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
     torch.nn.Module.__init__(pipeline)
     pipeline.audio_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
-    events: list[tuple[Any, ...]] = []
+    events = []
 
     @contextmanager
     def record_component(value):
@@ -577,26 +539,20 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
         finally:
             events.append(("offload", value))
 
-    def encode_embedded(*args, **kwargs):
-        events.append(("embedded",))
+    def encode_waveform(waveform, sample_rate):
+        events.append(("audio", int(waveform[0]), sample_rate))
         if encode_fails:
             raise RuntimeError("audio encode failed")
-        return torch.ones(1, 2), [1]
+        return torch.ones(80, 32), 80
 
-    def encode_standalone(*args, **kwargs):
-        events.append(("standalone",))
-        return torch.ones(1, 2), [1]
-
-    pipeline._encode_video_audio_conditions_resident = encode_embedded
-    pipeline._encode_audio_conditions_resident = encode_standalone
+    pipeline.audio_vae.encode_waveform.side_effect = encode_waveform
     monkeypatch.setattr(module, "sequential_offload_component", record_component)
 
     def call():
-        return pipeline._encode_reference_audio_conditions(
-            [{"input_has_audio": True}],
-            has_audio=[True],
-            standalone_audios=[(torch.ones(1), 16_000)],
-            max_duration_seconds=1.0,
+        return _encode_media(
+            pipeline,
+            video_audios=(),
+            audios=((torch.ones(32_000), 16_000),),
         )
 
     if encode_fails:
@@ -604,10 +560,61 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
             call()
     else:
         result = call()
-        assert result[0].shape == (1, 2)
-        assert result[2].shape == (1, 2)
+        assert result.audio_condition.shape == (80, 32)
+        assert result.audio_condition_lengths == (80,)
+        assert events[1:-1] == [("audio", 1, 16_000)]
 
     assert events[0] == ("activate", pipeline.audio_vae)
     assert events[-1] == ("offload", pipeline.audio_vae)
     assert events.count(("activate", pipeline.audio_vae)) == 1
     assert events.count(("offload", pipeline.audio_vae)) == 1
+
+
+@pytest.mark.parametrize("distributed", [False, True])
+def test_manual_component_offload_wins_over_requested_model_offload(distributed):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=not distributed,
+        enable_distributed_layerwise_offload=distributed,
+    )
+    pipeline._model_cpu_offload_modules = []
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+    component = Mock()
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+    with pipeline._component_on_device(component):
+        component.load_to_device.assert_called_once_with()
+        component.offload_to_cpu.assert_not_called()
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+    component.offload_to_cpu.assert_called_once_with()
+
+
+def test_requested_model_offload_without_backend_uses_resident_components():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+    )
+    pipeline._model_cpu_offload_modules = []
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.encode_ids.assert_called_once()

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py

@@ -186,9 +186,8 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
     common_kwargs = {
         "model": model_ref,
         "enforce_eager": True,
-        "tensor_parallel_size": 2,
         # The fused BF16 baseline only fits on H100-80GB with encoder TP.
-        "text_encoder_tp_size": 2,
+        "stage_0_tensor_parallel_size": 2,
         "vae_use_tiling": True,
     }
 

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -110,6 +110,8 @@ def _step_pipeline(model, *, packed_batch_supported: bool = True):
     from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline.load_text_encoder = False
+    pipeline.load_vae_encoder = False
     pipeline.transformer = model
     pipeline.device = torch.device("cpu")
     pipeline._transformer_for_task = lambda task: model
@@ -304,8 +306,17 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames):
     }
 
     pipeline = _step_pipeline(_SegmentMeanModel())
-    monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {})))
-    monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context)
+    conditioning = object()
+    monkeypatch.setattr(
+        mod.MiniMaxH3Pipeline,
+        "_extract_encoder_conditioning",
+        staticmethod(lambda _: conditioning),
+    )
+    monkeypatch.setattr(
+        mod.MiniMaxH3Pipeline,
+        "_prepare_encoder_conditioning_inputs",
+        lambda self, value, sampling: context,
+    )
     monkeypatch.setattr(
         mod.MiniMaxH3Pipeline,
         "_build_denoise_inputs",
@@ -439,50 +450,3 @@ def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention):
     from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
 
     assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False
-
-
-def test_broadcast_rank0_exception_single_rank_reraises():
-    """Single-rank execution has no group; the helper just reraises."""
-    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
-        _broadcast_rank0_exception,
-    )
-    from vllm_omni.errors import OmniClientError
-
-    _broadcast_rank0_exception(None)
-    with pytest.raises(OmniClientError, match="bad ref"):
-        _broadcast_rank0_exception(OmniClientError("bad ref"))
-
-
-def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch):
-    """A rank-0 error becomes a matching client error on every other DiT rank."""
-    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
-    from vllm_omni.errors import OmniClientError
-
-    def fake_rank_world(rank):
-        return lambda: (object(), rank, 4)
-
-    def make_broadcast(rank0_payload):
-        def fake_broadcast(payload_list, *, src, group):
-            payload_list[0] = rank0_payload
-
-        return fake_broadcast
-
-    monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0))
-    err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError")
-    rank0_payload = {
-        "type": type(err).__name__,
-        "message": str(err),
-        "status_code": err.status_code,
-        "error_type": err.error_type,
-    }
-    monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload))
-    with pytest.raises(OmniClientError) as rank0_info:
-        mod._broadcast_rank0_exception(err)
-    assert rank0_info.value is err
-
-    monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2))
-    with pytest.raises(OmniClientError) as rank2_info:
-        mod._broadcast_rank0_exception(None)
-    assert rank2_info.value.status_code == 422
-    assert rank2_info.value.error_type == "UnprocessableEntityError"
-    assert "invalid reference-video file" in str(rank2_info.value)

tests/diffusion/models/minimax_h3/test_vsa_layout.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import math
+import sys
+import types
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    VideoTokenLayout,
+    VideoTokenSpan,
+)
+from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import (
+    MiniMaxH3VSAImpl,
+    _get_h3_tile_metadata,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
+    extra: dict[str, Any] = {
+        "vsa_h3_prefix_segments": prefix_segments,
+    }
+    if gate is not None:
+        extra["gate_compress"] = gate
+    return AttentionMetadata(
+        packed_padding=packed_padding,
+        video_layout=VideoTokenLayout(
+            used_len=sum(prefix_segments) + math.prod(video_shape),
+            video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
+        ),
+        extra=extra,
+    )
+
+
+def _fake_block_sparse_kernel(monkeypatch, calls):
+    """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
+
+    def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
+        calls["block_map"] = block_map.clone()
+        calls["variable_block_sizes"] = variable_block_sizes.tolist()
+        calls["q_shape"] = tuple(q.shape)
+        blocks = block_map.shape[-1]
+        assert blocks * 64 == q.shape[2]
+        # Expand the per-block map to a per-token mask so the stand-in honours
+        # the sparsity contract the real kernel implements, and drop the rows
+        # each partial block pads with, which the real kernel skips through
+        # variable_block_sizes.
+        token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
+        within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
+        token_mask = token_mask & within.reshape(1, 1, 1, -1)
+        out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
+        return out.to(q.dtype), None
+
+    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    setattr(module, "block_sparse_attn", block_sparse_attn)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
+
+
+def _h3_impl(**backend_kwargs):
+    kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
+    kwargs.update(backend_kwargs)
+    return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
+
+
+def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5, 70), (2, 4, 4)
+    seq_len = sum(prefix_segments) + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    assert output.shape == query.shape
+    # topk covers every video block, so the block map is dense and the tiled
+    # kernel must reproduce plain attention over the same rows.
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
+    # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
+    assert calls["variable_block_sizes"] == [5, 64, 6, 32]
+
+
+def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=1)
+    prefix_segments, video_shape = (5,), (8, 4, 4)
+    seq_len = 5 + 8 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    # 1 prefix block + 2 video blocks. The odd count is padded to an even one
+    # for the paired-CTA kernel contract, and the Triton route slices that
+    # transport-only partner off again before launching.
+    block_map = calls["block_map"]
+    assert block_map.shape[-1] == 3
+    assert calls["variable_block_sizes"] == [5, 64, 64]
+    assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
+    assert block_map[..., :1].all(), "prefix keys are exempt from selection"
+    assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
+
+    # The same odd geometry must still untile back onto the original rows.
+    dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
+
+
+def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    seq_len = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    zero_gate = torch.zeros_like(query)
+    with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
+    live_gate = torch.full_like(query, 0.5)
+    with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
+
+    torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
+    assert not torch.allclose(with_live_gate.float(), without_gate.float())
+
+
+def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    valid = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
+    padding = PackedPaddingMetadata(
+        q_length=valid,
+        kv_length=valid,
+        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
+        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
+    )
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
+
+    assert output.shape == query.shape
+    assert torch.count_nonzero(output[:, valid:]) == 0
+
+
+def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
+    partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
+        (5, 70, 9), (5, 6, 6), torch.device("cpu")
+    )
+    assert sizes[:4].tolist() == [5, 64, 6, 9]
+    assert prefix_blocks == 4
+    assert video_blocks == 8
+    assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
+    assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
+
+
+@pytest.mark.parametrize("fallback_on_error", [False, True])
+def test_h3_kernel_failure_preserves_fallback_policy(monkeypatch, fallback_on_error):
+    impl = _h3_impl(fallback_on_error=fallback_on_error)
+    query = torch.randn(1, 37, 2, 8)
+    metadata = _h3_metadata((5,), (2, 4, 4))
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("kernel failure")
+
+    monkeypatch.setattr(impl, "_forward_h3", fail)
+    monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
+    if fallback_on_error:
+        output = impl.forward_cuda(query, query, query, metadata)
+        expected = impl.sdpa_fallback.forward_cuda(query, query, query, metadata)
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+    else:
+        with pytest.raises(RuntimeError, match="kernel failure"):
+            impl.forward_cuda(query, query, query, metadata)
+
+
+def test_h3_accelerator_failure_is_not_swallowed(monkeypatch):
+    impl = _h3_impl(fallback_on_error=True)
+    query = torch.randn(1, 37, 2, 8)
+
+    def fail(*args, **kwargs):
+        raise torch.AcceleratorError("device failure")
+
+    monkeypatch.setattr(impl, "_forward_h3", fail)
+    with pytest.raises(torch.AcceleratorError, match="device failure"):
+        impl.forward_cuda(query, query, query, _h3_metadata((5,), (2, 4, 4)))
+
+
+def test_h3_without_layout_delegates_to_shared_backend(monkeypatch):
+    from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSAImpl
+
+    impl = _h3_impl()
+    query = torch.randn(1, 37, 2, 8)
+    metadata = AttentionMetadata()
+    expected = torch.empty_like(query)
+    calls = []
+
+    def forward(self, q, k, v, attn_metadata):
+        calls.append((self, q, k, v, attn_metadata))
+        return expected
+
+    monkeypatch.setattr(FastVideoVSAImpl, "forward_cuda", forward)
+    assert impl.forward_cuda(query, query, query, metadata) is expected
+    assert len(calls) == 1
+    assert all(actual is wanted for actual, wanted in zip(calls[0], (impl, query, query, query, metadata)))

tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.qwen_image import qwen_image_transformer
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenImageTransformerBlock
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _StaticOutput(nn.Module):
+    def __init__(self, output: torch.Tensor):
+        super().__init__()
+        self.output = output
+
+    def forward(self, _: torch.Tensor) -> torch.Tensor:
+        return self.output
+
+
+class _RecordingNorm(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+    def forward(self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
+        self.calls.append((x, scale, shift))
+        return x
+
+
+class _ZeroAttention(nn.Module):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        **_: object,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.zeros_like(hidden_states), torch.zeros_like(encoder_hidden_states)
+
+
+class _ZeroMLP(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.zeros_like(x)
+
+
+def test_non_cuda_select01_uses_adalayernorm_dispatch(monkeypatch):
+    hidden_size = 4
+    block = QwenImageTransformerBlock.__new__(QwenImageTransformerBlock)
+    nn.Module.__init__(block)
+    block.zero_cond_t = True
+    block.img_mod = _StaticOutput(torch.zeros(2, 6 * hidden_size))
+    block.txt_mod = _StaticOutput(torch.zeros(1, 6 * hidden_size))
+    block.img_norm1 = _RecordingNorm()
+    block.img_norm2 = _RecordingNorm()
+    block.txt_norm1 = _RecordingNorm()
+    block.txt_norm2 = _RecordingNorm()
+    block.attn = _ZeroAttention()
+    block.img_mlp = _ZeroMLP()
+    block.txt_mlp = _ZeroMLP()
+
+    monkeypatch.setattr(qwen_image_transformer, "can_use_qwen_select01_triton", lambda _: False)
+
+    def fail_if_called(*args, **kwargs):
+        pytest.fail("non-CUDA select01 must preserve the AdaLayerNorm dispatch")
+
+    monkeypatch.setattr(qwen_image_transformer, "fused_layernorm_select01", fail_if_called)
+    monkeypatch.setattr(qwen_image_transformer, "fused_residual_layernorm_select01", fail_if_called)
+
+    hidden_states = torch.randn(1, 3, hidden_size)
+    encoder_hidden_states = torch.randn(1, 2, hidden_size)
+    modulate_index = torch.tensor([[0, 1, 0]])
+    block(
+        hidden_states,
+        encoder_hidden_states,
+        torch.ones(1, 2),
+        torch.zeros(2, hidden_size),
+        (torch.empty(0), torch.empty(0)),
+        modulate_index=modulate_index,
+    )
+
+    assert len(block.img_norm1.calls) == 1
+    assert len(block.img_norm2.calls) == 1
+    for _, scale, shift in (*block.img_norm1.calls, *block.img_norm2.calls):
+        assert scale.shape == hidden_states.shape
+        assert shift.shape == hidden_states.shape

tests/diffusion/models/wan2_2/conftest.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from contextlib import contextmanager
@@ -33,7 +36,7 @@ class StubScheduler:
         self.config = SimpleNamespace(num_train_timesteps=1000)
         self.set_timesteps_calls: list[tuple[int, torch.device]] = []
 
-    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
+    def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
         self.set_timesteps_calls.append((num_steps, device))
 
 

tests/diffusion/models/wan2_2/fixtures/native_unipc_shift12.json

@@ -0,0 +1,97 @@
+{
+  "source_url": "https://github.com/Wan-Video/Wan2.2/blob/42bf4cfaa384bc21833865abc2f9e6c0e67233dc/wan/utils/fm_solvers_unipc.py",
+  "source_sha256": "0dec8c7ed17f6f2049275c6848113314da6ccec1c8db5bdc89df43c05c6038d9",
+  "git_head": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc",
+  "num_train_timesteps": 1000,
+  "use_dynamic_shifting": false,
+  "reproduction": "Load FlowUniPCMultistepScheduler from source_url; construct with num_train_timesteps=1000, shift=1.0, use_dynamic_shifting=False; call set_timesteps(40, device='cpu', shift=12.0); record timesteps.tolist() and sigmas.tolist().",
+  "constructor_shift": 1.0,
+  "runtime_shift": 12.0,
+  "steps": 40,
+  "timesteps": [
+    999,
+    997,
+    995,
+    993,
+    990,
+    988,
+    985,
+    982,
+    979,
+    976,
+    972,
+    969,
+    965,
+    961,
+    956,
+    952,
+    947,
+    941,
+    936,
+    929,
+    922,
+    915,
+    907,
+    898,
+    888,
+    877,
+    865,
+    852,
+    837,
+    819,
+    799,
+    776,
+    749,
+    717,
+    678,
+    631,
+    571,
+    492,
+    386,
+    235
+  ],
+  "sigmas": [
+    0.9999166131019592,
+    0.9977826476097107,
+    0.9955461621284485,
+    0.9931996464729309,
+    0.9907346963882446,
+    0.988142192363739,
+    0.9854119420051575,
+    0.9825326800346375,
+    0.9794917702674866,
+    0.9762753844261169,
+    0.9728676676750183,
+    0.9692512154579163,
+    0.9654061794281006,
+    0.961310088634491,
+    0.9569376707077026,
+    0.95225989818573,
+    0.9472436308860779,
+    0.9418508410453796,
+    0.9360373020172119,
+    0.9297519326210022,
+    0.9229347705841064,
+    0.9155154228210449,
+    0.9074103832244873,
+    0.8985199332237244,
+    0.888724148273468,
+    0.8778772950172424,
+    0.865800678730011,
+    0.8522725701332092,
+    0.8370144367218018,
+    0.8196719288825989,
+    0.7997865080833435,
+    0.7767546772956848,
+    0.7497654557228088,
+    0.717703104019165,
+    0.6789888143539429,
+    0.6313128471374512,
+    0.5711563229560852,
+    0.49288034439086914,
+    0.38684695959091187,
+    0.23510953783988953,
+    0.0
+  ],
+  "implementation": "Wan-Video/Wan2.2 wan/utils/fm_solvers_unipc.py"
+}

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py

@@ -1,15 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
     get_wan22_i2v_post_process_func,
@@ -50,7 +51,7 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
 
 
 def _make_i2v_sampling(**overrides):
-    values = {
+    values: dict[str, object] = {
         "height": 16,
         "width": 16,
         "num_frames": 5,
@@ -275,16 +276,18 @@ def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
     assert first_frame_mask.shape == (2, 1, 2, 2, 2)
 
 
-def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
+@pytest.mark.parametrize("solver", ["unipc", "euler"])
+@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
+def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch, solver: str, shift: float) -> None:
     monkeypatch.setattr(
         "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
         SimpleNamespace(is_available=lambda: False),
     )
     pipeline = _make_i2v_pipeline(expand_timesteps=True)
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 5.0
+    pipeline.scheduler = build_wan_scheduler(solver, shift)
+    pipeline.od_config = SimpleNamespace(flow_shift=shift)
+    pipeline._sample_solver = solver
+    pipeline._flow_shift = shift
     pipeline.boundary_ratio = 0.875
     pipeline.has_image_encoder = False
     pipeline._guidance_scale = None
@@ -327,6 +330,8 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
                     generator=gen_a,
                     latents=latents_a,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
             SimpleNamespace(
@@ -336,12 +341,25 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
                     generator=gen_b,
                     latents=latents_b,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
         ]
     )
 
     outputs = pipeline.forward(batch)
+    if solver == "unipc":
+        sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
+        sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
+        torch.testing.assert_close(
+            pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
+        )
+        assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
+    else:
+        reference = build_wan_scheduler("euler", shift)
+        reference.set_timesteps(50, device="cpu")
+        torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
 
     assert prepare_call["batch_size"] == 4
     assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -10,7 +10,7 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, build_wan_scheduler
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -19,6 +19,12 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.fixture(autouse=True)
+def _cpu_platform(monkeypatch):
+    module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
+    monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
+
+
 class _StubTransformer(nn.Module):
     @property
     def dtype(self) -> torch.dtype:
@@ -51,10 +57,10 @@ class _StubScheduler:
     def __init__(self, timesteps: list[int]) -> None:
         self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
         self.config = SimpleNamespace(num_train_timesteps=1000)
-        self.set_timesteps_calls: list[tuple[int, torch.device]] = []
+        self.set_timesteps_calls: list[tuple[int, torch.device, float | None]] = []
 
-    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
-        self.set_timesteps_calls.append((num_steps, device))
+    def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
+        self.set_timesteps_calls.append((num_steps, device, shift))
 
 
 @contextmanager
@@ -184,7 +190,24 @@ def test_forward_delegates_denoising_to_diffuse(
     assert captured["boundary_timestep"] == pytest.approx(875.0)
     assert captured["latent_condition"] is None
     assert captured["first_frame_mask"] is None
-    assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
+    assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"), 5.0)]
+
+
+@pytest.mark.parametrize("solver", ["unipc", "euler"])
+def test_forward_passes_request_shift_without_mutating_scheduler_config(solver: str) -> None:
+    pipeline = _make_pipeline()
+    pipeline.diffuse = lambda **kwargs: kwargs["latents"]
+    for shift in (3.0, 12.0, 5.0):
+        sampling = _make_sampling(num_inference_steps=5, extra_args={"sample_solver": solver, "flow_shift": shift})
+        request = OmniDiffusionRequest(prompt="prompt", request_id="schedule", sampling_params=sampling)
+        pipeline.forward(DiffusionRequestBatch(requests=[request]))
+        reference = build_wan_scheduler(solver, shift)
+        if solver == "unipc":
+            reference.set_timesteps(5, device="cpu", shift=shift)
+            assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
+        else:
+            reference.set_timesteps(5, device="cpu")
+        torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
 
 
 def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
@@ -7,6 +10,7 @@ import pytest
 import torch
 import torch.nn as nn
 
+from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
     Wan22S2VPipeline,
     _make_clip_generators,
@@ -14,11 +18,32 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
 from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("configured_shift", [None, 12.0])
+def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | None) -> None:
+    module = "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v"
+    config = SimpleNamespace(model="unused", flow_shift=configured_shift, enable_diffusion_pipeline_profiler=False)
+
+    def init_components(pipeline, *args):
+        pipeline.vae = SimpleNamespace(config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8))
+
+    with (
+        patch(f"{module}.get_local_device", return_value=torch.device("cpu")),
+        patch(f"{module}._resolve_model_path", return_value="unused"),
+        patch(f"{module}._is_diffusers_format", return_value=True),
+        patch.object(Wan22S2VPipeline, "_init_diffusers_format", init_components),
+        patch.object(Wan22S2VPipeline, "setup_diffusion_pipeline_profiler"),
+    ):
+        pipeline = Wan22S2VPipeline(od_config=config)
+    assert pipeline._flow_shift == (3.0 if configured_shift is None else configured_shift)
+    assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
+    assert pipeline.scheduler.sigma_max == float(np.float32(0.999))
 
 
 def _make_s2v_sampling(**overrides):
-    values = {
+    values: dict[str, object] = {
         "height": 16,
         "width": 16,
         "num_frames": 8,
@@ -211,7 +236,8 @@ def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
         pipeline.forward(batch)
 
 
-def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
+@pytest.mark.parametrize("shift", [3.0, 12.0])
+def test_s2v_forward_batches_request_local_inputs_and_splits_outputs(shift: float) -> None:
     pipeline = object.__new__(Wan22S2VPipeline)
     nn.Module.__init__(pipeline)
     pipeline.device = torch.device("cpu")
@@ -227,7 +253,8 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
         enable_cpu_offload=False,
         parallel_config=SimpleNamespace(use_hsdp=False),
     )
-    pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
+    pipeline._flow_shift = shift
+    pipeline.scheduler = FlowUniPCMultistepScheduler(shift=1.0)
     pipeline.vae_scale_factor_spatial = 8
     pipeline.resolution_divisor = 16
     pipeline.motion_frames = 7
@@ -262,6 +289,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
                     "multi_modal_data": {"image": image, "audio": audio_a},
                 },
                 sampling_params=_make_s2v_sampling(
+                    num_inference_steps=5,
                     num_outputs_per_prompt=2,
                     generator=generators_a,
                     latents=latents_a,
@@ -275,6 +303,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
                     "multi_modal_data": {"image": image, "audio": audio_b},
                 },
                 sampling_params=_make_s2v_sampling(
+                    num_inference_steps=5,
                     num_outputs_per_prompt=2,
                     generator=generators_b,
                     latents=latents_b,
@@ -287,6 +316,12 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
         platform.is_available.return_value = False
         outputs = pipeline.forward(batch)
 
+    expected = np.linspace(float(np.float32(0.999)), 0.0, 6)[:-1]
+    expected = shift * expected / (1.0 + (shift - 1.0) * expected)
+    torch.testing.assert_close(
+        pipeline.scheduler.sigmas, torch.tensor(np.append(expected, 0.0), dtype=torch.float32), rtol=0, atol=0
+    )
+    assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
     assert len(outputs) == 2
     assert outputs[0].output[0].shape[0] == 2
     assert outputs[1].output[0].shape[0] == 2
@@ -519,6 +554,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
     mock_scheduler.timesteps = torch.linspace(999, 0, 5)
     mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),))
     pipeline.scheduler = mock_scheduler
+    pipeline._flow_shift = 3.0
 
     # -- Bind methods from the real class --
     pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline)

tests/diffusion/models/wan2_2/test_wan22_scheduler.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare actual Wan schedules to native Wan outputs and independent sigma arithmetic."""
+
+import json
+from pathlib import Path
+
+import numpy as np
+import numpy.typing as npt
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    FASTWAN_DMD_SCHEDULER_SHIFT,
+    build_wan_scheduler,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_wan_unipc_matches_captured_native_full_schedule() -> None:
+    # The fixture pins the official Wan source and records how to regenerate it.
+    golden = json.loads((Path(__file__).parent / "fixtures/native_unipc_shift12.json").read_text())
+    scheduler = build_wan_scheduler("unipc", golden["runtime_shift"])
+    for steps in (golden["steps"], 20, golden["steps"]):
+        scheduler.set_timesteps(steps, device="cpu", shift=golden["runtime_shift"])
+        if steps == golden["steps"]:
+            assert scheduler.timesteps.tolist() == golden["timesteps"]
+            torch.testing.assert_close(scheduler.sigmas, torch.tensor(golden["sigmas"]), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("shift", [1.0, 3.0, 5.0, 12.0])
+@pytest.mark.parametrize("steps", [20, 40])
+def test_wan_unipc_shifts_unmodified_training_endpoints(shift: float, steps: int) -> None:
+    # Native Wan: float32 training sigma_max=(1-1/1000), sigma_min=0;
+    # interpolate in float64, apply the runtime shift, then cast sigma to FP32.
+    unshifted = np.linspace(float(np.float32(0.999)), 0.0, steps + 1)[:-1]
+    shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted)
+    expected_timesteps = (shifted * 1000).astype(np.int64)
+    expected_sigmas: npt.NDArray[np.float32] = np.append(shifted, 0.0).astype(np.float32)
+    scheduler = build_wan_scheduler("unipc", shift)
+    scheduler.set_timesteps(steps, device="cpu", shift=shift)
+    assert scheduler.config.shift == scheduler.config["shift"] == 1.0
+    torch.testing.assert_close(scheduler.timesteps, torch.from_numpy(expected_timesteps), rtol=0, atol=0)
+    torch.testing.assert_close(scheduler.sigmas, torch.from_numpy(expected_sigmas), rtol=0, atol=0)
+
+
+def test_wan_euler_keeps_dmd_shift_semantics() -> None:
+    scheduler = build_wan_scheduler("euler", FASTWAN_DMD_SCHEDULER_SHIFT)
+    scheduler.set_timesteps(3, device="cpu")
+    # The unchanged Euler branch starts at sigma=1, with shift=8 for DMD.
+    expected = torch.tensor([1.0, 16.0 / 17.0, 0.8, 0.0])
+    torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=1e-7)
+    torch.testing.assert_close(scheduler.timesteps, expected[:-1] * 1000, rtol=0, atol=1e-4)

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -1,15 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from dataclasses import dataclass, field
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import build_wan_scheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
     Wan22VACEPipeline,
     create_vace_transformer_from_config,
@@ -33,6 +35,8 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
         )
         enable_layerwise_offload: bool = True
         post_init_calls: int = 0
+        _resolved_diffusion_offload: object = None
+        _diffusion_offload_flags_materialized: bool = False
 
         def __post_init__(self) -> None:
             self.post_init_calls += 1
@@ -41,7 +45,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
     resolved_offload = object()
     config._resolved_diffusion_offload = resolved_offload
     config._diffusion_offload_flags_materialized = True
-    captured = {}
+    captured: dict[str, object] = {}
 
     def capture_base_init(self, *, od_config, prefix="") -> None:
         captured.update(od_config=od_config, prefix=prefix)
@@ -51,6 +55,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
     Wan22VACEPipeline(od_config=config, prefix="stage")
 
     forwarded = captured["od_config"]
+    assert isinstance(forwarded, RuntimeConfig)
     assert forwarded is not config
     assert config.flow_shift is None
     assert forwarded.flow_shift == 3.0
@@ -75,7 +80,7 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
 
 
 def _make_vace_sampling(**overrides):
-    values = {
+    values: dict[str, object] = {
         "height": 16,
         "width": 16,
         "num_frames": 5,
@@ -123,7 +128,7 @@ def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
 
 
 def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None:
-    captured = {}
+    captured: dict[str, object] = {}
 
     class FakeVACETransformer:
         def __init__(self, **kwargs) -> None:
@@ -203,7 +208,9 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
     torch.testing.assert_close(result, torch.ones_like(latents))
 
 
-def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
+@pytest.mark.parametrize("solver", ["unipc", "euler"])
+@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
+def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch, solver: str, shift: float) -> None:
     monkeypatch.setattr(
         "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
         SimpleNamespace(is_available=lambda: False),
@@ -211,10 +218,10 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
     pipeline = _make_vace_pipeline()
     pipeline.transformer.vace_patch_embedding = None
     pipeline.transformer_2 = None
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 3.0
+    pipeline.scheduler = build_wan_scheduler(solver, shift)
+    pipeline.od_config = SimpleNamespace(flow_shift=shift)
+    pipeline._sample_solver = solver
+    pipeline._flow_shift = shift
     pipeline.boundary_ratio = None
     pipeline._guidance_scale = None
     pipeline._num_timesteps = None
@@ -247,6 +254,8 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
                     generator=gen_a,
                     latents=latents_a,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
             SimpleNamespace(
@@ -256,12 +265,25 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
                     generator=gen_b,
                     latents=latents_b,
                     num_outputs_per_prompt=2,
+                    num_inference_steps=50,
+                    extra_args={"sample_solver": solver, "flow_shift": shift},
                 ),
             ),
         ]
     )
 
     outputs = pipeline.forward(batch)
+    if solver == "unipc":
+        sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
+        sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
+        torch.testing.assert_close(
+            pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
+        )
+        assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
+    else:
+        reference = build_wan_scheduler("euler", shift)
+        reference.set_timesteps(50, device="cpu")
+        torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
 
     assert prepare_call["batch_size"] == 4
     assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]

tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py

@@ -39,7 +39,8 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
 from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
 from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage
 from vllm_omni.engine.stage_init_utils import StageMetadata
-from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver
 from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py

@@ -42,7 +42,8 @@ from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
 from vllm_omni.engine.stage_client import StagePoolClient
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.engine.stage_pool import StagePool
-from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
 from vllm_omni.outputs import OmniRequestOutput
 

tests/diffusion/test_diffusion_engine.py

@@ -15,7 +15,12 @@ from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import (
+    DIFFUSION_REQUEST_LIFECYCLE_KEY,
+    DIFFUSION_REQUEST_STARTED,
+    DiffusionOutput,
+    OmniDiffusionConfig,
+)
 from vllm_omni.diffusion.diffusion_engine import (
     DiffusionEngine,
     DiffusionExecutionMode,
@@ -183,6 +188,28 @@ def _make_request_mode_sched_output(*request_ids: str) -> RealDiffusionScheduler
     )
 
 
+@pytest.mark.cpu
+def test_request_started_output_is_emitted_only_for_opted_in_requests() -> None:
+    sched_output = _make_request_mode_sched_output("tracked", "untracked")
+    sched_output.scheduled_new_reqs[0].req.sampling_params.emit_request_lifecycle = True
+    engine = object.__new__(DiffusionEngine)
+    emitted = []
+    engine._put_output = lambda request_id, output: emitted.append((request_id, output))
+
+    engine._emit_request_started_outputs(sched_output)
+
+    assert len(emitted) == 1
+    request_id, output = emitted[0]
+    assert request_id == "tracked"
+    assert output.request_started is True
+    assert output.finished is False
+    [formatted] = engine.postprocess_output(sched_output.scheduled_new_reqs[0].req, output)
+    assert formatted.custom_output == {
+        DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED,
+    }
+    assert formatted.finished is False
+
+
 class TestRequestBatchCapability:
     pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -7,6 +7,7 @@ from unittest.mock import Mock
 import pytest
 
 from vllm_omni.diffusion import io_support
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
@@ -16,6 +17,22 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.mark.parametrize("model_class_name", ["DreamZeroPipeline", "LingBotWorldCausalDMDPipeline"])
+def test_observation_conditioned_startup_skips_generic_warmup(model_class_name: str) -> None:
+    """Observation-conditioned pipelines must not receive generic text requests."""
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = OmniDiffusionConfig.__new__(OmniDiffusionConfig)
+    engine.od_config.model_class_name = model_class_name
+    engine.od_config.diffusion_load_format = "default"
+    engine.add_req_and_wait_for_response = Mock(side_effect=AssertionError("generic text warmup submitted"))
+    engine.close = Mock()
+
+    engine.run_startup_warmup()
+
+    engine.add_req_and_wait_for_response.assert_not_called()
+    engine.close.assert_not_called()
+
+
 @pytest.mark.parametrize("step_execution", [False, True])
 def test_dummy_run_uses_enough_steps_for_execution_mode(
     monkeypatch: pytest.MonkeyPatch,

tests/diffusion/test_diffusion_engine_metrics.py

@@ -90,6 +90,7 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) ->
     Args:
         source: Full file source code.
         class_name: Enclosing class name, or None for module-level functions.
+            The method may be inherited from a base class defined in the same file.
         func_name: Function/method name.
 
     Returns:
@@ -97,13 +98,24 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) ->
     """
     tree = ast.parse(source)
     if class_name is not None:
-        for node in ast.walk(tree):
-            if isinstance(node, ast.ClassDef) and node.name == class_name:
-                for item in node.body:
-                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name:
-                        result = ast.get_source_segment(source, item)
-                        assert result is not None, f"{class_name}.{func_name} source not found"
-                        return result
+        # Follow the bases: ``Orchestrator`` inherits its loops and helpers from
+        # ``OrchestratorBase`` in the same module, and the lookup stays anchored
+        # on the public class rather than on where the split happens to fall.
+        classes = {node.name: node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)}
+        pending = [class_name]
+        seen: set[str] = set()
+        while pending:
+            name = pending.pop(0)
+            if name in seen or name not in classes:
+                continue
+            seen.add(name)
+            node = classes[name]
+            for item in node.body:
+                if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name:
+                    result = ast.get_source_segment(source, item)
+                    assert result is not None, f"{class_name}.{func_name} source not found"
+                    return result
+            pending.extend(base.id for base in node.bases if isinstance(base, ast.Name))
     else:
         for node in ast.walk(tree):
             if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:

tests/diffusion/test_diffusion_output_formatter.py

@@ -443,3 +443,14 @@ def test_format_empty_diffusion_outputs_preserves_empty_response_shape() -> None
     assert [result.prompt for result in results] == ["prompt-0"]
     assert [result.images for result in results] == [[]]
     assert [result.metrics for result in results] == [{}]
+
+
+def test_format_empty_diffusion_outputs_preserves_custom_output() -> None:
+    [result] = format_empty_diffusion_outputs(
+        _request("prompt-0"),
+        finished=False,
+        custom_output={"lifecycle": "started"},
+    )
+
+    assert result.custom_output == {"lifecycle": "started"}
+    assert result.finished is False

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import asyncio
@@ -7,7 +10,11 @@ from unittest.mock import MagicMock, patch
 import pytest
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import (
+    DIFFUSION_REQUEST_LIFECYCLE_KEY,
+    DIFFUSION_REQUEST_STARTED,
+    OmniDiffusionConfig,
+)
 from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -80,6 +87,43 @@ async def test_inline_dispatch_request_success(client, mock_engine):
     assert output.request_id == "req-1"
 
 
+@pytest.mark.asyncio
+async def test_inline_non_streaming_dispatches_lifecycle_before_final(client, mock_engine):
+    lifecycle = OmniRequestOutput.from_diffusion(
+        request_id="req-lifecycle",
+        images=[],
+        custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED},
+        finished=False,
+    )
+    intermediate = OmniRequestOutput.from_diffusion(
+        request_id="req-lifecycle",
+        images=[],
+        custom_output={"chunk": 0},
+        finished=False,
+    )
+    final = OmniRequestOutput.from_diffusion(request_id="req-lifecycle", images=[MagicMock()])
+
+    async def _step_streaming(_request):
+        yield [lifecycle]
+        yield [intermediate]
+        yield [final]
+
+    mock_engine.step_streaming = _step_streaming
+
+    await client.add_request_async("req-lifecycle", "A test prompt", OmniDiffusionSamplingParams())
+
+    outputs = []
+    for _ in range(20):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            outputs.append(output)
+            if output.finished:
+                break
+        await asyncio.sleep(0.01)
+
+    assert outputs == [lifecycle, final]
+
+
 @pytest.mark.asyncio
 async def test_inline_requests_clone_shared_sampling_params(client, mock_engine):
     requests = []

tests/diffusion/test_stage_diffusion_proc.py

@@ -9,6 +9,7 @@ import pytest
 
 import vllm_omni.diffusion.stage_diffusion_proc as stage_diffusion_proc
 import vllm_omni.plugins as omni_plugins
+from vllm_omni.diffusion.data import DIFFUSION_REQUEST_LIFECYCLE_KEY, DIFFUSION_REQUEST_STARTED
 from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -103,6 +104,49 @@ async def test_proc_streaming_request_yields_each_engine_chunk():
     assert captured["request"].kv_sender_info == {0: {"host": "127.0.0.1"}}
 
 
+@pytest.mark.asyncio
+async def test_proc_non_streaming_forwards_lifecycle_before_final_output():
+    lifecycle = OmniRequestOutput.from_diffusion(
+        request_id="",
+        images=[],
+        custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED},
+        finished=False,
+    )
+    intermediate = OmniRequestOutput.from_diffusion(
+        request_id="",
+        images=[],
+        custom_output={"chunk": 0},
+        finished=False,
+    )
+    final = OmniRequestOutput.from_diffusion(request_id="", images=[], finished=True)
+
+    class _LifecycleEngine:
+        async def step_streaming(self, request):
+            del request
+            yield [lifecycle]
+            yield [intermediate]
+            yield [final]
+
+    stage_proc = object.__new__(StageDiffusionProc)
+    stage_proc._engine = _LifecycleEngine()
+    intermediate_outputs = []
+
+    async def _capture(output):
+        intermediate_outputs.append(output)
+
+    result = await stage_proc._process_request(
+        request_id="req-lifecycle",
+        prompt="prompt",
+        sampling_params_dict=asdict(OmniDiffusionSamplingParams()),
+        on_request_started=_capture,
+    )
+
+    assert intermediate_outputs == [lifecycle]
+    assert lifecycle.request_id == "req-lifecycle"
+    assert result is final
+    assert result.request_id == "req-lifecycle"
+
+
 @pytest.mark.asyncio
 async def test_proc_process_request_with_batching_async_output():
     stage_proc = object.__new__(StageDiffusionProc)

tests/e2e/features/comfyui/test_comfyui_integration.py

@@ -16,9 +16,11 @@ import time
 import traceback
 from collections.abc import Iterable, Sequence
 from enum import Enum
+from pathlib import Path
 from types import SimpleNamespace
 from typing import Any, NamedTuple
 
+import av
 import pytest
 import requests
 import torch
@@ -82,6 +84,7 @@ class SamplingKind(str, Enum):
     VIDEO_FASTH3 = "video_fasth3"
     VIDEO_REF2VA_IMAGE_AUDIO = "video_ref2va_image_audio"
     VIDEO_REF2VA_MULTI_VIDEO = "video_ref2va_multi_video"
+    VIDEO_REF2VA_MIXED = "video_ref2va_mixed"
 
 
 # Pre-defined arguments to be used in function calls during the tests
@@ -153,6 +156,13 @@ H3_MODEL_PARAMS = MiniMaxH3ModelSpecificParams(
     }
 )
 
+H3_STAGE_CONFIG = {
+    "stage_type": "diffusion",
+    "final_output": True,
+    "final_output_type": "video",
+    "engine_args": {"model_class_name": "MiniMaxH3Pipeline"},
+}
+
 LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
 
 
@@ -442,17 +452,30 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
                     "audio_flow_shift": 3.0,
                 },
             )
-        elif sampling_case.kind is SamplingKind.VIDEO_REF2VA_MULTI_VIDEO:
+        elif sampling_case.kind in (SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, SamplingKind.VIDEO_REF2VA_MIXED):
             assert len(received_sampling_params_list) == 1
             assert isinstance(prompt, dict)
             multi_modal_data = prompt.get("multi_modal_data")
             assert isinstance(multi_modal_data, dict)
             input_videos = multi_modal_data.get("video")
             assert isinstance(input_videos, list)
-            assert len(input_videos) == 2
             assert all(isinstance(path, str) for path in input_videos)
-            assert "image" not in multi_modal_data
-            assert "audio" not in multi_modal_data
+            if sampling_case.kind is SamplingKind.VIDEO_REF2VA_MULTI_VIDEO:
+                assert len(input_videos) == 2
+                assert "image" not in multi_modal_data
+                assert "audio" not in multi_modal_data
+            else:
+                assert [Path(path).read_bytes() for path in input_videos] == [b"video-1", b"video-2", b"video-3"]
+                images = multi_modal_data.get("image")
+                assert isinstance(images, list)
+                assert len(images) == 6
+                assert [image.getpixel((0, 0)) for image in images] == [(i * 40,) * 3 for i in range(1, 7)]
+                audios = multi_modal_data.get("audio")
+                assert isinstance(audios, list)
+                assert len(audios) == 3
+                for path, sample_rate in zip(audios, (24000, 32000, 48000)):
+                    with av.open(path) as container:
+                        assert container.streams.audio[0].sample_rate == sample_rate
             _assert_sampling_param_values(
                 received_sampling_params_list[0],
                 {
@@ -970,7 +993,7 @@ async def test_fast_h3_deployment_node(api_server: str, sampling_case: SamplingC
             ServerCase(
                 served_model="MiniMaxAI/MiniMax-H3",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
+                stage_configs=[H3_STAGE_CONFIG],
                 outputs=[_build_diffusion_video_output()],
             ),
             SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
@@ -981,13 +1004,24 @@ async def test_fast_h3_deployment_node(api_server: str, sampling_case: SamplingC
             ServerCase(
                 served_model="MiniMaxAI/MiniMax-H3",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
+                stage_configs=[H3_STAGE_CONFIG],
                 outputs=[_build_diffusion_video_output()],
             ),
             SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
             "multi_video",
             id="ref2va-multi-video",
         ),
+        pytest.param(
+            ServerCase(
+                served_model="MiniMaxAI/MiniMax-H3",
+                stage_list=["diffusion"],
+                stage_configs=[H3_STAGE_CONFIG],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MIXED, sampling_params=None),
+            "mixed",
+            id="ref2va-mixed-12-references",
+        ),
     ],
     indirect=["server_case", "sampling_case"],
 )
@@ -1017,6 +1051,17 @@ async def test_video_generation_node_minimax_h3_ref2va(
         prompt = (
             "Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2."
         )
+    elif ref_mode == "mixed":
+        inputs = {f"image_{i}": torch.full((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), i * 40 / 255) for i in range(1, 7)}
+        inputs.update({f"video_{i}": VideoInput(f"video-{i}".encode()) for i in range(1, 4)})
+        inputs.update(
+            {
+                f"audio_{i}": {"waveform": torch.zeros(1, 1, sample_rate), "sample_rate": sample_rate}
+                for i, sample_rate in enumerate((24000, 32000, 48000), start=1)
+            }
+        )
+        (references,) = refs_node.get_references(**inputs)
+        prompt = "Combine the subjects from the images with the reference videos and audio clips."
     else:
         raise AssertionError(f"Unknown ref_mode: {ref_mode}")
 

tests/e2e/features/comfyui/test_h3_reference_workflow.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import json
+from graphlib import TopologicalSorter
+from pathlib import Path
+
+import pytest
+from comfyui_vllm_omni import nodes as omni_nodes
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+WORKFLOW = (
+    Path(__file__).resolve().parents[4]
+    / "apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax-H3 Reference to Video.json"
+)
+
+
+@pytest.fixture
+def workflow():
+    return json.loads(WORKFLOW.read_text())
+
+
+def test_reference_workflow_connections(workflow):
+    nodes = {node["id"]: node for node in workflow["nodes"]}
+    links = {link[0]: link for link in workflow["links"]}
+    assert len(nodes) == len(workflow["nodes"])
+    assert len(links) == len(workflow["links"])
+    graph: dict[int, set[int]] = {node_id: set() for node_id in nodes}
+    for link_id, source, output_slot, target, input_slot, kind in links.values():
+        output = nodes[source]["outputs"][output_slot]
+        input_ = nodes[target]["inputs"][input_slot]
+        assert output["type"] == input_["type"] == kind
+        assert link_id in output["links"]
+        assert input_["link"] == link_id
+        graph[target].add(source)
+    assert len(tuple(TopologicalSorter(graph).static_order())) == len(nodes)
+    for node in nodes.values():
+        for slot, input_ in enumerate(node.get("inputs", [])):
+            if input_["link"] is not None:
+                assert links[input_["link"]][3:5] == [node["id"], slot]
+        for slot, output in enumerate(node.get("outputs", [])):
+            for link_id in output.get("links") or []:
+                assert links[link_id][1:3] == [node["id"], slot]
+
+
+def test_reference_workflow_matches_omni_node_interfaces(workflow):
+    for node in workflow["nodes"]:
+        if not node["type"].startswith("VLLMOmni"):
+            continue
+        cls = getattr(omni_nodes, node["type"])
+        schema = cls.INPUT_TYPES()
+        inputs = {**schema.get("required", {}), **schema.get("optional", {})}
+        for input_ in node["inputs"]:
+            assert input_["type"] == inputs[input_["name"]][0]
+        assert tuple(output["type"] for output in node["outputs"]) == cls.RETURN_TYPES
+        expected_widgets = len(schema.get("required", {}))
+        if "seed" in schema.get("required", {}):
+            expected_widgets += 1
+        assert len(node["widgets_values"]) == expected_widgets
+    refs = next(node for node in workflow["nodes"] if node["type"] == "VLLMOmniVideoReferences")
+    assert [input_["name"] for input_ in refs["inputs"]] == list(
+        omni_nodes.VLLMOmniVideoReferences.INPUT_TYPES()["optional"]
+    )
+
+
+def test_reference_workflow_default_route_and_output(workflow):
+    nodes = {node["id"]: node for node in workflow["nodes"]}
+    links = {link[0]: link for link in workflow["links"]}
+    generate = next(node for node in nodes.values() if node["type"] == "VLLMOmniGenerateVideo")
+    inputs = {input_["name"]: input_["link"] for input_ in generate["inputs"]}
+    assert inputs["frame"] is None
+    assert inputs["lora"] is None
+    for name, kind in (
+        ("references", "VLLMOmniVideoReferences"),
+        ("sampling_params", "VLLMOmniDiffusionSampling"),
+        ("model_params", "VLLMOmniMiniMaxH3Params"),
+    ):
+        assert nodes[links[inputs[name]][1]]["type"] == kind
+    save = next(node for node in nodes.values() if node["type"] == "SaveVideo")
+    assert links[save["inputs"][0]["link"]][1] == generate["id"]
+    assert save["widgets_values"] == ["video/MiniMax-H3-Ref2VA", "mp4", "h264"]
+
+
+def test_reference_workflow_defaults_and_portable_assets(workflow):
+    by_type = {node["type"]: node for node in workflow["nodes"]}
+    values = by_type["VLLMOmniGenerateVideo"]["widgets_values"]
+    assert values[0] == "http://localhost:8000/v1"
+    assert values[1] == "MiniMaxAI/MiniMax-H3"
+    assert "<Picture 1>" in values[2]
+    assert values[4:8] == [1344, 768, 24, 5.167]
+    num_frames = round(values[7] * values[6])
+    assert num_frames == 124
+    assert values[3] == ""
+    assert (num_frames - 5) % 17 == 0
+    assert 4 <= values[7] <= 15
+    assert by_type["VLLMOmniDiffusionSampling"]["widgets_values"] == [1, 50, 1.0, 1.0, False, False, 42, "fixed"]
+    assert by_type["VLLMOmniMiniMaxH3Params"]["widgets_values"] == [3.0, 12.0]
+    assert by_type["VLLMOmniRemoteLoRA"]["widgets_values"][0] == ""
+    for kind in ("LoadImage", "LoadVideo", "LoadAudio"):
+        filename = by_type[kind]["widgets_values"][0]
+        assert filename and Path(filename).name == filename
+        assert ":" not in filename and "\\" not in filename
+    assert set(by_type) == {
+        "LoadImage",
+        "LoadVideo",
+        "LoadAudio",
+        "SaveVideo",
+        "MarkdownNote",
+        "VLLMOmniVideoReferences",
+        "VLLMOmniGenerateVideo",
+        "VLLMOmniDiffusionSampling",
+        "VLLMOmniMiniMaxH3Params",
+        "VLLMOmniRemoteLoRA",
+    }

tests/e2e/features/comfyui/test_video_references.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from io import BytesIO
+from unittest.mock import AsyncMock
+
+import av
+import pytest
+import torch
+from comfy_api.input import VideoInput
+from comfyui_vllm_omni.nodes import VLLMOmniVideoReferences
+from comfyui_vllm_omni.utils import api_client
+from PIL import Image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def api_calls(monkeypatch):
+    calls = AsyncMock(return_value={"id": "test-video", "status": "completed"})
+    monkeypatch.setattr(api_client, "url_json", calls)
+    monkeypatch.setattr(api_client, "url_bytes", AsyncMock(return_value=b"generated-video"))
+    monkeypatch.setattr(api_client, "bytes_to_video", lambda data: data)
+    return calls
+
+
+async def _generate(references, **kwargs):
+    return await api_client.VLLMOmniClient("http://localhost/v1").generate_video(
+        model="MiniMaxAI/MiniMax-H3",
+        prompt="Use the reference subjects and sounds.",
+        width=896,
+        height=512,
+        num_frames=90,
+        fps=24,
+        references=references,
+        **kwargs,
+    )
+
+
+def _references(image_count, video_count, audio_count):
+    refs = {f"image_{i}": torch.full((1, 16, 16, 3), i / 10) for i in range(1, image_count + 1)}
+    refs.update({f"video_{i}": VideoInput(f"video-{i}".encode()) for i in range(1, video_count + 1)})
+    refs.update(
+        {f"audio_{i}": {"waveform": torch.zeros(1, 1, 24000), "sample_rate": 24000} for i in range(1, audio_count + 1)}
+    )
+    return refs
+
+
+def test_reference_node_preserves_existing_ports_and_exposes_h3_limits():
+    ports = VLLMOmniVideoReferences.INPUT_TYPES()["optional"]
+    assert list(ports)[:6] == ["image_1", "image_2", "audio_1", "audio_2", "video_1", "video_2"]
+    assert {name for name, kind in ports.items() if kind == ("IMAGE",)} == {f"image_{i}" for i in range(1, 10)}
+    assert {name for name, kind in ports.items() if kind == ("VIDEO",)} == {"video_1", "video_2", "video_3"}
+    assert {name for name, kind in ports.items() if kind == ("AUDIO",)} == {"audio_1", "audio_2", "audio_3"}
+    inputs = _references(6, 3, 3)
+    (refs,) = VLLMOmniVideoReferences().get_references(**inputs)
+    assert refs.keys() == inputs.keys()
+    assert all(refs[name] is value for name, value in inputs.items())
+
+
+@pytest.mark.parametrize(
+    "counts", [(1, 0, 0), (9, 0, 0), (0, 3, 0), (1, 0, 3), (0, 1, 3), (6, 3, 3), (9, 3, 0), (9, 0, 3)]
+)
+async def test_reference_uploads_keep_numeric_order_and_media_types(api_calls, counts):
+    refs = _references(*counts)
+    result = await _generate(dict(reversed(list(refs.items()))))
+    assert result == b"generated-video"
+    fields = api_calls.call_args_list[0].kwargs["data"]._fields
+    uploads = [
+        (options["filename"], headers["Content-Type"], value)
+        for options, headers, value in fields
+        if options["name"] == "input_references"
+    ]
+    expected = [
+        (f"{kind}_{i}.{extension}", content_type)
+        for kind, count, extension, content_type in (
+            ("image", counts[0], "png", "image/png"),
+            ("video", counts[1], "mp4", "video/mp4"),
+            ("audio", counts[2], "mp3", "audio/mpeg"),
+        )
+        for i in range(1, count + 1)
+    ]
+    assert [(name, content_type) for name, content_type, _ in uploads] == expected
+    for name, content_type, payload in uploads:
+        data = payload.getvalue()
+        if content_type == "image/png":
+            index = int(name.removeprefix("image_").removesuffix(".png"))
+            assert Image.open(BytesIO(data)).getpixel((0, 0)) == (int(index / 10 * 255),) * 3
+        elif content_type == "video/mp4":
+            assert data == name.removesuffix(".mp4").replace("_", "-").encode()
+        else:
+            with av.open(BytesIO(data)) as container:
+                assert container.streams.audio[0].sample_rate == 24000
+                assert sum(frame.samples for frame in container.decode(audio=0)) > 0
+
+
+async def test_sparse_reference_slots_are_not_dropped(api_calls):
+    inputs = _references(9, 3, 3)
+    inputs = {name: inputs[name] for name in ("audio_3", "image_9", "video_3", "image_2")}
+    (refs,) = VLLMOmniVideoReferences().get_references(image_1=None, **inputs)
+    await _generate(refs)
+    fields = api_calls.call_args_list[0].kwargs["data"]._fields
+    assert [options["filename"] for options, _, _ in fields if options["name"] == "input_references"] == [
+        "image_2.png",
+        "image_9.png",
+        "video_3.mp4",
+        "audio_3.mp3",
+    ]
+
+
+@pytest.mark.parametrize(
+    "counts, message",
+    [
+        ((0, 0, 0), "at least one image or video"),
+        ((0, 0, 3), "at least one image or video"),
+        ((9, 3, 1), "at most 12"),
+        ((9, 3, 3), "at most 12"),
+    ],
+)
+async def test_invalid_reference_combinations_fail_before_upload(api_calls, counts, message):
+    with pytest.raises(ValueError, match=message):
+        await _generate(_references(*counts))
+    api_calls.assert_not_called()
+
+
+@pytest.mark.parametrize("name", ["image_10", "video_4", "audio_4", "image_0"])
+async def test_unsupported_reference_slots_are_not_silently_ignored(api_calls, name):
+    refs = _references(1, 0, 0)
+    refs[name] = object()
+    with pytest.raises(ValueError, match="Unsupported reference input"):
+        await _generate(refs)
+    api_calls.assert_not_called()
+
+
+async def test_frame_and_references_are_mutually_exclusive(api_calls):
+    with pytest.raises(ValueError, match="only one of frame or references"):
+        await _generate(_references(1, 0, 0), frame=torch.zeros(1, 16, 16, 3))
+    api_calls.assert_not_called()

tests/e2e/features/fullduplex/test_core_contracts.py

@@ -14,7 +14,7 @@ import dataclasses
 
 import pytest
 
-from vllm_omni.engine.duplex.messages import DuplexFence
+from vllm_omni.engine.duplex.contracts import DuplexFence
 from vllm_omni.experimental.fullduplex.core import protocol as ev
 from vllm_omni.experimental.fullduplex.core.adapter import (
     DuplexAdapter,
@@ -32,12 +32,15 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_fence_is_frozen_value_identity():
+    # The fence carries (session_id, epoch, turn_id). response_seq was
+    # dropped with the incarnation concept: response identity is the session's
+    # active_response_id, not a field of the fence.
     f = DuplexFence(session_id="s1")
-    assert (f.epoch, f.turn_id, f.response_seq) == (0, 0, 0)
+    assert (f.epoch, f.turn_id) == (0, 0)
     with pytest.raises(dataclasses.FrozenInstanceError):
         f.epoch = 1  # type: ignore[misc]
-    assert DuplexFence("s1", 1, 2, 3) == DuplexFence("s1", 1, 2, 3)
-    assert DuplexFence("s1", 1, 2, 3) != DuplexFence("s1", 2, 2, 3)
+    assert DuplexFence("s1", 1, 2) == DuplexFence("s1", 1, 2)
+    assert DuplexFence("s1", 1, 2) != DuplexFence("s1", 2, 2)
 
 
 def test_session_identity_rules():

tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py

@@ -71,15 +71,8 @@ from vllm_omni.experimental.fullduplex.video_stacking import (  # noqa: E402
 from vllm_omni.metrics.definitions import compute_audio_rtf  # noqa: E402
 
 
-def _url_with_model(url, model, *, autostart=None, session_id=None):
-    # This driver speaks the MiniCPM-o native duplex wire contract.
-    return build_realtime_url(
-        url,
-        model,
-        autostart=autostart,
-        session_id=session_id,
-        extra_query={"native_duplex": "1"},
-    )
+def _url_with_model(url, model, *, autostart=None):
+    return build_realtime_url(url, model, autostart=autostart)
 
 
 _read_wav_pcm16 = read_pcm16_wav
@@ -630,7 +623,6 @@ def _session_update_event(args: DemoArgs) -> dict[str, object]:
         "playback_commit_policy": "ack_only",
         "extra_body": {
             "auto_response": True,
-            "native_duplex": True,
             "force_listen_count": 0,
         },
     }
@@ -646,9 +638,6 @@ def _session_update_event(args: DemoArgs) -> dict[str, object]:
         "type": "session.update",
         "session": session_payload,
     }
-    session_id = getattr(args, "session_id", None)
-    if session_id:
-        session_payload["session_id"] = session_id
     return event
 
 
@@ -1462,7 +1451,6 @@ async def run_demo(args: DemoArgs) -> dict[str, object]:
         args.url,
         args.model,
         autostart=False if getattr(args, "ref_audio", None) else None,
-        session_id=getattr(args, "session_id", None),
     )
     state = DemoState()
     stop = asyncio.Event()
@@ -1674,6 +1662,10 @@ async def run_demo(args: DemoArgs) -> dict[str, object]:
     )
     unexpected_error_events = _unexpected_error_events(state)
     continuous_input_ok = not continuous_input or state.count("input_audio_buffer.committed") == 0
+    # The server allocates every session id; label metrics with the one it announced.
+    created = state.events[0] if state.events else {}
+    created_session = created.get("session") if isinstance(created, dict) else None
+    server_session_id = created_session.get("id") if isinstance(created_session, dict) else None
     result = {
         "ok": terminal_activity_ok
         and state.count("session.closed") > 0
@@ -1711,8 +1703,8 @@ async def run_demo(args: DemoArgs) -> dict[str, object]:
         "input_transcription_ok": input_transcription_ok,
         "completed_response_ids": completed_response_ids,
         "response_timings": state.response_timing_summaries(),
-        "request_metrics": state.session_request_metrics(session_id=args.session_id),
-        "session_metrics": state.session_metric_summary(session_id=args.session_id),
+        "request_metrics": state.session_request_metrics(session_id=server_session_id),
+        "session_metrics": state.session_metric_summary(session_id=server_session_id),
         "lifecycle_counts_ok": lifecycle_counts_ok,
         "validation_mode": validation_mode,
         "scenario": scenario,
@@ -1773,7 +1765,6 @@ def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
     parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
-    parser.add_argument("--session-id", help="Use an explicit public session ID, including for close/reopen tests.")
     parser.add_argument("--input-wav", required=True)
     parser.add_argument("--ref-audio", help="Optional WAV used as the MiniCPM-o voice prompt.")
     parser.add_argument(

tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py

@@ -17,7 +17,6 @@ import hashlib
 import io
 import json
 import sys
-import uuid
 import wave
 from pathlib import Path
 
@@ -266,8 +265,6 @@ async def _run_demo_process(
         str(args.chunk_ms),
         "--timeout-s",
         str(args.timeout_s),
-        "--session-id",
-        f"duplex-pair-{label}-{uuid.uuid4().hex}",
     ]
     if args.no_realtime_pacing:
         command.append("--no-realtime-pacing")

tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py

@@ -21,7 +21,7 @@ from websockets.asyncio.client import ClientConnection
 from websockets.exceptions import ConnectionClosed
 
 SCRIPT_DIR = Path(__file__).resolve().parent
-REPO_ROOT = SCRIPT_DIR.parents[3]
+REPO_ROOT = SCRIPT_DIR.parents[2]
 if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
 
@@ -42,6 +42,15 @@ def _url_with_model(*args, **kwargs) -> str:
     return _scenario_module()._url_with_model(*args, **kwargs)
 
 
+def _created_session_id(created: dict[str, object]) -> str:
+    """The server-allocated session id announced by ``session.created``."""
+    session = created.get("session")
+    session_id = session.get("id") if isinstance(session, dict) else None
+    if not isinstance(session_id, str) or not session_id:
+        raise RuntimeError("session.created omitted the server-allocated session id")
+    return session_id
+
+
 async def run_demo(args):
     return await _scenario_module().run_demo(args)
 
@@ -217,7 +226,6 @@ async def _open_admission_session(
         args.url,
         args.model,
         autostart=False if getattr(args, "ref_audio", None) else None,
-        session_id=session_id,
     )
     ws = await websockets.connect(url, max_size=64 * 1024 * 1024)
     await ws.send(
@@ -225,10 +233,8 @@ async def _open_admission_session(
             {
                 "type": "session.update",
                 "session": {
-                    "session_id": session_id,
                     "model": args.model,
                     "modalities": ["audio", "text"],
-                    "extra_body": {"native_duplex": True},
                     **({"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) else {}),
                 },
             }
@@ -255,12 +261,10 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
         for index in range(limit):
             accepted.append(await _open_admission_session(args, f"{prefix}-accepted-{index}"))
 
-        overflow_id = f"{prefix}-overflow"
         overflow_url = _url_with_model(
             args.url,
             args.model,
             autostart=False if getattr(args, "ref_audio", None) else None,
-            session_id=overflow_id,
         )
         async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow:
             await overflow.send(
@@ -268,10 +272,8 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
                     {
                         "type": "session.update",
                         "session": {
-                            "session_id": overflow_id,
                             "model": args.model,
                             "modalities": ["audio", "text"],
-                            "extra_body": {"native_duplex": True},
                             **(
                                 {"ref_audio": _ref_audio_data_url(args.ref_audio)}
                                 if getattr(args, "ref_audio", None)
@@ -326,7 +328,6 @@ async def _resume_probe(
         args.url,
         args.model,
         autostart=False if getattr(args, "ref_audio", None) else None,
-        session_id=session_id,
     )
     async with websockets.connect(url, max_size=64 * 1024 * 1024) as first:
         await first.send(
@@ -334,10 +335,8 @@ async def _resume_probe(
                 {
                     "type": "session.update",
                     "session": {
-                        "session_id": session_id,
                         "model": args.model,
                         "modalities": ["audio", "text"],
-                        "extra_body": {"native_duplex": True},
                         **(
                             {"ref_audio": _ref_audio_data_url(args.ref_audio)}
                             if getattr(args, "ref_audio", None)
@@ -348,10 +347,10 @@ async def _resume_probe(
             )
         )
         created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s)
+        session_id = _created_session_id(created)
         token = created.get("resume_token")
-        incarnation = created.get("incarnation")
         generation = created.get("attachment_generation")
-        if not isinstance(token, str) or not isinstance(incarnation, int):
+        if not isinstance(token, str):
             raise RuntimeError("session.created omitted resumable credentials")
         last_seq = max(_server_event_sequences(first_events), default=0)
 
@@ -366,7 +365,6 @@ async def _resume_probe(
                 {
                     "type": "session.resume",
                     "session_id": session_id,
-                    "incarnation": incarnation,
                     "resume_token": token,
                     "last_received_server_event_seq": last_seq,
                 }
@@ -429,7 +427,6 @@ async def _takeover_probe(
         args.url,
         args.model,
         autostart=False if getattr(args, "ref_audio", None) else None,
-        session_id=session_id,
     )
     first = await websockets.connect(url, max_size=64 * 1024 * 1024)
     second = None
@@ -439,10 +436,8 @@ async def _takeover_probe(
                 {
                     "type": "session.update",
                     "session": {
-                        "session_id": session_id,
                         "model": args.model,
                         "modalities": ["audio", "text"],
-                        "extra_body": {"native_duplex": True},
                         **(
                             {"ref_audio": _ref_audio_data_url(args.ref_audio)}
                             if getattr(args, "ref_audio", None)
@@ -453,10 +448,10 @@ async def _takeover_probe(
             )
         )
         created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s)
+        session_id = _created_session_id(created)
         token = created.get("resume_token")
-        incarnation = created.get("incarnation")
         generation = created.get("attachment_generation")
-        if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int):
+        if not isinstance(token, str) or not isinstance(generation, int):
             raise RuntimeError("session.created omitted takeover credentials")
         last_seq = max(_server_event_sequences(first_events), default=0)
 
@@ -466,7 +461,6 @@ async def _takeover_probe(
                 {
                     "type": "session.resume",
                     "session_id": session_id,
-                    "incarnation": incarnation,
                     "resume_token": token,
                     "last_received_server_event_seq": last_seq,
                 }

tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py

@@ -18,7 +18,6 @@ import hashlib
 import io
 import json
 import sys
-import uuid
 import wave
 from pathlib import Path
 
@@ -412,8 +411,6 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]:
         str(args.chunk_ms),
         "--timeout-s",
         str(args.timeout_s),
-        "--session-id",
-        f"duplex-soft-interrupt-{uuid.uuid4().hex}",
     ]
     command.extend(["--ref-audio", str(_canonical_path(args.ref_audio))])
     temperature = getattr(args, "temperature", None)

tests/e2e/online_serving/test_duplex_client_live.py

@@ -12,7 +12,6 @@ incremental playback acks, a forced transport drop recovered by automatic
 from __future__ import annotations
 
 import asyncio
-import uuid
 from pathlib import Path
 
 import pytest
@@ -47,7 +46,6 @@ async def _run_live_session(*, url: str, model: str, ref_audio: Path, input_wav:
         url,
         model=model,
         config=config,
-        session_id=f"duplex-client-live-{uuid.uuid4().hex}",
         reconnect=ReconnectPolicy(max_attempts=5, backoff_s=(0.25, 1.0)),
         # Short interval so the periodic heartbeat path runs during the
         # session — including across the forced transport drop below.

tests/e2e/online_serving/test_gepard_tts.py

@@ -0,0 +1,524 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""E2E online serving tests for Gepard-1.0 via ``/v1/audio/speech``.
+
+Zero-shot default voice. Needs a GPU and the NeMo NanoCodec; the weekly
+``TTS · Gepard-1.0 · L4`` step installs NeMo then runs this file.
+"""
+
+from __future__ import annotations
+
+import array
+import os
+import statistics
+import struct
+import time
+import wave
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from io import BytesIO
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+import pytest
+import requests
+import soundfile as sf
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import convert_audio_bytes_to_text
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytest.importorskip("nemo.collections.tts.models")
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+MODEL = "nineninesix/gepard-1.0"
+SAMPLE_RATE = 22050
+SAMPLES_PER_FRAME = 1024
+PCM16_BYTES_PER_FRAME = SAMPLES_PER_FRAME * 2
+DEFAULT_TIMEOUT_S = 180.0
+
+_DISTINGUISHABLE_PROMPTS = {
+    "He drinks coffee every morning.": "coffee",
+    "Machine learning is interesting.": "learning",
+    "Please close the window before leaving.": "window",
+    "My favorite color is purple.": "purple",
+}
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("gepard.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+            stage_init_timeout=900,
+        ),
+        id="gepard",
+    )
+]
+
+
+def _base_config(omni_server, text: str, **extra) -> dict:
+    cfg = {
+        "model": omni_server.model,
+        "input": text,
+        "voice": "default",
+        "timeout": DEFAULT_TIMEOUT_S,
+        "seed": 7,
+        # NanoCodec is quieter than the 24 kHz models the PCM HNR helper was
+        # calibrated on; keep the catastrophic-failure check, not a quality gate.
+        "min_hnr_db": -2.0,
+        "transcript_escalation_model": "large-v3",
+    }
+    cfg.update(extra)
+    return cfg
+
+
+def _wav_sample_rate(wav_bytes: bytes) -> int:
+    return struct.unpack_from("<I", wav_bytes, 24)[0]
+
+
+def _wav_pcm_payload_len(wav_bytes: bytes) -> int:
+    data, _sr = sf.read(BytesIO(wav_bytes), dtype="int16")
+    return int(data.size) * 2
+
+
+def _pcm16le_to_wav(pcm: bytes, sample_rate: int = SAMPLE_RATE) -> bytes:
+    """Whisper helpers expect a container; streamed speech is raw s16le."""
+    buf = BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(pcm)
+    return buf.getvalue()
+
+
+def _pcm16_samples(pcm: bytes) -> array.array:
+    samples = array.array("h")
+    samples.frombytes(pcm)
+    return samples
+
+
+def _pcm16_diag(label: str, left: bytes, right: bytes) -> None:
+    extra_frames = (len(left) - len(right)) // PCM16_BYTES_PER_FRAME
+    overlap = min(len(left), len(right))
+    first_frame_match = (
+        overlap >= PCM16_BYTES_PER_FRAME and left[:PCM16_BYTES_PER_FRAME] == right[:PCM16_BYTES_PER_FRAME]
+    )
+    worst = None
+    corr = None
+    first_diff_sample = None
+    if overlap >= 2 and overlap % 2 == 0:
+        a = _pcm16_samples(left[:overlap])
+        b = _pcm16_samples(right[:overlap])
+        worst = max(abs(x - y) for x, y in zip(a, b, strict=True))
+        for i, (x, y) in enumerate(zip(a, b, strict=True)):
+            if x != y:
+                first_diff_sample = i
+                break
+        try:
+            corr = round(statistics.correlation(a, b), 6)
+        except statistics.StatisticsError:
+            corr = None
+    print(
+        f"[Gepard pcm-diag] {label} left={len(left)} right={len(right)} "
+        f"extra_frames={extra_frames} first_frame_match={first_frame_match} "
+        f"first_diff_sample={first_diff_sample} max_abs={worst} corr={corr}"
+    )
+
+
+def _scan_logs_for_preemption(omni_server) -> None:
+    log_paths = getattr(omni_server, "_stage_log_paths", {}) or {}
+    hits: list[str] = []
+    for path in log_paths.values():
+        if path is None or not path.exists():
+            continue
+        text = path.read_text(encoding="utf-8", errors="replace").lower()
+        if "preempt" in text:
+            hits.append(str(path))
+    assert not hits, f"preemption mentioned in stage logs: {hits}"
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+@pytest.mark.parametrize("response_format", ["wav", "pcm"])
+def test_text_to_audio_basic(omni_server, online_client, response_format: str, run_level: str) -> None:
+    text = "Hello, this is Gepard speaking."
+    # whisper-small/large-v3 mishear 22.05 kHz WAV of this prompt as
+    # "Jeb Ard" / "Jeff Bard" (cosine 0.80 / 0.75). Offline uses keyword
+    # containment; the PCM twin below still goes through the 0.9 ASR gate.
+    if response_format == "wav":
+        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
+        payload = {
+            "model": omni_server.model,
+            "input": text,
+            "voice": "default",
+            "seed": 7,
+            "stream": False,
+            "response_format": "wav",
+        }
+        r = requests.post(url, json=payload, timeout=DEFAULT_TIMEOUT_S)
+        r.raise_for_status()
+        audio = r.content
+        assert _wav_sample_rate(audio) == SAMPLE_RATE
+        pcm_len = _wav_pcm_payload_len(audio)
+    else:
+        [resp] = online_client.send_audio_speech_request(
+            _base_config(omni_server, text, stream=False, response_format=response_format)
+        )
+        assert resp.audio_bytes
+        pcm_len = len(resp.audio_bytes)
+    assert pcm_len % PCM16_BYTES_PER_FRAME == 0
+    if run_level in {"advanced_model", "full_model"}:
+        duration = pcm_len / 2 / SAMPLE_RATE
+        assert 0.5 < duration < 30.0, f"implausible duration {duration:.2f}s"
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_seed_determinism(omni_server, online_client) -> None:
+    text = "Hello, this is Gepard speaking."
+    [a] = online_client.send_audio_speech_request(_base_config(omni_server, text, seed=7, response_format="pcm"))
+    [b] = online_client.send_audio_speech_request(_base_config(omni_server, text, seed=7, response_format="pcm"))
+    [c] = online_client.send_audio_speech_request(_base_config(omni_server, text, seed=11, response_format="pcm"))
+    assert a.audio_bytes == b.audio_bytes
+    assert a.audio_bytes != c.audio_bytes
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_streaming_pcm_incremental(omni_server, online_client) -> None:
+    text = "Hello, this is Gepard speaking."
+    url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
+    payload = {
+        "model": omni_server.model,
+        "input": text,
+        "voice": "default",
+        "seed": 7,
+        "stream": True,
+        "stream_format": "audio",
+        "response_format": "pcm",
+    }
+    start = time.perf_counter()
+    arrivals: list[tuple[float, int]] = []
+    streamed = bytearray()
+    with requests.post(url, json=payload, stream=True, timeout=DEFAULT_TIMEOUT_S) as resp:
+        resp.raise_for_status()
+        content_type = resp.headers.get("content-type") or ""
+        assert "pcm" in content_type or "octet-stream" in content_type
+        for chunk in resp.iter_content(chunk_size=None):
+            if not chunk:
+                continue
+            arrivals.append((time.perf_counter(), len(chunk)))
+            streamed.extend(chunk)
+    assert len(arrivals) >= 2, f"expected incremental chunks, got {len(arrivals)}"
+    assert arrivals[1][0] > arrivals[0][0]
+    ttfa_ms = (arrivals[0][0] - start) * 1000.0
+    cadence_ms = [(arrivals[i][0] - arrivals[i - 1][0]) * 1000.0 for i in range(1, len(arrivals))]
+    print(f"[Gepard streaming] TTFA_ms={ttfa_ms:.1f} chunk_cadence_ms={cadence_ms}")
+
+    [non_stream] = online_client.send_audio_speech_request(
+        _base_config(omni_server, text, seed=7, stream=False, response_format="pcm")
+    )
+    assert bytes(streamed) == non_stream.audio_bytes
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_seeded_streaming_stops(omni_server) -> None:
+    url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
+    payload = {
+        "model": omni_server.model,
+        "input": "Hello, this is Gepard speaking.",
+        "voice": "default",
+        "seed": 7,
+        "stream": True,
+        "stream_format": "audio",
+        "response_format": "pcm",
+    }
+    with requests.post(url, json=payload, stream=True, timeout=DEFAULT_TIMEOUT_S) as resp:
+        resp.raise_for_status()
+        audio = b"".join(chunk for chunk in resp.iter_content(chunk_size=None) if chunk)
+    duration = len(audio) / 2 / SAMPLE_RATE
+    assert duration < 30.0, f"seeded stream ran {duration:.1f}s — stop constraint dropped?"
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_sse_stream_emits_delta_then_done(omni_server) -> None:
+    url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
+    payload = {
+        "model": omni_server.model,
+        "input": "Hello, this is Gepard speaking.",
+        "voice": "default",
+        "seed": 7,
+        "stream": True,
+        "stream_format": "sse",
+        "response_format": "pcm",
+    }
+    with requests.post(url, json=payload, stream=True, timeout=DEFAULT_TIMEOUT_S) as resp:
+        resp.raise_for_status()
+        assert resp.headers.get("content-type", "").startswith("text/event-stream")
+        body = b"".join(resp.iter_content(chunk_size=None)).decode("utf-8", errors="replace")
+    assert "speech.audio.delta" in body
+    assert "speech.audio.done" in body
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_concurrent_requests_stay_isolated(omni_server, online_client, run_level: str) -> None:
+    def _one(text: str):
+        return online_client.send_audio_speech_request(_base_config(omni_server, text, seed=7, response_format="wav"))[
+            0
+        ]
+
+    texts = list(_DISTINGUISHABLE_PROMPTS)
+    with ThreadPoolExecutor(max_workers=4) as pool:
+        responses = list(pool.map(_one, texts))
+
+    assert all(r.success for r in responses)
+    for i in range(len(responses)):
+        for j in range(i + 1, len(responses)):
+            assert responses[i].audio_bytes != responses[j].audio_bytes
+
+    if run_level in {"advanced_model", "full_model"}:
+        for text, resp in zip(texts, responses, strict=True):
+            keyword = _DISTINGUISHABLE_PROMPTS[text]
+            transcript = convert_audio_bytes_to_text(resp.audio_bytes, language="en").lower()
+            assert keyword in transcript, f"expected {keyword!r} in {transcript!r} for {text!r}"
+
+    _scan_logs_for_preemption(omni_server)
+
+
+def _speech_url(omni_server) -> str:
+    return f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
+
+
+def _stream_payload(omni_server, text: str, **extra) -> dict:
+    payload = {
+        "model": omni_server.model,
+        "input": text,
+        "voice": "default",
+        "seed": 7,
+        "stream": True,
+        "stream_format": "audio",
+        "response_format": "pcm",
+    }
+    payload.update(extra)
+    return payload
+
+
+@dataclass
+class _StreamCapture:
+    pcm: bytes
+    arrivals: list[tuple[float, int]]
+
+    @property
+    def duration_s(self) -> float:
+        return len(self.pcm) / 2 / SAMPLE_RATE
+
+
+def _capture_stream(
+    url: str,
+    payload: dict,
+    timeout_s: float,
+    *,
+    sleep_after_first_s: float = 0.0,
+) -> _StreamCapture:
+    start = time.perf_counter()
+    arrivals: list[tuple[float, int]] = []
+    pcm = bytearray()
+    with requests.post(url, json=payload, stream=True, timeout=timeout_s) as resp:
+        resp.raise_for_status()
+        first = True
+        for chunk in resp.iter_content(chunk_size=None):
+            if not chunk:
+                continue
+            arrivals.append((time.perf_counter() - start, len(chunk)))
+            pcm.extend(chunk)
+            if first and sleep_after_first_s > 0:
+                time.sleep(sleep_after_first_s)
+            first = False
+    return _StreamCapture(bytes(pcm), arrivals)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_concurrent_streaming_stays_isolated(omni_server, run_level: str) -> None:
+    url = _speech_url(omni_server)
+    texts = list(_DISTINGUISHABLE_PROMPTS)
+
+    def _one(text: str) -> tuple[str, _StreamCapture]:
+        return text, _capture_stream(url, _stream_payload(omni_server, text), DEFAULT_TIMEOUT_S)
+
+    with ThreadPoolExecutor(max_workers=4) as pool:
+        results = list(pool.map(_one, texts))
+
+    captures = [cap for _, cap in results]
+    assert all(len(c.arrivals) >= 2 for c in captures)
+    for i in range(len(captures)):
+        for j in range(i + 1, len(captures)):
+            assert captures[i].pcm != captures[j].pcm
+
+    if run_level in {"advanced_model", "full_model"}:
+        for text, cap in results:
+            keyword = _DISTINGUISHABLE_PROMPTS[text]
+            transcript = convert_audio_bytes_to_text(_pcm16le_to_wav(cap.pcm), language="en").lower()
+            assert keyword in transcript, f"expected {keyword!r} in {transcript!r} for {text!r}"
+
+    _scan_logs_for_preemption(omni_server)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_slow_client_still_completes(omni_server, online_client) -> None:
+    """Slow reader must finish, and seed=7 must match a prior non-stream.
+
+    Runs before ``test_client_disconnect_releases_capacity`` so an abort
+    storm cannot contaminate this comparison. Prints overlap diagnostics
+    before the bit-exact asserts.
+    """
+    text = "Hello, this is Gepard speaking."
+    [baseline] = online_client.send_audio_speech_request(
+        _base_config(omni_server, text, seed=7, stream=False, response_format="pcm")
+    )
+    cap = _capture_stream(
+        url=_speech_url(omni_server),
+        payload=_stream_payload(omni_server, text),
+        timeout_s=DEFAULT_TIMEOUT_S,
+        sleep_after_first_s=3.0,
+    )
+    assert len(cap.arrivals) >= 2
+    assert len(cap.pcm) % PCM16_BYTES_PER_FRAME == 0
+    assert 0.5 < cap.duration_s < 30.0
+    print(f"[Gepard slow-client] chunks={len(cap.arrivals)} duration_s={cap.duration_s:.3f}")
+    _pcm16_diag("slow-vs-prior-nonstream", cap.pcm, baseline.audio_bytes)
+    [after] = online_client.send_audio_speech_request(
+        _base_config(omni_server, text, seed=7, stream=False, response_format="pcm")
+    )
+    _pcm16_diag("after-nonstream-vs-prior", after.audio_bytes, baseline.audio_bytes)
+    assert cap.pcm == baseline.audio_bytes
+    assert after.audio_bytes == baseline.audio_bytes
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_client_disconnect_releases_capacity(omni_server, online_client) -> None:
+    """A follow-up request must complete after in-flight streams are closed.
+
+    Logs PCM diagnostics against a pre-abort ``seed=7`` baseline. Bit-exact
+    equality after abort is not asserted here: a dying stream can still share
+    a GPU batch with the next generate after ``abort()`` returns, and Gepard's
+    in-model stop head / bf16 GEMM are not row-stable in that batch.
+    """
+    url = _speech_url(omni_server)
+    text = "Hello, this is Gepard speaking."
+    [baseline] = online_client.send_audio_speech_request(
+        _base_config(omni_server, text, seed=7, stream=False, response_format="pcm")
+    )
+    long_payload = _stream_payload(
+        omni_server,
+        "Please keep talking about the weather, the window, coffee, and purple flowers.",
+        max_new_tokens=80,
+    )
+    for _ in range(4):
+        resp = requests.post(url, json=long_payload, stream=True, timeout=DEFAULT_TIMEOUT_S)
+        resp.raise_for_status()
+        try:
+            chunk = next(c for c in resp.iter_content(chunk_size=None) if c)
+            assert chunk
+        finally:
+            resp.close()
+
+    t0 = time.perf_counter()
+    follow = requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": text,
+            "voice": "default",
+            "seed": 7,
+            "stream": False,
+            "response_format": "pcm",
+        },
+        timeout=DEFAULT_TIMEOUT_S,
+    )
+    follow.raise_for_status()
+    elapsed = time.perf_counter() - t0
+    pcm_len = len(follow.content)
+    assert pcm_len % PCM16_BYTES_PER_FRAME == 0
+    print(f"[Gepard disconnect] followup_s={elapsed:.2f} pcm_bytes={pcm_len}")
+    _pcm16_diag("disconnect-followup-vs-prior", follow.content, baseline.audio_bytes)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_mixed_stream_and_nonstream_isolated(omni_server, online_client, run_level: str) -> None:
+    url = _speech_url(omni_server)
+    stream_text = "He drinks coffee every morning."
+    nonstream_text = "My favorite color is purple."
+
+    def _stream():
+        return _capture_stream(url, _stream_payload(omni_server, stream_text), DEFAULT_TIMEOUT_S)
+
+    def _nonstream():
+        return online_client.send_audio_speech_request(
+            _base_config(omni_server, nonstream_text, seed=7, response_format="wav")
+        )[0]
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        stream_future = pool.submit(_stream)
+        nonstream_future = pool.submit(_nonstream)
+        streamed = stream_future.result()
+        nonstream = nonstream_future.result()
+
+    assert streamed.pcm
+    assert nonstream.success
+    assert streamed.pcm != nonstream.audio_bytes
+    if run_level in {"advanced_model", "full_model"}:
+        stream_tr = convert_audio_bytes_to_text(_pcm16le_to_wav(streamed.pcm), language="en").lower()
+        nonstream_tr = convert_audio_bytes_to_text(nonstream.audio_bytes, language="en").lower()
+        assert "coffee" in stream_tr, stream_tr
+        assert "purple" in nonstream_tr, nonstream_tr
+    _scan_logs_for_preemption(omni_server)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_max_new_tokens_caps_whole_frames(omni_server) -> None:
+    frames = 16
+    url = _speech_url(omni_server)
+    r = requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": "Hello, this is Gepard speaking.",
+            "voice": "default",
+            "seed": 7,
+            "stream": False,
+            "response_format": "pcm",
+            "max_new_tokens": frames,
+        },
+        timeout=DEFAULT_TIMEOUT_S,
+    )
+    assert r.status_code == 200, r.text
+    pcm_len = len(r.content)
+    assert pcm_len % PCM16_BYTES_PER_FRAME == 0
+    assert pcm_len == frames * PCM16_BYTES_PER_FRAME, f"got {pcm_len} bytes, want {frames} frames"
+    duration = pcm_len / 2 / SAMPLE_RATE
+    assert duration < 1.0, f"budget cap should stop near 16 frames, got {duration:.2f}s"
+
+    over = requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": "Hello, this is Gepard speaking.",
+            "voice": "default",
+            "max_new_tokens": 4097,
+        },
+        timeout=DEFAULT_TIMEOUT_S,
+    )
+    assert over.status_code == 400
+    assert "max_new_tokens" in over.text

tests/e2e/online_serving/test_minicpmo_4_5_duplex.py

@@ -6,8 +6,8 @@
 from __future__ import annotations
 
 import asyncio
+import base64
 import json
-import uuid
 from pathlib import Path
 
 import pytest
@@ -90,25 +90,16 @@ async def _receive_protocol_events(ws, required_types: set[str], *, timeout_s: f
 
 
 async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[dict[str, object]]:
-    session_id = f"duplex-ci-protocol-{uuid.uuid4().hex}"
-    websocket_url = build_realtime_url(
-        url,
-        model,
-        autostart=False,
-        session_id=session_id,
-        extra_query={"native_duplex": "1"},
-    )
+    websocket_url = build_realtime_url(url, model, autostart=False)
     async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws:
         await ws.send(
             json.dumps(
                 {
                     "type": "session.update",
                     "session": {
-                        "session_id": session_id,
                         "model": model,
                         "modalities": ["audio", "text"],
                         "ref_audio": _ref_audio_data_url(str(ref_audio)),
-                        "extra_body": {"native_duplex": True},
                     },
                 }
             )
@@ -123,6 +114,170 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[
     return events
 
 
+async def _run_text_only_response_create(
+    *,
+    url: str,
+    model: str,
+    ref_audio: Path,
+    text: str,
+    timeout_s: float = 120.0,
+) -> dict[str, object]:
+    """Drive a text prompt with no audio and return the server's answer.
+
+    A Realtime client asks for speech from text with a user message item
+    followed by ``response.create``. A model-native duplex session generates
+    from audio units, so this returns whatever the server answers rather than
+    assuming a response.
+    """
+    websocket_url = build_realtime_url(url, model, autostart=False)
+    async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws:
+        await ws.send(
+            json.dumps(
+                {
+                    "type": "session.update",
+                    "session": {
+                        "model": model,
+                        "modalities": ["audio", "text"],
+                        "output_audio_format": "pcm16",
+                        "turn_detection": None,
+                        "temperature": 0.0,
+                        "ref_audio": _ref_audio_data_url(str(ref_audio)),
+                        "extra_body": {"auto_response": False},
+                    },
+                }
+            )
+        )
+        await _receive_protocol_events(ws, {"session.created", "session.updated"}, timeout_s=60)
+        await ws.send(
+            json.dumps(
+                {
+                    "type": "conversation.item.create",
+                    "item": {
+                        "type": "message",
+                        "role": "user",
+                        "content": [{"type": "input_text", "text": text}],
+                    },
+                }
+            )
+        )
+        await ws.send(json.dumps({"type": "response.create"}))
+
+        async def receive_outcome() -> dict[str, object]:
+            while True:
+                raw = await ws.recv()
+                if not isinstance(raw, str):
+                    continue
+                event = json.loads(raw)
+                if isinstance(event, dict) and event.get("type") in {"error", "response.done"}:
+                    return event
+
+        return await asyncio.wait_for(receive_outcome(), timeout=timeout_s)
+
+
+async def _run_seeded_text_to_audio(
+    *,
+    url: str,
+    model: str,
+    ref_audio: Path | None,
+    text: str,
+    modalities: tuple[str, ...] = ("audio", "text"),
+    silence_seconds: float = 12.0,
+    timeout_s: float = 180.0,
+) -> dict[str, object]:
+    """Speak a seeded text: the duplex route's text-to-speech shape.
+
+    A model-native session takes its text once, in the session context
+    (``duplex_initial_user_text``), and then generates per audio unit. Silence
+    carries no content of its own, so it only advances the clock and lets the
+    model answer the seeded turn.
+    """
+    websocket_url = build_realtime_url(url, model, autostart=False)
+    audio_bytes = 0
+    transcript: list[str] = []
+    output_text: list[str] = []
+    seen: list[str] = []
+    session_payload: dict[str, object] = {
+        "model": model,
+        "modalities": list(modalities),
+        "input_audio_format": "pcm16",
+        "output_audio_format": "pcm16",
+        "turn_detection": None,
+        "temperature": 0.0,
+        "extra_body": {
+            "auto_response": True,
+            "force_listen_count": 0,
+            "duplex_initial_user_text": text,
+        },
+    }
+    if ref_audio is not None:
+        session_payload["ref_audio"] = _ref_audio_data_url(str(ref_audio))
+    async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws:
+        await ws.send(json.dumps({"type": "session.update", "session": session_payload}))
+
+        async def reader() -> None:
+            nonlocal audio_bytes
+            while True:
+                raw = await ws.recv()
+                if not isinstance(raw, str):
+                    continue
+                event = json.loads(raw)
+                if not isinstance(event, dict):
+                    continue
+                event_type = event.get("type")
+                if isinstance(event_type, str):
+                    seen.append(event_type)
+                delta = event.get("delta")
+                if event_type == "response.output_audio.delta" and isinstance(delta, str):
+                    audio_bytes += len(base64.b64decode(delta))
+                elif event_type == "response.output_audio_transcript.delta" and isinstance(delta, str):
+                    transcript.append(delta)
+                elif event_type == "response.output_text.delta" and isinstance(delta, str):
+                    output_text.append(delta)
+                elif event_type == "error":
+                    raise AssertionError(f"seeded text turn received an error: {event}")
+
+        reader_task = asyncio.create_task(reader())
+        silence = bytes(2 * 16_000 * 200 // 1000)
+        sent_ms = 0
+        try:
+            while sent_ms < silence_seconds * 1000 and "response.done" not in seen:
+                sent_ms += 200
+                await ws.send(
+                    json.dumps(
+                        {
+                            "type": "input_audio_buffer.append",
+                            "audio": base64.b64encode(silence).decode("ascii"),
+                            "input_audio_format": "pcm16",
+                            "sample_rate_hz": 16_000,
+                            "duration_ms": 200,
+                            "audio_end_ms": sent_ms,
+                        }
+                    )
+                )
+                await asyncio.sleep(0.2)
+            loop = asyncio.get_running_loop()
+            deadline = loop.time() + timeout_s
+            while loop.time() < deadline and "response.done" not in seen:
+                await asyncio.sleep(0.5)
+            if reader_task.done():
+                await reader_task
+            # Close explicitly. Dropping the socket parks the session in its
+            # disconnect grace, where it keeps holding an admission slot and
+            # starves the tests that follow.
+            await ws.send(json.dumps({"type": "session.close"}))
+            deadline = loop.time() + 30
+            while loop.time() < deadline and "session.closed" not in seen:
+                await asyncio.sleep(0.25)
+        finally:
+            reader_task.cancel()
+    return {
+        "audio_bytes": audio_bytes,
+        "transcript": "".join(transcript),
+        "output_text": "".join(output_text),
+        "event_types": seen,
+    }
+
+
 @pytest.mark.core_model
 @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
@@ -205,6 +360,215 @@ def test_duplex_single_session_video_input(omni_server, tmp_path: Path) -> None:
     _assert_session_metrics(result["session_metrics"], expected_count=2)
 
 
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+@pytest.mark.parametrize(
+    ("locale", "text"),
+    [
+        ("en", "Please say exactly: the quick brown fox jumps over the lazy dog."),
+        ("zh", "请朗读：今天天气很好，我们一起去公园散步。"),
+    ],
+)
+def test_duplex_seeded_text_to_audio(omni_server, locale: str, text: str) -> None:
+    """text -> audio over the duplex route, in both locales.
+
+    This is the case the deleted turn-based ``text -> audio`` tests covered.
+    It is also the shape the Seed-TTS Realtime backend uses, so it guards the
+    benchmark path as well as the modality.
+    """
+    result = asyncio.run(
+        _run_seeded_text_to_audio(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            ref_audio=resolve_ref_audio(),
+            text=text,
+        )
+    )
+
+    assert "response.done" in result["event_types"], result["event_types"]
+    assert int(result["audio_bytes"]) > 0, f"{locale} seeded text produced no audio"
+    assert str(result["transcript"]).strip(), f"{locale} seeded text produced audio with no transcript"
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_seeded_text_to_text_needs_no_reference_voice(omni_server) -> None:
+    """text -> text, the one duplex session that opens without a reference voice.
+
+    ``ref_audio`` is required only when the session asks for audio output, so a
+    ``modalities: ["text"]`` session is the duplex equivalent of the deleted
+    turn-based ``text -> text`` case. The model is model-native and still
+    speaks its answer, so this asserts the text side and the absence of the
+    ``ref_audio_required`` rejection, not the absence of audio.
+    """
+    result = asyncio.run(
+        _run_seeded_text_to_audio(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            ref_audio=None,
+            text="What is the capital of France? Answer in one short sentence.",
+            modalities=("text",),
+        )
+    )
+
+    assert "response.done" in result["event_types"], result["event_types"]
+    produced_text = str(result["output_text"]) or str(result["transcript"])
+    assert produced_text.strip(), f"text-only session produced nothing: {result['event_types']}"
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_seeded_text_to_long_audio_output(omni_server) -> None:
+    """A longer answer, so Code2Wav runs across many frames rather than one.
+
+    Replaces the deleted ``text_to_audio_long_output`` case. The assertion is
+    on the amount of audio rather than its content: the point is that a long
+    answer streams to completion, not that the model says any given word.
+    """
+    result = asyncio.run(
+        _run_seeded_text_to_audio(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            ref_audio=resolve_ref_audio(),
+            text="What is the capital of China? Answer in about 40 words.",
+            silence_seconds=30.0,
+            timeout_s=240.0,
+        )
+    )
+
+    assert "response.done" in result["event_types"], result["event_types"]
+    # 24 kHz mono pcm16: 2 s of speech is 96000 bytes, comfortably more than a
+    # single Code2Wav frame and well under a 40-word answer.
+    assert int(result["audio_bytes"]) > 96_000, f"expected a long answer, got {result['audio_bytes']} bytes"
+    assert str(result["transcript"]).strip()
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_seeded_long_form_generation(omni_server) -> None:
+    """Long-form Chinese generation, replacing the deleted long-form case."""
+    result = asyncio.run(
+        _run_seeded_text_to_audio(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            ref_audio=resolve_ref_audio(),
+            text="帮我讲一个100字的故事。",
+            silence_seconds=30.0,
+            timeout_s=240.0,
+        )
+    )
+
+    assert "response.done" in result["event_types"], result["event_types"]
+    assert int(result["audio_bytes"]) > 96_000, f"expected long-form audio, got {result['audio_bytes']} bytes"
+    assert str(result["transcript"]).strip()
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_sequential_sessions_are_independent(omni_server) -> None:
+    """A second session must answer its own prompt, not replay the first one.
+
+    Replaces the deleted ``sequential_requests_independent`` case. Sessions are
+    engine-resident now, so this is the check that one session's context does
+    not leak into the next one on the same stages.
+    """
+    first = asyncio.run(
+        _run_seeded_text_to_audio(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            ref_audio=resolve_ref_audio(),
+            text="What is the capital of France? Answer in one short sentence.",
+        )
+    )
+    second = asyncio.run(
+        _run_seeded_text_to_audio(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            ref_audio=resolve_ref_audio(),
+            text="What is the capital of China? Answer in one short sentence.",
+        )
+    )
+
+    for result in (first, second):
+        assert "response.done" in result["event_types"], result["event_types"]
+        assert int(result["audio_bytes"]) > 0
+        assert str(result["transcript"]).strip()
+
+    # Different prompts must not produce the same answer: that would mean the
+    # second session inherited the first one's context.
+    assert str(first["transcript"]).strip() != str(second["transcript"]).strip()
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_text_only_response_create_is_rejected(omni_server) -> None:
+    """A duplex session generates from audio: a text-only prompt is refused, not ignored.
+
+    Every session on this route is model-native, and such a model decides per
+    audio unit whether to speak: conversation items are context, not a turn.
+    Refusing is the point -- opening a response the model never fills costs the
+    caller the session's whole idle timeout to discover. A text prompt reaches
+    the model through the session's seeded opening turn
+    (``initial_user_text``) instead, which is what ``/v1/chat/completions``
+    uses and what the Seed-TTS Realtime backend has to be ported to.
+    """
+    outcome = asyncio.run(
+        _run_text_only_response_create(
+            url=realtime_url(omni_server),
+            model=omni_server.model,
+            ref_audio=resolve_ref_audio(),
+            text="Please say: the quick brown fox jumps over the lazy dog.",
+        )
+    )
+
+    assert outcome.get("type") == "error", outcome
+    error = outcome.get("error")
+    assert isinstance(error, dict), outcome
+    assert error.get("code") == "text_only_turn_unsupported", outcome
+    assert error.get("type") == "invalid_request_error", outcome
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_single_session_still_image_input(omni_server, tmp_path: Path) -> None:
+    """A spoken turn with one still image attached.
+
+    The still-image case is a one-element frame list, which the client repeats
+    on every unit of the turn: unlike the camera track it never advances, so
+    this covers the image path the turn-based ``image -> text + audio`` request
+    used to cover.
+    """
+    args = demo_args(
+        omni_server=omni_server,
+        input_wav=validated_input_wav(),
+        ref_audio=resolve_ref_audio(),
+        output_dir=tmp_path / "still_image",
+    )
+    args.turns = 2
+    args.turn_duration_ms = [args.first_turn_ms] * args.turns
+    args.video_frames_b64 = duplex_camera_frames(seconds=1, cache_dir=tmp_path / "still")[:1]
+
+    result = asyncio.run(run_demo(args))
+
+    assert result["ok"] is True
+    assert result["video_frame_count"] == 1
+    assert result["video_stacked_frame_count"] == 0
+    _assert_positive_int(result["audio_delta_count"])
+    assert result["done_count"] == 2
+    assert result["error_count"] == 0
+    assert result["all_audio_responses_have_transcript"] is True
+    assert result["transcript_delta_done_ok"] is True
+    _assert_request_metrics(result["request_metrics"], expected_count=2)
+    _assert_session_metrics(result["session_metrics"], expected_count=2)
+
+
 @pytest.mark.advanced_model
 @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)

tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py

@@ -30,7 +30,7 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import
     run_soft_interrupt,
 )
 from tests.helpers.mark import hardware_test
-from vllm_omni.entrypoints.duplex.server_vad import (
+from vllm_omni.engine.duplex.vad import (
     SILERO_VAD_FILENAME,
     SILERO_VAD_REPO_ID,
     SILERO_VAD_REVISION,

tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py

@@ -152,7 +152,6 @@ def test_realtime_duplex_demo_pair_launches_demo_processes_concurrently(tmp_path
                 "parser.add_argument('--output-dir')",
                 "parser.add_argument('--chunk-ms')",
                 "parser.add_argument('--timeout-s')",
-                "parser.add_argument('--session-id')",
                 "parser.add_argument('--ref-audio')",
                 "parser.add_argument('--require-audio', action='store_true')",
                 "parser.add_argument('--no-realtime-pacing', action='store_true')",
@@ -1027,7 +1026,7 @@ def test_realtime_duplex_demo_resolves_distinct_turn_inputs():
         demo._turn_input_paths(primary, ["second.wav"], turns=3)
 
 
-def test_realtime_duplex_demo_explicitly_enables_native_runtime_before_connect():
+def test_realtime_duplex_demo_url_carries_no_model_opt_in_or_session_id():
     demo = _load_demo_module()
 
     url = demo._url_with_model(
@@ -1036,23 +1035,12 @@ def test_realtime_duplex_demo_explicitly_enables_native_runtime_before_connect()
     )
 
     query = parse_qs(urlsplit(url).query)
-    assert query["native_duplex"] == ["1"]
+    assert query["duplex"] == ["1"]
+    assert "native_duplex" not in query
+    assert "session_id" not in query
 
 
-def test_realtime_duplex_demo_explicit_session_id_reaches_autostart_query():
-    demo = _load_demo_module()
-
-    url = demo._url_with_model(
-        "ws://localhost:8099/v1/realtime?duplex=1",
-        "openbmb/MiniCPM-o-4_5",
-        session_id="reopen-e2e",
-    )
-
-    query = parse_qs(urlsplit(url).query)
-    assert query["session_id"] == ["reopen-e2e"]
-
-
-def test_realtime_duplex_demo_session_update_uses_explicit_session_id():
+def test_realtime_duplex_demo_session_update_leaves_session_id_to_the_server():
     demo = _load_demo_module()
 
     event = demo._session_update_event(
@@ -1066,8 +1054,8 @@ def test_realtime_duplex_demo_session_update_uses_explicit_session_id():
 
     assert event["type"] == "session.update"
     assert "session_id" not in event
-    assert event["session"]["session_id"] == "reopen-e2e"
-    assert event["session"]["extra_body"]["native_duplex"] is True
+    assert "session_id" not in event["session"]
+    assert "native_duplex" not in event["session"]["extra_body"]
 
 
 def test_realtime_duplex_demo_response_required_uses_deterministic_sampling():

tests/e2e/online_serving/test_voxtral_tts_expansion.py

@@ -12,11 +12,14 @@ import os
 os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
 
 import pytest
+from openai import BadRequestError as OpenAIBadRequestError
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
+# Note this model is the public checkpoint without the encoder weights required for voice cloning
 MODEL = "mistralai/Voxtral-4B-TTS-2603"
 STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml")
 EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"]
@@ -173,3 +176,21 @@ def test_speech_batches(omni_server, online_client) -> None:
             "timeout": 120.0,
         }
     )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_speech_ref_audio(omni_server, online_client) -> None:
+    """Test ref_audio failure is handled as a BadRequest instead of a RuntimeError"""
+    audio_data = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+    with pytest.raises(OpenAIBadRequestError) as exc_info:
+        online_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "how are you",
+                "ref_audio": audio_data,
+                "language": "English",
+                "response_format": "wav",
+            }
+        )
+
+    assert exc_info.value.status_code == 400

tests/engine/duplex/test_duplex_audio.py

@@ -6,7 +6,7 @@ import base64
 import numpy as np
 import pytest
 
-from vllm_omni.entrypoints.duplex.audio import convert_input_audio_with_rate
+from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 

tests/engine/duplex/test_duplex_commands.py

@@ -0,0 +1,453 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OpenAI Realtime client events -> typed ``DuplexCommand`` objects and their mailbox payloads."""
+
+from __future__ import annotations
+
+import base64
+
+import numpy as np
+import pytest
+
+from vllm_omni.engine.duplex.commands import (
+    REALTIME_COMMAND_TYPES,
+    AckPlayback,
+    AppendAudio,
+    AppendText,
+    BargeIn,
+    CancelInput,
+    CancelResponse,
+    ClearInput,
+    ClearOutputAudio,
+    CloseSession,
+    Commit,
+    CreateItem,
+    CreateResponse,
+    DeleteItem,
+    DuplexCommand,
+    DuplexCommandError,
+    Heartbeat,
+    SignalTurn,
+    TruncateItem,
+    UpdateSession,
+    command_from_realtime,
+)
+from vllm_omni.engine.duplex.realtime_commands import (
+    RealtimeInputDefaults,
+    translate_realtime_command,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_LOUD_PCM16 = base64.b64encode(b"\x00\x10" * 8).decode("ascii")
+_SILENT_PCM16 = base64.b64encode(b"\x00\x00" * 8).decode("ascii")
+_JPEG_FRAME = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 12).decode("ascii")
+
+_MINIMAL_PAYLOADS: dict[str, tuple[dict[str, object], type[DuplexCommand]]] = {
+    "input_audio_buffer.append": ({"audio": _LOUD_PCM16}, AppendAudio),
+    "input_audio_buffer.commit": ({}, Commit),
+    "input_audio_buffer.clear": ({}, ClearInput),
+    "output_audio_buffer.clear": ({}, ClearOutputAudio),
+    "response.create": ({}, CreateResponse),
+    "response.cancel": ({}, CancelResponse),
+    "conversation.item.create": ({"item": {"type": "message", "role": "user", "content": []}}, CreateItem),
+    "conversation.item.delete": ({"item_id": "item_1"}, DeleteItem),
+    "conversation.item.truncate": ({"item_id": "item_1", "audio_end_ms": 100}, TruncateItem),
+    "session.update": ({"session": {"instructions": "hi"}}, UpdateSession),
+    "playback.ack": ({"played_ms": 10}, AckPlayback),
+    "session.heartbeat": ({}, Heartbeat),
+    "session.close": ({}, CloseSession),
+    "turn.signal": ({"event": "user_started"}, SignalTurn),
+    "input.text.append": ({"text": "hello"}, AppendText),
+    "input.cancel": ({}, CancelInput),
+    "barge_in": ({}, BargeIn),
+}
+
+
+def test_every_realtime_command_type_has_a_mapping_case():
+    assert set(_MINIMAL_PAYLOADS) == set(REALTIME_COMMAND_TYPES)
+
+
+@pytest.mark.parametrize("event_type", sorted(REALTIME_COMMAND_TYPES))
+def test_command_from_realtime_maps_each_type_to_its_dataclass_and_propagates_event_id(event_type: str):
+    body, expected_cls = _MINIMAL_PAYLOADS[event_type]
+
+    command = command_from_realtime({"type": event_type, "event_id": f"evt-{event_type}", **body})
+
+    assert type(command) is expected_cls
+    assert command.event_id == f"evt-{event_type}"
+    payload = command.payload()
+    assert payload["type"] == expected_cls.type
+    assert payload["realtime_event_id"] == f"evt-{event_type}"
+    assert "event_id" not in payload
+
+    anonymous = command_from_realtime({"type": event_type, **body})
+    assert anonymous.event_id is None
+    assert "realtime_event_id" not in anonymous.payload()
+
+
+@pytest.mark.parametrize(
+    ("alias", "expected_cls", "body"),
+    [
+        ("push_text", AppendText, {"text": "hello"}),
+        ("input_text.append", AppendText, {"text": "hello"}),
+        ("signal_turn", SignalTurn, {"event": "user_started"}),
+        ("close_session", CloseSession, {}),
+        ("close", CloseSession, {}),
+        ("audio.playback_ack", AckPlayback, {"played_ms": 10}),
+        ("input.commit", Commit, {}),
+    ],
+)
+def test_pre_realtime_client_event_aliases_map_to_their_canonical_command(
+    alias: str, expected_cls: type[DuplexCommand], body: dict[str, object]
+):
+    assert type(command_from_realtime({"type": alias, **body})) is expected_cls
+
+
+def test_append_audio_decodes_and_converts_pcm16_to_16k_float32():
+    command = command_from_realtime(
+        {"type": "input_audio_buffer.append", "audio": _LOUD_PCM16, "format": "pcm16", "sample_rate_hz": 16000}
+    )
+
+    assert isinstance(command, AppendAudio)
+    assert command.format == "pcm_f32le"
+    assert command.sample_rate_hz == 16000
+    samples = np.frombuffer(command.audio, dtype="<f4")
+    assert samples.shape == (8,)
+    assert np.allclose(samples, 4096 / 32768.0)
+    assert command.is_speech is True
+    assert command.video_frames == ()
+    assert command.duration_ms is None
+    assert command.audio_end_ms is None
+
+
+def test_append_audio_defaults_come_from_the_session_payload():
+    defaults = RealtimeInputDefaults().with_session_payload({"input_audio_format": "pcm16", "sample_rate_hz": 8000})
+    assert defaults.input_audio_format == "pcm16"
+    assert defaults.input_sample_rate_hz == 8000
+
+    command = translate_realtime_command({"type": "input_audio_buffer.append", "audio": _LOUD_PCM16}, defaults=defaults)
+
+    assert isinstance(command, AppendAudio)
+    assert command.format == "pcm_f32le"
+    # 8 kHz input is resampled to the model's 16 kHz float stream.
+    assert command.sample_rate_hz == 16000
+    assert len(command.audio) > 8 * 4
+
+    passthrough = translate_realtime_command(
+        {"type": "input_audio_buffer.append", "audio": base64.b64encode(np.zeros(4, dtype="<f4").tobytes()).decode()},
+        defaults=RealtimeInputDefaults(input_audio_format="pcm_f32le", input_sample_rate_hz=16000),
+    )
+    assert passthrough.format == "pcm_f32le"
+    assert passthrough.audio == np.zeros(4, dtype="<f4").tobytes()
+
+
+def test_append_audio_speech_classification_uses_hints_then_rms():
+    silent = command_from_realtime({"type": "input_audio_buffer.append", "audio": _SILENT_PCM16})
+    assert silent.is_speech is False
+
+    forced = command_from_realtime({"type": "input_audio_buffer.append", "audio": _SILENT_PCM16, "is_speech": True})
+    assert forced.is_speech is True
+    assert forced.hints == {"is_speech": True}
+
+    vad = command_from_realtime(
+        {"type": "input_audio_buffer.append", "audio": _LOUD_PCM16, "vad": {"speech_probability": 0.1}}
+    )
+    assert vad.is_speech is False
+    assert vad.hints["vad"] == {"speech_probability": 0.1}
+
+
+def test_append_audio_carries_wire_hints_and_video_frames():
+    command = command_from_realtime(
+        {
+            "type": "input_audio_buffer.append",
+            "event_id": "evt-append",
+            "audio": _LOUD_PCM16,
+            "transcript": "hi there",
+            "duration_ms": 1000,
+            "audio_end_ms": 3000,
+            "video_frames": [_JPEG_FRAME],
+            "unknown_key": "dropped",
+        }
+    )
+
+    assert isinstance(command, AppendAudio)
+    assert command.duration_ms == 1000
+    assert command.audio_end_ms == 3000
+    assert command.video_frames == (_JPEG_FRAME,)
+    assert dict(command.hints) == {"transcript": "hi there", "duration_ms": 1000, "audio_end_ms": 3000}
+
+    payload = command.payload()
+    assert payload["type"] == "input_audio_buffer.append"
+    assert payload["realtime_event_id"] == "evt-append"
+    assert payload["audio"] == base64.b64encode(command.audio).decode("ascii")
+    assert payload["format"] == "pcm_f32le"
+    assert payload["sample_rate_hz"] == 16000
+    assert payload["is_speech"] is True
+    assert payload["transcript"] == "hi there"
+    assert payload["duration_ms"] == 1000
+    assert payload["audio_end_ms"] == 3000
+    assert payload["video_frames"] == [_JPEG_FRAME]
+    assert "hints" not in payload
+    assert "unknown_key" not in payload
+
+
+def test_append_audio_payload_re_encodes_bytes_and_flattens_hints():
+    command = AppendAudio(audio=b"\x01\x02", hints={"vad": {"is_speech": True}}, is_speech=True)
+
+    assert command.payload() == {
+        "type": "input_audio_buffer.append",
+        "audio": "AQI=",
+        "format": "pcm16",
+        "is_speech": True,
+        "vad": {"is_speech": True},
+    }
+
+
+@pytest.mark.parametrize(
+    ("body", "code"),
+    [
+        ({"audio": _LOUD_PCM16, "format": "mp3"}, "unsupported_audio_format"),
+        ({"audio": "@@@@", "format": "pcm_f32le"}, "bad_audio"),
+        ({"audio": _LOUD_PCM16, "video_frames": ["aGVsbG8="]}, "invalid_video_frames"),
+        ({"audio": _LOUD_PCM16, "video_frames": [_JPEG_FRAME], "max_slice_nums": 4}, "invalid_video_frames"),
+        ({"audio": _LOUD_PCM16, "format": "pcm16", "sample_rate_hz": 3}, "bad_event"),
+    ],
+)
+def test_malformed_append_raises_command_error_with_code(body: dict[str, object], code: str):
+    with pytest.raises(DuplexCommandError) as excinfo:
+        command_from_realtime({"type": "input_audio_buffer.append", "event_id": "evt-bad", **body})
+
+    assert excinfo.value.code == code
+    assert excinfo.value.event_id == "evt-bad"
+
+
+# ---- other malformed payloads ----
+
+
+@pytest.mark.parametrize(
+    ("payload", "code"),
+    [
+        ({}, "bad_event"),
+        ({"type": 7}, "bad_event"),
+        ({"type": "session.resume", "session_id": "s", "resume_token": "t"}, "unknown_event"),
+        ({"type": "session.event_ack", "server_event_seq": 1}, "unknown_event"),
+        ({"type": "conversation.item.frobnicate", "item_id": "item_1"}, "unknown_event"),
+        ({"type": "conversation.item.create"}, "bad_event"),
+        (
+            {
+                "type": "conversation.item.create",
+                "item": {"type": "message", "role": "user", "content": [{"type": "input_audio", "format": "mp3"}]},
+            },
+            "unsupported_audio_format",
+        ),
+        ({"type": "conversation.item.delete"}, "missing_item_id"),
+        ({"type": "conversation.item.truncate", "audio_end_ms": 1}, "missing_item_id"),
+        ({"type": "conversation.item.truncate", "item_id": "item_1"}, "bad_event"),
+        ({"type": "session.update", "session": {"output_audio_format": "opus"}}, "unsupported_audio_format"),
+        (
+            {"type": "session.update", "session": {"turn_detection": {"type": "semantic_vad"}}},
+            "unsupported_turn_detection",
+        ),
+        ({"type": "response.create", "response": {"output_audio_format": "opus"}}, "unsupported_audio_format"),
+        ({"type": "playback.ack"}, "bad_event"),
+        ({"type": "input.text.append"}, "bad_event"),
+        ({"type": "turn.signal"}, "bad_event"),
+        ({"type": "turn.signal", "event": "conversation.item.delete", "payload": {}}, "missing_item_id"),
+        ({"type": "turn.signal", "event": "conversation.item.truncate", "payload": {"item_id": "i"}}, "bad_event"),
+    ],
+)
+def test_malformed_payloads_raise_command_error_with_code(payload: dict[str, object], code: str):
+    with pytest.raises(DuplexCommandError) as excinfo:
+        command_from_realtime({"event_id": "evt-bad", **payload})
+
+    assert excinfo.value.code == code
+    assert excinfo.value.event_id == "evt-bad"
+
+
+# ---- payload() rendering ----
+
+
+def test_commit_maps_create_response_to_response_create():
+    command = command_from_realtime(
+        {"type": "input_audio_buffer.commit", "event_id": "evt-commit", "response_create": True, "final": False}
+    )
+
+    assert command == Commit(event_id="evt-commit", final=False, create_response=True)
+    assert command.payload() == {
+        "type": "input_audio_buffer.commit",
+        "realtime_event_id": "evt-commit",
+        "final": False,
+        "response_create": True,
+    }
+    assert Commit().payload() == {"type": "input_audio_buffer.commit", "final": True}
+    assert Commit(is_speech=False, realtime_item_id="item_1").payload() == {
+        "type": "input_audio_buffer.commit",
+        "final": True,
+        "is_speech": False,
+        "realtime_item_id": "item_1",
+    }
+    assert (
+        command_from_realtime({"type": "input_audio_buffer.commit", "create_response": False}).create_response is False
+    )
+
+
+def test_create_response_renders_options_as_response_object():
+    command = command_from_realtime({"type": "response.create", "response": {"modalities": ["text"]}})
+
+    assert isinstance(command, CreateResponse)
+    assert dict(command.options) == {"modalities": ["text"]}
+    assert command.payload() == {"type": "response.create", "response": {"modalities": ["text"]}}
+    assert command_from_realtime({"type": "response.create"}).payload() == {"type": "response.create", "response": {}}
+
+
+def test_cancel_and_clear_commands_keep_optional_response_id():
+    assert command_from_realtime({"type": "response.cancel", "response_id": ""}).response_id is None
+    cancel = command_from_realtime({"type": "response.cancel", "response_id": "resp_1"})
+    assert cancel == CancelResponse(response_id="resp_1")
+    assert cancel.payload() == {"type": "response.cancel", "response_id": "resp_1"}
+    clear = command_from_realtime({"type": "output_audio_buffer.clear", "response_id": "resp_2"})
+    assert clear == ClearOutputAudio(response_id="resp_2")
+    assert ClearInput().payload() == {"type": "input_audio_buffer.clear"}
+
+
+def test_update_session_renders_as_turn_signal_payload():
+    command = command_from_realtime(
+        {"type": "session.update", "session": {"instructions": "be brief", "voice": "alloy"}}
+    )
+
+    assert command == UpdateSession(patch={"instructions": "be brief", "voice": "alloy"})
+    assert command.payload() == {
+        "type": "turn.signal",
+        "event": "session.update",
+        "payload": {"instructions": "be brief", "voice": "alloy"},
+    }
+
+
+def test_create_item_normalizes_item_and_renders_as_turn_signal_payload():
+    command = command_from_realtime(
+        {
+            "type": "conversation.item.create",
+            "previous_item_id": "item_0",
+            "item": {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]},
+        }
+    )
+
+    assert isinstance(command, CreateItem)
+    assert command.previous_item_id == "item_0"
+    assert command.item["id"].startswith("item_")
+    assert command.item["object"] == "realtime.item"
+    assert command.item["status"] == "completed"
+    payload = command.payload()
+    assert payload["type"] == "turn.signal"
+    assert payload["event"] == "conversation.item.create"
+    assert payload["payload"] == {"item": dict(command.item), "previous_item_id": "item_0"}
+    assert "item" not in payload
+    assert "previous_item_id" not in payload
+
+    without_previous = command_from_realtime({"type": "conversation.item.create", "item": {"content": "x"}})
+    assert without_previous.item["role"] == "user"
+    assert without_previous.item["content"] == []
+    assert without_previous.payload()["payload"] == {"item": dict(without_previous.item)}
+
+
+def test_delete_and_truncate_render_as_turn_signal_payloads():
+    delete = command_from_realtime({"type": "conversation.item.delete", "item_id": "item_1"})
+    assert delete == DeleteItem(item_id="item_1")
+    assert delete.payload() == {
+        "type": "turn.signal",
+        "event": "conversation.item.delete",
+        "payload": {"item_id": "item_1"},
+    }
+
+    truncate = command_from_realtime(
+        {"type": "conversation.item.truncate", "item_id": "item_1", "audio_end_ms": 1500.0, "content_index": 1}
+    )
+    assert truncate == TruncateItem(item_id="item_1", audio_end_ms=1500, content_index=1)
+    assert truncate.payload() == {
+        "type": "turn.signal",
+        "event": "conversation.item.truncate",
+        "payload": {"item_id": "item_1", "audio_end_ms": 1500, "content_index": 1},
+    }
+
+
+def test_playback_ack_close_and_text_append_payloads():
+    ack = command_from_realtime(
+        {"type": "playback.ack", "played_ms": 1200.0, "committed_ms": 1000, "response_id": "resp_1", "item_id": ""}
+    )
+    assert ack == AckPlayback(played_ms=1200, committed_ms=1000, response_id="resp_1")
+    assert ack.payload() == {
+        "type": "playback.ack",
+        "played_ms": 1200,
+        "committed_ms": 1000,
+        "response_id": "resp_1",
+    }
+
+    assert command_from_realtime({"type": "session.close"}) == CloseSession(reason="client_close")
+    assert command_from_realtime({"type": "session.close", "reason": "done"}).payload() == {
+        "type": "session.close",
+        "reason": "done",
+    }
+    assert command_from_realtime({"type": "input.text.append", "text": "hello"}).payload() == {
+        "type": "input.text.append",
+        "text": "hello",
+    }
+    assert Heartbeat().payload() == {"type": "session.heartbeat"}
+
+
+def test_turn_signal_renders_payload_only_when_present_and_dispatches_known_events():
+    signal = command_from_realtime({"type": "turn.signal", "event": "user_started", "payload": {"source": "client"}})
+    assert signal == SignalTurn(event="user_started", signal_payload={"source": "client"})
+    assert signal.payload() == {"type": "turn.signal", "event": "user_started", "payload": {"source": "client"}}
+    assert command_from_realtime({"type": "turn.signal", "event": "user_started"}).payload() == {
+        "type": "turn.signal",
+        "event": "user_started",
+    }
+
+    assert command_from_realtime({"type": "turn.signal", "event": "barge_in"}) == BargeIn()
+    assert command_from_realtime({"type": "turn.signal", "event": "input.cancel"}) == CancelInput()
+    assert command_from_realtime(
+        {"type": "turn.signal", "event": "response.cancel", "payload": {"response_id": "resp_1"}}
+    ) == CancelResponse(response_id="resp_1")
+    assert command_from_realtime(
+        {"type": "turn.signal", "event": "session.update", "payload": {"voice": "alloy"}}
+    ) == UpdateSession(patch={"voice": "alloy"})
+    assert command_from_realtime(
+        {"type": "turn.signal", "event": "conversation.item.delete", "payload": {"item_id": "item_1"}}
+    ) == DeleteItem(item_id="item_1")
+    assert command_from_realtime(
+        {"type": "turn.signal", "event": "conversation.item.truncate", "payload": {"item_id": "i", "audio_end_ms": 5}}
+    ) == TruncateItem(item_id="i", audio_end_ms=5)
+    created = command_from_realtime(
+        {"type": "turn.signal", "event": "conversation.item.create", "payload": {"item": {"role": "system"}}}
+    )
+    assert isinstance(created, CreateItem)
+    assert created.item["role"] == "system"
+
+
+def test_raw_wire_hints_cannot_override_the_normalized_typed_fields():
+    """``build_append_audio`` normalizes; a client's raw value must not undo it.
+
+    A wire ``"is_speech": 0`` used to land in ``hints`` and overwrite the
+    computed ``bool | None`` in the rendered payload, so the runner's
+    silent-commit fast path (``event.get("is_speech") is False``) missed and
+    overlap classification followed the unvalidated value.
+    """
+    command = AppendAudio(
+        audio=b"\x00\x00" * 8,
+        is_speech=False,
+        hints={"is_speech": 0, "rms": 0.001},
+    )
+
+    payload = command.payload()
+
+    assert payload["is_speech"] is False, "the normalized field wins over the raw hint"
+    assert payload["rms"] == 0.001, "a hint with no typed counterpart still comes through"
+
+
+def test_a_hint_survives_when_its_typed_field_is_unset():
+    """Unset typed fields are absent from the payload, so the hint is the only value."""
+    command = AppendAudio(audio=b"\x00\x00" * 8, hints={"is_speech": True})
+
+    assert command.payload()["is_speech"] is True

tests/engine/duplex/test_duplex_control_client.py

@@ -1,110 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.engine.duplex.control_client import (
-    DuplexControlClient,
-    DuplexControlRequestError,
-)
-from vllm_omni.engine.duplex.lease import DuplexLeaseActivity
-from vllm_omni.engine.duplex.messages import (
-    DuplexControlError,
-    DuplexControlResultMessage,
-    DuplexFence,
-    ResumeDuplexSessionMessage,
-    TouchDuplexSessionMessage,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _Transport:
-    def __init__(self) -> None:
-        self.calls = []
-
-    def execute(self, key, message, **kwargs):
-        self.calls.append((key, message, kwargs))
-        return DuplexControlResultMessage(
-            control_id=message.control_id,
-            fence=message.fence,
-            operation="touch" if isinstance(message, TouchDuplexSessionMessage) else "resume",
-            session_id=message.session_id,
-            ok=True,
-            stage_results=[],
-        )
-
-
-class _UnsupportedTransport:
-    def execute(self, key, message, **kwargs):
-        del key, kwargs
-        return DuplexControlResultMessage(
-            control_id=message.control_id,
-            fence=message.fence,
-            operation="touch",
-            session_id=message.session_id,
-            ok=False,
-            stage_results=[{"result": {"supported": False}}],
-            unsupported_count=1,
-            error=DuplexControlError(code="invalid_capability", message="unsupported", retryable=False),
-            accepted_fence=DuplexFence(message.session_id, epoch=2),
-            lease_generation=7,
-        )
-
-
-def test_control_client_routes_touch_and_resume_by_control_id() -> None:
-    transport = _Transport()
-    control_ids = iter(("touch-id", "resume-id"))
-    client = DuplexControlClient(transport, control_id_factory=lambda: next(control_ids))
-    fence = DuplexFence("sid-client")
-
-    assert (
-        client.touch(
-            fence.session_id,
-            fence=fence,
-            activity=DuplexLeaseActivity.PLAYBACK_ACK,
-            timeout=2.0,
-        )["ok"]
-        is True
-    )
-    assert (
-        client.resume(
-            fence.session_id,
-            fence=fence,
-            expected_lease_generation=7,
-            timeout=3.0,
-        )["ok"]
-        is True
-    )
-
-    touch_key, touch_message, touch_kwargs = transport.calls[0]
-    assert touch_key == ("duplex", "touch-id")
-    assert isinstance(touch_message, TouchDuplexSessionMessage)
-    assert touch_message.activity == DuplexLeaseActivity.PLAYBACK_ACK.value
-    assert touch_kwargs["timeout"] == 2.0
-
-    resume_key, resume_message, resume_kwargs = transport.calls[1]
-    assert resume_key == ("duplex", "resume-id")
-    assert isinstance(resume_message, ResumeDuplexSessionMessage)
-    assert resume_message.expected_lease_generation == 7
-    assert resume_kwargs["timeout"] == 3.0
-
-
-def test_control_client_treats_ok_false_as_authoritative_failure() -> None:
-    client = DuplexControlClient(_UnsupportedTransport(), control_id_factory=lambda: "unsupported")
-    fence = DuplexFence("sid-unsupported")
-
-    with pytest.raises(DuplexControlRequestError, match="duplex touch failed") as exc_info:
-        client.touch(
-            fence.session_id,
-            fence=fence,
-            activity=DuplexLeaseActivity.HEARTBEAT,
-            timeout=1.0,
-        )
-
-    assert exc_info.value.code == "invalid_capability"
-    assert exc_info.value.retryable is False
-    assert exc_info.value.accepted_fence == DuplexFence(fence.session_id, epoch=2)
-    assert exc_info.value.lease_generation == 7

tests/engine/duplex/test_duplex_events.py

@@ -0,0 +1,386 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Typed duplex events: pure ``to_realtime()`` rendering and the stateful Realtime projection."""
+
+from __future__ import annotations
+
+import base64
+import inspect
+
+import pytest
+
+from vllm_omni.engine.duplex import events as events_module
+from vllm_omni.engine.duplex.events import (
+    REALTIME_ERROR_TYPES_BY_CODE,
+    AudioDelta,
+    DuplexEvent,
+    DuplexRawEvent,
+    ErrorEvent,
+    InputCommitted,
+    ItemAdded,
+    ItemCreated,
+    ItemDone,
+    Listen,
+    OutputItemAdded,
+    OverlapDecision,
+    ResponseDone,
+    SessionClosed,
+    SessionCreated,
+    SessionExpired,
+    SessionHeartbeatAck,
+    SessionReplaced,
+    SessionResumed,
+    SessionResyncRequired,
+    TextDelta,
+    TranscriptDone,
+    TurnEvent,
+    error_event,
+)
+from vllm_omni.engine.duplex.realtime_events import (
+    RealtimeProjectionState,
+    project_internal_event,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+#: Classes whose wire object carries ``session_id`` by contract.
+_SESSION_ID_ON_WIRE = {
+    SessionClosed,
+    SessionExpired,
+    SessionHeartbeatAck,
+    Listen,
+    TurnEvent,
+    OverlapDecision,
+    SessionResumed,
+    SessionReplaced,
+    SessionResyncRequired,
+}
+#: Classes whose wire object carries ``epoch`` by contract.
+_EPOCH_ON_WIRE = {Listen, TurnEvent, OverlapDecision}
+
+
+def _public_event_classes() -> list[type[DuplexEvent]]:
+    classes = [
+        getattr(events_module, name)
+        for name in events_module.__all__
+        if inspect.isclass(getattr(events_module, name)) and issubclass(getattr(events_module, name), DuplexEvent)
+    ]
+    assert len(classes) > 40
+    return [cls for cls in classes if cls is not DuplexEvent]
+
+
+@pytest.mark.parametrize("event_cls", _public_event_classes(), ids=lambda cls: cls.__name__)
+def test_to_realtime_renders_wire_type_event_id_and_documented_fields(event_cls: type[DuplexEvent]):
+    event = event_cls(session_id="duplex-hidden", epoch=7)
+
+    wire = event.to_realtime()
+
+    assert wire["type"] == event.type
+    assert wire["event_id"] == event.event_id
+    assert wire["event_id"].startswith("event_")
+    if event_cls is DuplexRawEvent:
+        assert wire["type"] == "duplex."
+    else:
+        assert wire["type"] == event_cls.wire_type
+    if event_cls in _SESSION_ID_ON_WIRE:
+        assert wire["session_id"] == "duplex-hidden"
+    else:
+        assert "session_id" not in wire
+    if event_cls in _EPOCH_ON_WIRE:
+        assert wire["epoch"] == 7
+    else:
+        assert "epoch" not in wire
+    # Every value is plain JSON (mappings / sequences are copied into dict / list).
+    for value in wire.values():
+        assert isinstance(value, str | int | float | bool | dict | list) or value is None
+
+
+def test_optional_wire_fields_are_omitted_when_none():
+    created = SessionCreated(session={"id": "sid"})
+    assert created.optional_wire_fields == frozenset({"attachment_generation", "resume_token"})
+    wire = created.to_realtime()
+    assert "attachment_generation" not in wire
+    assert "resume_token" not in wire
+    assert wire["session"] == {"id": "sid"}
+
+    with_resume = SessionCreated(session={"id": "sid"}, attachment_generation=2, resume_token="tok").to_realtime()
+    assert with_resume["attachment_generation"] == 2
+    assert with_resume["resume_token"] == "tok"
+
+    delta = AudioDelta(response_id="resp_1", item_id="item_resp_1", delta="AAAA").to_realtime()
+    assert "sample_rate_hz" not in delta
+    assert "metadata" not in delta
+    assert delta["format"] == "pcm16"
+    assert delta["output_index"] == 0
+    assert delta["content_index"] == 0
+
+    anonymous_listen = Listen(session_id="sid", epoch=0, details={"reason": "silence"}).to_realtime()
+    assert "response_id" not in anonymous_listen
+    assert "id" not in anonymous_listen["response"]
+    assert anonymous_listen["response"]["status"] == "listening"
+    assert anonymous_listen["response"]["metadata"] == {"reason": "silence"}
+
+    bound_listen = Listen(session_id="sid", epoch=0, response_id="resp_1").to_realtime()
+    assert bound_listen["response_id"] == "resp_1"
+    assert bound_listen["response"]["id"] == "resp_1"
+
+
+def test_error_event_maps_codes_to_openai_error_types_and_echoes_client_event_id():
+    assert REALTIME_ERROR_TYPES_BY_CODE["resource_exhausted"] == "rate_limit_error"
+    assert REALTIME_ERROR_TYPES_BY_CODE["internal_error"] == "server_error"
+    assert REALTIME_ERROR_TYPES_BY_CODE["bad_event"] == "invalid_request_error"
+
+    error = ErrorEvent(
+        code="resource_exhausted",
+        message="too many sessions",
+        related_event_id="evt_client",
+        param="session",
+        extra={"retryable": True},
+    )
+
+    assert error.error_type == "rate_limit_error"
+    assert error.text == "too many sessions"
+    assert error.error == {
+        "type": "rate_limit_error",
+        "code": "resource_exhausted",
+        "message": "too many sessions",
+        "event_id": "evt_client",
+        "param": "session",
+        "retryable": True,
+    }
+    wire = error.to_realtime()
+    assert wire["type"] == "error"
+    assert wire["error"] == error.error
+    # The server event id and the echoed client event id are distinct.
+    assert wire["event_id"] != "evt_client"
+    assert ErrorEvent(code="never-seen-code").error_type == "invalid_request_error"
+
+
+def test_error_event_helper_normalizes_optional_client_event_id_and_param():
+    event = error_event("bad_event", "nope", event_id="evt_1", param="item", extra={"retryable": False})
+    assert isinstance(event, ErrorEvent)
+    assert event.related_event_id == "evt_1"
+    assert event.param == "item"
+    assert event.error["retryable"] is False
+
+    anonymous = error_event("bad_event", "nope", event_id=42, param="")
+    assert anonymous.related_event_id is None
+    assert anonymous.param is None
+    assert "event_id" not in anonymous.error
+    assert "param" not in anonymous.error
+
+
+def test_audio_delta_decodes_base64_and_invalid_payloads_yield_none():
+    raw = bytes(range(16))
+    delta = AudioDelta(delta=base64.b64encode(raw).decode("ascii"), format="pcm16", sample_rate_hz=24000)
+
+    assert delta.audio == raw
+    assert delta.to_realtime()["sample_rate_hz"] == 24000
+    assert AudioDelta(delta="@@not base64@@").audio is None
+
+
+def test_session_closed_and_expired_are_the_only_terminal_events():
+    closed = SessionClosed(session_id="sid", reason="client_close", details={"type": "session.closed"})
+    expired = SessionExpired(session_id="sid", reason="idle_timeout")
+
+    assert closed.is_terminal is True
+    assert expired.is_terminal is True
+    assert closed.to_realtime() == {
+        "type": "session.closed",
+        "event_id": closed.event_id,
+        "session_id": "sid",
+        "reason": "client_close",
+        "event": {"type": "session.closed"},
+    }
+    assert expired.to_realtime() == {
+        "type": "session.expired",
+        "event_id": expired.event_id,
+        "session_id": "sid",
+        "reason": "idle_timeout",
+    }
+    non_terminal = [cls for cls in _public_event_classes() if cls().is_terminal]
+    assert non_terminal == [SessionClosed, SessionExpired]
+
+
+def test_generic_accessors_expose_response_item_text_and_status():
+    assert TextDelta(response_id="resp_1", item_id="item_resp_1", delta="hi").text == "hi"
+    assert TranscriptDone(transcript="done").text == "done"
+    assert OutputItemAdded(item={"id": "item_1"}).item_id == "item_1"
+    assert ItemAdded(item={"id": 3}).item_id is None
+    assert ResponseDone(response={"status": "cancelled"}).status == "cancelled"
+    assert ResponseDone().status is None
+    assert TurnEvent(event="user_started", turn_state="user_speaking").response_id is None
+    assert TurnEvent().audio is None
+
+
+def test_raw_event_prefixes_internal_type_and_wraps_details():
+    raw = DuplexRawEvent(internal_type="model.custom", details={"type": "model.custom", "value": (1, 2)})
+
+    assert raw.type == "duplex.model.custom"
+    assert raw.to_realtime() == {
+        "type": "duplex.model.custom",
+        "event_id": raw.event_id,
+        "event": {"type": "model.custom", "value": [1, 2]},
+    }
+
+
+def test_turn_event_and_overlap_decision_carry_session_identity_on_the_wire():
+    turn = TurnEvent(session_id="sid", epoch=2, event="barge_in", turn_state="barge_in").to_realtime()
+    assert turn == {
+        "type": "turn.event",
+        "event_id": turn["event_id"],
+        "session_id": "sid",
+        "event": "barge_in",
+        "turn_state": "barge_in",
+        "epoch": 2,
+    }
+    decision = OverlapDecision(
+        session_id="sid", epoch=2, policy="listen_only", action="ignore", reason="short", details={"ms": 300}
+    ).to_realtime()
+    assert decision["policy"] == "listen_only"
+    assert decision["action"] == "ignore"
+    assert decision["metadata"] == {"ms": 300}
+    assert decision["session_id"] == "sid"
+
+
+# ---- RealtimeProjectionState / project_internal_event ----
+
+
+def _types(events: list[DuplexEvent]) -> list[str]:
+    return [event.type for event in events]
+
+
+def test_projection_of_the_main_internal_event_sequence():
+    state = RealtimeProjectionState(session_id="duplex-proj", model="test-model", initial_session_update=True)
+    pcm = base64.b64encode(b"\x00\x10" * 8).decode("ascii")
+
+    created = project_internal_event(state, {"type": "session.created", "session": {"voice": "alloy"}})
+    assert _types(created) == ["session.created", "session.updated"]
+    assert isinstance(created[0], SessionCreated)
+    session_payload = created[0].session
+    assert session_payload["id"] == "duplex-proj"
+    assert session_payload["model"] == "test-model"
+    assert session_payload["voice"] == "alloy"
+    assert session_payload["input_audio_format"] == "pcm16"
+    assert session_payload["output_audio_format"] == "pcm16"
+    assert state.initial_session_update is False
+    # The extra session.updated is a one-shot for the opening session.update.
+    assert _types(project_internal_event(state, {"type": "session.created", "session": {}})) == ["session.created"]
+
+    response_created = project_internal_event(
+        state, {"type": "response.created", "response_id": "resp_1", "modalities": ["audio", "text"]}
+    )
+    assert _types(response_created) == [
+        "response.created",
+        "conversation.item.added",
+        "conversation.item.created",
+        "response.output_item.added",
+        "response.content_part.added",
+    ]
+    assert state.active_response_id == "resp_1"
+    assert response_created[0].response_id == "resp_1"
+    assert response_created[0].to_realtime()["response"]["status"] == "in_progress"
+    assert isinstance(response_created[1], ItemAdded)
+    assert isinstance(response_created[2], ItemCreated)
+    assert response_created[1].item_id == "item_resp_1"
+    assert response_created[3].item_id == "item_resp_1"
+    assert response_created[4].to_realtime()["part"] == {"type": "audio", "transcript": ""}
+
+    audio_delta = project_internal_event(
+        state,
+        {"type": "response.output_audio.delta", "response_id": "resp_1", "audio": pcm, "text": "hi", "format": "pcm16"},
+    )
+    assert _types(audio_delta) == ["response.output_audio.delta", "response.output_audio_transcript.delta"]
+    assert isinstance(audio_delta[0], AudioDelta)
+    assert audio_delta[0].delta == pcm
+    assert audio_delta[0].format == "pcm16"
+    assert audio_delta[0].item_id == "item_resp_1"
+    assert audio_delta[1].text == "hi"
+    assert state.conversation_items["item_resp_1"]["content"][0]["transcript"] == "hi"
+
+    done = project_internal_event(state, {"type": "response.done", "response_id": "resp_1"})
+    assert _types(done) == [
+        "response.output_audio.done",
+        "response.output_audio_transcript.done",
+        "response.content_part.done",
+        "response.output_item.done",
+        "conversation.item.done",
+        "response.done",
+        "rate_limits.updated",
+    ]
+    assert done[1].text == "hi"
+    assert isinstance(done[4], ItemDone)
+    assert done[4].item["status"] == "completed"
+    assert isinstance(done[5], ResponseDone)
+    assert done[5].status == "completed"
+    assert done[5].response["output"][0]["content"][0] == {"type": "output_audio", "transcript": "hi"}
+    assert state.active_response_id is None
+    # A duplicate terminal for the same response is idempotent.
+    assert project_internal_event(state, {"type": "response.done", "response_id": "resp_1"}) == []
+
+    committed = project_internal_event(
+        state, {"type": "input.committed", "message": {"role": "user", "content": "hello"}, "input_commit_seq": 1}
+    )
+    assert _types(committed) == [
+        "conversation.item.added",
+        "conversation.item.created",
+        "input_audio_buffer.committed",
+        "conversation.item.done",
+    ]
+    assert isinstance(committed[2], InputCommitted)
+    assert committed[2].previous_item_id == "item_resp_1"
+    assert committed[2].item_id == committed[0].item_id
+    assert committed[2].to_realtime()["event"]["input_commit_seq"] == 1
+    assert committed[0].item["content"] == [{"type": "input_text", "text": "hello"}]
+
+    project_internal_event(state, {"type": "response.created", "response_id": "resp_2"})
+    project_internal_event(state, {"type": "response.output_audio.delta", "response_id": "resp_2", "audio": pcm})
+    cancelled = project_internal_event(
+        state, {"type": "audio.cancelled", "response_id": "resp_2", "reason": "barge_in", "committed_ms": 0}
+    )
+    assert _types(cancelled) == [
+        "response.output_audio.done",
+        "response.content_part.done",
+        "response.output_item.done",
+        "conversation.item.done",
+        "response.done",
+        "rate_limits.updated",
+    ]
+    assert cancelled[4].status == "cancelled"
+    assert cancelled[4].response["status_details"] == {"type": "cancelled", "reason": "barge_in"}
+    assert state.active_response_id is None
+
+    closed = project_internal_event(state, {"type": "session.closed", "reason": "client_close"})
+    assert _types(closed) == ["session.closed"]
+    assert closed[0].is_terminal is True
+    assert closed[0].to_realtime()["reason"] == "client_close"
+
+
+def test_projection_of_output_audio_buffer_clear_emits_cleared_before_terminals():
+    state = RealtimeProjectionState(session_id="duplex-clear")
+    pcm = base64.b64encode(b"\x00\x10" * 8).decode("ascii")
+    project_internal_event(state, {"type": "response.created", "response_id": "resp_1"})
+    project_internal_event(state, {"type": "response.output_audio.delta", "response_id": "resp_1", "audio": pcm})
+
+    cleared = project_internal_event(
+        state, {"type": "audio.cancelled", "reason": "output_audio_buffer_clear", "committed_ms": 250}
+    )
+
+    assert _types(cleared)[:2] == ["output_audio_buffer.cleared", "response.output_audio.done"]
+    assert cleared[0].response_id == "resp_1"
+    assert _types(cleared)[-2:] == ["response.done", "rate_limits.updated"]
+    assert state.item_truncation_cursors["item_resp_1"] == (0, 250)
+
+
+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.
+    state = RealtimeProjectionState(session_id="duplex-err")
+
+    projected = project_internal_event(state, {"type": "error", "code": "bad_event", "error": "nope"})
+
+    assert len(projected) == 1
+    assert isinstance(projected[0], DuplexRawEvent)
+    assert projected[0].type == "duplex.error"
+    assert projected[0].to_realtime()["event"]["code"] == "bad_event"

tests/engine/duplex/test_duplex_lease.py

@@ -1,19 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+"""Lease semantics of the engine-resident duplex session (idle TTL, disconnect grace, resume CAS)."""
+
 from __future__ import annotations
 
 from dataclasses import FrozenInstanceError
 
 import pytest
 
-from vllm_omni.engine.duplex.lease import (
+from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.engine.duplex.contracts import DuplexFence
+from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError
+from vllm_omni.engine.duplex.session.lease import (
     DuplexLeaseActivity,
     DuplexLeaseConfig,
+    DuplexLeaseState,
 )
-from vllm_omni.engine.duplex.messages import DuplexFence
-from vllm_omni.engine.duplex.runtime import DuplexInputMode
-from vllm_omni.engine.duplex.session import DuplexSessionRuntimeManager
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -30,195 +33,148 @@ class FakeMonotonicClock:
 
 
 def _lease_config(*, idle_ttl_s: float | None = 300.0) -> DuplexLeaseConfig:
-    return DuplexLeaseConfig(
-        idle_ttl_s=idle_ttl_s,
-        disconnect_grace_s=30.0,
+    return DuplexLeaseConfig(idle_ttl_s=idle_ttl_s, disconnect_grace_s=30.0)
+
+
+def _session(session_id: str, clock: FakeMonotonicClock, *, idle_ttl_s: float | None = 300.0) -> DuplexEngineSession:
+    return DuplexEngineSession(
+        session_id=session_id,
+        config=DuplexSessionConfig(),
+        lease=DuplexLeaseState(config=_lease_config(idle_ttl_s=idle_ttl_s), generation=0, last_activity=clock()),
+        _clock=clock,
     )
 
 
-def test_open_touch_detach_and_idle_expiry_use_monotonic_time() -> None:
+def test_open_touch_detach_and_expiry_use_monotonic_time() -> None:
     clock = FakeMonotonicClock(100.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    fence = DuplexFence("sid-expiry")
-    session = manager.open_session(fence, lease_config=_lease_config())
+    session = _session("sid-expiry", clock)
 
     assert session.lease.last_activity == 100.0
     clock.advance(10.0)
-    session.touch(fence, DuplexLeaseActivity.HEARTBEAT)
-    session.detach(fence)
+    session.touch_lease(DuplexLeaseActivity.HEARTBEAT)
+    assert session.lease.last_activity == 110.0
+    session.detach_lease()
     clock.advance(29.0)
     assert session.lease.disconnect_grace_expired(clock()) is False
-    assert manager.collect_expired() == []
+    assert session.lease.idle_expired(clock()) is False
 
-    clock.advance(272.0)
-    expired = manager.collect_expired()
+    clock.advance(1.0)
+    assert session.lease.disconnect_grace_expired(clock()) is True
 
-    assert [item.session_id for item in expired] == ["sid-expiry"]
-    assert expired[0].reason == "disconnect_grace_expired"
-    assert manager.get("sid-expiry") is session
-    manager.finalize_close_session(session)
-    assert manager.get("sid-expiry") is None
+    clock.advance(271.0)
+    assert session.lease.idle_expired(clock()) is True
 
 
-def test_detach_grace_and_resume_advance_token_independent_lease_generation() -> None:
+def test_detach_grace_and_resume_advance_the_lease_generation() -> None:
     clock = FakeMonotonicClock(10.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    fence = DuplexFence("sid-resume")
-    session = manager.open_session(fence, lease_config=_lease_config())
+    session = _session("sid-resume", clock)
 
-    session.detach(fence)
+    session.detach_lease()
     clock.advance(29.0)
     assert session.lease.disconnect_grace_expired(clock()) is False
     clock.advance(2.0)
     assert session.lease.disconnect_grace_expired(clock()) is True
 
-    generation = session.resume(fence, expected_lease_generation=0)
+    generation = session.resume_lease(expected_lease_generation=0)
 
     assert generation == 1
+    assert session.lease_generation == 1
     assert session.lease.detached_at is None
     assert session.lease.disconnect_grace_expired(clock()) is False
     with pytest.raises(ValueError, match="lease generation mismatch"):
-        session.resume(fence, expected_lease_generation=0)
+        session.resume_lease(expected_lease_generation=0)
 
 
 def test_active_operation_prevents_mid_transaction_expiry() -> None:
     clock = FakeMonotonicClock(0.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    fence = DuplexFence("sid-operation")
-    session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=5.0))
+    session = _session("sid-operation", clock, idle_ttl_s=5.0)
 
-    session.begin_operation(fence, "append-1")
+    session.begin_lease_operation(session.fence, "append-1")
     clock.advance(10.0)
-    assert manager.collect_expired() == []
+    assert session.lease.idle_expired(clock()) is False
 
-    session.end_operation(fence, "append-1")
-    assert manager.collect_expired() == []
+    session.end_lease_operation("append-1")
+    assert session.lease.idle_expired(clock()) is False
     clock.advance(6.0)
-    assert [item.session_id for item in manager.collect_expired()] == ["sid-operation"]
+    assert session.lease.idle_expired(clock()) is True
 
 
 def test_sessions_have_independent_activity_deadlines_and_resources() -> None:
     clock = FakeMonotonicClock(0.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    fence_a = DuplexFence("sid-a")
-    fence_b = DuplexFence("sid-b")
-    session_a = manager.open_session(fence_a, lease_config=_lease_config(idle_ttl_s=10.0))
-    session_b = manager.open_session(fence_b, lease_config=_lease_config(idle_ttl_s=10.0))
-    session_a.reserve_stage_request(0, "req-a-reserved", fence=fence_a)
-    session_a.bind_stage_request(1, "req-a-submitted", fence=fence_a)
-    session_b.bind_stage_request(0, "req-b", fence=fence_b)
+    session_a = _session("sid-a", clock, idle_ttl_s=10.0)
+    session_b = _session("sid-b", clock, idle_ttl_s=10.0)
+    session_a.reserve_stage_request(0, "req-a-reserved", fence=session_a.fence)
+    session_a.bind_stage_request(1, "req-a-submitted", fence=session_a.fence)
+    session_b.bind_stage_request(0, "req-b", fence=session_b.fence)
 
     clock.advance(6.0)
-    session_b.touch(fence_b, DuplexLeaseActivity.MODEL_OUTPUT)
+    session_b.touch_lease(DuplexLeaseActivity.MODEL_OUTPUT)
     clock.advance(5.0)
-    expired = manager.collect_expired()
 
-    assert len(expired) == 1
-    assert expired[0].session_id == "sid-a"
-    assert expired[0].reserved_request_ids == ("req-a-reserved",)
-    assert expired[0].submitted_request_ids == ("req-a-submitted",)
-    assert manager.get("sid-b") is session_b
+    assert session_a.lease.idle_expired(clock()) is True
+    assert session_b.lease.idle_expired(clock()) is False
+    assert session_a.resource_request_ids(submitted=False) == ["req-a-reserved"]
+    assert session_a.resource_request_ids(submitted=True) == ["req-a-submitted"]
     assert session_b.resource_request_ids() == ["req-b"]
 
 
-def test_close_and_reaper_select_exactly_one_terminal_transition() -> None:
+def test_close_is_a_single_terminal_transition() -> None:
     clock = FakeMonotonicClock(0.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    fence = DuplexFence("sid-race")
-    manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=1.0))
+    session = _session("sid-race", clock, idle_ttl_s=1.0)
     clock.advance(2.0)
+    assert session.lease.idle_expired(clock()) is True
 
-    assert len(manager.collect_expired()) == 1
-    assert manager.close_session(fence, reason="explicit_close") is None
-    assert manager.collect_expired() == []
-
-    fence_2 = DuplexFence("sid-race-2")
-    manager.open_session(fence_2, lease_config=_lease_config(idle_ttl_s=1.0))
-    assert manager.close_session(fence_2, reason="explicit_close") is not None
-    clock.advance(2.0)
-    assert manager.collect_expired() == []
+    assert session.begin_close(reason="explicit_close") is True
+    assert session.lease.terminal_reason == "explicit_close"
+    # A terminal lease is never reaped again and a second close is a no-op.
+    assert session.lease.idle_expired(clock()) is False
+    assert session.lease.disconnect_grace_expired(clock()) is False
+    assert session.begin_close(reason="idle_ttl_expired") is True
+    assert session.lease.terminal_reason == "explicit_close"
+    with pytest.raises(RuntimeError, match="terminal"):
+        session.touch_lease(DuplexLeaseActivity.HEARTBEAT)
 
 
-def test_stale_fence_cannot_touch_detach_or_resume_lease() -> None:
+def test_stale_fence_cannot_start_a_lease_operation() -> None:
     clock = FakeMonotonicClock(0.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    current = DuplexFence("sid-fence", epoch=1)
+    session = _session("sid-fence", clock)
+    session.epoch = 1
+    session.sync_fence()
     stale = DuplexFence("sid-fence", epoch=0)
-    session = manager.open_session(current, lease_config=_lease_config())
 
-    with pytest.raises(RuntimeError, match="fence mismatch"):
-        session.touch(stale, DuplexLeaseActivity.APPEND)
-    with pytest.raises(RuntimeError, match="fence mismatch"):
-        session.detach(stale)
-    with pytest.raises(RuntimeError, match="fence mismatch"):
-        session.resume(stale, expected_lease_generation=0)
+    with pytest.raises(DuplexFenceMismatchError, match="fence mismatch"):
+        session.begin_lease_operation(stale, "append-1")
+    with pytest.raises(DuplexFenceMismatchError, match="fence mismatch"):
+        session.begin_lease_operation(DuplexFence("other-session", epoch=1), "append-2")
 
 
-def test_disabled_idle_expiry_never_collects_session() -> None:
+def test_disabled_idle_expiry_never_expires_session() -> None:
     clock = FakeMonotonicClock(0.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    manager.open_session(DuplexFence("sid-no-expiry"), lease_config=_lease_config(idle_ttl_s=None))
+    session = _session("sid-no-expiry", clock, idle_ttl_s=None)
 
     clock.advance(1_000_000.0)
 
-    assert manager.collect_expired() == []
+    assert session.lease.expires_at is None
+    assert session.lease.idle_expired(clock()) is False
 
 
 def test_detached_session_expires_at_disconnect_grace_when_idle_ttl_is_disabled() -> None:
     clock = FakeMonotonicClock(0.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    fence = DuplexFence("sid-disconnect-grace")
-    session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=None))
+    session = _session("sid-disconnect-grace", clock, idle_ttl_s=None)
 
-    session.detach(fence)
+    session.detach_lease()
     clock.advance(29.0)
-    assert manager.collect_expired() == []
+    assert session.lease.disconnect_grace_expired(clock()) is False
 
     clock.advance(1.0)
-    expired = manager.collect_expired()
-
-    assert [item.session_id for item in expired] == [fence.session_id]
-    assert expired[0].reason == "disconnect_grace_expired"
+    assert session.lease.disconnect_grace_expired(clock()) is True
 
 
-def test_runtime_manager_enforces_server_owned_session_admission() -> None:
-    manager = DuplexSessionRuntimeManager(max_sessions=2)
-    manager.open_session(DuplexFence("sid-admission-a"))
-    manager.open_session(DuplexFence("sid-admission-b"))
-
-    with pytest.raises(RuntimeError, match="duplex_session_capacity_exhausted"):
-        manager.open_session(DuplexFence("sid-admission-c"))
-
-
-def test_completed_append_cache_is_bounded() -> None:
-    manager = DuplexSessionRuntimeManager(completed_append_limit=2)
-    fence = DuplexFence("sid-completed-cache")
-    session = manager.open_session(fence)
-
-    for index in range(3):
-        session.record_completed_append(
-            f"operation-{index}",
-            fence=fence,
-            mode=DuplexInputMode.TURN_COMMIT_ONLY,
-            final=True,
-            stage_results=[{"index": index}],
-        )
-
-    assert list(session.completed_appends) == ["operation-1", "operation-2"]
-
-    session.accept_fence(DuplexFence(fence.session_id, epoch=1))
-    assert session.completed_appends == {}
-
-
-def test_lease_config_and_expiry_record_are_immutable() -> None:
+def test_lease_config_is_immutable_and_validated() -> None:
     config = _lease_config()
     with pytest.raises(FrozenInstanceError):
         config.idle_ttl_s = 1.0  # type: ignore[misc]
-
-    clock = FakeMonotonicClock(0.0)
-    manager = DuplexSessionRuntimeManager(clock=clock)
-    manager.open_session(DuplexFence("sid-record"), lease_config=_lease_config(idle_ttl_s=1.0))
-    clock.advance(2.0)
-    record = manager.collect_expired()[0]
-    with pytest.raises(FrozenInstanceError):
-        record.reason = "changed"  # type: ignore[misc]
+    with pytest.raises(ValueError, match="idle_ttl_s"):
+        DuplexLeaseConfig(idle_ttl_s=0.0)
+    with pytest.raises(ValueError, match="disconnect_grace_s"):
+        DuplexLeaseConfig(disconnect_grace_s=0.0)

tests/engine/duplex/test_duplex_plugin.py

@@ -0,0 +1,324 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""``DuplexModelPlugin`` loading/validation helpers and the ``PipelineConfig.duplex_plugin`` binding."""
+
+from __future__ import annotations
+
+from dataclasses import fields
+from pathlib import Path
+from typing import Any
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, PipelineConfig
+from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+from vllm_omni.engine.duplex.contracts import DuplexAppendPlan
+from vllm_omni.engine.duplex.plugin import (
+    DuplexDataPlane,
+    DuplexModelPlugin,
+    DuplexModelSessionState,
+    DuplexRuntimeConfigError,
+    coerce_int,
+    load_duplex_plugin,
+    payload_turn_id,
+    reject_changed_runtime_value,
+    validate_duplex_plugin_sampling,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _encode_audio(audio: object, sample_rate_hz: int, fmt: str, speed: float | None) -> str | None:
+    del audio, sample_rate_hz, fmt, speed
+    return None
+
+
+class FakeDataPlane(DuplexDataPlane):
+    def begin_request(self, request_id: str) -> None:
+        pass
+
+    def is_terminal(self, request_id: str | None) -> bool:
+        return False
+
+    def mark_terminal(self, request_id: str) -> None:
+        pass
+
+    def close_stream(self, request_id: str) -> None:
+        pass
+
+    def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None:
+        pass
+
+    def project(self, result, *, context=None):
+        del result, context
+        return []
+
+
+class FakePlugin(DuplexModelPlugin):
+    """Minimal plugin; referenced by dotted path from ``load_duplex_plugin`` tests."""
+
+    plugin_id = "fake-duplex"
+
+    def __init__(self, encode_audio) -> None:
+        super().__init__(encode_audio)
+        self.data_plane = FakeDataPlane()
+
+    def configure_sampling_params(self, *, runtime_config, defaults):
+        del runtime_config
+        return tuple(defaults)
+
+    def plan_append(self, **kwargs) -> DuplexAppendPlan:
+        del kwargs
+        return DuplexAppendPlan(prompt={"prompt_token_ids": [1]})
+
+    def decide_output(self, **kwargs):
+        del kwargs
+        return None
+
+    def create_session_state(self) -> DuplexModelSessionState:
+        raise NotImplementedError
+
+    def capabilities(self, *, max_sessions: int) -> DuplexCapabilities:
+        del max_sessions
+        return DuplexCapabilities()
+
+    def validate_client_extra_body(self, extra_body: object) -> None:
+        pass
+
+    async def prepare_runtime_config(self, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]:
+        del config, model_config
+        return {}
+
+    def runtime_config_for_update(self, config, current):
+        del config
+        return dict(current)
+
+    def data_plane_context(self, **kwargs) -> object:
+        return dict(kwargs)
+
+
+class PluginWithoutId(FakePlugin):
+    plugin_id = ""
+
+
+class PluginWithoutDataPlane(FakePlugin):
+    def __init__(self, encode_audio) -> None:
+        DuplexModelPlugin.__init__(self, encode_audio)
+
+
+class PluginWithWrongDataPlane(FakePlugin):
+    def __init__(self, encode_audio) -> None:
+        DuplexModelPlugin.__init__(self, encode_audio)
+        self.data_plane = object()  # type: ignore[assignment]
+
+
+class NotAPlugin:
+    plugin_id = "lookalike"
+    data_plane = FakeDataPlane()
+
+    def __init__(self, encode_audio) -> None:
+        self.encode_audio = encode_audio
+
+
+def _path(cls: type) -> str:
+    return f"{__name__}.{cls.__name__}"
+
+
+# --------------------------------------------------------------------------- #
+# load_duplex_plugin                                                          #
+# --------------------------------------------------------------------------- #
+
+
+def test_load_duplex_plugin_instantiates_the_class_at_the_dotted_path() -> None:
+    plugin = load_duplex_plugin(_path(FakePlugin), _encode_audio)
+
+    assert type(plugin) is FakePlugin
+    assert plugin.plugin_id == "fake-duplex"
+    assert isinstance(plugin.data_plane, FakeDataPlane)
+    # Defaults of the optional class-level policy.
+    assert plugin.private_runtime_config_keys == frozenset()
+    assert plugin.silence_continuation_samples == 16000
+    assert plugin.runtime_config_for_function_output(DuplexSessionConfig(), {}, {"name": "tool"}) is None
+
+
+def test_load_duplex_plugin_rejects_a_non_plugin_class() -> None:
+    with pytest.raises(TypeError, match="not a DuplexModelPlugin"):
+        load_duplex_plugin(_path(NotAPlugin), _encode_audio)
+
+
+def test_load_duplex_plugin_requires_a_plugin_id() -> None:
+    with pytest.raises(TypeError, match="plugin_id"):
+        load_duplex_plugin(_path(PluginWithoutId), _encode_audio)
+
+
+@pytest.mark.parametrize("plugin_cls", [PluginWithoutDataPlane, PluginWithWrongDataPlane])
+def test_load_duplex_plugin_requires_a_data_plane(plugin_cls: type) -> None:
+    with pytest.raises(TypeError, match="DuplexDataPlane as data_plane"):
+        load_duplex_plugin(_path(plugin_cls), _encode_audio)
+
+
+def test_load_duplex_plugin_rejects_paths_without_a_module() -> None:
+    with pytest.raises(ValueError, match="Invalid duplex plugin path"):
+        load_duplex_plugin("FakePlugin", _encode_audio)
+
+
+def test_load_duplex_plugin_propagates_missing_attributes() -> None:
+    with pytest.raises(AttributeError):
+        load_duplex_plugin(f"{__name__}.NoSuchPlugin", _encode_audio)
+
+
+# --------------------------------------------------------------------------- #
+# validate_duplex_plugin_sampling                                             #
+# --------------------------------------------------------------------------- #
+
+
+def test_validate_duplex_plugin_sampling_accepts_one_parameter_per_stage() -> None:
+    plugin = FakePlugin(_encode_audio)
+
+    validate_duplex_plugin_sampling(plugin, sampling_defaults=(SamplingParams(), SamplingParams()))
+    validate_duplex_plugin_sampling(plugin, sampling_defaults=())
+
+
+def test_validate_duplex_plugin_sampling_rejects_stage_count_mismatch() -> None:
+    class WrongStageCountPlugin(FakePlugin):
+        def configure_sampling_params(self, *, runtime_config, defaults):
+            del runtime_config
+            return defaults[:1]
+
+    with pytest.raises(ValueError, match="one sampling parameter per stage"):
+        validate_duplex_plugin_sampling(
+            WrongStageCountPlugin(_encode_audio),
+            sampling_defaults=(SamplingParams(), SamplingParams()),
+        )
+
+
+def test_validate_duplex_plugin_sampling_rejects_type_mismatch() -> None:
+    class WrongSamplingTypePlugin(FakePlugin):
+        def configure_sampling_params(self, *, runtime_config, defaults):
+            del runtime_config
+            return tuple(object() for _ in defaults)
+
+    with pytest.raises(TypeError, match="sampling parameter type mismatch for stage 0"):
+        validate_duplex_plugin_sampling(
+            WrongSamplingTypePlugin(_encode_audio),
+            sampling_defaults=(SamplingParams(), SamplingParams()),
+        )
+
+
+def test_validate_duplex_plugin_sampling_skips_type_check_for_stages_without_defaults() -> None:
+    class ObjectSamplingPlugin(FakePlugin):
+        def configure_sampling_params(self, *, runtime_config, defaults):
+            del runtime_config
+            return tuple(object() for _ in defaults)
+
+    validate_duplex_plugin_sampling(ObjectSamplingPlugin(_encode_audio), sampling_defaults=(None, None))
+
+
+def test_validate_duplex_plugin_sampling_rejects_non_tuple_results() -> None:
+    class ListSamplingPlugin(FakePlugin):
+        def configure_sampling_params(self, *, runtime_config, defaults):
+            del runtime_config
+            return list(defaults)
+
+    with pytest.raises(TypeError, match="as a tuple"):
+        validate_duplex_plugin_sampling(ListSamplingPlugin(_encode_audio), sampling_defaults=(SamplingParams(),))
+
+
+# --------------------------------------------------------------------------- #
+# Runtime-config helpers                                                      #
+# --------------------------------------------------------------------------- #
+
+
+def test_reject_changed_runtime_value_only_raises_on_change() -> None:
+    reject_changed_runtime_value("voice-a", "voice-a", message="voice is fixed", code="voice_changed")
+
+    with pytest.raises(DuplexRuntimeConfigError, match="voice is fixed") as excinfo:
+        reject_changed_runtime_value("voice-b", "voice-a", message="voice is fixed", code="voice_changed")
+
+    assert excinfo.value.code == "voice_changed"
+    assert isinstance(excinfo.value, ValueError)
+
+
+def test_reject_changed_runtime_value_uses_the_requested_error_class() -> None:
+    class CustomRuntimeConfigError(DuplexRuntimeConfigError):
+        pass
+
+    with pytest.raises(CustomRuntimeConfigError) as excinfo:
+        reject_changed_runtime_value(1, 2, message="changed", code="custom", error_cls=CustomRuntimeConfigError)
+
+    assert excinfo.value.code == "custom"
+
+
+def test_duplex_runtime_config_error_defaults_its_code() -> None:
+    assert DuplexRuntimeConfigError("bad").code == "invalid_duplex_runtime_config"
+
+
+@pytest.mark.parametrize(
+    ("payload", "expected"),
+    [
+        ({"duplex_turn_id": 3}, 3),
+        ({"duplex_turn_id": "7"}, 7),
+        ({"model_turn_id": 2}, 2),
+        ({"duplex_turn_id": 5, "model_turn_id": 9}, 5),
+        ({"duplex_turn_id": "not-a-number"}, None),
+        ({}, None),
+        ("turn", None),
+        (None, None),
+    ],
+)
+def test_payload_turn_id_prefers_duplex_turn_id_and_coerces(payload: object, expected: int | None) -> None:
+    assert payload_turn_id(payload) == expected
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [(4, 4), ("11", 11), (2.9, 2), (True, 1), (None, None), ("x", None), (object(), None), ([1], None)],
+)
+def test_coerce_int(value: object, expected: int | None) -> None:
+    assert coerce_int(value) == expected
+
+
+# --------------------------------------------------------------------------- #
+# PipelineConfig binding                                                      #
+# --------------------------------------------------------------------------- #
+
+
+def test_pipeline_config_binds_one_duplex_plugin_path() -> None:
+    field_defaults = {f.name: f.default for f in fields(PipelineConfig)}
+
+    assert "duplex_plugin" in field_defaults
+    assert field_defaults["duplex_plugin"] is None
+
+    # The pre-framework fields survive only so that the pipelines of the models
+    # that are not ported yet still construct (RFC vllm-omni#7181 ports them).
+    # Nothing in the duplex framework may read them: a pipeline is a duplex
+    # model iff it declares ``duplex_plugin``.
+    legacy_fields = ("duplex_runtime_extension", "duplex_serving_adapter", "duplex_control_enabled")
+    for legacy_field in legacy_fields:
+        assert legacy_field in field_defaults
+
+    framework_sources = [
+        *(Path("vllm_omni/engine/duplex").glob("*.py")),
+        *(Path("vllm_omni/entrypoints/duplex").glob("*.py")),
+        Path("vllm_omni/engine/duplex_omni_engine.py"),
+        Path("vllm_omni/engine/duplex_orchestrator.py"),
+        Path("vllm_omni/entrypoints/duplex_omni.py"),
+    ]
+    readers = [
+        f"{source}:{legacy_field}"
+        for source in framework_sources
+        for legacy_field in legacy_fields
+        if legacy_field in source.read_text(encoding="utf-8")
+    ]
+    assert readers == []
+
+    # Same rule for the runtime-config knob the framework stopped reading: the
+    # append-retry table it bounded went away with the correlated append RPC.
+    assert "completed_append_cache_size" in {f.name for f in fields(DuplexSessionRuntimeConfig)}
+    assert [
+        str(source)
+        for source in framework_sources
+        if "completed_append_cache_size" in source.read_text(encoding="utf-8")
+    ] == []

tests/engine/duplex/test_duplex_runtime.py

@@ -1,563 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from dataclasses import FrozenInstanceError
-from types import SimpleNamespace
-
-import pytest
-from vllm.sampling_params import SamplingParams
-
-from vllm_omni.engine.duplex import runtime as duplex_runtime
-from vllm_omni.engine.duplex.contracts import (
-    duplex_data_plane_request_info,
-    duplex_resource_request_belongs_to_session,
-    duplex_resource_request_id,
-)
-from vllm_omni.engine.duplex.messages import DuplexFence
-from vllm_omni.engine.duplex.runtime import (
-    DuplexInputMode,
-    DuplexOutputAction,
-    DuplexOutputDecision,
-    DuplexRuntimeCapabilities,
-    DuplexSessionRuntimeManager,
-)
-from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import (
-    MiniCPMO45DuplexRuntimeExtension,
-    _duplex_hd_slice_count,
-    build_duplex_data_plane_prompt,
-    duplex_scheduler_token_budget,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_duplex_fence_is_immutable():
-    fence = DuplexFence("session")
-
-    with pytest.raises(FrozenInstanceError):
-        fence.epoch = 1  # type: ignore[misc]
-
-    assert not hasattr(fence, "__dict__")
-
-
-def test_duplex_runtime_tracks_stage_bindings_and_barge_in_epoch():
-    manager = DuplexSessionRuntimeManager()
-    session = manager.open_session(
-        DuplexFence("sid-1"),
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_TOKENS},
-        ),
-    )
-    session.bind_stage_request(stage_id=0, request_id="req-stage0", fence=session.fence)
-    session.bind_stage_request(stage_id=1, request_id="req-stage1", fence=session.fence)
-
-    update = session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence)
-    next_fence = DuplexFence("sid-1", epoch=1)
-    stale_request_ids = session.release_fence(session.fence)
-    session.accept_fence(next_fence)
-
-    assert update.seq == 1
-    assert session.fence == next_fence
-    assert stale_request_ids == ["req-stage0", "req-stage1"]
-    assert session.stage_bindings == {}
-
-
-def test_duplex_runtime_tracks_same_request_id_for_each_pipeline_stage():
-    manager = DuplexSessionRuntimeManager()
-    fence = DuplexFence("sid-shared-pipeline-request")
-    session = manager.open_session(fence)
-
-    session.reserve_stage_request(0, "req-shared", fence=fence)
-    session.bind_stage_request(0, "req-shared", fence=fence)
-    session.bind_stage_request(1, "req-shared", fence=fence)
-
-    assert session.stage_bindings[0].request_id == "req-shared"
-    assert session.stage_bindings[1].request_id == "req-shared"
-    assert session.resource_request_ids() == ["req-shared"]
-    assert session.input_seq == 0
-
-
-def test_duplex_runtime_extension_validation_rejects_missing_methods():
-    class IncompleteExtension:
-        def configure_sampling_params(self, *, runtime_config, defaults):
-            del runtime_config
-            return defaults
-
-    with pytest.raises(TypeError, match="plan_append"):
-        duplex_runtime.validate_duplex_runtime_extension(IncompleteExtension())
-
-
-def test_duplex_runtime_extension_validation_rejects_stage_count_mismatch():
-    class WrongStageCountExtension(MiniCPMO45DuplexRuntimeExtension):
-        def configure_sampling_params(self, *, runtime_config, defaults):
-            del runtime_config
-            return defaults[:1]
-
-    with pytest.raises(ValueError, match="one sampling parameter per stage"):
-        duplex_runtime.validate_duplex_runtime_extension(
-            WrongStageCountExtension(),
-            sampling_defaults=(SamplingParams(), SamplingParams()),
-        )
-
-
-def test_duplex_runtime_extension_validation_rejects_sampling_type_mismatch():
-    class WrongSamplingTypeExtension(MiniCPMO45DuplexRuntimeExtension):
-        def configure_sampling_params(self, *, runtime_config, defaults):
-            del runtime_config
-            return tuple(object() for _ in defaults)
-
-    with pytest.raises(TypeError, match="sampling parameter type"):
-        duplex_runtime.validate_duplex_runtime_extension(
-            WrongSamplingTypeExtension(),
-            sampling_defaults=(SamplingParams(), SamplingParams()),
-        )
-
-
-def test_duplex_output_decision_metadata_is_immutable():
-    decision = DuplexOutputDecision(
-        action=DuplexOutputAction.DIRECT_RESPONSE,
-        metadata={"model_listen": True},
-    )
-
-    with pytest.raises(TypeError):
-        decision.metadata["model_listen"] = False
-
-
-def _decide_minicpmo_output(
-    output: object,
-    *,
-    segment_token_ids: tuple[int, ...] = (),
-    segment_output_metadata: dict | None = None,
-):
-    return MiniCPMO45DuplexRuntimeExtension().decide_output(
-        stage_id=0,
-        final_stage_id=1,
-        segment_finished=True,
-        segment_token_ids=segment_token_ids,
-        segment_output_metadata=segment_output_metadata or {},
-        output=output,
-    )
-
-
-def test_minicpmo_extension_owns_stage_sampling_overrides():
-    defaults = (
-        SamplingParams(max_tokens=4),
-        SamplingParams(max_tokens=8),
-    )
-
-    configured = MiniCPMO45DuplexRuntimeExtension().configure_sampling_params(
-        runtime_config={
-            "duplex_stage_max_tokens": {"0": 20},
-            "duplex_stage_sampling_params": {"1": {"stop_token_ids": [151645]}},
-        },
-        defaults=defaults,
-    )
-
-    assert configured[0].max_tokens == 20
-    assert configured[1].stop_token_ids == [151645]
-    assert defaults[0].max_tokens == 4
-    assert 151645 not in (defaults[1].stop_token_ids or [])
-
-
-def test_minicpmo_output_decision_uses_raw_streaming_token_snapshot():
-    decision = _decide_minicpmo_output(
-        SimpleNamespace(outputs=[SimpleNamespace()]),
-        segment_token_ids=(151705,),
-        segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}},
-    )
-
-    assert decision is not None
-    assert decision.action is DuplexOutputAction.DIRECT_RESPONSE
-    assert decision.metadata["duplex_native_decision"] == "listen"
-    assert decision.metadata["model_listen"] is True
-
-
-@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"])
-def test_minicpmo_output_decision_ignores_output_level_token_history(attr):
-    output = SimpleNamespace(
-        multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
-        outputs=[SimpleNamespace()],
-        **{attr: [42, 151705]},
-    )
-
-    assert _decide_minicpmo_output(output) is None
-
-
-@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"])
-def test_minicpmo_output_decision_uses_completion_token_ids(attr):
-    output = SimpleNamespace(
-        multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
-        outputs=[SimpleNamespace(**{attr: [42, 151705]})],
-    )
-
-    assert _decide_minicpmo_output(output) is not None
-
-
-def test_minicpmo_output_decision_uses_completion_stop_reason():
-    output = SimpleNamespace(
-        multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
-        outputs=[SimpleNamespace(stop_reason=151705)],
-    )
-
-    assert _decide_minicpmo_output(output) is not None
-
-
-def test_duplex_runtime_cancel_fence_rejects_late_append_and_accepts_next_epoch():
-    manager = DuplexSessionRuntimeManager()
-    cancelled_fence = DuplexFence("sid-cancel-race")
-    next_fence = DuplexFence("sid-cancel-race", epoch=1)
-    session = manager.open_session(
-        cancelled_fence,
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-    session.bind_stage_request(
-        stage_id=0,
-        request_id="req-cancelled",
-        fence=cancelled_fence,
-    )
-
-    stale_request_ids = session.cancel_fence(cancelled_fence, next_fence)
-
-    assert stale_request_ids == ["req-cancelled"]
-    assert session.fence == next_fence
-    assert session.stage_bindings == {}
-    with pytest.raises(RuntimeError, match="fence mismatch"):
-        session.append_input(
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-            fence=cancelled_fence,
-        )
-    update = session.append_input(
-        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-        fence=next_fence,
-    )
-    assert update.seq == 1
-
-
-def test_duplex_runtime_stale_close_preserves_live_session_and_bindings():
-    manager = DuplexSessionRuntimeManager()
-    current_fence = DuplexFence("sid-stale-close", epoch=1)
-    session = manager.open_session(
-        current_fence,
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-    session.bind_stage_request(0, "req-live", fence=current_fence)
-
-    with pytest.raises(RuntimeError, match="fence mismatch"):
-        manager.close_session(DuplexFence("sid-stale-close"))
-
-    assert manager.get("sid-stale-close") is session
-    assert session.stage_request_ids() == ["req-live"]
-
-
-def test_duplex_runtime_reopen_rejects_late_append_from_old_incarnation():
-    manager = DuplexSessionRuntimeManager()
-    old_fence = DuplexFence("sid-reopen", incarnation=0)
-    old_session = manager.open_session(
-        old_fence,
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-    manager.close_session(old_fence)
-
-    new_fence = DuplexFence("sid-reopen", incarnation=1)
-    new_session = manager.open_session(
-        new_fence,
-        capabilities=old_session.capabilities,
-    )
-
-    with pytest.raises(RuntimeError, match="fence mismatch"):
-        new_session.append_input(
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-            fence=old_fence,
-        )
-    assert (
-        new_session.append_input(
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-            fence=new_fence,
-        ).seq
-        == 1
-    )
-
-
-def test_duplex_prompt_expands_incarnation_metadata():
-    fence = DuplexFence("sid-incarnation", incarnation=3)
-
-    prompt = build_duplex_data_plane_prompt(
-        request_id=duplex_resource_request_id(fence, "stage0"),
-        fence=fence,
-        session_config={},
-        runtime_config={},
-        seq=1,
-        turn_seq=1,
-        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-        payload={"is_speech": True},
-        final=False,
-    )
-
-    assert prompt["model_intermediate_buffer"]["duplex"]["incarnation"] == 3
-
-
-def test_duplex_runtime_tracks_turn_local_append_sequence():
-    manager = DuplexSessionRuntimeManager()
-    session = manager.open_session(
-        DuplexFence("sid-turn-seq"),
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-
-    first = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence)
-    second = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence)
-    next_turn = DuplexFence("sid-turn-seq", turn_id=1, response_seq=1)
-    third = session.append_input(
-        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-        fence=next_turn,
-    )
-    fourth = session.append_input(
-        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-        fence=next_turn,
-    )
-
-    assert [first.seq, second.seq, third.seq, fourth.seq] == [1, 2, 3, 4]
-    assert [first.turn_seq, second.turn_seq, third.turn_seq, fourth.turn_seq] == [1, 2, 1, 2]
-    assert [first.turn_id, second.turn_id, third.turn_id, fourth.turn_id] == [0, 0, 1, 1]
-
-
-def test_duplex_runtime_rejects_unsupported_append_mode():
-    manager = DuplexSessionRuntimeManager()
-    session = manager.open_session(
-        DuplexFence("sid-2"),
-        capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.TURN_COMMIT_ONLY}),
-    )
-
-    with pytest.raises(ValueError, match="not supported"):
-        session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence)
-
-
-def test_duplex_data_plane_request_info_extracts_structured_stage_result():
-    request_id, response_stage_id = duplex_data_plane_request_info(
-        {
-            "stage_results": [
-                {"result": {"supported": True}},
-                {
-                    "result": {
-                        "data_plane_append": True,
-                        "request_id": "duplex-sid-e0-stage0-s1",
-                        "response_stage_id": 1,
-                    }
-                },
-            ]
-        }
-    )
-
-    assert request_id == "duplex-sid-e0-stage0-s1"
-    assert response_stage_id == 1
-
-
-def test_duplex_data_plane_request_info_rejects_missing_request_id():
-    assert duplex_data_plane_request_info(
-        {
-            "stage_results": [
-                {
-                    "result": {
-                        "data_plane_append": True,
-                        "request_id": "",
-                        "response_stage_id": 1,
-                    }
-                }
-            ]
-        }
-    ) == (None, None)
-
-
-def test_duplex_scheduler_token_budget_estimates_pcm_slots():
-    assert (
-        duplex_scheduler_token_budget(
-            {
-                "audio": "AAAAAA==",
-                "format": "pcm_f32le",
-                "sample_rate_hz": 16000,
-            }
-        )
-        == 16
-    )
-
-
-def test_duplex_scheduler_token_budget_ignores_client_budget_fields():
-    assert (
-        duplex_scheduler_token_budget(
-            {
-                "audio": "AAAAAA==",
-                "format": "pcm_f32le",
-                "duplex_num_input_tokens": 999,
-                "num_input_tokens": 999,
-            }
-        )
-        == 16
-    )
-
-
-def test_resource_state_rejects_fence_regression_and_requires_explicit_fence():
-    current = DuplexFence("sid", epoch=2, turn_id=3, response_seq=4)
-    manager = DuplexSessionRuntimeManager()
-    session = manager.open_session(
-        current,
-        capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}),
-    )
-
-    for stale in (
-        DuplexFence("sid", epoch=1, turn_id=99, response_seq=99),
-        DuplexFence("sid", epoch=2, turn_id=2, response_seq=4),
-        DuplexFence("sid", epoch=2, turn_id=3, response_seq=3),
-    ):
-        with pytest.raises(RuntimeError, match="fence mismatch"):
-            session.accept_fence(stale)
-        assert session.fence == current
-
-    with pytest.raises(TypeError, match="fence"):
-        session.bind_stage_request(0, "request")
-    with pytest.raises(TypeError, match="fence"):
-        session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK)
-    with pytest.raises(TypeError, match="DuplexFence"):
-        manager.open_session("legacy-session")
-    with pytest.raises(TypeError, match="DuplexFence"):
-        manager.close_session("legacy-session")
-
-
-def test_resource_request_id_is_derived_from_fence_and_role():
-    fence = DuplexFence("sid-with-dashes", epoch=7, turn_id=11, response_seq=13)
-
-    assert duplex_resource_request_id(fence, "stage0") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage0"
-    assert duplex_resource_request_id(fence, "stage1") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage1"
-
-
-def test_resource_request_id_codec_separates_session_id_from_incarnation():
-    embedded_incarnation = duplex_resource_request_id(
-        DuplexFence("foo-i1", incarnation=0),
-        "stage0",
-    )
-    actual_incarnation = duplex_resource_request_id(
-        DuplexFence("foo", incarnation=1),
-        "stage0",
-    )
-
-    assert embedded_incarnation != actual_incarnation
-
-
-def test_resource_request_id_session_membership_uses_encoded_identity():
-    request_id = duplex_resource_request_id(
-        DuplexFence("sid-with-dashes", incarnation=2, epoch=7),
-        "stage0",
-    )
-
-    assert duplex_resource_request_belongs_to_session(request_id, "sid-with-dashes") is True
-    assert duplex_resource_request_belongs_to_session(request_id, "sid") is False
-    assert duplex_resource_request_belongs_to_session("duplex-s.invalid.i.x.e.7.r.stage0", "sid") is False
-
-
-def test_placeholder_budget_is_planned_inside_omni_engine_boundary():
-    fence = DuplexFence("sid", turn_id=1, response_seq=1)
-    prompt = build_duplex_data_plane_prompt(
-        request_id=duplex_resource_request_id(fence, "stage0"),
-        fence=fence,
-        session_config={},
-        runtime_config={},
-        seq=2,
-        turn_seq=1,
-        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-        payload={
-            "audio": "AAAAAA==",
-            "format": "pcm_f32le",
-            "duplex_num_input_tokens": 999,
-            "num_input_tokens": 999,
-        },
-        final=False,
-    )
-
-    assert len(prompt["prompt_token_ids"]) == 16
-    assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence
-    assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16
-
-
-def _jpeg_b64(width: int, height: int) -> str:
-    import base64
-    from io import BytesIO
-
-    from PIL import Image
-
-    buffer = BytesIO()
-    Image.new("RGB", (width, height), (32, 64, 96)).save(buffer, format="JPEG")
-    return base64.b64encode(buffer.getvalue()).decode("ascii")
-
-
-def _pcm_unit_payload(frames: list[str]) -> dict[str, object]:
-    import base64
-
-    return {
-        "audio": base64.b64encode(b"\x00" * (16000 * 4)).decode("ascii"),
-        "format": "pcm_f32le",
-        "sample_rate_hz": 16000,
-        "video_frames": frames,
-    }
-
-
-def test_duplex_scheduler_token_budget_one_frame_is_one_block():
-    audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([]))
-    one_frame = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(960, 540)]))
-
-    assert one_frame - audio_only == 66
-
-
-@pytest.mark.parametrize(
-    ("base_size", "expected_blocks"),
-    [
-        ((960, 540), 4),  # official camera frame: source + 2 HD patches + composite
-        ((540, 960), 4),  # portrait phone capture
-        ((640, 360), 4),
-        ((448, 448), 2),  # fits one tile: the processor does not slice it
-        ((320, 240), 2),
-    ],
-)
-def test_duplex_scheduler_token_budget_matches_stage0_hd_slicing(base_size, expected_blocks):
-    """A stacked pair is processed with max_slice_nums=[2, 1]; the budget must
-    equal the blocks Stage0 builds for that base-frame size."""
-    audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([]))
-    pair = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(*base_size), _jpeg_b64(448, 448)]))
-
-    assert pair - audio_only == expected_blocks * 66
-
-
-@pytest.mark.parametrize(
-    ("size", "expected_patches"),
-    [
-        # multiple = min(ceil(w * h / 448**2), max_slice_nums); <= 1 is never sliced.
-        ((448, 448), 0),  # exactly one tile
-        ((320, 240), 0),  # smaller than one tile
-        ((447, 449), 0),  # just under one tile
-        ((449, 449), 2),  # just over: the processor cuts a 2-tile grid
-        ((640, 360), 2),
-        ((960, 540), 2),  # official camera frame
-        ((540, 960), 2),  # same frame rotated
-        ((1920, 1080), 2),  # capped by max_slice_nums=2, not by area
-    ],
-)
-def test_duplex_hd_slice_count_matches_official_grid(size, expected_patches):
-    """Lock the port of ``MiniCPMVImageProcessor.get_sliced_grid``.
-
-    Stage0 slices a stacked unit with ``max_slice_nums=[2, 1]``, so the patch
-    count for the base frame decides the scheduler budget. A frame that fits one
-    ``scale_resolution=448`` tile is not sliced at all.
-    """
-    assert _duplex_hd_slice_count(size, 2) == expected_patches
-
-
-def test_duplex_scheduler_token_budget_unreadable_base_frame_keeps_hd_fallback():
-    audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([]))
-    pair = duplex_scheduler_token_budget(_pcm_unit_payload(["not-an-image", _jpeg_b64(448, 448)]))
-
-    assert pair - audio_only == 4 * 66

tests/engine/duplex/test_engine_session.py

@@ -0,0 +1,612 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Session-state semantics of the engine-resident duplex session (``DuplexEngineSession``)."""
+
+from __future__ import annotations
+
+from dataclasses import FrozenInstanceError
+
+import pytest
+
+from vllm_omni.engine.duplex.config import (
+    DuplexCapabilities,
+    DuplexConfigError,
+    DuplexOverlapPolicy,
+    DuplexSessionConfig,
+    DuplexTurnEventType,
+    DuplexTurnState,
+    ResponseCreateOptions,
+)
+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.model_executor.models.minicpmo_4_5.duplex.capabilities import (
+    minicpmo45_native_capabilities,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _session(session_id: str = "duplex-test", config: DuplexSessionConfig | None = None) -> DuplexEngineSession:
+    return DuplexEngineSession(session_id=session_id, config=config or DuplexSessionConfig(model="test-model"))
+
+
+def test_commit_audio_input_does_not_advance_model_turn_identity():
+    session = _session()
+
+    first = session.commit_audio_input(transcript="first chunk")
+    second = session.commit_audio_input(transcript="second chunk")
+
+    assert session.input_commit_seq == 2
+    assert first.input_commit_seq == 1
+    assert second.input_commit_seq == 2
+    assert first.turn_id == second.turn_id == session.turn_id == 0
+    assert first.message["transcript"] == "first chunk"
+    assert first.message["content"] == [
+        {"type": "audio_url", "audio_url": {"url": "native-duplex:input-audio"}, "transcript": "first chunk"}
+    ]
+    assert len(session.history) == 2
+    assert session.turn_state == DuplexTurnState.USER_COMMITTED
+
+
+def test_commit_audio_input_accepts_explicit_turn_id():
+    session = _session()
+
+    committed = session.commit_audio_input(turn_id=4)
+
+    assert committed.turn_id == 4
+    assert session.turn_id == 0
+    assert "transcript" not in committed.message
+
+
+# ---- response options ----
+
+
+def test_response_options_apply_to_one_response_without_mutating_session_defaults():
+    session = _session(config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64))
+    session.reserve_response_options(
+        ResponseCreateOptions(instructions="one response", voice="override-voice", max_tokens=8)
+    )
+
+    assert session.config.instructions == "base"
+    assert session.config.voice == "base-voice"
+    assert session.response_config is session.config
+    session.begin_response()
+    assert session.config.instructions == "base"
+    assert session.config.voice == "base-voice"
+    assert session.response_config.instructions == "one response"
+    assert session.response_config.voice == "override-voice"
+    assert session.response_config.max_tokens == 8
+
+    session.end_response()
+    assert session.response_config is session.config
+    assert session.config.instructions == "base"
+    assert session.config.voice == "base-voice"
+    assert session.config.max_tokens == 64
+
+
+def test_response_options_cannot_overwrite_an_unconsumed_reservation():
+    session = _session(config=DuplexSessionConfig(instructions="base"))
+    session.reserve_response_options(ResponseCreateOptions(instructions="first"))
+
+    with pytest.raises(RuntimeError, match="already reserved"):
+        session.reserve_response_options(ResponseCreateOptions(instructions="second"))
+
+    session.begin_response()
+    assert session.response_config.instructions == "first"
+
+
+def test_response_options_cannot_be_reserved_while_response_is_active():
+    session = _session(config=DuplexSessionConfig(instructions="base"))
+    session.begin_response()
+
+    with pytest.raises(RuntimeError, match="active"):
+        session.reserve_response_options(ResponseCreateOptions(instructions="too late"))
+
+
+def test_discarded_response_options_do_not_apply():
+    session = _session(config=DuplexSessionConfig(instructions="base"))
+    session.reserve_response_options(ResponseCreateOptions(instructions="dropped"))
+    session.discard_response_options()
+
+    session.begin_response()
+
+    assert session.response_config.instructions == "base"
+
+
+# ---- response / overlap identity ----
+
+
+def test_session_owns_response_and_overlap_identity():
+    session = _session()
+
+    response_id = session.begin_response()
+    assert session.turn_state == DuplexTurnState.ASSISTANT_GENERATING
+    session.accumulate_overlap_speech(320)
+    session.accumulate_overlap_speech(180)
+    session.end_response()
+
+    assert session.active_response_id is None
+    assert session.last_response_id == response_id
+    assert session.overlap_speech_ms == 500
+    assert session.reset_overlap_speech() == 500
+    assert session.overlap_speech_ms == 0
+    assert session.turn_state == DuplexTurnState.IDLE
+
+
+def test_session_composes_single_owner_ledgers_with_immutable_views():
+    session = _session()
+
+    session.bind_request("req-1")
+    response_id = session.begin_response(turn_id=3)
+    session.mark_audio_sent(duration_ms=240)
+
+    assert session.active_request_id == "req-1"
+    assert session.active_response_id == response_id
+    assert session.active_response_turn_id == 3
+    assert session.playback.sent_ms == 240
+    assert session.turn_state == DuplexTurnState.ASSISTANT_PLAYING
+    with pytest.raises(FrozenInstanceError):
+        session.playback.sent_ms = 480  # type: ignore[misc]
+    assert session.playback.sent_ms == 240
+
+
+def test_barge_in_advances_epoch_and_drops_uncommitted_assistant_text():
+    session = _session()
+    response_id = session.begin_response()
+    session.bind_request("chatcmpl-duplex-test")
+    session.append_assistant_text("unplayed answer")
+
+    new_epoch = session.barge_in()
+
+    assert response_id is not None
+    assert new_epoch == 1
+    assert session.epoch == 1
+    assert session.accepted_fence == DuplexFence(session.session_id, epoch=1, turn_id=0)
+    assert session.active_request_id is None
+    assert session.active_response_id is None
+    assert session.assistant_text_buffer == ()
+    assert session.history == ()
+    assert session.turn_state == DuplexTurnState.BARGE_IN
+
+
+# ---- playback ledger ----
+
+
+def test_playback_ack_tracks_committed_cursor_separately():
+    session = _session()
+    session.mark_audio_sent(duration_ms=10_000)
+
+    session.acknowledge_playback(played_ms=2_000)
+
+    assert session.playback.generated_ms == 10_000
+    assert session.playback.sent_ms == 10_000
+    assert session.playback.played_ms == 2_000
+    assert session.playback.committed_ms == 2_000
+    assert session.playback.as_dict() == {
+        "generated_ms": 10_000,
+        "sent_ms": 10_000,
+        "played_ms": 2_000,
+        "committed_ms": 2_000,
+    }
+
+
+def test_playback_ack_is_scoped_per_response():
+    session = _session()
+    first = session.begin_response()
+    session.mark_audio_sent(duration_ms=3_000)
+    session.end_response()
+    second = session.begin_response()
+    session.mark_audio_sent(duration_ms=1_000)
+
+    session.acknowledge_playback(played_ms=2_500, response_id=first)
+
+    assert session.playback_for_response(first).played_ms == 2_500
+    assert session.playback_for_response(second).played_ms == 0
+    assert session.playback.played_ms == 0
+
+
+def test_history_commit_uses_audio_text_alignment_marks():
+    session = _session()
+    session.begin_response()
+
+    session.append_assistant_text("hello ")
+    session.mark_audio_sent(duration_ms=1_000, text_chars=6)
+    session.append_assistant_text("world")
+    session.mark_audio_sent(duration_ms=2_000, text_chars=11)
+    assert [(mark.text_chars, mark.audio_end_ms) for mark in session.assistant_audio_text_marks] == [
+        (6, 1_000),
+        (11, 2_000),
+    ]
+    session.acknowledge_playback(played_ms=1_200, committed_ms=1_200)
+
+    committed = session.end_response(commit_text=True)
+
+    assert committed == {"role": "assistant", "content": "hello w"}
+    assert session.history[-1] == committed
+    assert session.last_assistant_full_message == {"role": "assistant", "content": "hello world"}
+
+
+def test_history_commit_with_ack_only_policy_defers_unacknowledged_text():
+    session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only"))
+    response_id = session.begin_response()
+    session.append_assistant_text("never played")
+    session.mark_audio_sent(duration_ms=1_000, text_chars=12)
+
+    committed = session.end_response(commit_text=True)
+
+    assert committed is None
+    assert session.history == ()
+    assert f"item_{response_id}" in session.pending_history_item_ids
+
+
+def test_truncate_history_item_uses_response_alignment_marks():
+    session = _session()
+    response_id = session.begin_response()
+    session.append_assistant_text("hello ")
+    session.mark_audio_sent(duration_ms=1_000, text_chars=6)
+    session.append_assistant_text("world")
+    session.mark_audio_sent(duration_ms=2_000, text_chars=11)
+    session.acknowledge_playback(played_ms=2_000)
+    committed = session.end_response(commit_text=True)
+    assert committed == {"role": "assistant", "content": "hello world"}
+    item_id = f"item_{response_id}"
+    session.register_history_item(item_id, committed)
+
+    assert session.truncate_history_item(item_id, audio_end_ms=1_000) is True
+
+    assert session.history[-1] == {"role": "assistant", "content": "hello"}
+    assert session.history_item_ids[item_id] == {"role": "assistant", "content": "hello"}
+    assert session.delete_history_item(item_id) is True
+    assert session.history == ()
+
+
+# ---- turn signals / model turns ----
+
+
+@pytest.mark.parametrize(
+    ("event_type", "turn_state"),
+    [
+        (DuplexTurnEventType.USER_STARTED, DuplexTurnState.USER_SPEAKING),
+        (DuplexTurnEventType.USER_COMMITTED, DuplexTurnState.USER_COMMITTED),
+        (DuplexTurnEventType.ASSISTANT_STARTED, DuplexTurnState.ASSISTANT_GENERATING),
+        (DuplexTurnEventType.ASSISTANT_DONE, DuplexTurnState.IDLE),
+        (DuplexTurnEventType.BARGE_IN, DuplexTurnState.BARGE_IN),
+    ],
+)
+def test_signal_turn_transitions_and_returns_typed_turn_event(event_type, turn_state):
+    session = _session()
+
+    event = session.signal_turn(event_type.value)
+
+    assert isinstance(event, TurnEvent)
+    assert session.turn_state == turn_state
+    wire = event.to_realtime()
+    assert wire["type"] == "turn.event"
+    assert wire["event"] == event_type.value
+    assert wire["turn_state"] == turn_state.value
+
+
+def test_signal_turn_playback_ack_updates_cursor_and_close_marks_session_closing():
+    session = _session()
+    session.mark_audio_sent(duration_ms=5_000)
+
+    ack = session.signal_turn("playback_ack", {"played_ms": 1_500, "committed_ms": 1_000})
+    assert ack.event == "playback_ack"
+    assert session.playback.played_ms == 1_500
+    assert session.playback.committed_ms == 1_000
+
+    close = session.signal_turn(DuplexTurnEventType.CLOSE.value)
+    assert close.event == "close"
+    assert session.state.value == "closing"
+
+
+def test_complete_model_turn_advances_turn_id_and_fence():
+    session = _session()
+    assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=0)
+
+    session.complete_model_turn(0)
+
+    assert session.turn_id == 1
+    assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=1)
+    assert session.accepted_fence == session.fence
+
+    # A stale terminal for an already-completed turn does not move identity backwards.
+    session.complete_model_turn(0)
+    assert session.turn_id == 1
+    session.complete_model_turn(3)
+    assert session.turn_id == 4
+    assert session.accepted_fence.turn_id == 4
+
+
+# ---- fence validation ----
+
+
+def test_accept_fence_rejects_stale_or_foreign_fences():
+    session = _session("sid-fence")
+    session.complete_model_turn(0)
+
+    with pytest.raises(DuplexFenceMismatchError):
+        session.accept_fence(DuplexFence("sid-fence", epoch=0, turn_id=0))
+    with pytest.raises(DuplexFenceMismatchError):
+        session.accept_fence(DuplexFence("sid-other", epoch=0, turn_id=1))
+
+    session.accept_fence(DuplexFence("sid-fence", epoch=1, turn_id=0))
+    assert session.accepted_fence == DuplexFence("sid-fence", epoch=1, turn_id=0)
+
+
+def test_prepare_and_commit_append_sequence_chunks_per_turn_and_epoch():
+    session = _session("sid-append")
+    fence = session.fence
+
+    first = session.commit_append(session.prepare_append(fence))
+    second = session.commit_append(session.prepare_append(fence))
+    assert (first.seq, first.turn_seq, first.turn_id) == (1, 1, 0)
+    assert (second.seq, second.turn_seq, second.turn_id) == (2, 2, 0)
+
+    session.complete_model_turn(0)
+    third = session.commit_append(session.prepare_append(session.fence))
+    assert (third.seq, third.turn_seq, third.turn_id) == (3, 1, 1)
+
+    # Barge-in advances the epoch (turn_id is kept) and restarts the append sequence.
+    session.barge_in()
+    assert session.fence == DuplexFence("sid-append", epoch=1, turn_id=1)
+    fourth = session.commit_append(session.prepare_append(session.fence))
+    assert (fourth.seq, fourth.turn_seq, fourth.turn_id) == (1, 1, 1)
+    assert session.input_seq == 1
+
+
+def test_commit_append_rejects_stale_reservation():
+    session = _session("sid-append-stale")
+    fence = session.fence
+    reservation = session.prepare_append(fence)
+    session.commit_append(session.prepare_append(fence))
+
+    with pytest.raises(RuntimeError, match="stale"):
+        session.commit_append(reservation)
+
+    with pytest.raises(DuplexFenceMismatchError):
+        session.prepare_append(DuplexFence("sid-other", epoch=0, turn_id=0))
+
+
+def test_cancel_fence_releases_stage_requests_and_advances_identity():
+    session = _session("sid-cancel")
+    cancelled = session.fence
+    session.reserve_stage_request(0, "req-a", fence=cancelled)
+    session.bind_stage_request(1, "req-b", fence=cancelled)
+    assert session.stage_request_submitted(1, "req-b") is True
+    assert session.stage_request_submitted(0, "req-a") is False
+    next_fence = DuplexFence("sid-cancel", epoch=1, turn_id=0)
+
+    stale = session.cancel_fence(cancelled, next_fence)
+
+    assert stale == ["req-a", "req-b"]
+    assert session.resource_request_ids() == []
+    assert session.accepted_fence == next_fence
+
+    with pytest.raises(DuplexFenceMismatchError):
+        session.cancel_fence(next_fence, DuplexFence("sid-cancel", epoch=1, turn_id=1))
+    with pytest.raises(DuplexFenceMismatchError):
+        session.cancel_fence(DuplexFence("sid-other", epoch=1, turn_id=0), DuplexFence("sid-cancel", epoch=2))
+
+
+# ---- public view ----
+
+
+def test_as_public_dict_exposes_identity_capabilities_and_playback():
+    session = _session("duplex-public", config=DuplexSessionConfig(model="m", voice="alloy"))
+    session.mark_audio_sent(duration_ms=100)
+
+    payload = session.as_public_dict()
+
+    assert payload["id"] == "duplex-public"
+    assert payload["model"] == "m"
+    assert payload["voice"] == "alloy"
+    assert payload["state"] == "open"
+    assert payload["turn_state"] == "assistant_playing"
+    assert payload["epoch"] == 0
+    assert payload["turn_id"] == 0
+    assert payload["capabilities"] == DuplexCapabilities().as_dict()
+    assert payload["playback"] == {"generated_ms": 100, "sent_ms": 100, "played_ms": 0, "committed_ms": 0}
+
+
+# ---- DuplexSessionConfig.from_realtime ----
+
+
+def test_from_realtime_rejects_unsupported_audio_formats():
+    with pytest.raises(DuplexConfigError) as excinfo:
+        DuplexSessionConfig.from_realtime({"input_audio_format": "mp3"})
+    assert excinfo.value.code == "unsupported_audio_format"
+
+    with pytest.raises(DuplexConfigError) as excinfo:
+        DuplexSessionConfig.from_realtime({"audio": {"output": {"format": {"type": "audio/opus"}}}})
+    assert excinfo.value.code == "unsupported_audio_format"
+
+
+def test_from_realtime_validates_turn_detection():
+    with pytest.raises(DuplexConfigError) as excinfo:
+        DuplexSessionConfig.from_realtime({"turn_detection": {"type": "semantic_vad"}})
+    assert excinfo.value.code == "unsupported_turn_detection"
+    assert excinfo.value.param == "turn_detection"
+
+    with pytest.raises(DuplexConfigError) as excinfo:
+        DuplexSessionConfig.from_realtime({"overlap_policy": "barge_in_on_speech"})
+    assert excinfo.value.code == "unsupported_turn_detection"
+
+    server_vad = DuplexSessionConfig.from_realtime({"turn_detection": {"type": "server_vad"}})
+    assert server_vad.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value
+    assert server_vad.extra_body["realtime_turn_detection"]["threshold"] == 0.5
+
+    model_owned = DuplexSessionConfig.from_realtime({"turn_detection": None})
+    assert model_owned.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
+    assert model_owned.extra_body["realtime_turn_detection"] is None
+
+
+def test_from_realtime_maps_wire_fields_and_stores_realtime_keys_in_extra_body():
+    config = DuplexSessionConfig.from_realtime(
+        {
+            "model": "openbmb/MiniCPM-o-4_5",
+            "instructions": "be brief",
+            "output_audio_format": "pcm16",
+            "max_response_output_tokens": "inf",
+            "audio": {"output": {"voice": "alloy", "speed": 1.25}},
+            "tools": [{"type": "function", "name": "lookup"}],
+            "tool_choice": "auto",
+            "metadata": {"tenant": "t1"},
+            "include": ["item.input_audio_transcription.logprobs"],
+            "input_audio_transcription": {"model": "whisper-1"},
+            "extra_body": {"custom": True},
+        }
+    )
+
+    assert config.model == "openbmb/MiniCPM-o-4_5"
+    assert config.instructions == "be brief"
+    assert config.voice == "alloy"
+    assert config.speed == 1.25
+    assert config.max_tokens is None
+    assert config.response_format == "pcm"
+    assert config.modalities == ["text", "audio"]
+    assert config.extra_body["custom"] is True
+    assert config.extra_body["realtime_tools"] == [{"type": "function", "name": "lookup"}]
+    assert config.extra_body["realtime_tool_choice"] == "auto"
+    assert config.extra_body["realtime_metadata"] == {"tenant": "t1"}
+    assert config.extra_body["realtime_include"] == ["item.input_audio_transcription.logprobs"]
+    assert config.extra_body["realtime_input_audio_transcription"] == {"model": "whisper-1"}
+    assert config.extra_body["realtime_audio"] == {"output": {"voice": "alloy", "speed": 1.25}}
+    assert config.extra_body["realtime_output_audio_format"] == "pcm16"
+    assert "extra_body" not in config.extra_body["realtime_session_payload"]
+    assert config.extra_body["realtime_session_payload"]["model"] == "openbmb/MiniCPM-o-4_5"
+
+
+def test_from_realtime_uses_served_model_when_payload_has_none():
+    config = DuplexSessionConfig.from_realtime({}, model="served-model")
+
+    assert config.model == "served-model"
+    assert config.response_format == "pcm"
+    assert config.idle_timeout_s == 300.0
+
+
+def test_from_realtime_ignores_client_chosen_session_ids():
+    config = DuplexSessionConfig.from_realtime({"id": "client-chosen", "session_id": "also-client-chosen"})
+
+    assert not hasattr(config, "session_id")
+    assert "id" not in config.as_dict()
+    session = _session("duplex-engine-allocated", config=config)
+    assert session.as_public_dict()["id"] == "duplex-engine-allocated"
+
+
+# ---- DuplexSessionConfig.apply_realtime_update ----
+
+
+def test_apply_realtime_update_rejects_model_change():
+    config = DuplexSessionConfig(model="served-model")
+
+    with pytest.raises(DuplexConfigError) as excinfo:
+        config.apply_realtime_update({"model": "other-model"}, session_id="sid")
+
+    assert excinfo.value.code == "model_update_unsupported"
+    assert config.model == "served-model"
+
+
+def test_apply_realtime_update_rejects_voice_change_after_audio_started():
+    config = DuplexSessionConfig(model="m", voice="alloy")
+
+    config.apply_realtime_update({"voice": "verse"}, audio_started=False)
+    assert config.voice == "verse"
+    with pytest.raises(DuplexConfigError) as excinfo:
+        config.apply_realtime_update({"audio": {"output": {"voice": "alloy"}}}, audio_started=True)
+    assert excinfo.value.code == "voice_update_after_audio_unsupported"
+    assert config.voice == "verse"
+
+
+def test_apply_realtime_update_rejects_ref_audio_change():
+    config = DuplexSessionConfig(model="m")
+
+    with pytest.raises(DuplexConfigError) as excinfo:
+        config.apply_realtime_update({"ref_audio": "/tmp/voice.wav"})
+
+    assert excinfo.value.code == "ref_audio_update_unsupported"
+    assert config.ref_audio is None
+
+
+def test_apply_realtime_update_patches_fields_and_realtime_keys():
+    config = DuplexSessionConfig(model="m", instructions="old", extra_body={"realtime_tools": [{"name": "x"}]})
+
+    config.apply_realtime_update(
+        {
+            "instructions": "new",
+            "output_audio_format": "g711_ulaw",
+            "max_output_tokens": 32,
+            "tools": None,
+            "metadata": {"k": "v"},
+            "overlap_policy": "not-a-policy",
+        }
+    )
+
+    assert config.instructions == "new"
+    assert config.response_format == "pcm"
+    assert config.max_tokens == 32
+    assert "realtime_tools" not in config.extra_body
+    assert config.extra_body["realtime_metadata"] == {"k": "v"}
+    assert config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
+    assert config.extra_body["realtime_session_payload"]["instructions"] == "new"
+
+
+# ---- overlap policy / capabilities ----
+
+
+def test_overlap_policy_defaults_and_invalid_values_to_listen_only():
+    assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
+    assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value
+    assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value
+    assert (
+        DuplexSessionConfig._normalize_overlap_policy(" Barge_In_On_Speech ")
+        == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value
+    )
+    assert DuplexSessionConfig.from_event({"session": {"overlap_policy": "bogus"}}).overlap_policy == "listen_only"
+
+
+def test_capabilities_as_dict_reports_model_native_duplex_constants():
+    caps = DuplexCapabilities().as_dict()
+
+    assert caps["implementation_level"] == "model_native_duplex"
+    assert caps["input_modes"] == ["append_audio_chunk"]
+    assert caps["supports_kv_lease"] is False
+    assert caps["supports_core_kv_lease"] is False
+    assert "client_event" in caps["signal_sources"]
+    assert set(caps) == set(minicpmo45_native_capabilities().as_dict())
+
+
+def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
+    caps = minicpmo45_native_capabilities(max_sessions=2).as_dict()
+
+    assert caps["implementation_level"] == "model_native_duplex"
+    assert caps["supports_input_append"] is True
+    assert caps["input_modes"] == ["append_audio_chunk"]
+    assert caps["adapter_patterns"] == ["scheduler_data_plane"]
+    assert caps["supports_model_internal_state"] is True
+    assert caps["requires_model_runner_kv"] is True
+    assert caps["requires_native_stage_role"] is True
+    assert caps["supports_kv_lease"] is False
+    assert caps["supports_core_kv_lease"] is False
+    assert caps["supports_stage_resumption"] is True
+    assert caps["supports_scheduler_native_append"] is False
+    assert caps["supports_core_resumable_request"] is True
+    assert caps["supports_stage_connector_handoff"] is True
+    assert caps["supports_audio_truncate"] is True
+    assert caps["supports_barge_in"] is True
+    assert caps["target_barge_in_latency_ms"] is None
+    assert caps["supports_multi_session"] is True
+    assert caps["supports_multi_session_same_replica"] is True
+    assert caps["supports_session_lease"] is True
+    assert caps["supports_session_resume"] is True
+    assert caps["session_admission_mode"] == "engine_managed"
+    assert caps["stage_handoff_transport"] == "scheduler_data_plane"
+
+
+def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
+    caps = minicpmo45_native_capabilities(max_sessions=1).as_dict()
+
+    assert caps["supports_multi_session"] is False
+    assert caps["supports_multi_session_same_replica"] is False

tests/engine/duplex/test_overlap_policy.py

@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""The overlap decision, reached directly.
+
+These rules decide what happens when input arrives while the model is speaking.
+They were unreachable except by driving a whole append through the runner, so
+none of them had a test that named the rule it was checking.
+"""
+
+from __future__ import annotations
+
+import base64
+
+import numpy as np
+import pytest
+
+from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.engine.duplex.session import overlap_policy
+from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _session(**config: object) -> DuplexEngineSession:
+    values: dict[str, object] = {"model": "fake-model"}
+    values.update(config)
+    return DuplexEngineSession(session_id="sid", config=DuplexSessionConfig(**values))
+
+
+def _pcm16(rms: float, samples: int = 1600) -> str:
+    data = np.full(samples, int(rms * 32768), dtype=np.int16)
+    return base64.b64encode(data.tobytes()).decode("ascii")
+
+
+# --------------------------------------------------------------------------- #
+# Speech detection                                                            #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("flag", [True, False])
+def test_an_explicit_is_speech_flag_wins_over_the_audio(flag: bool) -> None:
+    session = _session()
+    payload = {"format": "pcm16", "audio": _pcm16(0.9 if not flag else 0.0)}
+    assert overlap_policy.input_looks_like_speech(session, {"is_speech": flag}, payload) is flag
+
+
+def test_vad_probability_decides_when_no_flag_is_given() -> None:
+    session = _session()
+    payload = {"format": "pcm16", "audio": _pcm16(0.0)}
+    assert overlap_policy.input_looks_like_speech(session, {"vad": {"speech_probability": 0.9}}, payload) is True
+    assert overlap_policy.input_looks_like_speech(session, {"vad": {"speech_probability": 0.1}}, payload) is False
+
+
+def test_quiet_audio_below_the_configured_rms_is_not_speech() -> None:
+    session = _session(overlap_silence_rms=0.05)
+    loud = {"format": "pcm16", "audio": _pcm16(0.5)}
+    quiet = {"format": "pcm16", "audio": _pcm16(0.001)}
+    assert overlap_policy.input_looks_like_speech(session, {}, loud) is True
+    assert overlap_policy.input_looks_like_speech(session, {}, quiet) is False
+
+
+def test_undecodable_audio_is_treated_as_speech() -> None:
+    """Fail toward listening: dropping real speech is worse than a spurious turn."""
+    session = _session()
+    assert overlap_policy.input_looks_like_speech(session, {}, {"format": "pcm16", "audio": "!!not base64!!"}) is True
+
+
+# --------------------------------------------------------------------------- #
+# Barge-in                                                                    #
+# --------------------------------------------------------------------------- #
+
+
+def test_an_explicit_client_overlap_action_requests_barge_in() -> None:
+    session = _session()
+    payload = {"format": "pcm16", "audio": _pcm16(0.5)}
+    decision = overlap_policy.decide(session, {"overlap_action": "barge_in"}, payload, auto_responds=True)
+    assert decision["action"] == "barge_in"
+    assert decision["reason"] == "client_overlap_action"
+
+
+def test_barge_in_is_deferred_when_the_session_cannot_support_it() -> None:
+    session = _session()
+    session.capabilities = type(session.capabilities)(supports_barge_in=False)
+    payload = {"format": "pcm16", "audio": _pcm16(0.5)}
+    decision = overlap_policy.decide(session, {"force_barge_in": True}, payload, auto_responds=True)
+    assert decision["action"] != "barge_in"
+
+
+# --------------------------------------------------------------------------- #
+# Force-listen                                                                #
+# --------------------------------------------------------------------------- #
+
+
+def test_a_short_response_create_is_forced_to_listen() -> None:
+    """A sub-threshold ack should not take the turn from the model."""
+    session = _session(overlap_short_ack_ms=700)
+    short = {"response_create": True, "duration_ms": 200}
+    long = {"response_create": True, "duration_ms": 1500}
+    assert overlap_policy.should_force_listen_for_short_commit(session, short, {}) is True
+    assert overlap_policy.should_force_listen_for_short_commit(session, long, {}) is False
+
+
+def test_force_barge_in_overrides_the_short_ack_rule() -> None:
+    session = _session(overlap_short_ack_ms=700)
+    event = {"response_create": True, "duration_ms": 200, "force_barge_in": True}
+    assert overlap_policy.should_force_listen_for_short_commit(session, event, {}) is False
+
+
+def test_auto_response_overlap_forces_listen_only_when_auto_responding() -> None:
+    event = {"force_listen": True}
+    assert overlap_policy.should_force_listen_for_auto_response_overlap(event, {}, auto_responds=True) is True
+    assert overlap_policy.should_force_listen_for_auto_response_overlap(event, {}, auto_responds=False) is False
+
+
+# --------------------------------------------------------------------------- #
+# Payload merging                                                             #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("speech_first", [True, False], ids=["speech-then-silence", "silence-then-speech"])
+def test_merging_two_appends_keeps_speech_if_either_half_was_speech(speech_first: bool) -> None:
+    """Matching rates and decodable audio, so the halves really are concatenated.
+
+    Without a ``sample_rate_hz`` on both, ``merge_audio_payloads`` falls back
+    to the newer chunk before merging, and the speech flag it returned was the
+    second chunk's own -- the test passed without exercising the merge.
+    """
+    speech = {
+        "format": "pcm_f32le",
+        "audio": base64.b64encode(b"\x01" * 8).decode("ascii"),
+        "sample_rate_hz": 16000,
+        "is_speech": True,
+    }
+    silence = {
+        "format": "pcm_f32le",
+        "audio": base64.b64encode(b"\x00" * 8).decode("ascii"),
+        "sample_rate_hz": 16000,
+        "is_speech": False,
+    }
+    first, second = (speech, silence) if speech_first else (silence, speech)
+    merged = overlap_policy.merge_audio_payloads(first, second)
+    assert base64.b64decode(merged["audio"]) == base64.b64decode(first["audio"]) + base64.b64decode(second["audio"])
+    assert merged["sample_rate_hz"] == 16000
+    assert merged["is_speech"] is True
+
+
+@pytest.mark.parametrize(
+    ("first", "second"),
+    [
+        ({"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000}, {"format": "pcm16", "audio": ""}),
+        (
+            {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000},
+            {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 24000},
+        ),
+        (
+            {"format": "pcm_f32le", "audio": "!!not base64!!", "sample_rate_hz": 16000},
+            {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000},
+        ),
+    ],
+    ids=["format-mismatch", "rate-mismatch", "undecodable"],
+)
+def test_unmergeable_payloads_fall_back_to_the_newer_one(first: dict, second: dict) -> None:
+    """Whatever cannot be concatenated keeps the newer chunk and drops the older.
+
+    Worth pinning: the fallback discards ``first`` silently, so a bug that makes
+    two mergeable payloads look mismatched loses buffered audio rather than
+    raising.
+    """
+    assert overlap_policy.merge_audio_payloads(first, second) == second
+
+
+def test_mergeable_payloads_are_concatenated_in_order() -> None:
+    head = bytes([1, 2, 3, 4])
+    tail = bytes([5, 6, 7, 8])
+    first = {"format": "pcm_f32le", "audio": base64.b64encode(head).decode(), "sample_rate_hz": 16000}
+    second = {"format": "pcm_f32le", "audio": base64.b64encode(tail).decode(), "sample_rate_hz": 16000}
+    merged = overlap_policy.merge_audio_payloads(first, second)
+    assert base64.b64decode(merged["audio"]) == head + tail
+    assert merged["sample_rate_hz"] == 16000
+
+
+# --------------------------------------------------------------------------- #
+# Server VAD + barge_in_on_speech                                             #
+# --------------------------------------------------------------------------- #
+
+
+def test_server_vad_speech_started_reaches_a_barge_in_decision() -> None:
+    """The configuration this PR introduces must decide, not raise.
+
+    Both call sites read the VAD verdict through the module-level
+    ``vad_speech_started`` helper while binding its result to a local of the
+    same name, which made the name local for the whole scope and raised
+    ``UnboundLocalError`` on the call itself. Every ``server_vad`` +
+    ``barge_in_on_speech`` append hit it, and the client saw ``internal_error``
+    instead of a barge-in. The MiniCPM-o preset is ``listen_only``, so no demo
+    exercised this branch.
+    """
+    session = _session(overlap_policy="barge_in_on_speech")
+    payload = {"format": "pcm16", "audio": _pcm16(0.5)}
+    event = {"vad": {"speech_started": True}}
+
+    decision = overlap_policy.decide(session, event, payload, auto_responds=True)
+
+    assert decision["action"] == "barge_in"
+    assert decision["reason"] == "server_vad_speech_started"
+    assert decision["cancel_reason"] == "turn_detected"
+
+
+def test_server_vad_mid_utterance_keeps_listening_instead_of_barging_in() -> None:
+    """``speech_started`` false means the utterance is already running: not a new turn."""
+    session = _session(overlap_policy="barge_in_on_speech")
+    payload = {"format": "pcm16", "audio": _pcm16(0.5)}
+    event = {"vad": {"speech_started": False}}
+
+    decision = overlap_policy.decide(session, event, payload, auto_responds=True)
+
+    assert decision["action"] != "barge_in"
+
+
+def test_barge_in_on_speech_without_vad_still_barges_in_on_speech() -> None:
+    """No VAD verdict: the policy falls back to its own speech classification."""
+    session = _session(overlap_policy="barge_in_on_speech")
+    payload = {"format": "pcm16", "audio": _pcm16(0.5)}
+
+    decision = overlap_policy.decide(session, {}, payload, auto_responds=True)
+
+    assert decision["action"] == "barge_in"
+    assert decision["reason"] == "barge_in_on_speech"

tests/engine/duplex/test_vad_backend.py

@@ -0,0 +1,379 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""The engine-side Silero VAD: backend selection and endpoint rules.
+
+The endpoint rules are the part worth pinning hardest. They were reimplemented
+when turn detection moved into the engine, and drifted from upstream's
+``ThresholdEndpointPolicy`` in two ways that only show up mid-utterance: a loud
+frame used to cancel an in-progress silence timer, and ``audio_end_ms`` used to
+exclude the trailing silence OpenAI says it should include. The parity cases
+below hold upstream's own outputs, so neither can come back unnoticed.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import sys
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from vllm_omni.engine.duplex.turn_detection import validate_realtime_turn_detection
+from vllm_omni.engine.duplex.vad import (
+    SILERO_VAD_SHA256,
+    ServerVADUnavailableError,
+    SileroStreamingVAD,
+    SileroVADBackend,
+    SileroVADBackendProvider,
+    SileroVADConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+FRAME = 512
+SAMPLE_RATE_HZ = 16_000
+
+
+def _drive(config: SileroVADConfig, probabilities: list[float]) -> list[tuple[str, int | None]]:
+    """Feed one probability per frame and collect the endpoint events."""
+    scores = iter(probabilities)
+    vad = SileroStreamingVAD(config, frame_scorer=lambda _frame: next(scores))
+    events: list[tuple[str, int | None]] = []
+    for _ in probabilities:
+        result = vad.process(np.zeros(FRAME, dtype=np.float32))
+        if result.speech_started:
+            events.append(("start", result.speech_start_ms))
+        if result.speech_stopped:
+            events.append(("stop", result.speech_end_ms))
+    return events
+
+
+# --------------------------------------------------------------------------- #
+# Parity with the serving-side policy this was ported from                    #
+# --------------------------------------------------------------------------- #
+
+#: Probability sequences with the endpoint events upstream's
+#: ``ThresholdEndpointPolicy`` produced for them, captured from
+#: ``entrypoints/duplex/server_vad.py`` at `e2d2617f` before that module was
+#: removed. Golden rather than computed because the oracle no longer ships: the
+#: serving layer is transport only, so the VAD it used to carry is gone. Any
+#: drift in the engine's rules still breaks these.
+_PARITY_CASES = {
+    "one utterance": (
+        [0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+        {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 100, "min_speech_duration_ms": 32},
+        [("start", 32), ("stop", 256)],
+    ),
+    "prefix padding reaches back": (
+        [0.0, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0],
+        {"threshold": 0.5, "prefix_padding_ms": 300, "silence_duration_ms": 64, "min_speech_duration_ms": 32},
+        [("start", 0), ("stop", 160)],
+    ),
+    "loud frames inside the silence timer": (
+        [0.9, 0.9, 0.1, 0.45, 0.1, 0.1, 0.1],
+        {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 96, "min_speech_duration_ms": 32},
+        [("start", 0), ("stop", 160)],
+    ),
+    "min speech duration rejects a blip": (
+        [0.9, 0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0],
+        {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 64, "min_speech_duration_ms": 96},
+        [("start", 64), ("stop", 224)],
+    ),
+    "negative threshold floor at a low threshold": (
+        [0.9, 0.9, 0.1, 0.1, 0.1, 0.1],
+        {"threshold": 0.16, "prefix_padding_ms": 0, "silence_duration_ms": 64, "min_speech_duration_ms": 32},
+        [("start", 0)],
+    ),
+}
+
+
+@pytest.mark.parametrize(("probabilities", "options", "expected"), _PARITY_CASES.values(), ids=list(_PARITY_CASES))
+def test_endpoint_decisions_match_the_serving_side_policy(
+    probabilities: list[float],
+    options: dict[str, object],
+    expected: list[tuple[str, int | None]],
+) -> None:
+    """The engine detector reproduces the policy server VAD shipped with upstream."""
+    assert _drive(SileroVADConfig(**options), probabilities) == expected
+
+
+# --------------------------------------------------------------------------- #
+# The two rules that had drifted                                              #
+# --------------------------------------------------------------------------- #
+
+
+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.
+
+    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
+    who trails off unevenly never gets their turn committed.
+    """
+    config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=96, min_speech_duration_ms=32)
+    assert _drive(config, [0.9, 0.9, 0.1, 0.45, 0.1, 0.1]) == [("start", 0), ("stop", 160)]
+
+
+def test_audio_end_ms_includes_the_trailing_silence() -> None:
+    """OpenAI defines audio_end_ms as the end of the audio sent to the model.
+
+    That includes the silence spent deciding the turn was over, so it is the
+    frame boundary the detector stopped on, not where speech last was.
+    """
+    config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=100, min_speech_duration_ms=32)
+    events = _drive(config, [0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
+    stop_ms = next(ms for kind, ms in events if kind == "stop")
+    # 8 frames consumed when the timer expires: 8 * 512 / 16 kHz = 256 ms.
+    assert stop_ms == 256
+
+
+def test_reset_keeps_the_session_clock_and_clamps_the_prefix() -> None:
+    """A barge-in resets the detector but not the session's timeline.
+
+    Timestamps have to stay comparable with the rest of the session, and prefix
+    padding must not reach back into audio the reset discarded.
+    """
+    config = SileroVADConfig(threshold=0.5, prefix_padding_ms=300, silence_duration_ms=64, min_speech_duration_ms=32)
+    scores = iter([0.0] * 10 + [0.9, 0.9, 0.9])
+    vad = SileroStreamingVAD(config, frame_scorer=lambda _frame: next(scores))
+    for _ in range(10):
+        vad.process(np.zeros(FRAME, dtype=np.float32))
+
+    stream_start_ms = round(10 * FRAME * 1000 / SAMPLE_RATE_HZ)
+    vad.reset()
+    result = vad.process(np.zeros(FRAME * 3, dtype=np.float32))
+
+    assert result.speech_started
+    # Without the clock the start would land at 0; without the clamp, 300 ms of
+    # prefix would reach back before the reset.
+    assert result.speech_start_ms == stream_start_ms
+
+
+# --------------------------------------------------------------------------- #
+# Input handling                                                              #
+# --------------------------------------------------------------------------- #
+
+
+def _pcm16_b64(samples: np.ndarray) -> str:
+    return base64.b64encode((samples * 32767).astype("<i2").tobytes()).decode("ascii")
+
+
+def test_pcm16_input_is_accepted_alongside_float32() -> None:
+    config = SileroVADConfig(threshold=0.5)
+    tone = np.zeros(FRAME * 2, dtype=np.float32)
+    seen: list[int] = []
+    vad = SileroStreamingVAD(config, frame_scorer=lambda frame: seen.append(frame.size) or 0.0)
+    vad.process_base64(_pcm16_b64(tone), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ)
+    assert seen == [FRAME, FRAME]
+
+
+def test_an_odd_pcm16_byte_count_is_rejected() -> None:
+    vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=lambda _frame: 0.0)
+    with pytest.raises(ValueError, match="incomplete pcm16"):
+        vad.process_base64(base64.b64encode(b"\x00\x01\x02").decode(), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ)
+
+
+def test_resampling_keeps_the_frame_grid_from_drifting() -> None:
+    """The 24 kHz remainder carries across chunks, so frame boundaries cannot drift.
+
+    This is a weaker claim than upstream's ``StreamingAudioResampler``: the
+    interpolation is per chunk, so values near a chunk edge differ slightly.
+    What must hold is the frame *count* and alignment, which is what the
+    endpoint timestamps are derived from.
+    """
+    source = np.sin(np.linspace(0, 40 * np.pi, 24_000, dtype=np.float32))
+
+    def frames_for(chunk_sizes: list[int]) -> int:
+        frames = 0
+
+        def count(_frame: np.ndarray) -> float:
+            nonlocal frames
+            frames += 1
+            return 0.0
+
+        vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=count)
+        offset = 0
+        for size in chunk_sizes:
+            vad.process_base64(_pcm16_b64(source[offset : offset + size]), fmt="pcm16", sample_rate_hz=24_000)
+            offset += size
+        assert offset == source.size
+        return frames
+
+    whole = frames_for([source.size])
+    split = frames_for([7, 13, 511, 512, 513, source.size - 1_556])
+    assert whole == split == 16_000 // FRAME
+
+
+def test_a_sample_rate_change_mid_stream_is_rejected() -> None:
+    vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=lambda _frame: 0.0)
+    chunk = _pcm16_b64(np.zeros(1_024, dtype=np.float32))
+    vad.process_base64(chunk, fmt="pcm16", sample_rate_hz=24_000)
+    with pytest.raises(ValueError, match="sample rate cannot change"):
+        vad.process_base64(chunk, fmt="pcm16", sample_rate_hz=48_000)
+
+
+# --------------------------------------------------------------------------- #
+# Backend selection                                                           #
+# --------------------------------------------------------------------------- #
+
+
+def test_a_configured_model_path_that_does_not_exist_is_an_error() -> None:
+    """An operator who names an artifact gets told it is missing, not a fallback."""
+    provider = SileroVADBackendProvider(model_path="/nonexistent/silero_vad.onnx")
+    with pytest.raises(ServerVADUnavailableError, match="does not exist"):
+        provider.get()
+
+
+def test_a_configured_model_with_the_wrong_checksum_is_refused(tmp_path) -> None:
+    """A different Silero revision changes endpointing, so the pin is enforced."""
+    impostor = tmp_path / "silero_vad.onnx"
+    impostor.write_bytes(b"not the pinned model")
+    provider = SileroVADBackendProvider(model_path=str(impostor))
+    with pytest.raises(ServerVADUnavailableError, match="checksum mismatch"):
+        provider.get()
+    assert hashlib.sha256(impostor.read_bytes()).hexdigest() != SILERO_VAD_SHA256
+
+
+def test_the_onnx_backend_is_shared_by_every_session() -> None:
+    """One ONNX session for the whole engine: its state is passed in, not held."""
+    provider = SileroVADBackendProvider()
+    try:
+        backend = provider.get()
+    except ServerVADUnavailableError as exc:
+        pytest.skip(f"no Silero backend available here: {exc}")
+    if not isinstance(backend, SileroVADBackend):
+        pytest.skip("no local ONNX artifact; the torch fallback is per-session by design")
+
+    assert provider.get() is backend
+    # State is the caller's, so two sessions on one backend cannot interfere.
+    first, second = backend.new_state(), backend.new_state()
+    assert first is not second
+    probability, next_state = backend.infer(np.zeros(FRAME, dtype=np.float32), first)
+    assert 0.0 <= probability <= 1.0
+    assert next_state is not first
+
+
+def test_the_onnx_backend_holds_the_upstream_silero_v62_contract(monkeypatch, tmp_path) -> None:
+    """Pin the ONNX call shape without needing the real artifact or runtime.
+
+    The context carry is the part worth pinning: Silero v6.2 is fed 64 samples
+    of the previous frame ahead of the current 512, and that tail travels in the
+    caller's state rather than inside the session -- which is exactly what lets
+    one session serve every duplex session at once.
+    """
+    sessions: list[FakeInferenceSession] = []
+
+    class FakeSessionOptions:
+        inter_op_num_threads = 0
+        intra_op_num_threads = 0
+
+    class FakeInferenceSession:
+        def __init__(self, path: str, *, providers: list[str], sess_options: object) -> None:
+            self.providers = providers
+            self.sess_options = sess_options
+            self.calls: list[dict[str, np.ndarray]] = []
+            self.run_barrier: threading.Barrier | None = None
+            sessions.append(self)
+
+        def get_inputs(self) -> list[SimpleNamespace]:
+            return [SimpleNamespace(name=name) for name in ("input", "state", "sr")]
+
+        def run(self, _outputs: object, inputs: dict[str, np.ndarray]) -> list[np.ndarray]:
+            self.calls.append({name: np.array(value, copy=True) for name, value in inputs.items()})
+            if self.run_barrier is not None:
+                self.run_barrier.wait(timeout=5)
+            return [
+                np.asarray([[0.75]], dtype=np.float32),
+                np.full((2, 1, 128), len(self.calls), dtype=np.float32),
+            ]
+
+    monkeypatch.setitem(
+        sys.modules,
+        "onnxruntime",
+        SimpleNamespace(SessionOptions=FakeSessionOptions, InferenceSession=FakeInferenceSession),
+    )
+
+    model_path = tmp_path / "silero_vad.onnx"
+    model_path.write_bytes(b"fake-model")
+    backend = SileroVADBackend(model_path)
+    session = sessions[0]
+
+    # Single-threaded on purpose: many sessions share this one ORT session, so
+    # it must not fan out per call.
+    assert session.providers == ["CPUExecutionProvider"]
+    assert session.sess_options.inter_op_num_threads == 1
+    assert session.sess_options.intra_op_num_threads == 1
+    # Warm-up already ran: 64 context + 512 frame.
+    assert session.calls[0]["input"].shape == (1, 576)
+    assert session.calls[0]["state"].shape == (2, 1, 128)
+    assert session.calls[0]["sr"].item() == SAMPLE_RATE_HZ
+
+    state = backend.new_state()
+    first_frame = np.arange(backend.frame_samples, dtype=np.float32)
+    probability, state = backend.infer(first_frame, state)
+    first_call = session.calls[1]
+
+    assert probability == pytest.approx(0.75)
+    np.testing.assert_array_equal(
+        first_call["input"][:, : backend.context_samples],
+        np.zeros((1, backend.context_samples), dtype=np.float32),
+    )
+    np.testing.assert_array_equal(first_call["input"][:, backend.context_samples :], first_frame[None, :])
+
+    second_frame = -first_frame
+    backend.infer(second_frame, state)
+    second_call = session.calls[2]
+    # The previous frame's tail becomes this frame's context.
+    np.testing.assert_array_equal(
+        second_call["input"][:, : backend.context_samples],
+        first_frame[None, -backend.context_samples :],
+    )
+    np.testing.assert_array_equal(second_call["state"], np.full((2, 1, 128), 2, dtype=np.float32))
+
+    # Two sessions may enter the shared ORT session at the same time.
+    session.run_barrier = threading.Barrier(2)
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [
+            executor.submit(backend.infer, first_frame, backend.new_state()),
+            executor.submit(backend.infer, second_frame, backend.new_state()),
+        ]
+        assert [future.result()[0] for future in futures] == pytest.approx([0.75, 0.75])
+
+
+# --------------------------------------------------------------------------- #
+# Session payload validation                                                  #
+# --------------------------------------------------------------------------- #
+
+
+def test_an_unknown_server_vad_field_is_refused() -> None:
+    """A misspelled knob that silently does nothing is worse than a refusal.
+
+    The endpointing would still look configured while running on defaults.
+    """
+    error = validate_realtime_turn_detection({"turn_detection": {"type": "server_vad", "semantic_eagerness": "high"}})
+    assert error is not None
+    assert "semantic_eagerness" in error
+
+
+def test_the_documented_server_vad_fields_are_all_accepted() -> None:
+    assert (
+        validate_realtime_turn_detection(
+            {
+                "turn_detection": {
+                    "type": "server_vad",
+                    "threshold": 0.6,
+                    "prefix_padding_ms": 200,
+                    "silence_duration_ms": 400,
+                    "create_response": True,
+                    "interrupt_response": True,
+                    "min_speech_duration_ms": 120,
+                },
+                "overlap_policy": "barge_in_on_speech",
+            }
+        )
+        is None
+    )

tests/engine/test_async_omni_engine_abort_ack.py

@@ -54,7 +54,7 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture):
     rpc_q: queue.Queue = queue.Queue()
     engine = _make_engine(request_q, rpc_q)
     mocker.patch(
-        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        "vllm_omni.engine.omni_engine_base.uuid.uuid4",
         return_value=SimpleNamespace(hex="abort-rpc-1"),
     )
 
@@ -81,7 +81,7 @@ def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture):
     rpc_q: queue.Queue = queue.Queue()
     engine = _make_engine(request_q, rpc_q)
     mocker.patch(
-        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        "vllm_omni.engine.omni_engine_base.uuid.uuid4",
         return_value=SimpleNamespace(hex="abort-rpc-err"),
     )
 
@@ -109,7 +109,7 @@ def test_abort_async_times_out_without_result(mocker: MockerFixture):
     rpc_q: queue.Queue = queue.Queue()
     engine = _make_engine(request_q, rpc_q)
     mocker.patch(
-        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        "vllm_omni.engine.omni_engine_base.uuid.uuid4",
         return_value=SimpleNamespace(hex="abort-rpc-timeout"),
     )
 
@@ -139,7 +139,7 @@ def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture)
     rpc_q: queue.Queue = queue.Queue()
     engine = _make_engine(request_q, rpc_q)
     mocker.patch(
-        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        "vllm_omni.engine.omni_engine_base.uuid.uuid4",
         return_value=SimpleNamespace(hex="blocked-abort"),
     )
 

tests/engine/test_async_omni_engine_outputs.py

@@ -7,7 +7,6 @@ Focuses on the critical behavior: when the orchestrator thread dies,
 subsequent attempts to collect output raise RuntimeError.
 """
 
-import asyncio
 import queue
 import threading
 from concurrent.futures import ThreadPoolExecutor
@@ -18,16 +17,6 @@ from pytest_mock import MockerFixture
 
 from vllm_omni.engine.async_engine_utils import weak_shutdown_async_omni_engine
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-from vllm_omni.engine.duplex.control_client import (
-    DuplexControlClient,
-    DuplexControlRequestError,
-)
-from vllm_omni.engine.duplex.messages import (
-    AppendDuplexInputMessage,
-    DuplexControlResultMessage,
-    DuplexFence,
-    SignalDuplexTurnMessage,
-)
 from vllm_omni.engine.messages import (
     CollectiveRPCResultMessage,
     ErrorMessage,
@@ -166,9 +155,9 @@ def test_output_remains_on_shared_output_path(mocker: MockerFixture):
     raw_queue = queue.Queue()
     raw_queue.put_nowait(
         OutputMessage(
-            request_id="legacy-duplex-request",
+            request_id="shared-path-request",
             stage_id=0,
-            engine_outputs=OmniRequestOutput(request_id="legacy-duplex-request"),
+            engine_outputs=OmniRequestOutput(request_id="shared-path-request"),
             finished=False,
         )
     )
@@ -176,293 +165,7 @@ def test_output_remains_on_shared_output_path(mocker: MockerFixture):
 
     output = engine.try_get_output(timeout=0.01)
 
-    assert output.request_id == "legacy-duplex-request"
-
-
-def test_duplex_control_client_uses_engine_owned_correlated_transport():
-    engine = object.__new__(AsyncOmniEngine)
-    engine._duplex_control_client = None
-    engine._correlated_rpc_client = CorrelatedRpcClient(queue.Queue(), queue.Queue())
-
-    try:
-        client = engine._get_duplex_control_client()
-
-        assert isinstance(client, DuplexControlClient)
-        assert client._transport is engine._correlated_rpc_client
-        assert not hasattr(engine, "_submit_rpc_and_wait")
-    finally:
-        engine._correlated_rpc_client.close()
-
-
-def test_duplex_control_api_requires_explicit_fence():
-    engine = object.__new__(AsyncOmniEngine)
-
-    with pytest.raises(TypeError, match="fence"):
-        engine.open_duplex_session("sid")
-    with pytest.raises(TypeError, match="fence"):
-        engine.append_duplex_input("sid", mode="append_audio_chunk", payload={})
-    with pytest.raises(TypeError, match="fence"):
-        engine.signal_duplex_turn("sid", event="turn.end")
-    with pytest.raises(TypeError, match="fence"):
-        engine.close_duplex_session("sid")
-
-
-def test_open_duplex_session_waits_for_control_ack(mocker: MockerFixture):
-    request_q = queue.Queue()
-    rpc_q = queue.Queue()
-    control_result = DuplexControlResultMessage(
-        control_id="ctrl-1",
-        fence=DuplexFence("sid"),
-        operation="open",
-        session_id="sid",
-        ok=False,
-        stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False}}],
-        unsupported_count=1,
-        error_count=0,
-    )
-
-    engine = object.__new__(AsyncOmniEngine)
-    engine.request_queue = SimpleNamespace(sync_q=request_q)
-    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
-    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
-    engine._duplex_control_client = None
-    mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-1"))
-
-    with ThreadPoolExecutor(max_workers=1) as executor:
-        pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1)
-        msg = request_q.get(timeout=1)
-        assert msg.type == "open_duplex_session"
-        assert msg.control_id == "ctrl-1"
-        rpc_q.put(control_result)
-        with pytest.raises(DuplexControlRequestError) as exc_info:
-            pending.result(timeout=1)
-
-    assert exc_info.value.result["unsupported_count"] == 1
-    assert exc_info.value.result["stage_results"][0]["result"]["supported"] is False
-    engine._correlated_rpc_client.close()
-
-
-def test_signal_duplex_turn_message_carries_next_fence(mocker: MockerFixture):
-    request_q = queue.Queue()
-    rpc_q = queue.Queue()
-    cancelled_fence = DuplexFence("sid-cancel", epoch=4, turn_id=2)
-    next_fence = DuplexFence("sid-cancel", epoch=5, turn_id=2)
-    control_result = DuplexControlResultMessage(
-        control_id="ctrl-signal",
-        fence=cancelled_fence,
-        operation="signal",
-        session_id="sid-cancel",
-        ok=True,
-        stage_results=[],
-    )
-
-    engine = object.__new__(AsyncOmniEngine)
-    engine.request_queue = SimpleNamespace(sync_q=request_q)
-    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
-    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
-    engine._duplex_control_client = None
-    mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-signal"))
-
-    with ThreadPoolExecutor(max_workers=1) as executor:
-        pending = executor.submit(
-            engine.signal_duplex_turn,
-            "sid-cancel",
-            event="input.cancel",
-            fence=cancelled_fence,
-            next_fence=next_fence,
-            session_config={"temperature": 0.0},
-            timeout=1,
-        )
-        msg = request_q.get(timeout=1)
-        assert isinstance(msg, SignalDuplexTurnMessage)
-        assert msg.fence == cancelled_fence
-        assert msg.next_fence == next_fence
-        assert msg.session_config == {"temperature": 0.0}
-        rpc_q.put(control_result)
-        assert pending.result(timeout=1)["ok"] is True
-
-    engine._correlated_rpc_client.close()
-
-
-@pytest.mark.asyncio
-async def test_cancelled_async_append_does_not_stop_executor_callable(mocker: MockerFixture):
-    entered = threading.Event()
-    release = threading.Event()
-    completed = threading.Event()
-    engine = object.__new__(AsyncOmniEngine)
-
-    def blocking_append(*args, **kwargs):
-        del args, kwargs
-        entered.set()
-        release.wait(timeout=1)
-        completed.set()
-        return {}
-
-    mocker.patch.object(engine, "append_duplex_input", side_effect=blocking_append)
-    task = asyncio.create_task(
-        engine.append_duplex_input_async(
-            "sid-executor-cancel",
-            mode="append_audio_chunk",
-            payload=b"audio",
-            fence=DuplexFence("sid-executor-cancel"),
-        )
-    )
-    assert await asyncio.to_thread(entered.wait, 1)
-
-    task.cancel()
-    with pytest.raises(asyncio.CancelledError):
-        await task
-    release.set()
-
-    assert await asyncio.to_thread(completed.wait, 1)
-
-
-def test_duplex_controls_route_out_of_order_without_global_lock():
-    request_q = queue.Queue()
-    rpc_q = queue.Queue()
-    engine = object.__new__(AsyncOmniEngine)
-    engine.request_queue = SimpleNamespace(sync_q=request_q)
-    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
-    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
-    engine._duplex_control_client = None
-    first_fence = DuplexFence("sid", epoch=1)
-    second_fence = DuplexFence("sid", epoch=2)
-    first = AppendDuplexInputMessage(
-        control_id="first",
-        fence=first_fence,
-        session_id="sid",
-        expected_epoch=1,
-        mode="audio",
-        payload=b"first",
-    )
-    second = AppendDuplexInputMessage(
-        control_id="second",
-        fence=second_fence,
-        session_id="sid",
-        expected_epoch=2,
-        mode="audio",
-        payload=b"second",
-    )
-
-    with ThreadPoolExecutor(max_workers=2) as executor:
-        client = engine._get_duplex_control_client()
-        first_result = executor.submit(client.execute, first, timeout=1)
-        second_result = executor.submit(client.execute, second, timeout=1)
-        submitted = {request_q.get(timeout=1).control_id, request_q.get(timeout=1).control_id}
-        assert submitted == {"first", "second"}
-        rpc_q.put(
-            DuplexControlResultMessage(
-                control_id="second",
-                fence=second_fence,
-                operation="append",
-                session_id="sid",
-                ok=True,
-                stage_results=[],
-            )
-        )
-        rpc_q.put(
-            DuplexControlResultMessage(
-                control_id="first",
-                fence=first_fence,
-                operation="append",
-                session_id="sid",
-                ok=True,
-                stage_results=[],
-            )
-        )
-
-        assert first_result.result(timeout=1)["fence"] == first_fence
-        assert second_result.result(timeout=1)["fence"] == second_fence
-
-    assert not hasattr(engine, "_rpc_lock")
-    engine._correlated_rpc_client.close()
-
-
-def test_terminal_rpc_router_rejects_without_enqueuing_new_control():
-    request_q = queue.Queue()
-    rpc_q = queue.Queue()
-    engine = object.__new__(AsyncOmniEngine)
-    engine.request_queue = SimpleNamespace(sync_q=request_q)
-    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
-    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
-    engine._duplex_control_client = None
-    pending_message = AppendDuplexInputMessage(
-        control_id="pending",
-        fence=DuplexFence("sid"),
-        session_id="sid",
-        mode="audio",
-        payload=b"audio",
-    )
-    message = AppendDuplexInputMessage(
-        control_id="after-fatal",
-        fence=DuplexFence("sid"),
-        session_id="sid",
-        mode="audio",
-        payload=b"audio",
-    )
-
-    try:
-        with ThreadPoolExecutor(max_workers=1) as executor:
-            pending = executor.submit(engine._get_duplex_control_client().execute, pending_message, timeout=1)
-            assert request_q.get(timeout=1).control_id == "pending"
-            rpc_q.put(ErrorMessage(error="orchestrator failed", fatal=True))
-            with pytest.raises(RuntimeError, match="orchestrator failed"):
-                pending.result(timeout=1)
-
-        with pytest.raises(RuntimeError, match="orchestrator failed"):
-            engine._get_duplex_control_client().execute(message, timeout=1)
-
-        assert request_q.empty()
-    finally:
-        engine._correlated_rpc_client.close()
-
-
-def test_collective_and_duplex_results_share_router_without_swapping(mocker: MockerFixture):
-    request_q = queue.Queue()
-    rpc_q = queue.Queue()
-    engine = object.__new__(AsyncOmniEngine)
-    engine.request_queue = SimpleNamespace(sync_q=request_q)
-    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
-    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
-    engine._duplex_control_client = None
-    fence = DuplexFence("sid")
-    duplex = AppendDuplexInputMessage(
-        control_id="duplex-control",
-        fence=fence,
-        session_id="sid",
-        mode="audio",
-        payload=b"audio",
-    )
-    mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="collective-rpc"))
-
-    with ThreadPoolExecutor(max_workers=2) as executor:
-        duplex_result = executor.submit(engine._get_duplex_control_client().execute, duplex, timeout=1)
-        collective_result = executor.submit(engine.collective_rpc, "health", timeout=1)
-        submitted_types = {request_q.get(timeout=1).type, request_q.get(timeout=1).type}
-        assert submitted_types == {"append_duplex_input", "collective_rpc"}
-        rpc_q.put(
-            CollectiveRPCResultMessage(
-                rpc_id="collective-rpc",
-                method="health",
-                stage_ids=[0],
-                results=["healthy"],
-            )
-        )
-        rpc_q.put(
-            DuplexControlResultMessage(
-                control_id="duplex-control",
-                fence=fence,
-                operation="append",
-                session_id="sid",
-                ok=True,
-                stage_results=[],
-            )
-        )
-
-        assert collective_result.result(timeout=1) == ["healthy"]
-        assert duplex_result.result(timeout=1)["fence"] == fence
-
-    engine._correlated_rpc_client.close()
+    assert output.request_id == "shared-path-request"
 
 
 def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixture):
@@ -483,8 +186,7 @@ def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixtu
     engine.request_queue = SimpleNamespace(sync_q=request_q)
     engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
     engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
-    engine._duplex_control_client = None
-    mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc"))
+    mocker.patch("vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc"))
 
     with ThreadPoolExecutor(max_workers=1) as executor:
         pending = executor.submit(engine.collective_rpc, "health", timeout=1)
@@ -506,36 +208,6 @@ def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixtu
     engine._correlated_rpc_client.close()
 
 
-def test_open_duplex_session_raises_on_stage_control_error(mocker: MockerFixture):
-    request_q = queue.Queue()
-    rpc_q = queue.Queue()
-    control_result = DuplexControlResultMessage(
-        control_id="ctrl-error",
-        fence=DuplexFence("sid"),
-        operation="open",
-        session_id="sid",
-        ok=False,
-        stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False, "error": "boom"}}],
-        unsupported_count=1,
-        error_count=1,
-    )
-
-    engine = object.__new__(AsyncOmniEngine)
-    engine.request_queue = SimpleNamespace(sync_q=request_q)
-    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
-    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
-    engine._duplex_control_client = None
-    mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-error"))
-
-    with ThreadPoolExecutor(max_workers=1) as executor:
-        pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1)
-        request_q.get(timeout=1)
-        rpc_q.put(control_result)
-        with pytest.raises(RuntimeError, match="duplex open failed"):
-            pending.result(timeout=1)
-    engine._correlated_rpc_client.close()
-
-
 def test_shutdown_does_not_race_runtime_cleanup_with_live_orchestrator(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
     engine._shutdown_called = False

tests/engine/test_async_omni_engine_stage_init.py

@@ -12,7 +12,7 @@ import types
 import pytest
 
 from vllm_omni.diffusion.data import AttentionConfig
-from vllm_omni.engine import async_omni_engine as async_omni_engine_module
+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_init_utils import (
     LogicalStageInitPlan,
@@ -188,7 +188,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):
-    import vllm_omni.engine.async_omni_engine as engine_mod
+    import vllm_omni.engine.omni_engine_base as engine_mod
 
     engine = object.__new__(AsyncOmniEngine)
     engine.stage_configs = [types.SimpleNamespace()]

tests/engine/test_diffusion_streaming_output_kwarg.py

@@ -33,7 +33,7 @@ def _resolve(monkeypatch, kwargs: dict[str, Any]) -> dict[str, Any]:
         return "config.yaml", [], None
 
     monkeypatch.setattr(
-        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
+        "vllm_omni.engine.omni_engine_base.load_and_resolve_stage_configs",
         _fake_load,
     )
     engine = SimpleNamespace(

tests/engine/test_duplex_import_boundary.py

@@ -45,27 +45,62 @@ def test_isolated_import_ignores_conflicting_gpu_visibility(monkeypatch: pytest.
     _assert_isolated_import_succeeds("import vllm_omni.outputs")
 
 
-def test_stable_engine_imports_load_duplex_kernel_eagerly() -> None:
-    # The duplex kernel is imported eagerly by the stable engine modules.
-    # Model-specific duplex adapters must still load only via the
-    # dotted-string plugin paths, and the client-side vllm_omni.clients
-    # package must never enter the stable runtime import graph.
+def test_turn_based_stack_has_no_duplex_vocabulary() -> None:
+    # The generic bases and the turn-based siblings (AsyncOmni / AsyncOmniEngine /
+    # Orchestrator) carry no duplex code: importing them must not load the
+    # engine-side duplex package, the model plugins or the client library.
     _assert_isolated_import_succeeds("""
 import sys
 
 import vllm_omni.engine.async_omni_engine
+import vllm_omni.engine.omni_engine_base
 import vllm_omni.engine.orchestrator
 import vllm_omni.entrypoints.async_omni
+import vllm_omni.entrypoints.async_omni_base
+
+forbidden_prefixes = (
+    "vllm_omni.engine.duplex",
+    "vllm_omni.engine.duplex_omni_engine",
+    "vllm_omni.engine.duplex_orchestrator",
+    "vllm_omni.entrypoints.duplex",
+    "vllm_omni.entrypoints.duplex_omni",
+    "vllm_omni.clients",
+    "vllm_omni.model_executor.models.minicpmo_4_5.duplex",
+    "vllm_omni.model_executor.models.nemotron_voicechat.duplex",
+    "vllm_omni.model_executor.models.personaplex.duplex",
+)
+loaded = sorted(
+    name
+    for name in sys.modules
+    if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden_prefixes)
+)
+if loaded:
+    raise SystemExit("turn-based imports loaded duplex modules: " + ", ".join(loaded))
+""")
+
+
+def test_duplex_stack_loads_the_framework_but_not_the_model_plugins() -> None:
+    # The duplex siblings import the engine-side duplex framework eagerly;
+    # model plugins still load only through PipelineConfig.duplex_plugin, and
+    # the client-side package never enters the server import graph.
+    _assert_isolated_import_succeeds("""
+import sys
+
+import vllm_omni.engine.duplex_omni_engine
+import vllm_omni.engine.duplex_orchestrator
+import vllm_omni.entrypoints.duplex_omni
+import vllm_omni.entrypoints.duplex.serving
 
 expected_eager = (
-    "vllm_omni.engine.duplex.contracts",
-    "vllm_omni.engine.duplex.control_plane",
-    "vllm_omni.entrypoints.duplex_request_client",
-    "vllm_omni.outputs.duplex",
+    "vllm_omni.engine.duplex.commands",
+    "vllm_omni.engine.duplex.events",
+    "vllm_omni.engine.duplex.session.engine_session",
+    "vllm_omni.engine.duplex.session.manager",
+    "vllm_omni.engine.duplex.plugin",
 )
 missing = sorted(name for name in expected_eager if name not in sys.modules)
 if missing:
-    raise SystemExit("duplex kernel modules not imported eagerly: " + ", ".join(missing))
+    raise SystemExit("duplex framework modules not imported eagerly: " + ", ".join(missing))
 
 forbidden_prefixes = (
     "vllm_omni.clients",
@@ -79,7 +114,7 @@ loaded = sorted(
     if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden_prefixes)
 )
 if loaded:
-    raise SystemExit("stable imports loaded plugin-only duplex modules: " + ", ".join(loaded))
+    raise SystemExit("duplex imports loaded plugin-only or client modules: " + ", ".join(loaded))
 """)
 
 
@@ -190,16 +225,32 @@ def test_engine_duplex_uses_canonical_contract_module_names() -> None:
     core_dir = REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "core"
 
     for name in (
+        "commands.py",
         "contracts.py",
-        "lease.py",
+        "events.py",
         "messages.py",
-        "session.py",
-        "control_plane.py",
-        "control_client.py",
-        "runtime.py",
+        "plugin.py",
         "intermediate.py",
     ):
         assert (engine_dir / name).is_file()
+    # one session and everything that runs it lives in the session subpackage
+    for name in (
+        "__init__.py",
+        "engine_session.py",
+        "manager.py",
+        "runner.py",
+        "context.py",
+        "emitter.py",
+        "model_channel.py",
+        "control.py",
+        "append_task.py",
+        "helpers.py",
+        "lease.py",
+    ):
+        assert (engine_dir / "session" / name).is_file()
+    assert not (engine_dir / "session.py").exists()
+    for removed in ("control_plane.py", "control_client.py", "runtime.py"):
+        assert not (engine_dir / removed).exists()
     # the duplex_ prefix is dropped inside the duplex package
     assert not (engine_dir / "duplex_session.py").exists()
     assert not (engine_dir / "duplex_lease.py").exists()

tests/engine/test_duplex_omni_engine.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""``DuplexOmniEngine``: the session message surface over the generic engine base."""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vllm_omni.engine.duplex.commands import Heartbeat
+from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+from vllm_omni.engine.duplex.messages import (
+    CloseDuplexSessionMessage,
+    DuplexControlResultMessage,
+    DuplexSessionCommandMessage,
+    DuplexSessionError,
+    OpenDuplexSessionMessage,
+    ResumeDuplexSessionMessage,
+    TouchDuplexSessionMessage,
+)
+from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeRpcClient:
+    def __init__(self, result: Any) -> None:
+        self.result = result
+        self.calls: list[tuple[tuple[str, str], Any, float | None]] = []
+        #: Control RPCs must block rather than raise a raw queue.Full.
+        self.block_on_submit_flags: list[bool] = []
+
+    def execute(self, key, message, *, timeout=None, timeout_message=None, block_on_submit=False):
+        self.calls.append((key, message, timeout))
+        self.block_on_submit_flags.append(block_on_submit)
+        if isinstance(self.result, BaseException):
+            raise self.result
+        return self.result
+
+
+class _FakeSyncQueue:
+    def __init__(self, *, full: bool = False) -> None:
+        self.items: list[Any] = []
+        self.full = full
+
+    def put(self, item: Any, timeout: float | None = None) -> None:
+        if self.full:
+            raise queue.Full
+        self.items.append(item)
+
+
+def _engine(result: Any, *, alive: bool = True, full: bool = False) -> DuplexOmniEngine:
+    engine = object.__new__(DuplexOmniEngine)
+    engine._correlated_rpc_client = _FakeRpcClient(result)
+    engine.request_queue = SimpleNamespace(sync_q=_FakeSyncQueue(full=full))
+    engine.is_alive = lambda: alive
+    engine.plugin = None
+    return engine
+
+
+def _ok(operation: str, **fields: Any) -> DuplexControlResultMessage:
+    return DuplexControlResultMessage(control_id="c", operation=operation, session_id="sid", ok=True, **fields)
+
+
+@pytest.mark.asyncio
+async def test_open_session_runs_the_correlated_rpc_with_the_typed_config() -> None:
+    result = _ok("open", capabilities=DuplexCapabilities(), public_session={"id": "sid"}, lease_generation=0)
+    engine = _engine(result)
+    config = DuplexSessionConfig(model="m")
+
+    returned = await engine.open_session_async("sid", config, timeout=3.0)
+
+    assert returned is result
+    ((key, message, timeout),) = engine.rpc_client.calls
+    assert key == ("duplex", message.control_id)
+    assert isinstance(message, OpenDuplexSessionMessage)
+    assert message.session_id == "sid" and message.session_config is config
+    assert timeout == 3.0
+    # A momentarily full request queue is backpressure, not a failed open: the
+    # control RPC blocks rather than raising an untyped queue.Full at the caller.
+    assert engine.rpc_client.block_on_submit_flags == [True]
+
+
+@pytest.mark.asyncio
+async def test_close_resume_and_touch_build_their_messages_without_incarnation() -> None:
+    engine = _engine(_ok("close"))
+    await engine.close_session_async("sid", reason="bye")
+    await engine.resume_session_async("sid", expected_lease_generation=4)
+    await engine.touch_session_async("sid", activity="detach")
+
+    close, resume, touch = (call[1] for call in engine.rpc_client.calls)
+    assert isinstance(close, CloseDuplexSessionMessage) and close.reason == "bye"
+    assert isinstance(resume, ResumeDuplexSessionMessage) and resume.expected_lease_generation == 4
+    assert isinstance(touch, TouchDuplexSessionMessage) and touch.activity == "detach"
+    for message in (close, resume, touch):
+        assert message.session_id == "sid"
+        assert not hasattr(message, "incarnation")
+
+
+@pytest.mark.asyncio
+async def test_failed_control_result_maps_to_a_typed_session_error() -> None:
+    engine = _engine(
+        DuplexControlResultMessage(
+            control_id="c",
+            operation="open",
+            session_id="sid",
+            ok=False,
+            error_code="resource_exhausted",
+            error_message="duplex_session_capacity_exhausted: limit=1",
+            error_retryable=True,
+        )
+    )
+    with pytest.raises(DuplexSessionError) as excinfo:
+        await engine.open_session_async("sid", DuplexSessionConfig())
+    assert excinfo.value.code == "resource_exhausted"
+    assert excinfo.value.retryable is True
+    assert excinfo.value.session_id == "sid"
+
+
+@pytest.mark.asyncio
+async def test_engine_error_and_timeout_map_to_session_errors() -> None:
+    # The correlated RPC client raises RuntimeError for a closed router or a
+    # terminal orchestrator error; both become a typed engine_error.
+    engine = _engine(RuntimeError("orchestrator exploded"))
+    with pytest.raises(DuplexSessionError, match="orchestrator exploded") as excinfo:
+        await engine.close_session_async("sid")
+    assert excinfo.value.code == "engine_error"
+
+    engine = _engine(TimeoutError("duplex open timed out for session sid"))
+    with pytest.raises(DuplexSessionError) as excinfo:
+        await engine.open_session_async("sid", DuplexSessionConfig())
+    assert excinfo.value.code == "timeout" and excinfo.value.retryable is True
+
+
+@pytest.mark.asyncio
+async def test_submit_command_is_one_way_on_the_request_queue() -> None:
+    engine = _engine(None)
+    command = Heartbeat(event_id="evt-1")
+
+    await engine.submit_command_async("sid", command)
+
+    (message,) = engine.request_queue.sync_q.items
+    assert isinstance(message, DuplexSessionCommandMessage)
+    assert message.session_id == "sid" and message.command is command
+    assert engine.rpc_client.calls == []
+
+
+@pytest.mark.asyncio
+async def test_submit_command_rejects_dead_engine_and_full_queue() -> None:
+    with pytest.raises(DuplexSessionError) as excinfo:
+        await _engine(None, alive=False).submit_command_async("sid", Heartbeat())
+    assert excinfo.value.code == "engine_dead"
+
+    with pytest.raises(DuplexSessionError) as excinfo:
+        await _engine(None, full=True).submit_command_async("sid", Heartbeat())
+    assert excinfo.value.code == "engine_backpressure" and excinfo.value.retryable is True
+
+
+def test_duplex_capabilities_come_from_the_loaded_plugin() -> None:
+    engine = object.__new__(DuplexOmniEngine)
+    engine.plugin = None
+    with pytest.raises(RuntimeError, match="not loaded"):
+        engine.duplex_capabilities
+    capabilities = DuplexCapabilities(supports_barge_in=False)
+    engine.plugin = SimpleNamespace(capabilities=lambda *, max_sessions: capabilities)
+    engine.duplex_session_config = SimpleNamespace(max_sessions=3)
+    assert engine.duplex_capabilities is capabilities
+
+
+def test_validate_deployment_requires_a_plugin_and_duplex_session_mode() -> None:
+    engine = object.__new__(DuplexOmniEngine)
+    engine.model = "plain-model"
+    engine.pipeline_config = SimpleNamespace(duplex_plugin=None)
+    engine.deploy_config = SimpleNamespace(session_mode="duplex")
+    with pytest.raises(ValueError, match="not a duplex model"):
+        engine._validate_deployment()
+
+    engine.pipeline_config = SimpleNamespace(duplex_plugin="pkg.mod.Plugin")
+    engine.deploy_config = SimpleNamespace(session_mode="turn")
+    with pytest.raises(ValueError, match="session_mode: duplex"):
+        engine._validate_deployment()
+
+    engine.deploy_config = None
+    with pytest.raises(ValueError, match="none resolved"):
+        engine._validate_deployment()
+
+
+@pytest.mark.asyncio
+async def test_async_wrappers_run_the_blocking_calls_off_the_event_loop() -> None:
+    loop = asyncio.get_running_loop()
+    engine = _engine(_ok("touch"))
+    seen: list[bool] = []
+    original = engine.rpc_client.execute
+
+    def execute(*args, **kwargs):
+        # Called from the default executor thread, never on the event loop thread.
+        try:
+            seen.append(asyncio.get_running_loop() is loop)
+        except RuntimeError:
+            seen.append(False)
+        return original(*args, **kwargs)
+
+    engine.rpc_client.execute = execute
+    await engine.touch_session_async("sid", activity="heartbeat")
+    assert seen == [False]

tests/engine/test_duplex_orchestrator.py

@@ -0,0 +1,517 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""``DuplexOrchestrator``: the template seams and the stage port around ``OrchestratorBase``."""
+
+from __future__ import annotations
+
+import asyncio
+import struct
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+import pytest
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from tests.engine.test_orchestrator import (
+    FakeOutputProcessor,
+    FakeRunningCounter,
+    FakeStageClient,
+    _build_stage_pools,
+)
+from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
+from vllm_omni.engine.duplex import commands
+from vllm_omni.engine.duplex.config import DuplexSessionConfig, DuplexSessionState
+from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id
+from vllm_omni.engine.duplex.messages import (
+    CloseDuplexSessionMessage,
+    DuplexControlResultMessage,
+    DuplexSessionCommandMessage,
+    OpenDuplexSessionMessage,
+)
+from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator, DuplexOrchestratorRequestState
+from vllm_omni.engine.messages import AbortRequestMessage, ShutdownRequestMessage
+from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import MiniCPMO45DuplexPlugin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SESSION_ID = "duplex-orch"
+
+
+def _encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None:
+    del sample_rate_hz, response_format, speed
+    samples = int(np.asarray(audio, dtype=np.float32).size) if audio is not None else 0
+    return f"wav-{samples}" if samples > 0 else None
+
+
+def _stage_configs(count: int) -> list[object]:
+    return [SimpleNamespace(model_config=SimpleNamespace(max_model_len=4096)) for _ in range(count)]
+
+
+def _build(
+    *,
+    stages: int = 1,
+    running_counter: FakeRunningCounter | None = None,
+    runtime_config: DuplexSessionRuntimeConfig | None = None,
+) -> tuple[DuplexOrchestrator, list[FakeStageClient], asyncio.Queue, asyncio.Queue]:
+    clients = [FakeStageClient(stage_type="llm", final_output=index == stages - 1) for index in range(stages)]
+    pools = _build_stage_pools(
+        [[client] for client in clients],
+        output_processors=[FakeOutputProcessor() for _ in clients],
+        stage_vllm_configs=_stage_configs(stages),
+    )
+    rpc_q: asyncio.Queue = asyncio.Queue()
+    output_q: asyncio.Queue = asyncio.Queue()
+    orchestrator = DuplexOrchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=output_q,
+        rpc_async_queue=rpc_q,
+        stage_pools=pools,
+        running_counter=running_counter,
+        plugin=MiniCPMO45DuplexPlugin(_encode_audio),
+        duplex_session_config=runtime_config or DuplexSessionRuntimeConfig(reaper_interval_s=0.01),
+        model_config=None,
+    )
+    return orchestrator, clients, rpc_q, output_q
+
+
+def _open_message(extra_body: dict[str, object] | None = None) -> OpenDuplexSessionMessage:
+    return OpenDuplexSessionMessage(
+        control_id=f"open-{SESSION_ID}",
+        session_id=SESSION_ID,
+        session_config=DuplexSessionConfig(
+            model="openbmb/MiniCPM-o-4_5",
+            modalities=["text"],
+            instructions="public instructions",
+            extra_body={"auto_response": True, **(extra_body or {})},
+        ),
+    )
+
+
+async def _open(orchestrator: DuplexOrchestrator, rpc_q: asyncio.Queue, **kwargs: Any) -> DuplexControlResultMessage:
+    await orchestrator._dispatch_message(_open_message(**kwargs))
+    result = await asyncio.wait_for(rpc_q.get(), timeout=2.0)
+    assert isinstance(result, DuplexControlResultMessage)
+    return result
+
+
+async def _close(orchestrator: DuplexOrchestrator, rpc_q: asyncio.Queue) -> DuplexControlResultMessage:
+    await orchestrator._dispatch_message(
+        CloseDuplexSessionMessage(control_id=f"close-{SESSION_ID}", session_id=SESSION_ID, reason="test")
+    )
+    result = await asyncio.wait_for(rpc_q.get(), timeout=2.0)
+    await _settle(orchestrator)
+    return result
+
+
+def _stage0_request_id(epoch: int = 0) -> str:
+    return duplex_resource_request_id(DuplexFence(SESSION_ID, epoch=epoch), "stage0")
+
+
+def _append_audio(samples: int = 16000) -> commands.AppendAudio:
+    return commands.AppendAudio(
+        audio=struct.pack(f"<{samples}f", *([0.05] * samples)),
+        format="pcm_f32le",
+        sample_rate_hz=16000,
+        is_speech=True,
+    )
+
+
+async def _submit(orchestrator: DuplexOrchestrator, command: commands.DuplexCommand) -> None:
+    assert await orchestrator._dispatch_message(DuplexSessionCommandMessage(session_id=SESSION_ID, command=command))
+    await _settle(orchestrator)
+
+
+async def _settle(orchestrator: DuplexOrchestrator, *, timeout_s: float = 3.0) -> None:
+    """Let the session runner drain its mailbox and append tasks."""
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout_s
+    quiet = 0
+    while loop.time() < deadline:
+        runner = orchestrator.session_manager.runners.get(SESSION_ID)
+        busy = runner is not None and (
+            not runner._mailbox.empty()
+            or any(not task.done() for task in runner.tasks.append_tasks)
+            or any(not task.done() for task in runner._background_tasks)
+        )
+        busy = busy or any(not task.done() for task in orchestrator.session_manager._dispatched_control_tasks)
+        quiet = 0 if busy else quiet + 1
+        if quiet >= 5:
+            return
+        await asyncio.sleep(0.01)
+
+
+def _tts_output(request_id: str, *, samples: int = 24000, text: str = "hello") -> SimpleNamespace:
+    return SimpleNamespace(
+        request_id=request_id,
+        finished=False,
+        outputs=[SimpleNamespace(text=text, token_ids=[], multimodal_output={})],
+        multimodal_output={
+            "audio": np.zeros(samples, dtype=np.float32),
+            "sr": 24000,
+            "meta.duplex_turn_id": np.array([0], dtype=np.int32),
+            "meta.duplex_epoch": np.array([0], dtype=np.int32),
+        },
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Message dispatch and session lifecycle                                      #
+# --------------------------------------------------------------------------- #
+
+
+def test_turn_based_orchestrator_has_no_session_manager() -> None:
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=[],
+    )
+    assert not hasattr(orchestrator, "session_manager")
+
+
+@pytest.mark.asyncio
+async def test_duplex_messages_are_dispatched_and_generic_ones_are_not() -> None:
+    orchestrator, _, rpc_q, _ = _build()
+    assert await orchestrator._dispatch_message(AbortRequestMessage(request_ids=["x"])) is False
+    assert await orchestrator._dispatch_message(ShutdownRequestMessage()) is False
+    assert await orchestrator._dispatch_message(_open_message()) is True
+    assert (await asyncio.wait_for(rpc_q.get(), timeout=2.0)).ok is True
+    await orchestrator.session_manager.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_open_preregisters_the_stage0_request_and_close_releases_it() -> None:
+    counter = FakeRunningCounter()
+    orchestrator, clients, rpc_q, _ = _build(running_counter=counter)
+    result = await _open(orchestrator, rpc_q)
+    assert result.ok is True
+    assert result.public_session["id"] == SESSION_ID
+    assert result.capabilities is not None and result.capabilities.supports_input_append
+
+    request_id = _stage0_request_id()
+    request_state = orchestrator.request_states[request_id]
+    assert isinstance(request_state, DuplexOrchestratorRequestState)
+    assert request_state.session_owned is True
+    assert request_state.session_id == SESSION_ID
+    assert request_state.fence == DuplexFence(SESSION_ID)
+    assert request_state.streaming.enabled is True
+    # Preregistration reserves the id; nothing is running until an append submits.
+    assert counter.value == 0
+    assert clients[0].add_request_calls == []
+    assert orchestrator.session_manager.runner_for_request_id(request_id) is not None
+
+    closed = await _close(orchestrator, rpc_q)
+    assert closed.ok is True
+    assert request_id not in orchestrator.request_states
+    assert orchestrator.session_manager.runner_for_request_id(request_id) is None
+    assert counter.value == 0
+    assert orchestrator.session_manager.active_count() == 0
+
+
+@pytest.mark.asyncio
+async def test_open_failure_rolls_back_session_and_reserved_request() -> None:
+    orchestrator, _, rpc_q, _ = _build()
+    # Server-owned runtime keys may not come from the client.
+    result = await _open(orchestrator, rpc_q, extra_body={"duplex_stage_max_tokens": {"0": 99}})
+
+    assert result.ok is False
+    assert result.error_code == "invalid_duplex_runtime_config"
+    assert orchestrator.request_states == {}
+    assert orchestrator.session_manager.active_count() == 0
+
+
+@pytest.mark.asyncio
+async def test_open_fails_cleanly_when_a_stage_has_no_live_replica() -> None:
+    orchestrator, _, rpc_q, _ = _build(stages=2)
+    orchestrator.stage_pools[1].evict_replica(0)
+
+    result = await _open(orchestrator, rpc_q)
+
+    assert result.ok is False
+    assert "stage 1 has no live replica" in (result.error_message or "")
+    assert orchestrator.request_states == {}
+    assert orchestrator.session_manager.active_count() == 0
+
+
+@pytest.mark.asyncio
+async def test_bridge_state_keeps_public_and_runtime_config_separate() -> None:
+    orchestrator, _, rpc_q, _ = _build()
+    await _open(orchestrator, rpc_q)
+    request_state = orchestrator.request_states[_stage0_request_id()]
+    bridge = request_state.streaming.bridge_states["duplex"]
+
+    assert bridge["session_id"] == SESSION_ID
+    assert bridge["fence"] == DuplexFence(SESSION_ID)
+    assert (bridge["epoch"], bridge["turn_id"], bridge["model_turn_id"]) == (0, 0, 0)
+    assert "incarnation" not in bridge
+    assert bridge["session_config"]["instructions"] == "public instructions"
+    assert "duplex_stage_sampling_params" not in bridge["session_config"]
+    assert bridge["runtime_config"]["instructions"] == "public instructions"
+    assert bridge["runtime_config"]["duplex_stage_max_tokens"] == {"0": 20, "1": 8192}
+    # The plugin's sampling policy is applied to the preregistered request.
+    assert request_state.sampling_params_list[0].max_tokens == 20
+    await orchestrator.session_manager.shutdown()
+
+
+# --------------------------------------------------------------------------- #
+# Stage port                                                                  #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.asyncio
+async def test_append_submits_the_resumable_stage0_request_and_counts_it_running() -> None:
+    counter = FakeRunningCounter()
+    orchestrator, clients, rpc_q, _ = _build(running_counter=counter)
+    await _open(orchestrator, rpc_q)
+    request_id = _stage0_request_id()
+
+    await _submit(orchestrator, _append_audio())
+    assert len(clients[0].add_request_calls) == 1
+    submitted = clients[0].add_request_calls[0][0]
+    assert submitted.request_id == request_id
+    assert submitted.resumable is True
+    assert submitted.sampling_params.max_tokens == 20
+    request_state = orchestrator.request_states[request_id]
+    assert request_state.stage_fences[0] == DuplexFence(SESSION_ID)
+    assert 0 in request_state.stage_submit_ts
+    assert counter.value == 1
+    session = orchestrator.session_manager.get(SESSION_ID)
+    assert session is not None and session.stage_request_submitted(0, request_id)
+
+    # Later units update the same resumable request; the counter stays at one.
+    await _submit(orchestrator, _append_audio())
+    assert len(clients[0].add_request_calls) == 2
+    assert clients[0].add_request_calls[1][0].request_id == request_id
+    assert counter.value == 1
+
+    await _close(orchestrator, rpc_q)
+    assert clients[0].abort_calls == [[request_id]]
+    assert counter.value == 0
+
+
+@pytest.mark.asyncio
+async def test_session_update_refreshes_the_next_append_sampling_params() -> None:
+    orchestrator, clients, rpc_q, _ = _build()
+    await _open(orchestrator, rpc_q)
+    request_state = orchestrator.request_states[_stage0_request_id()]
+    assert request_state.sampling_params_list[0].max_tokens == 20
+
+    await _submit(orchestrator, commands.UpdateSession(patch={"max_output_tokens": 7}))
+    await _submit(orchestrator, _append_audio())
+
+    assert request_state.sampling_params_list[0].max_tokens == 7
+    assert clients[0].add_request_calls[-1][0].sampling_params.max_tokens == 7
+    await orchestrator.session_manager.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_forwarded_stage_requests_are_bound_and_barge_in_aborts_them() -> None:
+    orchestrator, clients, rpc_q, _ = _build(stages=2)
+    await _open(orchestrator, rpc_q)
+    request_id = _stage0_request_id()
+    await _submit(orchestrator, _append_audio())
+    request_state = orchestrator.request_states[request_id]
+    session = orchestrator.session_manager.get(SESSION_ID)
+    assert session is not None
+
+    # The base forwards Stage0 text to the TTS stage and reports the submission.
+    forwarded = SimpleNamespace(request_id=request_id, prompt_token_ids=[1, 2], resumable=True)
+    replica_id = await orchestrator.stage_pools[1].submit_initial(request_id, request_state, forwarded)
+    orchestrator._on_stage_submitted(1, request_id, replica_id, request_state)
+    assert request_state.stage_fences[1] == DuplexFence(SESSION_ID)
+    assert session.stage_request_submitted(1, request_id)
+
+    await _submit(orchestrator, commands.BargeIn())
+    assert session.epoch == 1
+    assert clients[0].abort_calls == [[request_id]]
+    assert clients[1].abort_calls == [[request_id]]
+    assert request_id not in orchestrator.request_states
+    assert session.resource_request_ids() == []
+    await orchestrator.session_manager.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue() -> None:
+    orchestrator, _, rpc_q, output_q = _build(stages=2)
+    await _open(orchestrator, rpc_q)
+    request_id = _stage0_request_id()
+    await _submit(orchestrator, _append_audio())
+    request_state = orchestrator.request_states[request_id]
+    while not output_q.empty():
+        output_q.get_nowait()
+
+    consumed = await orchestrator._intercept_stage_output(1, 0, _tts_output(request_id), request_state, None, None)
+    await _settle(orchestrator)
+
+    assert consumed is True
+    session = orchestrator.session_manager.get(SESSION_ID)
+    assert session is not None and session.active_response_id is not None
+    types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
+    assert "response.created" in types and "response.output_audio.delta" in types
+
+    orphan = DuplexOrchestratorRequestState(
+        request_id="duplex-s.b3RoZXI.e.0.r.stage0",
+        prompt=None,
+        sampling_params_list=[],
+        final_stage_id=1,
+        session_owned=True,
+        session_id="other",
+        fence=DuplexFence("other"),
+    )
+    assert await orchestrator._intercept_stage_output(1, 0, _tts_output(orphan.request_id), orphan, None, None)
+    await orchestrator.session_manager.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_forward_failure_closes_the_owning_session() -> None:
+    orchestrator, clients, rpc_q, output_q = _build(stages=2)
+    await _open(orchestrator, rpc_q)
+    request_id = _stage0_request_id()
+    await _submit(orchestrator, _append_audio())
+    request_state = orchestrator.request_states[request_id]
+    session = orchestrator.session_manager.get(SESSION_ID)
+    assert session is not None
+
+    absorbed = await orchestrator._handle_forward_failure(request_id, 1, request_state, ValueError("bad thinker"))
+    await _settle(orchestrator)
+
+    assert absorbed is True
+    assert request_id not in orchestrator.request_states
+    assert clients[0].abort_calls == [[request_id]]
+    assert SESSION_ID not in orchestrator.session_manager.runners
+    assert session.state == DuplexSessionState.CLOSED
+    assert orchestrator.session_manager.active_count() == 0
+    types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
+    # The stage failure is reported before the session expires, never after.
+    assert types.index("error") < types.index("session.expired")
+    assert types[-1] == "session.expired"
+
+
+@pytest.mark.asyncio
+async def test_request_cleanup_failure_is_deferred_and_retried_by_the_reaper(monkeypatch) -> None:
+    orchestrator, clients, rpc_q, _ = _build()
+    await _open(orchestrator, rpc_q)
+    request_id = _stage0_request_id()
+    await _submit(orchestrator, _append_audio())
+    manager = orchestrator.session_manager
+
+    async def failing_abort(request_ids: list[str]) -> None:
+        raise RuntimeError("stage abort failed")
+
+    monkeypatch.setattr(clients[0], "abort_requests_async", failing_abort)
+    with pytest.raises(RuntimeError, match="stage abort failed"):
+        await orchestrator._cleanup_request_ids([request_id], abort=True, release_owners=True)
+    await _settle(orchestrator)
+
+    assert SESSION_ID not in manager.runners
+    assert SESSION_ID in manager._closing  # admission slot retained
+    assert manager._pending_request_cleanups and not manager._request_cleanups_in_progress
+    assert manager.active_count() == 0 and manager._admission_count() == 1
+
+    monkeypatch.undo()
+    assert await manager.reap_expired() >= 1
+    assert manager._pending_request_cleanups == {}
+    assert manager._admission_count() == 0
+    assert request_id not in orchestrator.request_states
+
+
+@pytest.mark.asyncio
+async def test_reaper_runs_as_a_background_task_and_survives_one_failure(monkeypatch) -> None:
+    orchestrator, _, _, _ = _build(runtime_config=DuplexSessionRuntimeConfig(reaper_interval_s=0.01))
+    calls = {"count": 0}
+
+    async def reap_expired(now: float | None = None) -> int:
+        calls["count"] += 1
+        if calls["count"] == 1:
+            raise RuntimeError("transient cleanup failure")
+        return 0
+
+    monkeypatch.setattr(orchestrator.session_manager, "reap_expired", reap_expired)
+    tasks = orchestrator._background_tasks()
+    assert len(tasks) == 1
+    task = asyncio.create_task(tasks[0])
+    await asyncio.sleep(0.05)
+    orchestrator._shutdown_event.set()
+    await task
+    assert calls["count"] >= 2
+    await orchestrator._shutdown_extensions()
+
+
+# --------------------------------------------------------------------------- #
+# The hierarchy the two surfaces depend on                                    #
+# --------------------------------------------------------------------------- #
+
+
+def test_the_duplex_stack_extends_the_turn_based_one() -> None:
+    """One engine serves a session and an ordinary request, so duplex must extend turn-based.
+
+    Pinned because it is the kind of relationship a later tidy-up turns back
+    into siblings over a shared base. If it does, a duplex server silently
+    stops serving /v1/chat/completions: the chat service needs a real
+    EngineClient with generate(), which only the turn-based side provides.
+    """
+    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+    from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine
+    from vllm_omni.engine.orchestrator import Orchestrator
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.entrypoints.duplex_omni import DuplexOmni
+
+    assert issubclass(DuplexOrchestrator, Orchestrator)
+    assert issubclass(DuplexOmniEngine, AsyncOmniEngine)
+    assert issubclass(DuplexOmni, AsyncOmni)
+    # What the chat route actually needs off the engine client.
+    assert hasattr(DuplexOmni, "generate")
+
+
+@pytest.mark.asyncio
+async def test_an_unrecognised_message_reaches_the_turn_based_handler() -> None:
+    """The duplex orchestrator consumes session messages and passes the rest down.
+
+    Without the fall-through, an ordinary add_request on a duplex engine is
+    silently dropped and the chat request that sent it hangs forever.
+    """
+    orchestrator, _clients, _rpc_q, _output_q = _build()
+    handled: list[str] = []
+
+    async def _turn_based(msg: object) -> None:
+        handled.append(getattr(msg, "type", ""))
+
+    orchestrator._handle_add_request = _turn_based  # type: ignore[method-assign]
+    msg = SimpleNamespace(type="add_request")
+
+    assert await orchestrator._dispatch_message(msg) is True
+    assert handled == ["add_request"], "an add_request must reach Orchestrator, not be dropped"
+
+
+@pytest.mark.asyncio
+async def test_a_dead_replica_closes_the_sessions_it_was_serving() -> None:
+    """A session outlives its stage request, so replica death has to release the owner.
+
+    ``_handle_dead_replica`` cleaned the stage request but not the runner that
+    owned it: the session stayed alive holding an admission slot, and the client
+    was told nothing — the request-scoped ``ErrorMessage`` it emits has no
+    frontend ``request_states`` entry to land in for a session-owned request, so
+    the frontend drops it. Heartbeats would then keep an unusable session
+    occupying capacity.
+    """
+    orchestrator, clients, rpc_q, output_q = _build()
+    await _open(orchestrator, rpc_q)
+    request_id = _stage0_request_id()
+    await _submit(orchestrator, _append_audio())
+    session = orchestrator.session_manager.get(SESSION_ID)
+    assert session is not None
+    assert orchestrator.session_manager.active_count() == 1
+
+    await orchestrator._handle_dead_replica(0, 0, EngineDeadError("stage 0 replica 0 died"))
+    await _settle(orchestrator)
+
+    assert request_id not in orchestrator.request_states
+    assert SESSION_ID not in orchestrator.session_manager.runners
+    assert session.state == DuplexSessionState.CLOSED
+    assert orchestrator.session_manager.active_count() == 0, "the admission slot must come back"
+    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}"

tests/engine/test_orchestrator.py

@@ -21,20 +21,7 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason
 from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.metrics.stats import IterationStats
 
-from vllm_omni.engine.duplex.control_plane import DuplexControlPlane
-from vllm_omni.engine.duplex.messages import (
-    AppendDuplexInputMessage,
-    CloseDuplexSessionMessage,
-    DuplexFence,
-    OpenDuplexSessionMessage,
-    SignalDuplexTurnMessage,
-)
-from vllm_omni.engine.duplex.runtime import (
-    DuplexInputMode,
-    DuplexRuntimeCapabilities,
-    DuplexSessionRuntimeState,
-    duplex_resource_request_id,
-)
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
     AbortResultMessage,
@@ -54,7 +41,6 @@ from vllm_omni.engine.orchestrator import (
 )
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import MiniCPMO45DuplexRuntimeExtension
 from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -1118,778 +1104,6 @@ async def test_handle_abort_emits_error_result_on_failure() -> None:
     assert "cleanup exploded" in (result.error or "")
 
 
-def _duplex_open_message(
-    session_id: str,
-    *,
-    incarnation: int = 0,
-    session_config: dict[str, object] | None = None,
-    runtime_config: dict[str, object] | None = None,
-) -> OpenDuplexSessionMessage:
-    return OpenDuplexSessionMessage(
-        control_id=f"open-{session_id}",
-        fence=DuplexFence(session_id, incarnation=incarnation),
-        session_id=session_id,
-        capabilities={
-            "input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value],
-            "implementation_level": "model_native_duplex",
-        },
-        session_config=session_config or {},
-        runtime_config=runtime_config or {},
-    )
-
-
-async def _handle_duplex(orchestrator: Orchestrator, message: object) -> None:
-    await orchestrator._require_duplex_control_plane().handle(message)
-
-
-def _duplex_request_state(
-    orchestrator: Orchestrator,
-    session: DuplexSessionRuntimeState,
-    *,
-    stage_id: int,
-) -> OrchestratorRequestState | None:
-    context = orchestrator._require_duplex_control_plane().ensure_stage_request(
-        session,
-        stage_id=stage_id,
-    )
-    if context is None:
-        return None
-    return orchestrator.request_states.get(context.request_id)
-
-
-@pytest.mark.asyncio
-async def test_duplex_control_plane_keeps_public_and_runtime_config_separate() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=_build_stage_pools([[stage0]]),
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    open_message = _duplex_open_message(
-        "sid-config-channels",
-        incarnation=3,
-        session_config={
-            "instructions": "public instructions",
-            "extra_body": {"duplex_stage_max_tokens": {"0": 99}},
-        },
-        runtime_config={
-            "duplex_stage_max_tokens": {"0": 3},
-            "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151705]}},
-        },
-    )
-
-    await _handle_duplex(orchestrator, open_message)
-
-    assert rpc_q.get_nowait().ok is True
-    session = orchestrator.duplex_sessions.require(open_message.session_id)
-    request_state = _duplex_request_state(orchestrator, session, stage_id=0)
-    assert request_state.sampling_params_list[0].max_tokens == 3
-    assert request_state.sampling_params_list[0].stop_token_ids == [151705]
-    bridge = request_state.streaming.bridge_states["duplex"]
-    assert bridge["incarnation"] == open_message.fence.incarnation
-    assert bridge["session_config"] == open_message.session_config
-    assert bridge["runtime_config"] == open_message.runtime_config
-
-
-@pytest.mark.asyncio
-async def test_duplex_control_plane_preserves_turn_commit_without_model_extension() -> None:
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=[],
-        duplex_runtime_extension=None,
-        enable_duplex_control=True,
-    )
-    assert isinstance(orchestrator.duplex_control_plane, DuplexControlPlane)
-
-    message = OpenDuplexSessionMessage(
-        control_id="open-turn-commit",
-        fence=DuplexFence("sid-turn-commit"),
-        session_id="sid-turn-commit",
-        capabilities={"input_modes": [DuplexInputMode.TURN_COMMIT_ONLY.value]},
-        session_config={},
-    )
-    await _handle_duplex(orchestrator, message)
-
-    result = rpc_q.get_nowait()
-    assert result.ok is True
-    assert result.stage_results[0]["result"]["scheduler_request_context"] is False
-
-
-def test_ordinary_orchestrator_bypasses_duplex_control_plane() -> None:
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=asyncio.Queue(),
-        stage_pools=[],
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-    )
-
-    assert orchestrator.duplex_control_plane is None
-
-
-@pytest.mark.asyncio
-async def test_duplex_close_cleans_preregistered_request_without_append() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    running_counter = FakeRunningCounter()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=_build_stage_pools([[stage0]]),
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-        running_counter=running_counter,
-    )
-    open_message = _duplex_open_message("sid-preregister-close")
-
-    await _handle_duplex(orchestrator, open_message)
-    open_result = rpc_q.get_nowait()
-    request_id = open_result.stage_results[0]["result"]["request_id"]
-    assert request_id in orchestrator.request_states
-
-    await _handle_duplex(
-        orchestrator,
-        CloseDuplexSessionMessage(
-            control_id="close-preregistered",
-            fence=open_message.fence,
-            session_id=open_message.session_id,
-        ),
-    )
-
-    assert rpc_q.get_nowait().ok is True
-    assert request_id not in orchestrator.request_states
-    assert orchestrator.duplex_sessions.get(open_message.session_id) is None
-    assert running_counter.value == 0
-
-
-@pytest.mark.asyncio
-async def test_duplex_open_failure_rolls_back_session_and_reserved_request() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=_build_stage_pools([[stage0]]),
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    open_message = _duplex_open_message(
-        "sid-open-rollback",
-        runtime_config={
-            "duplex_stage_sampling_params": {
-                "0": {"stop_token_ids": 123},
-            }
-        },
-    )
-
-    await _handle_duplex(orchestrator, open_message)
-
-    assert rpc_q.get_nowait().ok is False
-    assert orchestrator.duplex_sessions.get(open_message.session_id) is None
-    assert orchestrator.request_states == {}
-
-
-@pytest.mark.asyncio
-async def test_duplex_running_counter_tracks_only_submitted_request() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    running_counter = FakeRunningCounter()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=_build_stage_pools([[stage0]]),
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-        running_counter=running_counter,
-    )
-    open_message = _duplex_open_message("sid-duplex-counter")
-    await _handle_duplex(orchestrator, open_message)
-    assert rpc_q.get_nowait().ok is True
-    assert running_counter.value == 0
-
-    await _handle_duplex(
-        orchestrator,
-        AppendDuplexInputMessage(
-            control_id="append-duplex-counter",
-            fence=open_message.fence,
-            session_id=open_message.session_id,
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-            payload={"is_speech": True},
-        ),
-    )
-    assert rpc_q.get_nowait().ok is True
-    assert running_counter.value == 1
-
-    await _handle_duplex(
-        orchestrator,
-        CloseDuplexSessionMessage(
-            control_id="close-duplex-counter",
-            fence=open_message.fence,
-            session_id=open_message.session_id,
-        ),
-    )
-    assert rpc_q.get_nowait().ok is True
-    assert running_counter.value == 0
-
-
-@pytest.mark.asyncio
-async def test_duplex_failed_append_does_not_advance_sequence_or_fence() -> None:
-    class FailingPlanExtension(MiniCPMO45DuplexRuntimeExtension):
-        def plan_append(self, **kwargs):
-            raise RuntimeError("planned append failure")
-
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=_build_stage_pools([[stage0]]),
-        duplex_runtime_extension=FailingPlanExtension(),
-        enable_duplex_control=True,
-    )
-    fence = DuplexFence("sid-append-rollback", turn_id=1)
-    session = orchestrator.duplex_sessions.open_session(
-        DuplexFence("sid-append-rollback"),
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-
-    await _handle_duplex(
-        orchestrator,
-        AppendDuplexInputMessage(
-            control_id="append-rollback",
-            fence=fence,
-            session_id=fence.session_id,
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-            payload={"is_speech": True},
-        ),
-    )
-
-    result = rpc_q.get_nowait()
-    assert result.ok is False
-    assert session.fence == DuplexFence("sid-append-rollback")
-    assert session.input_seq == 0
-    assert session.input_turn_seq == 0
-    assert stage0.add_request_calls == []
-
-
-@pytest.mark.asyncio
-async def test_duplex_duplicate_append_operation_is_submitted_once() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=_build_stage_pools([[stage0]]),
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    fence = DuplexFence("sid-idempotent-append")
-    session = orchestrator.duplex_sessions.open_session(
-        fence,
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-
-    for control_id in ("append-first", "append-timeout-retry"):
-        await _handle_duplex(
-            orchestrator,
-            AppendDuplexInputMessage(
-                control_id=control_id,
-                operation_id="physical-input-1",
-                fence=fence,
-                session_id=fence.session_id,
-                mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-                payload={"is_speech": True},
-            ),
-        )
-
-    first = rpc_q.get_nowait()
-    retry = rpc_q.get_nowait()
-    assert first.ok is True
-    assert retry.ok is True
-    assert first.stage_results == retry.stage_results
-    assert len(stage0.add_request_calls) == 1
-    assert session.input_seq == 1
-
-
-@pytest.mark.asyncio
-async def test_duplex_barge_in_aborts_bound_stage_requests_before_releasing_fence() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=False)
-    stage1 = FakeStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools(
-        [[stage0], [stage1]],
-        output_processors=[FakeOutputProcessor(), FakeOutputProcessor()],
-        stage_vllm_configs=[
-            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
-            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
-        ],
-    )
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=stage_pools,
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    session = orchestrator.duplex_sessions.open_session(
-        DuplexFence("sid-stage-signal"),
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-    session.bind_stage_request(0, "req-stage0", fence=session.fence)
-    session.bind_stage_request(1, "req-stage1", fence=session.fence)
-    assert stage_pools[0].select_replica_id("req-stage0") == 0
-    assert stage_pools[1].select_replica_id("req-stage1") == 0
-    session.append_input(
-        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-        fence=session.fence,
-    )
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="ctrl-signal",
-            fence=session.fence,
-            next_fence=DuplexFence("sid-stage-signal", epoch=1),
-            session_id="sid-stage-signal",
-            event="barge_in",
-        ),
-    )
-
-    assert stage0.abort_calls == [["req-stage0"]]
-    assert stage1.abort_calls == [["req-stage1"]]
-    assert stage0.collective_rpc_calls == []
-    assert stage1.collective_rpc_calls == []
-    assert session.stage_bindings == {}
-    assert session.fence == DuplexFence("sid-stage-signal", epoch=1)
-    result = rpc_q.get_nowait()
-    assert result.ok is True
-
-
-@pytest.mark.asyncio
-async def test_duplex_late_barge_in_releases_only_cancelled_fence_bindings() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=False)
-    stage1 = FakeStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools(
-        [[stage0], [stage1]],
-        output_processors=[FakeOutputProcessor(), FakeOutputProcessor()],
-        stage_vllm_configs=[
-            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
-            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
-        ],
-    )
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=stage_pools,
-        enable_duplex_control=True,
-    )
-    cancelled_fence = DuplexFence("sid-late-stage-signal")
-    session = orchestrator.duplex_sessions.open_session(
-        cancelled_fence,
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-    session.bind_stage_request(0, "req-stage0", fence=cancelled_fence)
-    session.bind_stage_request(1, "req-stage1", fence=cancelled_fence)
-    assert stage_pools[0].select_replica_id("req-stage0") == 0
-    assert stage_pools[1].select_replica_id("req-stage1") == 0
-    current_fence = DuplexFence("sid-late-stage-signal", epoch=1)
-    session.accept_fence(current_fence)
-
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="ctrl-late-signal",
-            fence=cancelled_fence,
-            next_fence=current_fence,
-            session_id="sid-late-stage-signal",
-            event="barge_in",
-        ),
-    )
-
-    assert stage0.abort_calls == [["req-stage0"]]
-    assert stage1.abort_calls == [["req-stage1"]]
-    assert session.stage_bindings == {}
-    assert session.fence == current_fence
-    result = rpc_q.get_nowait()
-    assert result.ok is True
-
-
-@pytest.mark.asyncio
-async def test_duplex_cancel_without_next_fence_is_rejected_without_releasing_bindings() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools(
-        [[stage0]],
-        output_processors=[FakeOutputProcessor()],
-        stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))],
-    )
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=stage_pools,
-        enable_duplex_control=True,
-    )
-    fence = DuplexFence("sid-cancel-contract")
-    session = orchestrator.duplex_sessions.open_session(fence)
-    session.bind_stage_request(0, "req-live", fence=fence)
-
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="cancel-without-next",
-            fence=fence,
-            session_id=fence.session_id,
-            event="input.cancel",
-        ),
-    )
-
-    result = rpc_q.get_nowait()
-    assert result.ok is False
-    assert result.error_count == 1
-    assert "next_fence" in result.stage_results[0]["result"]["error"]
-    assert session.fence == fence
-    assert session.stage_request_ids() == ["req-live"]
-    assert stage0.abort_calls == []
-
-
-@pytest.mark.asyncio
-async def test_duplex_session_update_replaces_runtime_config() -> None:
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=[],
-        enable_duplex_control=True,
-    )
-    fence = DuplexFence("sid-update-config")
-    session = orchestrator.duplex_sessions.open_session(
-        fence,
-        session_config={"temperature": 0.7},
-    )
-
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="update-config",
-            fence=fence,
-            session_id=fence.session_id,
-            event="session.update",
-            session_config={"temperature": 0.0, "instructions": "updated"},
-        ),
-    )
-
-    result = rpc_q.get_nowait()
-    assert result.ok is True
-    assert session.session_config == {"temperature": 0.0, "instructions": "updated"}
-
-
-@pytest.mark.asyncio
-async def test_duplex_session_update_refreshes_next_append_sampling_params() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools(
-        [[stage0]],
-        output_processors=[FakeOutputProcessor()],
-        stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))],
-    )
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=stage_pools,
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    fence = DuplexFence("sid-update-policy")
-    session = orchestrator.duplex_sessions.open_session(
-        fence,
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-        runtime_config={
-            "duplex_stage_max_tokens": {"0": 2},
-            "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151645]}},
-        },
-    )
-    req_state = _duplex_request_state(orchestrator, session, stage_id=0)
-    assert req_state is not None
-    assert req_state.sampling_params_list[0].max_tokens == 2
-
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="update-policy",
-            fence=fence,
-            session_id=fence.session_id,
-            event="session.update",
-            runtime_config={
-                "duplex_stage_max_tokens": {"0": 7},
-                "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151705]}},
-            },
-        ),
-    )
-    assert rpc_q.get_nowait().ok is True
-
-    await _handle_duplex(
-        orchestrator,
-        AppendDuplexInputMessage(
-            control_id="append-updated-policy",
-            fence=fence,
-            session_id=fence.session_id,
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-            payload={"is_speech": True},
-        ),
-    )
-
-    assert rpc_q.get_nowait().ok is True
-    assert req_state.sampling_params_list[0].max_tokens == 7
-    assert req_state.sampling_params_list[0].stop_token_ids == [151705]
-    submitted_request = stage0.add_request_calls[0][0]
-    assert submitted_request.sampling_params.stop_token_ids == [151705]
-    assert submitted_request.sampling_params.max_tokens == 7
-
-
-@pytest.mark.asyncio
-async def test_duplex_invalid_session_update_preserves_previous_config() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=_build_stage_pools([[stage0]]),
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    fence = DuplexFence("sid-invalid-update")
-    original_runtime_config = {
-        "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151645]}},
-    }
-    session = orchestrator.duplex_sessions.open_session(
-        fence,
-        runtime_config=original_runtime_config,
-    )
-
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="invalid-update",
-            fence=fence,
-            session_id=fence.session_id,
-            event="session.update",
-            runtime_config={
-                "duplex_stage_sampling_params": {"0": {"stop_token_ids": 123}},
-            },
-        ),
-    )
-
-    result = rpc_q.get_nowait()
-    assert result.ok is False
-    assert session.runtime_config == original_runtime_config
-    assert session.config_generation == 0
-
-
-@pytest.mark.asyncio
-async def test_duplex_arbitrary_non_cancel_signal_is_rejected() -> None:
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=[],
-        enable_duplex_control=True,
-    )
-    fence = DuplexFence("sid-unsupported-signal")
-    orchestrator.duplex_sessions.open_session(fence)
-
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="unsupported-signal",
-            fence=fence,
-            session_id=fence.session_id,
-            event="turn.end",
-        ),
-    )
-
-    result = rpc_q.get_nowait()
-    assert result.ok is False
-    assert "unsupported duplex runtime signal" in result.stage_results[0]["result"]["error"]
-
-
-@pytest.mark.asyncio
-async def test_duplex_cancel_rejects_late_old_append_and_accepts_next_epoch() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools(
-        [[stage0]],
-        output_processors=[FakeOutputProcessor()],
-        stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))],
-    )
-    rpc_q: asyncio.Queue = asyncio.Queue()
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=rpc_q,
-        stage_pools=stage_pools,
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    cancelled_fence = DuplexFence("sid-cancel-late-append")
-    next_fence = DuplexFence("sid-cancel-late-append", epoch=1)
-    session = orchestrator.duplex_sessions.open_session(
-        cancelled_fence,
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-    )
-
-    await _handle_duplex(
-        orchestrator,
-        SignalDuplexTurnMessage(
-            control_id="cancel-old-fence",
-            fence=cancelled_fence,
-            next_fence=next_fence,
-            session_id=session.session_id,
-            event="input.cancel",
-        ),
-    )
-
-    cancel_result = rpc_q.get_nowait()
-    assert cancel_result.ok is True
-    assert session.fence == next_fence
-
-    await _handle_duplex(
-        orchestrator,
-        AppendDuplexInputMessage(
-            control_id="late-old-append",
-            fence=cancelled_fence,
-            session_id=session.session_id,
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-            payload={"is_speech": True},
-        ),
-    )
-
-    stale_result = rpc_q.get_nowait()
-    assert stale_result.ok is False
-    assert stage0.add_request_calls == []
-    assert session.stage_bindings == {}
-
-    await _handle_duplex(
-        orchestrator,
-        AppendDuplexInputMessage(
-            control_id="next-epoch-append",
-            fence=next_fence,
-            session_id=session.session_id,
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-            payload={"is_speech": True},
-        ),
-    )
-
-    next_result = rpc_q.get_nowait()
-    assert next_result.ok is True
-    assert len(stage0.add_request_calls) == 1
-    assert session.stage_bindings[0].fence == next_fence
-
-
-@pytest.mark.asyncio
-async def test_duplex_append_updates_bridge_turn_id_on_long_lived_stage0_request() -> None:
-    stage0 = FakeStageClient(stage_type="llm", final_output=True)
-    stage_pools = _build_stage_pools(
-        [[stage0]],
-        output_processors=[FakeOutputProcessor()],
-        stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))],
-    )
-    orchestrator = Orchestrator(
-        request_async_queue=asyncio.Queue(),
-        output_async_queue=asyncio.Queue(),
-        rpc_async_queue=asyncio.Queue(),
-        stage_pools=stage_pools,
-        duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(),
-        enable_duplex_control=True,
-    )
-    session = orchestrator.duplex_sessions.open_session(
-        DuplexFence("sid-bridge-turn"),
-        capabilities=DuplexRuntimeCapabilities(
-            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
-        ),
-        session_config={
-            "voice": "test",
-        },
-        runtime_config={
-            "duplex_stage_sampling_params": {
-                "0": {"stop_token_ids": [151645]},
-            },
-        },
-    )
-
-    await _handle_duplex(
-        orchestrator,
-        AppendDuplexInputMessage(
-            control_id="append-1",
-            fence=DuplexFence("sid-bridge-turn"),
-            session_id="sid-bridge-turn",
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-            payload={"is_speech": True},
-        ),
-    )
-    req_state1 = _duplex_request_state(orchestrator, session, stage_id=0)
-    assert req_state1 is not None
-    duplex_state1 = req_state1.streaming.bridge_states["duplex"]
-    assert duplex_state1["session_id"] == "sid-bridge-turn"
-    assert duplex_state1["epoch"] == 0
-    assert duplex_state1["turn_id"] == 0
-    assert duplex_state1["session_config"] == session.session_config
-    assert req_state1.sampling_params_list[0].stop_token_ids == [151645]
-    submitted_request = stage0.add_request_calls[0][0]
-    assert submitted_request.sampling_params.stop_token_ids == [151645]
-    assert submitted_request.sampling_params.max_tokens == 1
-
-    await _handle_duplex(
-        orchestrator,
-        AppendDuplexInputMessage(
-            control_id="append-2",
-            fence=DuplexFence("sid-bridge-turn", turn_id=1, response_seq=1),
-            session_id="sid-bridge-turn",
-            mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value,
-            payload={"is_speech": True, "new_user_turn": True},
-        ),
-    )
-    req_state2 = _duplex_request_state(orchestrator, session, stage_id=0)
-    assert req_state2 is req_state1
-    duplex_state2 = req_state2.streaming.bridge_states["duplex"]
-    assert duplex_state2["turn_id"] == 1
-    assert duplex_state2["epoch"] == 0
-    expected_request_id = duplex_resource_request_id(DuplexFence("sid-bridge-turn"), "stage0")
-    assert [call[0].request_id for call in stage0.add_request_calls] == [
-        expected_request_id,
-        expected_request_id,
-    ]
-
-
 # ---------------------------------------------------------------------------
 # Multi-replica tests
 # ---------------------------------------------------------------------------
@@ -2641,61 +1855,6 @@ async def test_stage_pool_process_llm_raw_outputs_mutates_iteration_stats() -> N
     assert iteration_stats.num_generation_tokens == 3
 
 
-@pytest.mark.asyncio
-async def test_duplex_reaper_loop_waits_between_ticks():
-    class _Plane:
-        def __init__(self) -> None:
-            self.calls = 0
-
-        async def reap_expired(self) -> int:
-            self.calls += 1
-            return 0
-
-    orchestrator = object.__new__(Orchestrator)
-    orchestrator.duplex_control_plane = _Plane()
-    orchestrator._duplex_reaper_interval_s = 0.01
-    orchestrator._shutdown_event = asyncio.Event()
-
-    task = asyncio.create_task(orchestrator._duplex_reaper_loop())
-    await asyncio.sleep(0.035)
-    orchestrator._shutdown_event.set()
-    await task
-
-    assert 2 <= orchestrator.duplex_control_plane.calls <= 5
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("first_cleanup_delay", [0.0, 0.05], ids=["immediate", "delayed"])
-async def test_duplex_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: float) -> None:
-    recovered = asyncio.Event()
-
-    class _Plane:
-        def __init__(self) -> None:
-            self.calls = 0
-
-        async def reap_expired(self) -> int:
-            self.calls += 1
-            if self.calls == 1:
-                await asyncio.sleep(first_cleanup_delay)
-                raise RuntimeError("transient cleanup failure")
-            recovered.set()
-            return 0
-
-    orchestrator = object.__new__(Orchestrator)
-    orchestrator.duplex_control_plane = _Plane()
-    orchestrator._duplex_reaper_interval_s = 0.01
-    orchestrator._shutdown_event = asyncio.Event()
-
-    task = asyncio.create_task(orchestrator._duplex_reaper_loop())
-    try:
-        # Wait for recovery itself, including when cleanup exceeds the old 35 ms window.
-        await asyncio.wait_for(recovered.wait(), timeout=5.0)
-        assert orchestrator.duplex_control_plane.calls >= 2
-    finally:
-        orchestrator._shutdown_event.set()
-        await asyncio.wait_for(task, timeout=5.0)
-
-
 @pytest.mark.asyncio
 async def test_abort_retry_does_not_repeat_successful_stage_abort():
     class _Pool:
@@ -2736,45 +1895,58 @@ async def test_abort_retry_does_not_repeat_successful_stage_abort():
 
 
 @pytest.mark.asyncio
-async def test_request_cleanup_failure_is_deferred_to_control_plane():
-    class _FailingPool:
-        async def abort_requests(self, request_ids: list[str]) -> None:
-            raise RuntimeError("stage abort failed")
+async def test_duplex_session_request_error_finish_is_delivered_as_request_error() -> None:
+    """A session-owned request the scheduler finished with FinishReason.ERROR
+    (e.g. its prompt could not grow past max_model_len) must reach the
+    session's consumer as a request-scoped error, not vanish: its terminal
+    outputs are not processed like other requests, and the output processor
+    may already have dropped the request state."""
+    output_queue: asyncio.Queue = asyncio.Queue()
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=output_queue,
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=[],
+    )
+    duplex_state = OrchestratorRequestState(request_id="duplex-req", session_owned=True)
+    plain_state = OrchestratorRequestState(request_id="plain-req")
+    reason = "context_length_exceeded: streaming session prompt would grow to 8220 tokens, above max_model_len 8192"
 
-        def release_bindings(self, request_ids: list[str]) -> None:
-            pass
+    await orchestrator._report_duplex_session_request_error(
+        0,
+        0,
+        OmniEngineCoreOutput(
+            request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.ERROR, stop_reason=reason
+        ),
+        duplex_state,
+    )
+    error = output_queue.get_nowait()
+    assert isinstance(error, ErrorMessage)
+    assert error.request_id == "duplex-req"
+    assert error.stage_id == 0
+    assert error.fatal is False
+    assert error.error == reason
 
-    class _Plane:
-        def __init__(self) -> None:
-            self.deferred: list[str] = []
-            self.finalized: list[str] = []
-
-        def close_sessions_for_request_ids(self, request_ids: list[str], **kwargs):
-            assert kwargs == {"abort": True, "cleanup_in_progress": True}
-            return {"sid-cleanup": ["req-a"]}
-
-        def defer_request_cleanups(self, session_ids: list[str]) -> None:
-            self.deferred.extend(session_ids)
-
-        def finalize_closed_sessions(self, session_ids: list[str]) -> None:
-            self.finalized.extend(session_ids)
-
-    from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
-
-    orchestrator = object.__new__(Orchestrator)
-    orchestrator.stage_pools = [_FailingPool()]
-    orchestrator.duplex_control_plane = _Plane()
-    orchestrator._pd_kv_params = {}
-    orchestrator.request_states = {}
-    orchestrator._running_counter = None
-    orchestrator._cfg_tracker = CfgCompanionTracker()
-
-    with pytest.raises(RuntimeError, match="stage abort failed"):
-        await orchestrator._cleanup_request_ids(
-            ["req-a"],
-            abort=True,
-            close_duplex_sessions=True,
-        )
-
-    assert orchestrator.duplex_control_plane.deferred == ["sid-cleanup"]
-    assert orchestrator.duplex_control_plane.finalized == []
+    # An ordinary segment stop, a segment-finished error and a request that is
+    # not session-owned produce nothing.
+    await orchestrator._report_duplex_session_request_error(
+        0,
+        0,
+        OmniEngineCoreOutput(request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.STOP),
+        duplex_state,
+    )
+    await orchestrator._report_duplex_session_request_error(
+        0,
+        0,
+        OmniEngineCoreOutput(
+            request_id="duplex-req", new_token_ids=[], finish_reason=FinishReason.ERROR, is_segment_finished=True
+        ),
+        duplex_state,
+    )
+    await orchestrator._report_duplex_session_request_error(
+        0,
+        0,
+        OmniEngineCoreOutput(request_id="plain-req", new_token_ids=[], finish_reason=FinishReason.ERROR),
+        plain_state,
+    )
+    assert output_queue.empty()

tests/engine/test_orchestrator_error_handling.py

@@ -1054,7 +1054,7 @@ async def test_streaming_input_processor_client_error_does_not_forward_terminal_
         final_stage_id=1,
     )
     state.streaming.enabled = True
-    assert state.duplex_identity is None
+    assert state.session_owned is False
     orchestrator.request_states[state.request_id] = state
     try:
         await orchestrator._route_output(0, 0, _build_request_output(state.request_id), state, stage_metrics=None)
@@ -1074,34 +1074,14 @@ async def test_streaming_input_processor_client_error_does_not_forward_terminal_
             q.close()
 
 
-@pytest.mark.asyncio
-async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None:
-    class FailingStage(FakeStageClient):
-        def process_engine_inputs(self, *_args, **_kwargs):
-            raise ValueError("No latent or hidden_states found in thinker output")
-
-    fixture = orchestrator_factory(
-        [FakeStageClient(stage_type="llm"), FailingStage(stage_type="llm", final_output=True)]
-    )
-    state = OrchestratorRequestState(
-        request_id="bad",
-        prompt=SimpleNamespace(request_id="bad", prompt_token_ids=[1]),
-        sampling_params_list=[_sampling_params(), _sampling_params()],
-        final_stage_id=1,
-        duplex_identity=SimpleNamespace(),
-    )
-
-    async def no_cleanup(*_args, **_kwargs):
-        pass
-
-    monkeypatch.setattr(fixture.orchestrator, "_cleanup_request_ids", no_cleanup)
-    try:
-        await fixture.orchestrator._forward_to_next_stage_unguarded("bad", 0, _build_request_output("raw"), state)
-        error = await _wait_for_error_message(fixture, request_id="bad")
-        assert error.fatal is False
-    finally:
-        fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
-        fixture.thread.join(timeout=5)
+# ``test_duplex_input_processor_failure_is_request_scoped`` lived here upstream,
+# where one Orchestrator served both turn-based and duplex requests. Duplex
+# request ownership now belongs to DuplexOrchestrator, and the base
+# ``_handle_forward_failure`` deliberately absorbs nothing, so the equivalent
+# assertion is
+# ``tests/engine/test_duplex_orchestrator.py::test_forward_failure_closes_the_owning_session``
+# -- which also checks the stage abort, the session teardown and the
+# error-before-session.expired ordering.
 
 
 @pytest.mark.asyncio

tests/engine/test_orchestrator_segment_stage_keying.py

@@ -2,10 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Streaming segment-boundary state must be keyed by stage.
 
-``Orchestrator._orchestration_loop`` polls every stage into the same
+``OrchestratorBase._orchestration_loop`` polls every stage into the same
 ``OrchestratorRequestState``. While the segment-boundary fields lived in a single
 flat slot per request, whichever stage polled most recently overwrote the others,
 and the per-stage consumers then read back a different stage's boundary.
+
+The duplex consumer of the per-stage segment (``DuplexOrchestrator``'s output
+context) is covered in ``test_duplex_orchestrator.py``.
 """
 
 from __future__ import annotations
@@ -15,8 +18,6 @@ from dataclasses import dataclass, field
 import pytest
 from vllm.sampling_params import SamplingParams
 
-from vllm_omni.engine.duplex.contracts import DuplexRequestIdentity
-from vllm_omni.engine.duplex.messages import DuplexFence
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
@@ -54,22 +55,19 @@ class _RecordingPool:
         return _StageMetrics()
 
 
-def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState:
+def _streaming_req_state(request_id: str = "req-streaming") -> OrchestratorRequestState:
     req_state = OrchestratorRequestState(
         request_id=request_id,
         sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)],
         final_stage_id=TALKER_STAGE,
+        session_owned=True,
     )
     req_state.streaming.enabled = True
-    req_state.duplex_identity = DuplexRequestIdentity(
-        session_id="session-1",
-        fence=DuplexFence(session_id="session-1"),
-    )
     return req_state
 
 
-def test_duplex_output_context_reads_only_its_own_stage_segment() -> None:
-    req_state = _duplex_req_state()
+def test_segment_accessor_reads_only_its_own_stage() -> None:
+    req_state = _streaming_req_state()
     req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(
         finished=True,
         token_ids=[11, 22],
@@ -78,26 +76,28 @@ def test_duplex_output_context_reads_only_its_own_stage_segment() -> None:
     # The talker polls after the thinker and is still mid-segment.
     req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False)
 
-    thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE)
-    talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE)
-
-    assert thinker is not None
-    assert talker is not None
+    thinker = req_state.streaming.segment(THINKER_STAGE)
+    talker = req_state.streaming.segment(TALKER_STAGE)
 
     # These are the assertions that fail on the flat field: the talker is written
     # second, so the thinker loses its boundary to whichever stage polled last.
-    assert thinker.segment_finished is True
-    assert thinker.segment_token_ids == (11, 22)
-    assert thinker.segment_output_metadata == {"stage": "thinker"}
+    assert thinker.finished is True
+    assert thinker.token_ids == [11, 22]
+    assert thinker.output_metadata == {"stage": "thinker"}
 
-    assert talker.segment_finished is False
-    assert talker.segment_token_ids == ()
-    assert talker.segment_output_metadata == {}
+    assert talker.finished is False
+    assert talker.token_ids == []
+    assert talker.output_metadata == {}
+
+    # A stage that never reported a boundary gets a fresh, unshared segment.
+    unreported = req_state.streaming.segment(2)
+    unreported.output_metadata["poisoned"] = True
+    assert req_state.streaming.segment(2).output_metadata == {}
 
 
 @pytest.mark.asyncio
 async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None:
-    req_state = _duplex_req_state()
+    req_state = _streaming_req_state()
     req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True)
 
     talker_pool = _RecordingPool()

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -7,23 +7,17 @@ import asyncio
 import queue
 from types import SimpleNamespace
 from typing import Any
-from unittest.mock import ANY, AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import janus
 import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 
-from vllm_omni.engine.duplex.contracts import (
-    DuplexStageRequestContext,
-    DuplexStageSubmission,
-)
-from vllm_omni.engine.duplex.messages import DuplexFence
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
     StreamingSegmentState,
-    _OrchestratorDuplexStagePort,
 )
 from vllm_omni.engine.stage_pool import StagePool
 
@@ -131,64 +125,6 @@ class FakePrewarmPool:
         return self.stage_client if self.get_bound_replica_id(request_id) is not None else None
 
 
-def _duplex_stage_port_submission():
-    stage_pools = []
-    for stage_id in range(3):
-        pool = SimpleNamespace(
-            stage_client=SimpleNamespace(default_sampling_params=SamplingParams(max_tokens=1)),
-            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
-            submit_initial=AsyncMock(return_value=stage_id + 10),
-            submit_update=AsyncMock(return_value=stage_id + 20),
-        )
-        stage_pools.append(pool)
-    request_states: dict[str, OrchestratorRequestState] = {}
-    prewarm = AsyncMock()
-    port = _OrchestratorDuplexStagePort(
-        stage_pools=stage_pools,
-        request_states=request_states,
-        running_counter=None,
-        cleanup_request_ids=AsyncMock(),
-        async_chunk=True,
-        prewarm_async_chunk_stages=prewarm,
-    )
-    context = DuplexStageRequestContext(
-        request_id="req-duplex",
-        session_id="session-duplex",
-        fence=DuplexFence("session-duplex"),
-        stage_id=0,
-        final_stage_id=2,
-        config_generation=0,
-        sampling_params=tuple(SamplingParams(max_tokens=1) for _ in range(3)),
-    )
-    port.ensure_request(context)
-    submission = DuplexStageSubmission(
-        context=context,
-        prompt={"prompt_token_ids": [1, 2]},
-        already_submitted=False,
-    )
-    return port, stage_pools, request_states, prewarm, submission
-
-
-def test_duplex_bridge_state_catches_up_after_serving_turn_boundary() -> None:
-    port, _, request_states, _, submission = _duplex_stage_port_submission()
-    request_state = request_states[submission.context.request_id]
-    bridge_state = request_state.streaming.bridge_states["duplex"]
-    bridge_state["model_turn_id"] = 0
-
-    next_context = DuplexStageRequestContext(
-        request_id=submission.context.request_id,
-        session_id=submission.context.session_id,
-        fence=DuplexFence(submission.context.session_id, turn_id=1),
-        stage_id=submission.context.stage_id,
-        final_stage_id=submission.context.final_stage_id,
-        config_generation=submission.context.config_generation,
-        sampling_params=submission.context.sampling_params,
-    )
-    port.ensure_request(next_context)
-
-    assert bridge_state["model_turn_id"] == 1
-
-
 def _request_output(request_id: str) -> RequestOutput:
     completion = CompletionOutput(
         index=0,
@@ -279,13 +215,12 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N
         stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info)
     orchestrator.stage_pools = [stage0, stage1, stage2]
     orchestrator._emit_tx_edge = lambda **_kwargs: None
-    orchestrator._record_duplex_stage_submission = MagicMock()
+    orchestrator._on_stage_submitted = MagicMock()
     req_state = OrchestratorRequestState(
         request_id="req-prewarm",
         prompt={"prompt_token_ids": [1, 2]},
         sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)],
         final_stage_id=2,
-        duplex_identity=SimpleNamespace(),
     )
 
     prewarmed = await orchestrator._prewarm_async_chunk_stages(
@@ -306,7 +241,7 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N
     assert stage2.submitted[0].resumable is True
     assert 1 not in req_state.stage_submit_ts
     assert 2 in req_state.stage_submit_ts
-    orchestrator._record_duplex_stage_submission.assert_called_once_with(
+    orchestrator._on_stage_submitted.assert_called_once_with(
         2,
         "req-prewarm",
         0,
@@ -314,42 +249,6 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N
     )
 
 
-@pytest.mark.asyncio
-async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
-    port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
-    prewarm.return_value = True
-
-    result = await port.submit(submission)
-
-    assert result.stage_id == 0
-    stage_pools[0].submit_initial.assert_awaited_once()
-    prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"])
-
-
-@pytest.mark.asyncio
-async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None:
-    """A failed prewarm has already aborted the request and popped its state.
-
-    Returning a success result here would hand the control plane a replica for a
-    request that no longer exists, and the trailing bookkeeping would re-register
-    a running counter the cleanup just released -- a leak that never decrements.
-    """
-    port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
-    counter = MagicMock()
-    port._running_counter = counter
-    prewarm.return_value = False
-    request_state = request_states["req-duplex"]
-
-    with pytest.raises(RuntimeError, match="prewarm failed"):
-        await port.submit(submission)
-
-    prewarm.assert_awaited_once()
-    assert request_state.duplex_stage_fences == {}
-    assert request_state.stage_submit_ts == {}
-    assert request_state.running_counter_registered is False
-    counter.increment.assert_not_called()
-
-
 @pytest.mark.asyncio
 async def test_async_route_forwards_to_outgoing_only_stage() -> None:
     orchestrator = object.__new__(Orchestrator)

tests/engine/test_output_metadata_snapshots.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Snapshot metadata must not accumulate as if it were generated content."""
+
+import pytest
+import torch
+from vllm.outputs import PoolingRequestOutput
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.outputs.mm_outputs import MultimodalPayload
+from vllm_omni.outputs.output_modality import OutputModality
+from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+RATE_KEYS = ("sr", "sample_rate", "audio_sample_rate")
+
+
+def _state(kind: RequestOutputKind) -> OmniRequestState:
+    return OmniRequestState(
+        request_id="audio",
+        external_req_id="audio",
+        parent_req=None,
+        request_index=0,
+        lora_request=None,
+        prompt=None,
+        prompt_token_ids=[0],
+        prompt_embeds=None,
+        logprobs_processor=None,
+        detokenizer=None,
+        max_tokens_param=None,
+        arrival_time=0.0,
+        queue=None,
+        log_stats=False,
+        stream_interval=1,
+        output_kind=kind,
+    )
+
+
+@pytest.mark.parametrize("key", RATE_KEYS)
+@pytest.mark.parametrize("representation", ["integer", "scalar_tensor", "vector_tensor", "metadata_tensor"])
+def test_rate_is_latest_snapshot_before_consolidation(key, representation):
+    payload = MultimodalPayload()
+    for index, rate in enumerate((16000, 24000, 48000)):
+        value = (
+            rate if representation == "integer" else torch.tensor([rate] if representation == "vector_tensor" else rate)
+        )
+        incoming = (
+            MultimodalPayload(tensors={"audio": torch.full((2,), index)}, metadata={key: value})
+            if representation == "metadata_tensor"
+            else MultimodalPayload.from_dict({"audio": torch.full((2,), index), key: value})
+        )
+        assert incoming is not None
+        payload = payload.merged_with(incoming)
+        snapshot = payload[key]
+        assert not isinstance(snapshot, list)
+        assert int(snapshot) == rate
+        assert payload.get(key) is payload.to_dict()[key]
+    payload = payload.merged_with(MultimodalPayload(tensors={"audio": torch.full((2,), 3)}))
+    assert int(payload[key]) == 48000  # Missing metadata means retain, not clear.
+    payload.consolidate_tensors(OutputModality.AUDIO)
+    payload.consolidate_metadata()
+    torch.testing.assert_close(payload["audio"], torch.arange(4).repeat_interleave(2))
+
+
+@pytest.mark.parametrize("key", RATE_KEYS)
+@pytest.mark.parametrize("tensor_first", [False, True])
+def test_rate_representation_change_removes_stale_partition(key, tensor_first):
+    first = torch.tensor(16000) if tensor_first else 16000
+    last = 24000 if tensor_first else torch.tensor(24000)
+    payload = MultimodalPayload.from_dict({key: first})
+    incoming = MultimodalPayload.from_dict({key: last})
+    assert payload is not None and incoming is not None
+    incoming_keys = (set(incoming.tensors), set(incoming.metadata))
+    payload = payload.merged_with(incoming)
+    assert int(payload[key]) == 24000
+    assert payload[key] is payload.to_dict()[key]
+    assert (key in payload.tensors) != (key in payload.metadata)
+    assert incoming[key] is last
+    assert incoming_keys == (set(incoming.tensors), set(incoming.metadata))
+
+
+@pytest.mark.parametrize("key", RATE_KEYS)
+def test_self_merge_retains_snapshot_and_content(key):
+    payload = MultimodalPayload.from_dict({key: torch.tensor(24000), "audio": torch.tensor([1.0, 2.0])})
+    assert payload is not None
+    payload = payload.merged_with(payload)
+    assert int(payload[key]) == 24000
+    payload.consolidate_tensors(OutputModality.AUDIO)
+    torch.testing.assert_close(payload["audio"], torch.tensor([1.0, 2.0, 1.0, 2.0]))
+
+
+@pytest.mark.parametrize("key", RATE_KEYS)
+@pytest.mark.parametrize("kind", [RequestOutputKind.DELTA, RequestOutputKind.CUMULATIVE])
+def test_long_stream_keeps_rate_bounded_without_losing_audio(key, kind):
+    state = _state(kind)
+    for index in range(1000):
+        state.add_multimodal_tensor(
+            {"model_outputs": torch.full((3,), index / 1000), key: torch.tensor(24000)}, "audio"
+        )
+        output = state.make_request_output([], None, None, None)
+        assert output is not None and not isinstance(output, PoolingRequestOutput)
+        payload = output.outputs[0].multimodal_output
+        assert isinstance(payload[key], torch.Tensor) and payload[key].numel() == 1
+        audio = payload["audio"]
+        assert audio.numel() == (3 if kind == RequestOutputKind.DELTA else 3 * (index + 1))
+        torch.testing.assert_close(audio[-3:], torch.full((3,), index / 1000))
+        assert ("audio" in state.mm_accumulated) is (kind == RequestOutputKind.CUMULATIVE)
+    final = state.make_request_output([], None, FinishReason.STOP, None)
+    assert final is not None and final.finished
+    if kind == RequestOutputKind.DELTA:
+        assert "audio" not in final.outputs[0].multimodal_output
+
+
+@pytest.mark.parametrize("pending", [0, 3])
+def test_abort_flushes_only_pending_content_with_latest_rate(pending):
+    state = _state(RequestOutputKind.DELTA)
+    processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False, output_modality=OutputModality.AUDIO)
+    processor.request_states[state.request_id] = state
+    processor.external_req_ids[state.external_req_id] = [state.request_id]
+    for rate in (16000, 24000):
+        state.add_multimodal_tensor({"model_outputs": torch.ones(3), "sr": torch.tensor(rate)}, "audio")
+        output = state.make_request_output([], None, None, None)
+        assert output is not None
+    if pending:
+        state.add_multimodal_tensor({"model_outputs": torch.ones(pending)}, "audio")
+    aborted, outputs = processor.abort_requests_collecting_outputs([state.request_id], internal=True)
+    assert aborted == [state.request_id] and len(outputs) == 1 and outputs[0].finished
+    payload = outputs[0].outputs[0].multimodal_output
+    assert int(payload["sr"]) == 24000
+    assert (payload["audio"].numel() if "audio" in payload else 0) == pending
+    assert state.request_id not in processor.request_states
+
+
+def test_non_snapshot_tensors_still_accumulate_and_empty_merge_keeps_identity():
+    first = MultimodalPayload.from_dict({"latent": torch.tensor([1.0]), "audio": torch.tensor([2.0])})
+    second = MultimodalPayload.from_dict({"latent": torch.tensor([3.0]), "audio": torch.tensor([4.0])})
+    assert first is not None and second is not None
+    assert MultimodalPayload().merged_with(first) is first
+    merged = first.merged_with(second)
+    merged.consolidate_tensors(OutputModality.AUDIO)
+    torch.testing.assert_close(merged["latent"], torch.tensor([1.0, 3.0]))
+    torch.testing.assert_close(merged["audio"], torch.tensor([2.0, 4.0]))

tests/engine/test_output_modality.py

@@ -41,6 +41,7 @@ _mm_mod = _load_module(
 OutputModality = _om_mod.OutputModality
 TensorAccumulationStrategy = _om_mod.TensorAccumulationStrategy
 get_accumulation_strategy = _om_mod.get_accumulation_strategy
+register_key_accumulation_strategy = _om_mod.register_key_accumulation_strategy
 MultimodalPayload = _mm_mod.MultimodalPayload
 MultimodalCompletionOutput = _mm_mod.MultimodalCompletionOutput
 
@@ -156,3 +157,90 @@ def test_output_modality_printed_examples(capsys):
     captured = capsys.readouterr()
     assert "OutputModality Parsing" in captured.out
     assert "MultimodalPayload" in captured.out
+
+
+def test_get_accumulation_strategy_key_override_precedence():
+    """A registered per-key override wins over the modality-wide default;
+    any other key under the same modality keeps using that default."""
+    assert get_accumulation_strategy(OutputModality.AUDIO, "unregistered.key") == TensorAccumulationStrategy.CONCAT_LAST
+
+    register_key_accumulation_strategy("__test_override__.frames", TensorAccumulationStrategy.CONCAT_DIM0)
+    assert (
+        get_accumulation_strategy(OutputModality.AUDIO, "__test_override__.frames")
+        == TensorAccumulationStrategy.CONCAT_DIM0
+    )
+    # A plain waveform key under the same AUDIO modality is unaffected.
+    assert get_accumulation_strategy(OutputModality.AUDIO, "waveform") == TensorAccumulationStrategy.CONCAT_LAST
+    assert get_accumulation_strategy(OutputModality.AUDIO) == TensorAccumulationStrategy.CONCAT_LAST
+
+
+def _talker_codec_frame_payload(audio_key: str, ref_key: str) -> MultimodalPayload:
+    """Synthetic stand-in for a TTS talker's accumulated per-step output.
+
+    ``audio_key`` grows one ``[1, num_codebooks]`` row per decode step on top
+    of a ``[5, num_codebooks]`` prefill block (Qwen3-TTS's ``codes.audio``);
+    ``ref_key`` re-emits the same constant ``[100, num_codebooks]``
+    reference-context matrix unchanged at every step (Qwen3-TTS's
+    ``codes.ref``). Neither is a waveform chunk.
+    """
+    payload = MultimodalPayload()
+    payload.tensors[audio_key] = [
+        torch.zeros(5, 16),
+        torch.full((1, 16), 1.0),
+        torch.full((1, 16), 2.0),
+        torch.full((1, 16), 3.0),
+    ]
+    ref_block = torch.arange(100 * 16, dtype=torch.float32).reshape(100, 16)
+    payload.tensors[ref_key] = [ref_block.clone(), ref_block.clone(), ref_block.clone()]
+    return payload
+
+
+def test_consolidate_tensors_default_audio_strategy_corrupts_codec_frame_keys():
+    """Reproduces both failure modes of forcing codec-frame keys through the
+    AUDIO modality's waveform-tuned default (CONCAT_LAST) with no per-key
+    override registered: the audio-frames key's dim-0 growth makes
+    CONCAT_LAST's dim=-1 concat raise (mismatched dim 0), which the
+    flatten-chunks fallback turns into a 1-D blob instead of the intended
+    ``[8, 16]``; the ref-frames key's matching last dim lets CONCAT_LAST
+    "succeed" by silently concatenating three identical ``[100, 16]`` blocks
+    into ``[100, 48]`` instead of keeping a single copy.
+    """
+    payload = _talker_codec_frame_payload("unregistered.codes.audio", "unregistered.codes.ref")
+
+    payload.consolidate_tensors(OutputModality.AUDIO)
+
+    assert tuple(payload.tensors["unregistered.codes.audio"].shape) != (8, 16)
+    assert tuple(payload.tensors["unregistered.codes.ref"].shape) == (100, 48)
+
+
+def test_consolidate_tensors_with_qwen3_tts_key_overrides_is_correct():
+    """With the per-key overrides Qwen3-TTS's pipeline module registers
+    (``codes.audio`` -> CONCAT_DIM0, ``codes.ref`` -> REPLACE), both keys
+    consolidate correctly under the AUDIO modality despite not being
+    waveform chunks."""
+    register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0)
+    register_key_accumulation_strategy("codes.ref", TensorAccumulationStrategy.REPLACE)
+
+    payload = _talker_codec_frame_payload("codes.audio", "codes.ref")
+    expected_ref = payload.tensors["codes.ref"][0]
+
+    payload.consolidate_tensors(OutputModality.AUDIO)
+
+    assert tuple(payload.tensors["codes.audio"].shape) == (8, 16)
+    assert tuple(payload.tensors["codes.ref"].shape) == (100, 16)
+    assert torch.equal(payload.tensors["codes.ref"], expected_ref)
+
+
+def test_consolidate_tensors_raises_with_key_name_instead_of_silently_keeping_last():
+    """A concat failure under a non-CONCAT_LAST strategy must be surfaced
+    with the offending key name rather than resolved by silently keeping
+    only the last chunk. Previously the except-branch only special-cased
+    the literal key ``"audio"``, so any other mismatched key (e.g. a
+    codec-frame key routed to CONCAT_DIM0) was corrupted instead of
+    failing loudly.
+    """
+    payload = MultimodalPayload()
+    payload.tensors["latent.frames"] = [torch.zeros(2, 8), torch.zeros(3, 9)]
+
+    with pytest.raises(RuntimeError, match="latent.frames"):
+        payload.consolidate_tensors(OutputModality.LATENT)

tests/engine/test_single_stage_mode.py

@@ -390,7 +390,7 @@ class TestSingleStageModeDetection:
                 return_value=None,
             )
             mocker.patch(
-                "vllm_omni.engine.async_omni_engine.load_deploy_config",
+                "vllm_omni.engine.omni_engine_base.load_deploy_config",
                 return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()),
             )
 
@@ -438,7 +438,9 @@ class TestSingleStageModeDetection:
         )
         engine = self._make_engine_no_thread(mocker, resolved_pipeline=pipeline)
 
-        StageConfigFactory.get_pipeline_config.assert_not_called()
+        StageConfigFactory.get_pipeline_config.assert_called_once_with(
+            model="fake-model", trust_remote_code=False, deploy_config_path=None
+        )
         assert engine._event_driven_orch_default is expected
 
     def test_stage_id_kwarg_promotes_to_single_stage_mode(self, mocker: MockerFixture):
@@ -467,11 +469,11 @@ class TestSingleStageModeDetection:
     def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture):
         duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
         get_pipeline_config = mocker.patch(
-            "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
+            "vllm_omni.engine.omni_engine_base.StageConfigFactory.get_pipeline_config",
             return_value=None,
         )
         load_deploy_config = mocker.patch(
-            "vllm_omni.engine.async_omni_engine.load_deploy_config",
+            "vllm_omni.engine.omni_engine_base.load_deploy_config",
             return_value=SimpleNamespace(duplex_session=duplex_session),
         )
 
@@ -488,7 +490,8 @@ class TestSingleStageModeDetection:
             deploy_config_path="/fake/duplex.yaml",
         )
         load_deploy_config.assert_called_once_with("/fake/duplex.yaml")
-        assert engine.duplex_session_config is duplex_session
+        # The turn-based engine only keeps the resolved deploy profile for introspection.
+        assert engine.deploy_config.duplex_session is duplex_session
 
     def test_auto_discovered_deploy_loads_duplex_runtime_config(
         self,
@@ -497,11 +500,11 @@ class TestSingleStageModeDetection:
         deploy_path = "/resolved/qwen3_omni_moe.yaml"
         duplex_session = DuplexSessionRuntimeConfig(server_vad_model_path="/models/silero_vad.onnx")
         mocker.patch(
-            "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
+            "vllm_omni.engine.omni_engine_base.StageConfigFactory.get_pipeline_config",
             return_value=None,
         )
         load_deploy_config = mocker.patch(
-            "vllm_omni.engine.async_omni_engine.load_deploy_config",
+            "vllm_omni.engine.omni_engine_base.load_deploy_config",
             return_value=SimpleNamespace(duplex_session=duplex_session),
         )
 
@@ -578,7 +581,7 @@ class TestEndpointRestrictionsTrustRemoteCode:
 
     def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine:
         mocker.patch(
-            "vllm_omni.engine.async_omni_engine.load_deploy_config",
+            "vllm_omni.engine.omni_engine_base.load_deploy_config",
             return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()),
         )
         mocker.patch.object(

tests/entrypoints/duplex/test_duplex_serving.py

@@ -0,0 +1,609 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""The thin websocket handler: handshake, command translation, event pump, resume/takeover."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import json
+from contextlib import suppress
+from typing import Any
+
+import pytest
+from fastapi import WebSocketDisconnect
+
+from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
+from vllm_omni.engine.duplex import commands
+from vllm_omni.engine.duplex.config import DuplexCapabilities
+from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent, SessionClosed, SessionCreated
+from vllm_omni.engine.duplex.messages import DuplexSessionError
+from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request
+from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler
+from vllm_omni.entrypoints.duplex.websocket import MAX_EVENT_BYTES
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_DISCONNECT = object()
+
+
+class FakeWebSocket:
+    def __init__(self, query: dict[str, str] | None = None) -> None:
+        self.query_params = dict(query or {})
+        self.sent: list[dict[str, Any]] = []
+        self.accepted = False
+        self.closed: list[tuple[int, str]] = []
+        self._send_failure: str | None = None
+        self._inbound: asyncio.Queue[Any] = asyncio.Queue()
+
+    async def accept(self) -> None:
+        self.accepted = True
+
+    async def send_json(self, payload: dict[str, Any]) -> None:
+        if self.closed or self._send_failure is not None:
+            raise RuntimeError(
+                self._send_failure or "Unexpected ASGI message 'websocket.send', after sending 'websocket.close'."
+            )
+        self.sent.append(json.loads(json.dumps(payload)))
+
+    async def receive_text(self) -> str:
+        item = await self._inbound.get()
+        if item is _DISCONNECT:
+            raise WebSocketDisconnect(code=1000)
+        return item
+
+    async def close(self, code: int = 1000, reason: str = "") -> None:
+        self.closed.append((code, reason))
+        # A closed socket ends the reader exactly like a client disconnect.
+        self._inbound.put_nowait(_DISCONNECT)
+
+    # ---- test helpers ----
+
+    def feed(self, payload: dict[str, Any] | str) -> None:
+        self._inbound.put_nowait(payload if isinstance(payload, str) else json.dumps(payload))
+
+    def disconnect(self) -> None:
+        self._inbound.put_nowait(_DISCONNECT)
+
+    def break_sends(self) -> None:
+        """Kill the write half only: the reader stays parked, as it does in practice."""
+        self._send_failure = "Unexpected ASGI message 'websocket.send', after sending 'websocket.close'."
+
+    def types(self) -> list[str]:
+        return [payload["type"] for payload in self.sent]
+
+    async def wait_for(self, wire_type: str, *, timeout_s: float = 2.0) -> dict[str, Any]:
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + timeout_s
+        while loop.time() < deadline:
+            for payload in self.sent:
+                if payload["type"] == wire_type:
+                    return payload
+            await asyncio.sleep(0.005)
+        raise AssertionError(f"no {wire_type!r} in {self.types()}")
+
+
+class FakeHandle:
+    def __init__(self, session_id: str, capabilities: DuplexCapabilities, *, idle_timeout_s: float = 300) -> None:
+        self.session_id = session_id
+        self.capabilities = capabilities
+        self.public_session: dict[str, Any] = {"id": session_id, "idle_timeout_s": idle_timeout_s}
+        self.lease_generation = 0
+        self.closed = False
+        self.close_reasons: list[str] = []
+        self.commands: list[commands.DuplexCommand] = []
+        self._outbox: asyncio.Queue[DuplexEvent] = asyncio.Queue()
+
+    def deliver(self, event: DuplexEvent) -> None:
+        self._outbox.put_nowait(event)
+
+    async def submit(self, command: commands.DuplexCommand) -> None:
+        if self.closed:
+            raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id)
+        self.commands.append(command)
+
+    async def events(self):
+        while True:
+            event = await self._outbox.get()
+            yield event
+            if event.is_terminal:
+                return
+
+    async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None:
+        self.close_reasons.append(reason)
+        if self.closed:
+            return
+        self.closed = True
+        self.deliver(SessionClosed(session_id=self.session_id, reason=reason))
+
+
+class FakeOmni:
+    def __init__(
+        self, *, resumable: bool = True, replay_max_bytes: int = 64 * 1024, idle_timeout_s: float = 300
+    ) -> None:
+        self.duplex_session_config = DuplexSessionRuntimeConfig(
+            resume_replay_ttl_s=60.0, resume_replay_max_bytes_per_session=replay_max_bytes
+        )
+        self.capabilities = DuplexCapabilities(supports_session_resume=resumable)
+        self.idle_timeout_s = idle_timeout_s
+        self.opened: list[dict[str, Any]] = []
+        self.handles: dict[str, FakeHandle] = {}
+        self.resumed: list[tuple[str, int]] = []
+        self.detached: list[str] = []
+        self.open_error: DuplexSessionError | None = None
+
+    async def open_session(self, config: Any) -> FakeHandle:
+        self.opened.append(dict(config))
+        if self.open_error is not None:
+            raise self.open_error
+        session_id = f"duplex-{len(self.handles) + 1:032x}"
+        handle = FakeHandle(session_id, self.capabilities, idle_timeout_s=self.idle_timeout_s)
+        self.handles[session_id] = handle
+        handle.deliver(SessionCreated(session_id=session_id, session={"id": session_id, "model": config.get("model")}))
+        return handle
+
+    def get_session(self, session_id: str) -> FakeHandle | None:
+        return self.handles.get(session_id)
+
+    async def resume_session(self, session_id: str, *, expected_lease_generation: int) -> FakeHandle:
+        self.resumed.append((session_id, expected_lease_generation))
+        handle = self.handles[session_id]
+        handle.lease_generation = expected_lease_generation + 1
+        return handle
+
+    async def detach_session(self, session_id: str) -> None:
+        self.detached.append(session_id)
+
+
+def _handler(omni: FakeOmni, **kwargs: Any) -> OmniDuplexSessionHandler:
+    kwargs.setdefault("config_timeout_s", 1.0)
+    kwargs.setdefault("idle_timeout_s", 5.0)
+    return OmniDuplexSessionHandler(duplex_omni=omni, **kwargs)
+
+
+def _session_update(**session: Any) -> dict[str, Any]:
+    return {"type": "session.update", "session": {"model": "test-model", "modalities": ["audio", "text"], **session}}
+
+
+async def _open(
+    handler: OmniDuplexSessionHandler, omni: FakeOmni, **session: Any
+) -> tuple[FakeWebSocket, FakeHandle, asyncio.Task]:
+    ws = FakeWebSocket({"duplex": "1", "autostart": "0"})
+    task = asyncio.create_task(handler.handle_realtime_session(ws))
+    ws.feed(_session_update(**session))
+    created = await ws.wait_for("session.created")
+    handle = omni.handles[created["session"]["id"]]
+    return ws, handle, task
+
+
+# --------------------------------------------------------------------------- #
+# Handshake and commands                                                      #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.asyncio
+async def test_session_update_opens_a_session_and_announces_the_server_id() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni, session_id="mine", instructions="hi")
+
+    assert ws.accepted
+    created = ws.sent[0]
+    assert created["type"] == "session.created"
+    assert created["session"]["id"] == handle.session_id != "mine"
+    assert created["attachment_generation"] == 1
+    assert isinstance(created["resume_token"], str) and created["resume_token"]
+    assert "incarnation" not in created
+    assert "server_event_seq" not in created
+    # The whole session object went to open_session; the id inside it is ignored there.
+    assert omni.opened == [
+        {"model": "test-model", "modalities": ["audio", "text"], "session_id": "mine", "instructions": "hi"}
+    ]
+
+    ws.feed({"type": "input_audio_buffer.append", "audio": base64.b64encode(b"\x00\x01" * 8).decode("ascii")})
+    ws.feed({"type": "input_audio_buffer.commit", "event_id": "evt-c"})
+    await asyncio.sleep(0.05)
+    assert [type(command) for command in handle.commands] == [commands.AppendAudio, commands.Commit]
+    assert handle.commands[1].event_id == "evt-c"
+
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    # A resumable session is detached (engine-owned grace), not closed.
+    assert omni.detached == [handle.session_id]
+    assert handle.close_reasons == []
+
+
+@pytest.mark.asyncio
+async def test_session_events_are_journaled_and_sent_in_order() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    delta = await ws.wait_for("response.output_audio.delta")
+    assert delta["delta"] == "aGk=" and delta["server_event_seq"] == 1
+
+    ws.feed({"type": "session.event_ack", "server_event_seq": 1})
+    ws.feed({"type": "session.event_ack", "server_event_seq": 9, "event_id": "evt-ack"})
+    error = await ws.wait_for("error")
+    assert error["error"]["code"] == "invalid_event_ack" and error["error"]["event_id"] == "evt-ack"
+
+    await handle.close(reason="client_close")
+    closed = await ws.wait_for("session.closed")
+    assert closed["reason"] == "client_close"
+    await asyncio.wait_for(task, timeout=2.0)
+    assert ws.closed[0][1] == "client_close"
+
+
+@pytest.mark.asyncio
+async def test_envelope_errors_are_answered_locally_without_touching_the_session() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+
+    ws.feed("not json")
+    ws.feed("x" * (MAX_EVENT_BYTES + 1))
+    ws.feed({"type": "totally.unknown", "event_id": "evt-u"})
+    ws.feed({"type": "session.resume", "event_id": "evt-r"})
+    ws.feed({"type": "playback.ack", "event_id": "evt-p"})  # played_ms missing
+    await asyncio.sleep(0.1)
+
+    codes = [payload["error"]["code"] for payload in ws.sent if payload["type"] == "error"]
+    assert codes == ["invalid_json", "event_too_large", "unknown_event", "unsupported_session_resume", "bad_event"]
+    errors = [payload for payload in ws.sent if payload["type"] == "error"]
+    assert errors[2]["error"]["event_id"] == "evt-u"
+    assert errors[4]["error"]["event_id"] == "evt-p"
+    assert handle.commands == []
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_open_rejection_is_reported_and_the_socket_ends() -> None:
+    omni = FakeOmni()
+    omni.open_error = DuplexSessionError("no room", code="resource_exhausted", retryable=True)
+    handler = _handler(omni)
+    ws = FakeWebSocket({"duplex": "1"})
+    task = asyncio.create_task(handler.handle_realtime_session(ws))
+    ws.feed(_session_update())
+    await asyncio.wait_for(task, timeout=2.0)
+    assert ws.types() == ["error"]
+    assert ws.sent[0]["error"]["code"] == "resource_exhausted"
+
+
+@pytest.mark.asyncio
+async def test_config_timeout_and_first_message_validation() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni, config_timeout_s=0.05)
+    ws = FakeWebSocket({"duplex": "1"})
+    await asyncio.wait_for(handler.handle_realtime_session(ws), timeout=2.0)
+    assert ws.sent[0]["error"]["code"] == "config_timeout"
+
+    ws = FakeWebSocket({"duplex": "1"})
+    task = asyncio.create_task(handler.handle_realtime_session(ws))
+    ws.feed("{not json")
+    await asyncio.wait_for(task, timeout=2.0)
+    assert ws.sent[0]["error"]["code"] == "invalid_json"
+    assert omni.opened == []
+
+
+@pytest.mark.asyncio
+async def test_idle_timeout_detaches_a_resumable_session_like_a_disconnect() -> None:
+    omni = FakeOmni(idle_timeout_s=0.05)
+    handler = _handler(omni, idle_timeout_s=0.05)
+    ws, handle, task = await _open(handler, omni)
+
+    await asyncio.wait_for(task, timeout=2.0)
+    # Serving makes no session-lifetime decision: the engine lease decides
+    # whether a silent, detached session expires.
+    assert handle.close_reasons == []
+    assert omni.detached == [handle.session_id]
+
+
+@pytest.mark.asyncio
+async def test_idle_timeout_closes_a_non_resumable_session() -> None:
+    omni = FakeOmni(idle_timeout_s=0.05, resumable=False)
+    handler = _handler(omni, idle_timeout_s=0.05)
+    ws, handle, task = await _open(handler, omni)
+
+    await asyncio.wait_for(task, timeout=2.0)
+    assert handle.close_reasons == ["disconnect"]
+    assert omni.detached == []
+
+
+@pytest.mark.asyncio
+async def test_transport_send_failure_detaches_the_session_instead_of_closing_it() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+
+    # Only the write half dies, so the pump's send is the one and only report
+    # of the broken socket (the reader is still parked on receive).
+    ws.break_sends()
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    await asyncio.sleep(0.05)
+
+    # The engine session is alive and resumable, so this is a disconnect of
+    # the current attachment, not a session close; the pump keeps journaling.
+    assert handle.close_reasons == []
+    assert omni.detached == [handle.session_id]
+    assert not handler._pumps[handle.session_id].done()
+
+    # The reader sees the same broken socket a moment later. The attachment is
+    # already gone, so this must not detach a second time: another detach would
+    # restart the engine's disconnect grace window.
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    assert handle.close_reasons == []
+
+    handle.deliver(SessionClosed(session_id=handle.session_id, reason="client_close"))
+    await asyncio.sleep(0.05)
+    assert handle.session_id not in handler._pumps
+
+
+@pytest.mark.asyncio
+async def test_non_resumable_session_is_closed_on_disconnect() -> None:
+    omni = FakeOmni(resumable=False)
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    assert "resume_token" not in ws.sent[0]
+
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert handle.close_reasons == ["disconnect"]
+    assert omni.detached == []
+
+
+@pytest.mark.asyncio
+async def test_a_client_close_still_delivers_session_closed_before_the_socket_goes() -> None:
+    """The read loop must not outrun the pump's terminal event.
+
+    ``DuplexSessionHandle._deliver`` queues ``SessionClosed`` and marks the
+    handle closed in one synchronous step, so ``handle.closed`` is already true
+    while the event is still sitting in the outbox. ``_read_loop`` loops on
+    exactly that flag: it returned, the endpoint returned, and the ASGI server
+    tore the socket down with ``session.closed`` unsent -- the client saw an
+    abrupt close (no close frame) instead of the terminal event it was waiting
+    for.
+
+    The teardown is what makes this observable, so it is modelled here: a real
+    connection stops accepting writes the moment the endpoint returns, which is
+    why a pump that is merely *still scheduled* is already too late.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws = FakeWebSocket({"duplex": "1", "autostart": "0"})
+
+    async def serve() -> None:
+        try:
+            await handler.handle_realtime_session(ws)
+        finally:
+            # The ASGI server drops the connection when the endpoint returns;
+            # anything the pump writes after this never reaches the client.
+            ws.break_sends()
+
+    task = asyncio.create_task(serve())
+    ws.feed(_session_update())
+    created = await ws.wait_for("session.created")
+    handle = omni.handles[created["session"]["id"]]
+    submit = handle.submit
+
+    async def closing_submit(command: commands.DuplexCommand) -> None:
+        await submit(command)
+        if isinstance(command, commands.CloseSession):
+            # Same order as the real handle: queue the event, then flip the flag.
+            handle.deliver(SessionClosed(session_id=handle.session_id, reason="client_close"))
+            handle.closed = True
+
+    handle.submit = closing_submit  # type: ignore[method-assign]
+
+    ws.feed({"type": "session.close"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+    assert ws.types()[-1] == "session.closed", f"terminal event never reached the wire: {ws.types()}"
+    assert ws.closed and ws.closed[0][0] == 1000, "the socket must close with a normal close frame"
+
+
+# --------------------------------------------------------------------------- #
+# Resume and takeover                                                         #
+# --------------------------------------------------------------------------- #
+
+
+async def _resume(
+    handler: OmniDuplexSessionHandler,
+    session_id: str,
+    token: str,
+    *,
+    last_seq: int = 0,
+) -> tuple[FakeWebSocket, asyncio.Task]:
+    ws = FakeWebSocket({"duplex": "1", "resume": "1"})
+    task = asyncio.create_task(handler.handle_realtime_session(ws))
+    ws.feed(
+        {
+            "type": "session.resume",
+            "session_id": session_id,
+            "resume_token": token,
+            "last_received_server_event_seq": last_seq,
+        }
+    )
+    return ws, task
+
+
+@pytest.mark.asyncio
+async def test_resume_after_disconnect_replays_missed_events_and_rotates_the_token() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="one"))
+    await ws.wait_for("response.output_audio.delta")
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    # Emitted while detached: journaled for replay only.
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="two"))
+    await asyncio.sleep(0.05)
+
+    ws2, task2 = await _resume(handler, handle.session_id, token, last_seq=1)
+    resumed = await ws2.wait_for("session.resumed")
+    assert resumed["session_id"] == handle.session_id
+    assert resumed["attachment_generation"] == 2
+    assert resumed["resume_token"] != token
+    assert "incarnation" not in resumed
+    replayed = await ws2.wait_for("response.output_audio.delta")
+    assert replayed["delta"] == "two" and replayed["server_event_seq"] == 2
+    assert omni.resumed == [(handle.session_id, 0)]
+
+    # The old token is revoked by the rotation.
+    ws3, task3 = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(task3, timeout=2.0)
+    assert ws3.sent[0]["error"]["code"] == "invalid_resume_token"
+
+    ws2.feed({"type": "session.close"})
+    await asyncio.sleep(0.05)
+    ws2.disconnect()
+    await asyncio.wait_for(task2, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_resume_takes_over_a_live_attachment() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+
+    ws2, task2 = await _resume(handler, handle.session_id, token)
+    await ws2.wait_for("session.resumed")
+    replaced = await ws.wait_for("session.replaced")
+    assert replaced["attachment_generation"] == 1
+    assert ws.closed and ws.closed[0][1] == "session_replaced"
+    # The replaced socket's later input is ignored; the new one drives the session.
+    ws.feed({"type": "input_audio_buffer.clear"})
+    ws2.feed({"type": "input_audio_buffer.clear", "event_id": "evt-new"})
+    await asyncio.sleep(0.05)
+    assert [command.event_id for command in handle.commands] == ["evt-new"]
+    await asyncio.wait_for(task, timeout=2.0)
+
+    ws2.disconnect()
+    await asyncio.wait_for(task2, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+
+
+@pytest.mark.asyncio
+async def test_resume_validation_errors() -> None:
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws = FakeWebSocket({"duplex": "1", "resume": "1"})
+    task = asyncio.create_task(handler.handle_realtime_session(ws))
+    ws.feed({"type": "session.resume", "session_id": "duplex-x"})
+    await asyncio.wait_for(task, timeout=2.0)
+    assert ws.sent[0]["error"]["code"] == "invalid_session_resume"
+
+    ws, task = await _resume(handler, "duplex-unknown", "token")
+    await asyncio.wait_for(task, timeout=2.0)
+    assert ws.sent[0]["error"]["code"] == "session_resume_expired"
+
+
+@pytest.mark.asyncio
+async def test_journal_overflow_degrades_to_live_delivery_with_resync_required() -> None:
+    omni = FakeOmni(replay_max_bytes=256)
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="x" * 400))
+    resync = await ws.wait_for("session.resync_required")
+    assert resync["reason"] == "journal_overflow"
+    delta = await ws.wait_for("response.output_audio.delta")
+    assert "server_event_seq" not in delta
+    assert ws.types().index("session.resync_required") < ws.types().index("response.output_audio.delta")
+
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+# --------------------------------------------------------------------------- #
+# Envelope helpers                                                            #
+# --------------------------------------------------------------------------- #
+
+
+def test_realtime_envelope_query_rules() -> None:
+    envelope = RealtimeEnvelope.from_query_params({"model": "m"})
+    assert envelope.initial_open_payload() == {"model": "m"}
+    assert envelope.initial_open_payload() is None  # autostart happens once
+
+    envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"})
+    assert envelope.resume_only is True
+    assert envelope.initial_open_payload() is None
+    assert RealtimeEnvelope.from_query_params({"resume": "1"}).resume_only is True
+    assert RealtimeEnvelope.from_query_params({"session_id": "mine"}).default_session_payload() == {"model": None}
+
+
+def test_realtime_envelope_first_message_classification() -> None:
+    envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"})
+    resume = envelope.first_message({"type": "session.resume", "session_id": "s"})
+    assert resume.kind == "resume" and resume.resume_payload["session_id"] == "s"
+
+    envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"})
+    opened = envelope.first_message({"type": "session.update", "session": {"model": "x", "instructions": "hi"}})
+    assert opened.kind == "open" and opened.session_payload == {"model": "x", "instructions": "hi"}
+    assert opened.pending_command_payload is None
+
+    envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"})
+    autostarted = envelope.first_message({"type": "input_audio_buffer.commit"})
+    assert autostarted.kind == "open" and autostarted.session_payload == {"model": "m"}
+    assert autostarted.pending_command_payload == {"type": "input_audio_buffer.commit"}
+
+
+def test_parse_resume_request_requires_the_three_fields_only() -> None:
+    request = parse_resume_request({"session_id": "s", "resume_token": "t", "last_received_server_event_seq": 3})
+    assert request is not None
+    assert (request.session_id, request.resume_token, request.last_received_server_event_seq) == ("s", "t", 3)
+    assert parse_resume_request({"session_id": "s", "resume_token": "t"}).last_received_server_event_seq == 0
+    assert parse_resume_request({"session_id": "s", "resume_token": "t", "incarnation": 1}) is not None
+    assert parse_resume_request({"session_id": "s"}) is None
+    assert parse_resume_request({"session_id": "s", "resume_token": "t", "last_received_server_event_seq": -1}) is None
+
+
+@pytest.mark.asyncio
+async def test_a_resume_that_fails_to_activate_does_not_detach_the_live_attachment() -> None:
+    """A rejected resume must roll back only what it owns.
+
+    Two reconnects can both authenticate and both complete the engine resume;
+    only one activates. The loser's activation raises, and the rollback used to
+    detach the *session*, starting the engine's disconnect grace for the winner
+    -- a grace ordinary heartbeats do not clear.
+
+    The activation failure is injected rather than raced for: what is under test
+    is the rollback's precondition, not the window that produces it.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+    try:
+        assert await registry.has_attachment(handle.session_id), "the opener is attached"
+        detached_before = list(omni.detached)
+
+        async def failing_resume(*args, **kwargs):
+            raise RuntimeError("transport activation lost the race")
+
+        original_resume = registry.resume
+        registry.resume = failing_resume  # type: ignore[method-assign]
+        try:
+            ws_lose, task_lose = await _resume(handler, handle.session_id, token)
+            error = await ws_lose.wait_for("error")
+        finally:
+            registry.resume = original_resume  # type: ignore[method-assign]
+
+        assert error["error"]["code"] == "session_resume_conflict"
+        assert omni.detached == detached_before, "the loser must not detach the winner"
+        assert await registry.has_attachment(handle.session_id), "the winner is still attached"
+    finally:
+        for pending in (task, locals().get("task_lose")):
+            if pending is not None:
+                pending.cancel()
+                with suppress(asyncio.CancelledError, Exception):
+                    await pending

tests/entrypoints/duplex/test_runtime_adapter_boundary.py

@@ -1,145 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import subprocess
-import sys
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-
-from vllm_omni.entrypoints.duplex.runtime_adapter import (
-    validate_serving_runtime_adapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.mark.parametrize(
-    "module_name",
-    [
-        "vllm_omni.entrypoints.duplex.session_runner",
-        "vllm_omni.entrypoints.duplex.runtime_bridge",
-        "vllm_omni.entrypoints.duplex.serving",
-    ],
-)
-def test_generic_openai_runtime_import_does_not_load_model_adapters(module_name: str) -> None:
-    repo_root = Path(__file__).resolve().parents[3]
-    script = f"""
-import sys
-
-import {module_name}
-
-model_duplex_prefixes = (
-    "vllm_omni.model_executor.models.minicpmo_4_5.duplex",
-    "vllm_omni.model_executor.models.nemotron_voicechat.duplex",
-    "vllm_omni.model_executor.models.personaplex.duplex",
-)
-loaded = sorted(
-    name
-    for name in sys.modules
-    if any(
-        name == prefix or name.startswith(prefix + ".")
-        for prefix in model_duplex_prefixes
-    )
-)
-if loaded:
-    raise SystemExit("generic OpenAI runtime loaded model duplex modules: " + ", ".join(loaded))
-"""
-
-    result = subprocess.run(
-        [sys.executable, "-c", script],
-        cwd=repo_root,
-        capture_output=True,
-        text=True,
-        timeout=60,
-        check=False,
-    )
-
-    assert result.returncode == 0, result.stderr
-
-
-def test_generic_handler_supports_adapterless_chat_fallback() -> None:
-    repo_root = Path(__file__).resolve().parents[3]
-    script = """
-from types import SimpleNamespace
-
-from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler
-
-chat_service = SimpleNamespace(
-    engine_client=SimpleNamespace(),
-    model_config=SimpleNamespace(model="test-model"),
-)
-handler = OmniDuplexSessionHandler(chat_service=chat_service)
-if handler._serving_runtime_adapter is not None:
-    raise SystemExit("generic handler silently selected a model serving adapter")
-"""
-
-    result = subprocess.run(
-        [sys.executable, "-c", script],
-        cwd=repo_root,
-        capture_output=True,
-        text=True,
-        timeout=60,
-        check=False,
-    )
-
-    assert result.returncode == 0, result.stderr
-
-
-def _valid_runtime_adapter() -> SimpleNamespace:
-    data_plane = SimpleNamespace(
-        begin_request=lambda request_id: None,
-        is_terminal=lambda request_id: False,
-        mark_terminal=lambda request_id: None,
-        close_stream=lambda request_id: None,
-        close_session=lambda session_id, **kwargs: None,
-        project=lambda result, **kwargs: (),
-    )
-    return SimpleNamespace(
-        adapter_id="test-adapter",
-        session_states={},
-        data_plane=data_plane,
-        clean_response_done_prefix="clean",
-        interrupted_tts_prefix="interrupted",
-        private_runtime_config_keys=frozenset({"private_key"}),
-        create_session_state=lambda: object(),
-        session_state=lambda session_id: object(),
-        remove_session_state=lambda session_id: None,
-        is_enabled=lambda config: True,
-        capabilities=lambda **kwargs: object(),
-        validate_client_extra_body=lambda extra_body: None,
-        prepare_runtime_config=lambda config, **kwargs: {},
-        runtime_config_for_update=lambda config, current: {},
-        data_plane_context=lambda **kwargs: object(),
-    )
-
-
-@pytest.mark.parametrize(
-    ("attribute", "invalid_value"),
-    [
-        ("session_states", None),
-        ("clean_response_done_prefix", None),
-        ("interrupted_tts_prefix", None),
-        ("private_runtime_config_keys", {"private_key"}),
-    ],
-)
-def test_runtime_adapter_validator_rejects_invalid_protocol_attributes(
-    attribute: str,
-    invalid_value: object,
-) -> None:
-    adapter = _valid_runtime_adapter()
-    setattr(adapter, attribute, invalid_value)
-
-    with pytest.raises(TypeError, match=attribute):
-        validate_serving_runtime_adapter(adapter)
-
-
-def test_runtime_adapter_validator_rejects_incomplete_data_plane() -> None:
-    adapter = _valid_runtime_adapter()
-    adapter.data_plane.close_session = None
-
-    with pytest.raises(TypeError, match="data_plane.*close_session"):
-        validate_serving_runtime_adapter(adapter)

tests/entrypoints/duplex/test_runtime_control_redaction.py

@@ -1,115 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-"""Regression tests for duplex ``runtime_control`` payload redaction.
-
-The redacted result is emitted on the WebSocket (for example
-``session.created.runtime_control``), so it must always survive
-``json.dumps``. Control results carry typed values straight from the engine --
-``fence`` and ``accepted_fence`` are ``DuplexFence`` instances added by
-``DuplexControlClient.execute`` -- which previously reached
-``WebSocket.send_json`` unconverted and killed the session at handshake.
-"""
-
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass
-from enum import Enum
-
-import pytest
-
-from vllm_omni.engine.duplex.messages import DuplexFence
-from vllm_omni.entrypoints.duplex.runtime_bridge import NativeRuntimeBridgeMixin
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_redact = NativeRuntimeBridgeMixin._redact_runtime_control_result
-
-
-def _control_result() -> dict[str, object]:
-    """The shape ``DuplexControlClient.execute`` returns for an open."""
-    return {
-        "ok": True,
-        "operation": "open",
-        "session_id": "sess-1",
-        "fence": DuplexFence("sess-1", epoch=0, turn_id=0, response_seq=0, incarnation=0),
-        "accepted_fence": DuplexFence("sess-1", epoch=1, turn_id=2, response_seq=3, incarnation=4),
-        "stage_results": [{"stage_id": -1, "replica_id": -1, "result": {"supported": True}}],
-    }
-
-
-def test_redacted_control_result_is_json_serializable() -> None:
-    """The whole payload must serialize the way send_json does."""
-    redacted = _redact(_control_result())
-    # Must not raise.
-    json.dumps({"type": "session.created", "runtime_control": redacted})
-
-
-def test_fences_are_converted_to_plain_dicts() -> None:
-    redacted = _redact(_control_result())
-    assert redacted["fence"] == {
-        "session_id": "sess-1",
-        "epoch": 0,
-        "turn_id": 0,
-        "response_seq": 0,
-        "incarnation": 0,
-    }
-    # Field values are preserved, not just the type.
-    assert redacted["accepted_fence"]["epoch"] == 1
-    assert redacted["accepted_fence"]["turn_id"] == 2
-    assert redacted["accepted_fence"]["response_seq"] == 3
-    assert redacted["accepted_fence"]["incarnation"] == 4
-
-
-def test_nested_dataclasses_are_converted() -> None:
-    """Fences can appear inside stage results, not only at the top level."""
-    redacted = _redact(
-        {
-            "stage_results": [
-                {"stage_id": 0, "result": {"fence": DuplexFence("sess-2", epoch=7)}},
-            ]
-        }
-    )
-    json.dumps(redacted)
-    assert redacted["stage_results"][0]["result"]["fence"]["epoch"] == 7
-
-
-def test_enums_are_reduced_to_values() -> None:
-    class _Mode(str, Enum):
-        APPEND = "append_audio_chunk"
-
-    redacted = _redact({"mode": _Mode.APPEND})
-    json.dumps(redacted)
-    assert redacted["mode"] == "append_audio_chunk"
-
-
-def test_redaction_still_drops_heavy_payload_keys() -> None:
-    """The pre-existing contract must be preserved."""
-    redacted = _redact(
-        {
-            "ok": True,
-            "stage_handoff": object(),
-            "tts_handoff": object(),
-            "omni_payload": object(),
-            "tts_hidden_states": object(),
-            "tts_token_ids": object(),
-            "traceback": "long trace",
-        }
-    )
-    assert redacted == {"ok": True}
-
-
-def test_plain_values_pass_through_unchanged() -> None:
-    payload = {"ok": True, "count": 3, "name": "x", "missing": None, "items": [1, "two"]}
-    assert _redact(payload) == payload
-
-
-def test_dataclass_types_are_not_treated_as_instances() -> None:
-    """A class object is not a dataclass instance and must not be asdict'd."""
-
-    @dataclass
-    class _Shape:
-        x: int = 1
-
-    assert _redact(_Shape) is _Shape

tests/entrypoints/duplex/test_websocket_actor.py

@@ -1,145 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import asyncio
-
-import pytest
-
-from vllm_omni.entrypoints.duplex.websocket import (
-    DuplexWebSocketActor,
-    normalize_duplex_input_event,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class FakeWebSocket:
-    def __init__(self) -> None:
-        self.sent: list[dict[str, object]] = []
-
-    async def send_json(self, payload: dict[str, object]) -> None:
-        self.sent.append(dict(payload))
-
-
-@pytest.mark.parametrize(
-    ("event", "expected"),
-    [
-        ({"type": "signal_turn", "event": "barge_in"}, {"type": "turn.signal", "event": "barge_in"}),
-        ({"type": "close_session"}, {"type": "session.close"}),
-        ({"type": "audio.playback_ack", "played_ms": 1}, {"type": "playback.ack", "played_ms": 1}),
-        ({"type": "input_text.append", "text": "a"}, {"type": "input.text.append", "text": "a"}),
-        ({"type": "push_text", "text": "b"}, {"type": "input.text.append", "text": "b"}),
-        (
-            {"type": "input.audio.append", "audio": "wav"},
-            {"type": "input_audio_buffer.append", "audio": "wav", "format": "wav"},
-        ),
-        (
-            {"type": "push_chunk", "audio": "wav", "format": "pcm_f32le"},
-            {"type": "input_audio_buffer.append", "audio": "wav", "format": "pcm_f32le"},
-        ),
-    ],
-)
-def test_input_aliases_normalize_once_at_mailbox_boundary(event, expected):
-    assert normalize_duplex_input_event(event) == expected
-
-
-def test_actor_uses_one_fifo_mailbox_for_inbound_events():
-    actor = DuplexWebSocketActor(FakeWebSocket())
-
-    assert isinstance(actor.mailbox, asyncio.Queue)
-    assert not isinstance(actor.mailbox, asyncio.PriorityQueue)
-    assert not hasattr(actor, "control_queue")
-    assert not hasattr(actor, "input_queue")
-    assert not hasattr(actor, "event_queue")
-    assert not hasattr(actor, "lifecycle_state")
-    assert not hasattr(actor, "last_response_id")
-    assert not hasattr(actor, "overlap_speech_ms")
-    assert not hasattr(actor, "transition")
-    assert not hasattr(actor, "output_generation_in_flight")
-    assert not hasattr(actor, "session")
-    assert not hasattr(actor, "runtime_opened")
-    assert not hasattr(actor, "runtime_closed")
-    assert not hasattr(actor, "drain_input_queue")
-    assert not hasattr(actor, "control_events_seen")
-    assert not hasattr(actor, "input_events_seen")
-    assert not hasattr(actor, "cancel_count")
-
-
-@pytest.mark.asyncio
-async def test_control_event_does_not_overtake_earlier_audio_input():
-    actor = DuplexWebSocketActor(FakeWebSocket())
-    await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"})
-    await actor.enqueue_event({"type": "response.cancel"})
-
-    assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"}
-    assert await actor.next_event() == {"type": "response.cancel"}
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("event_type", ["input.cancel", "session.close", "close_session"])
-async def test_terminal_control_preserves_wire_order(event_type: str):
-    actor = DuplexWebSocketActor(FakeWebSocket())
-    await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"})
-    await actor.enqueue_event({"type": event_type})
-
-    assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"}
-    assert await actor.next_event() == {"type": event_type}
-
-
-@pytest.mark.asyncio
-async def test_mailbox_delivers_every_enqueued_event_exactly_once():
-    actor = DuplexWebSocketActor(FakeWebSocket())
-    events = [
-        {"type": "input_audio_buffer.append", "audio": str(index)}
-        if index % 2
-        else {"type": "response.created", "response_id": f"resp-{index}"}
-        for index in range(20)
-    ]
-
-    await asyncio.gather(*(actor.enqueue_event(dict(event)) for event in events))
-    received = [await actor.next_event() for _ in events]
-
-    assert sorted(received, key=repr) == sorted(events, key=repr)
-
-
-@pytest.mark.asyncio
-async def test_writer_is_single_owner_of_websocket_send():
-    websocket = FakeWebSocket()
-    actor = DuplexWebSocketActor(websocket)
-    writer = asyncio.create_task(actor.writer_loop())
-
-    await actor.send_json({"type": "one"})
-    await actor.send_json({"type": "two"})
-    await actor.close_writer()
-    await writer
-
-    assert websocket.sent == [{"type": "one"}, {"type": "two"}]
-
-
-@pytest.mark.asyncio
-async def test_stale_fence_payload_is_dropped_before_websocket_send():
-    websocket = FakeWebSocket()
-    actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2)
-    writer = asyncio.create_task(actor.writer_loop())
-
-    await actor.send_json({"type": "response.output_audio.delta", "epoch": 1})
-    await actor.send_json({"type": "response.output_audio.delta", "epoch": 2})
-    await actor.close_writer()
-    await writer
-
-    assert websocket.sent == [{"type": "response.output_audio.delta", "epoch": 2}]
-    assert actor.stale_output_dropped == 1
-
-
-@pytest.mark.asyncio
-async def test_writer_does_not_revoke_accepted_terminal_after_close_starts():
-    websocket = FakeWebSocket()
-    actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 1)
-    await actor.send_json({"type": "response.done", "epoch": 1, "response_id": "resp-1"})
-    actor.closing = True
-    await actor.close_writer()
-
-    await actor.writer_loop()
-
-    assert websocket.sent == [{"type": "response.done", "epoch": 1, "response_id": "resp-1"}]
-    assert actor.stale_output_dropped == 0

tests/entrypoints/openai/test_duplex_protocol.py

@@ -1,601 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import asyncio
-import json
-
-import pytest
-
-from vllm_omni.entrypoints.duplex.protocol import (
-    DuplexOverlapPolicy,
-    DuplexSession,
-    DuplexSessionConfig,
-    DuplexSessionRegistry,
-    DuplexTurnController,
-    DuplexTurnEventType,
-    DuplexTurnState,
-    ResponseCreateOptions,
-)
-from vllm_omni.entrypoints.duplex.realtime_session import (
-    NativeRealtimeSessionProtocol,
-)
-from vllm_omni.entrypoints.duplex.realtime_state import RealtimeStateOwner
-from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import (
-    minicpmo45_native_capabilities,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_realtime_state_owner_uses_explicit_bindings_not_dynamic_attribute_proxy():
-    assert "__getattr__" not in RealtimeStateOwner.__dict__
-    assert "__setattr__" not in RealtimeStateOwner.__dict__
-
-
-class _ProtocolWebSocket:
-    def __init__(self, *events: dict[str, object]) -> None:
-        self._events = asyncio.Queue()
-        for event in events:
-            self._events.put_nowait(json.dumps(event))
-
-    async def receive_text(self) -> str:
-        return await self._events.get()
-
-
-def test_duplex_session_commits_text_and_audio_as_one_turn():
-    registry = DuplexSessionRegistry()
-    session = registry.create(DuplexSessionConfig(model="test-model"))
-
-    session.append_text("hello")
-    session.append_audio("YWJj", fmt="wav", sample_rate_hz=16000)
-    committed = session.commit_user_input()
-
-    assert committed is not None
-    assert committed.turn_id == 0
-    assert committed.input_commit_seq == 1
-    assert committed.epoch == 0
-    assert len(session.history) == 1
-    content = session.history[0]["content"]
-    assert isinstance(content, list)
-    assert content[0] == {"type": "text", "text": "hello"}
-    assert content[1]["type"] == "audio_url"
-    assert content[1]["audio_url"]["url"] == "data:audio/wav;base64,YWJj"
-    assert session.turn_state == DuplexTurnState.USER_COMMITTED
-
-
-def test_response_options_apply_to_one_response_without_mutating_session_defaults():
-    session = DuplexSession(
-        session_id="sid-response-options",
-        config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64),
-    )
-    session.reserve_response_options(
-        ResponseCreateOptions(
-            instructions="one response",
-            voice="override-voice",
-            max_tokens=8,
-        )
-    )
-
-    assert session.config.instructions == "base"
-    assert session.config.voice == "base-voice"
-    session.begin_response()
-    assert session.config.instructions == "base"
-    assert session.config.voice == "base-voice"
-    assert session.response_config.instructions == "one response"
-    assert session.response_config.voice == "override-voice"
-    assert session.response_config.max_tokens == 8
-
-    session.end_response()
-    assert session.config.instructions == "base"
-    assert session.config.voice == "base-voice"
-    assert session.config.max_tokens == 64
-
-
-def test_response_options_cannot_overwrite_an_unconsumed_reservation():
-    session = DuplexSession(
-        session_id="sid-response-options-pending",
-        config=DuplexSessionConfig(instructions="base"),
-    )
-    session.reserve_response_options(ResponseCreateOptions(instructions="first"))
-
-    with pytest.raises(RuntimeError, match="already reserved"):
-        session.reserve_response_options(ResponseCreateOptions(instructions="second"))
-
-    session.begin_response()
-    assert session.response_config.instructions == "first"
-
-
-def test_response_options_cannot_be_reserved_while_response_is_active():
-    session = DuplexSession(
-        session_id="sid-response-options-active",
-        config=DuplexSessionConfig(instructions="base"),
-    )
-    session.begin_response()
-
-    with pytest.raises(RuntimeError, match="active"):
-        session.reserve_response_options(ResponseCreateOptions(instructions="too late"))
-
-
-def test_realtime_model_name_does_not_implicitly_enable_native_duplex():
-    protocol = NativeRealtimeSessionProtocol({})
-
-    event = protocol._session_create_from_realtime({"model": "openbmb/MiniCPM-o-4_5"})
-
-    assert "native_duplex" not in event["session"]["extra_body"]
-    assert "minicpmo45_native_duplex" not in event["session"]["extra_body"]
-
-
-def test_realtime_explicit_native_duplex_flag_is_preserved():
-    protocol = NativeRealtimeSessionProtocol({})
-
-    event = protocol._session_create_from_realtime(
-        {
-            "model": "openbmb/MiniCPM-o-4_5",
-            "extra_body": {"native_duplex": True},
-        }
-    )
-
-    assert event["session"]["extra_body"]["native_duplex"] is True
-
-
-def test_realtime_legacy_native_duplex_alias_is_normalized():
-    # The deprecated model-prefixed spelling still opts in, but everything
-    # downstream (and every session echo) sees only the canonical key.
-    protocol = NativeRealtimeSessionProtocol({})
-
-    event = protocol._session_create_from_realtime(
-        {
-            "model": "openbmb/MiniCPM-o-4_5",
-            "extra_body": {"minicpmo45_native_duplex": True},
-        }
-    )
-
-    assert event["session"]["extra_body"]["native_duplex"] is True
-    assert "minicpmo45_native_duplex" not in event["session"]["extra_body"]
-
-
-def test_realtime_explicit_query_native_duplex_flag_is_available_before_autostart():
-    protocol = NativeRealtimeSessionProtocol(
-        {
-            "model": "openbmb/MiniCPM-o-4_5",
-            "native_duplex": "1",
-        }
-    )
-
-    event = json.loads(asyncio.run(protocol.receive_internal_event_text(None)))
-
-    assert event["type"] == "session.create"
-    assert event["session"]["extra_body"]["native_duplex"] is True
-
-
-def test_realtime_legacy_query_native_duplex_alias_seeds_canonical_key():
-    protocol = NativeRealtimeSessionProtocol(
-        {
-            "model": "openbmb/MiniCPM-o-4_5",
-            "minicpmo45_native_duplex": "1",
-        }
-    )
-
-    event = json.loads(asyncio.run(protocol.receive_internal_event_text(None)))
-
-    assert event["type"] == "session.create"
-    assert event["session"]["extra_body"]["native_duplex"] is True
-    assert "minicpmo45_native_duplex" not in event["session"]["extra_body"]
-
-
-@pytest.mark.asyncio
-async def test_realtime_resume_heartbeat_and_event_ack_translate_without_session_autostart():
-    protocol = NativeRealtimeSessionProtocol(
-        {
-            "model": "openbmb/MiniCPM-o-4_5",
-            "native_duplex": "1",
-        }
-    )
-
-    resume = await protocol._to_duplex_event(
-        {
-            "type": "session.resume",
-            "session_id": "sid-resume",
-            "incarnation": 3,
-            "resume_token": "opaque-token",
-            "last_received_server_event_seq": 17,
-        }
-    )
-    heartbeat = await protocol._to_duplex_event(
-        {
-            "type": "session.heartbeat",
-            "event_id": "event-heartbeat",
-        }
-    )
-    event_ack = await protocol._to_duplex_event(
-        {
-            "type": "session.event_ack",
-            "server_event_seq": 19,
-        }
-    )
-
-    assert resume == {
-        "type": "session.resume",
-        "session_id": "sid-resume",
-        "incarnation": 3,
-        "resume_token": "opaque-token",
-        "last_received_server_event_seq": 17,
-    }
-    assert heartbeat == {
-        "type": "session.heartbeat",
-        "event_id": "event-heartbeat",
-    }
-    assert event_ack == {
-        "type": "session.event_ack",
-        "server_event_seq": 19,
-    }
-    assert protocol._opened is False
-
-    wire_protocol = NativeRealtimeSessionProtocol({})
-    first_wire_event = json.loads(
-        await wire_protocol.receive_internal_event_text(
-            _ProtocolWebSocket(
-                {
-                    "type": "session.resume",
-                    "session_id": "sid-resume",
-                    "incarnation": 3,
-                    "resume_token": "opaque-token",
-                    "last_received_server_event_seq": 17,
-                }
-            )
-        )
-    )
-
-    assert first_wire_event["type"] == "session.resume"
-
-
-@pytest.mark.asyncio
-async def test_realtime_resume_query_suppresses_model_autostart_on_same_url():
-    protocol = NativeRealtimeSessionProtocol(
-        {
-            "model": "openbmb/MiniCPM-o-4_5",
-            "native_duplex": "1",
-            "resume": "1",
-        }
-    )
-    resume = {
-        "type": "session.resume",
-        "session_id": "sid-resume",
-        "incarnation": 3,
-        "resume_token": "opaque-token",
-        "last_received_server_event_seq": 17,
-    }
-
-    first_wire_event = json.loads(await protocol.receive_internal_event_text(_ProtocolWebSocket(resume)))
-
-    assert first_wire_event == resume
-    assert protocol._autostarted_default_session is False
-
-
-def test_realtime_projects_resume_lifecycle_events_without_duplex_prefix():
-    protocol = NativeRealtimeSessionProtocol({})
-
-    created = protocol.encode_outbound_event(
-        {
-            "type": "session.created",
-            "session": {"id": "sid-resume"},
-            "incarnation": 3,
-            "attachment_generation": 1,
-            "resume_token": "first-token",
-        }
-    )[0]
-    resumed = protocol.encode_outbound_event(
-        {
-            "type": "session.resumed",
-            "session_id": "sid-resume",
-            "incarnation": 3,
-            "attachment_generation": 2,
-            "resume_token": "second-token",
-        }
-    )[0]
-    heartbeat = protocol.encode_outbound_event(
-        {
-            "type": "session.heartbeat_ack",
-            "session_id": "sid-resume",
-        }
-    )[0]
-    replaced = protocol.encode_outbound_event(
-        {
-            "type": "session.replaced",
-            "session_id": "sid-resume",
-            "attachment_generation": 1,
-        }
-    )[0]
-    resync = protocol.encode_outbound_event(
-        {
-            "type": "session.resync_required",
-            "session_id": "sid-resume",
-            "reason": "journal_gap",
-        }
-    )[0]
-
-    assert created["type"] == "session.created"
-    assert created["incarnation"] == 3
-    assert created["attachment_generation"] == 1
-    assert created["resume_token"] == "first-token"
-    assert resumed["type"] == "session.resumed"
-    assert resumed["resume_token"] == "second-token"
-    assert heartbeat["type"] == "session.heartbeat_ack"
-    assert replaced["type"] == "session.replaced"
-    assert resync == {
-        "type": "session.resync_required",
-        "session_id": "sid-resume",
-        "reason": "journal_gap",
-        "event_id": resync["event_id"],
-    }
-
-
-def test_realtime_projects_listen_response_identity():
-    protocol = NativeRealtimeSessionProtocol({})
-    protocol.encode_outbound_event({"type": "session.created", "session": {"id": "sid-listen"}})
-
-    bound = protocol.encode_outbound_event(
-        {
-            "type": "response.listen",
-            "session_id": "sid-listen",
-            "epoch": 0,
-            "reason": "model_listen",
-            "model_listen": True,
-            "response_id": "resp_1",
-        }
-    )[0]
-    # A listen that terminates a precreated response keeps the response
-    # identity on the Realtime dialect; clients demultiplex on response_id.
-    assert bound["type"] == "response.listen"
-    assert bound["response_id"] == "resp_1"
-    assert bound["response"]["id"] == "resp_1"
-    assert bound["response"]["status"] == "listening"
-
-    anonymous = protocol.encode_outbound_event(
-        {
-            "type": "response.listen",
-            "session_id": "sid-listen",
-            "epoch": 0,
-            "reason": "silence_or_noise",
-        }
-    )[0]
-    assert "response_id" not in anonymous
-    assert "id" not in anonymous["response"]
-
-
-def test_duplex_session_registry_advances_incarnation_when_id_is_reused():
-    registry = DuplexSessionRegistry()
-
-    first = registry.create(session_id="reused-session")
-    registry.close(first.session_id)
-    second = registry.create(session_id="reused-session")
-
-    assert first.incarnation == 0
-    assert second.incarnation == 1
-
-
-def test_native_input_commit_does_not_advance_model_turn_identity():
-    session = DuplexSession(
-        session_id="sid-model-owned-turn",
-        config=DuplexSessionConfig(extra_body={"auto_response": True}),
-    )
-
-    first = session.commit_native_audio_input(transcript="first chunk")
-    second = session.commit_native_audio_input(transcript="second chunk")
-
-    assert session.input_commit_seq == 2
-    assert first.input_commit_seq == 1
-    assert second.input_commit_seq == 2
-    assert first.turn_id == second.turn_id == session.turn_id == 0
-
-
-def test_duplex_session_owns_response_and_overlap_identity():
-    session = DuplexSession(session_id="sid-state-owner", config=DuplexSessionConfig())
-
-    response_id = session.begin_response()
-    session.accumulate_overlap_speech(320)
-    session.accumulate_overlap_speech(180)
-    session.end_response()
-
-    assert session.active_response_id is None
-    assert session.last_response_id == response_id
-    assert session.overlap_speech_ms == 500
-    assert session.reset_overlap_speech() == 500
-    assert session.overlap_speech_ms == 0
-
-
-def test_duplex_session_composes_single_owner_ledgers():
-    session = DuplexSession(session_id="sid-ledgers", config=DuplexSessionConfig())
-
-    session.append_text("hello")
-    session.bind_request("req-1")
-    response_id = session.begin_response(turn_id=3)
-    session.mark_audio_sent(duration_ms=240)
-
-    assert not hasattr(session, "input")
-    assert not hasattr(session, "response")
-    assert not hasattr(session, "playback_ledger")
-    assert not hasattr(session, "conversation")
-    assert session.pending_text == ("hello",)
-    assert session.active_request_id == "req-1"
-    assert session.active_response_id == response_id
-    assert session.active_response_turn_id == 3
-    assert session.playback.sent_ms == 240
-
-    try:
-        session.playback.sent_ms = 480
-    except (AttributeError, TypeError):
-        pass
-    else:
-        raise AssertionError("playback view must be immutable")
-
-
-def test_duplex_barge_in_advances_epoch_and_drops_uncommitted_assistant_text():
-    registry = DuplexSessionRegistry()
-    session = registry.create()
-    response_id = session.begin_response()
-    session.bind_request("chatcmpl-duplex-test")
-    session.append_assistant_text("unplayed answer")
-
-    new_epoch = session.barge_in()
-
-    assert response_id is not None
-    assert new_epoch == 1
-    assert session.epoch == 1
-    assert session.active_request_id is None
-    assert session.active_response_id is None
-    assert session.assistant_text_buffer == ()
-    assert session.history == ()
-    assert session.turn_state == DuplexTurnState.BARGE_IN
-
-
-def test_duplex_playback_ack_tracks_committed_cursor_separately():
-    registry = DuplexSessionRegistry()
-    session = registry.create()
-    session.mark_audio_sent(duration_ms=10_000)
-
-    session.acknowledge_playback(played_ms=2_000)
-
-    assert session.playback.generated_ms == 10_000
-    assert session.playback.sent_ms == 10_000
-    assert session.playback.played_ms == 2_000
-    assert session.playback.committed_ms == 2_000
-
-
-@pytest.mark.parametrize("sample_rate_hz,frame_count", [(22_050, 23), (24_000, 24)])
-def test_duplex_generated_audio_rounds_only_cumulative_frames(sample_rate_hz: int, frame_count: int):
-    session = DuplexSessionRegistry().create()
-    response_id = session.begin_response()
-
-    durations = [
-        session.record_generated_audio(response_id, frame_count=1, sample_rate_hz=sample_rate_hz)
-        for _ in range(frame_count)
-    ]
-
-    assert durations[:-1] == [0] * (frame_count - 1)
-    assert durations[-1] == 1
-    assert session.playback.generated_ms == 1
-    assert session.playback.sent_ms == 0
-    assert session.playback.played_ms == 0
-
-
-def test_duplex_generated_audio_rejects_rate_change_without_advancing():
-    session = DuplexSessionRegistry().create()
-    response_id = session.begin_response()
-    session.record_generated_audio(response_id, frame_count=2205, sample_rate_hz=22_050)
-
-    with pytest.raises(ValueError, match="sample rate changed"):
-        session.record_generated_audio(response_id, frame_count=2400, sample_rate_hz=24_000)
-
-    assert session.playback.generated_ms == 100
-    assert session.playback.sent_ms == 0
-
-
-def test_duplex_late_audio_accounting_cannot_change_new_response():
-    session = DuplexSessionRegistry().create()
-    old_response = session.begin_response()
-    session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050)
-    session.barge_in()
-    new_response = session.begin_response()
-    session.record_generated_audio(new_response, frame_count=2400, sample_rate_hz=48_000)
-    current_state = session.turn_state
-
-    assert session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) is None
-    session.mark_audio_sent(100, response_id=old_response, text_chars=99)
-
-    assert session.playback_for_response(old_response).sent_ms == 100
-    assert session.playback.generated_ms == 50
-    assert session.playback.sent_ms == 0
-    assert session.assistant_audio_text_marks == ()
-    assert session.turn_state == current_state
-
-
-def test_duplex_history_commit_uses_audio_text_alignment_marks():
-    registry = DuplexSessionRegistry()
-    session = registry.create()
-    session.begin_response()
-
-    session.append_assistant_text("hello ")
-    session.mark_audio_sent(duration_ms=1_000, text_chars=6)
-    session.append_assistant_text("world")
-    session.mark_audio_sent(duration_ms=2_000, text_chars=11)
-    session.acknowledge_playback(played_ms=1_200, committed_ms=1_200)
-
-    committed = session.end_response(commit_text=True)
-
-    assert committed == {"role": "assistant", "content": "hello w"}
-
-
-def test_turn_controller_accepts_external_signal_source():
-    registry = DuplexSessionRegistry()
-    session = registry.create()
-    controller = DuplexTurnController()
-
-    event = controller.signal(session, DuplexTurnEventType.USER_STARTED.value)
-
-    assert event["type"] == "turn.event"
-    assert event["event"] == "user_started"
-    assert event["turn_state"] == "user_speaking"
-    assert session.turn_state == DuplexTurnState.USER_SPEAKING
-
-
-def test_duplex_capabilities_do_not_claim_core_kv_or_input_append():
-    registry = DuplexSessionRegistry()
-    session = registry.create()
-    caps = session.capabilities.as_dict()
-
-    assert caps["implementation_level"] == "serving_session_adapter"
-    assert caps["supports_kv_lease"] is False
-    assert caps["supports_input_append"] is False
-    assert caps["supports_reencode_context"] is True
-    assert caps["adapter_patterns"] == ["chunk_group_append"]
-    assert "turn_commit_only" in caps["input_modes"]
-    assert "client_event" in caps["signal_sources"]
-
-
-def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
-    caps = minicpmo45_native_capabilities(max_sessions=2).as_dict()
-
-    assert caps["implementation_level"] == "model_native_duplex"
-    assert caps["supports_input_append"] is True
-    assert caps["input_modes"] == ["append_audio_chunk"]
-    assert caps["adapter_patterns"] == ["scheduler_data_plane"]
-    assert caps["supports_model_internal_state"] is True
-    assert caps["requires_model_runner_kv"] is True
-    assert caps["requires_native_stage_role"] is True
-    assert caps["supports_kv_lease"] is False
-    assert caps["supports_core_kv_lease"] is False
-    assert caps["supports_stage_resumption"] is True
-    assert caps["supports_scheduler_native_append"] is False
-    assert caps["supports_core_resumable_request"] is True
-    assert caps["supports_stage_connector_handoff"] is True
-    assert caps["supports_audio_truncate"] is True
-    assert caps["supports_barge_in"] is True
-    assert caps["target_barge_in_latency_ms"] is None
-    assert caps["supports_multi_session"] is True
-    assert caps["supports_multi_session_same_replica"] is True
-    assert caps["supports_session_lease"] is True
-    assert caps["supports_session_resume"] is True
-    assert caps["session_admission_mode"] == "engine_managed"
-    assert caps["stage_handoff_transport"] == "scheduler_data_plane"
-
-
-def test_duplex_overlap_policy_auto_falls_back_to_listen_only():
-    from vllm_omni.entrypoints.duplex.protocol import (
-        DuplexOverlapPolicy,
-        DuplexSessionConfig,
-    )
-
-    assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value
-
-
-def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
-    caps = minicpmo45_native_capabilities(max_sessions=1).as_dict()
-
-    assert caps["supports_multi_session"] is False
-    assert caps["supports_multi_session_same_replica"] is False
-
-
-def test_duplex_overlap_policy_defaults_and_invalid_values_to_listen_only():
-    assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
-    assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value

tests/entrypoints/openai/test_duplex_session_attachment.py

@@ -133,7 +133,7 @@ async def test_registry_failed_send_is_accepted_only_when_journaled(mocker, jour
         assert accepted.call_args_list == expected_calls
         raise RuntimeError("transport lost")
 
-    await registry.create("sid-send-failure", incarnation=0, send=failing_send, close=mocker.AsyncMock())
+    await registry.create("sid-send-failure", send=failing_send, close=mocker.AsyncMock())
 
     with pytest.raises(RuntimeError, match="transport lost"):
         await registry.send_event(
@@ -152,7 +152,7 @@ async def test_registry_detached_event_is_accepted_only_when_journaled(mocker, j
     registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
     send = mocker.AsyncMock()
     accepted = mocker.Mock()
-    created = await registry.create("sid-detached-acceptance", incarnation=0, send=send, close=mocker.AsyncMock())
+    created = await registry.create("sid-detached-acceptance", send=send, close=mocker.AsyncMock())
     await registry.detach("sid-detached-acceptance", attachment_generation=created.attachment_generation)
 
     entry = await registry.send_event(
@@ -175,7 +175,7 @@ async def test_registry_unjournaled_event_is_accepted_after_successful_send(mock
     async def send(payload):
         accepted.assert_not_called()
 
-    await registry.create("sid-live-acceptance", incarnation=0, send=send, close=mocker.AsyncMock())
+    await registry.create("sid-live-acceptance", send=send, close=mocker.AsyncMock())
 
     entry = await registry.send_event(
         "sid-live-acceptance",
@@ -193,7 +193,7 @@ async def test_registry_overflow_does_not_accept_or_send_event(mocker) -> None:
     registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=100)
     accepted = mocker.Mock()
     send = mocker.AsyncMock()
-    await registry.create("sid-overflow-acceptance", incarnation=0, send=send, close=mocker.AsyncMock())
+    await registry.create("sid-overflow-acceptance", send=send, close=mocker.AsyncMock())
 
     with pytest.raises(DuplexJournalOverflowError, match="byte limit"):
         await registry.send_event(
@@ -231,13 +231,12 @@ async def test_registry_resume_rotates_token_replays_and_atomically_replaces_att
     async def close_b(reason):
         closes_b.append(reason)
 
-    created = await registry.create("sid", incarnation=2, send=send_a, close=close_a)
+    created = await registry.create("sid", send=send_a, close=close_a)
     await registry.send_event("sid", {"type": "event-1"})
     await registry.send_event("sid", {"type": "event-2"})
 
     resumed = await registry.resume(
         "sid",
-        incarnation=2,
         resume_token=created.resume_token.plaintext,
         last_received_server_event_seq=1,
         send=send_b,
@@ -260,6 +259,34 @@ async def test_registry_resume_rotates_token_replays_and_atomically_replaces_att
     assert closes_a == [] and sends_b == [] and closes_b == []
 
 
+@pytest.mark.asyncio
+async def test_registry_detach_without_a_generation_targets_the_current_attachment_once() -> None:
+    registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
+    sends: list[dict] = []
+
+    async def send(payload):
+        sends.append(payload)
+
+    async def close(reason):
+        return None
+
+    await registry.create("sid-current", send=send, close=close)
+
+    # The outbound pump outlives any one connection and only knows the session:
+    # ``None`` means "whichever socket is attached right now".
+    assert await registry.detach("sid-current", attachment_generation=None) is True
+    # A second report of the same dead socket is not a new detach, so it cannot
+    # restart the engine's disconnect grace window.
+    assert await registry.detach("sid-current", attachment_generation=None) is False
+    assert await registry.detach("sid-current", attachment_generation=1) is False
+    assert await registry.is_current_attachment("sid-current", 1) is False
+
+    # A detached session still journals for a later resume.
+    await registry.send_event("sid-current", {"type": "event-1"})
+    assert sends == []
+    assert await registry.detach("missing-session", attachment_generation=None) is False
+
+
 @pytest.mark.asyncio
 async def test_registry_resume_sends_activation_then_replay_before_new_live_events() -> None:
     registry = DuplexSessionAttachmentRegistry(
@@ -282,13 +309,12 @@ async def test_registry_resume_sends_activation_then_replay_before_new_live_even
             activation_started.set()
             await release_activation.wait()
 
-    created = await registry.create("sid-order", incarnation=0, send=old_send, close=close)
+    created = await registry.create("sid-order", send=old_send, close=close)
     await registry.detach("sid-order", attachment_generation=1)
     await registry.send_event("sid-order", {"type": "replayed"})
     resume_task = asyncio.create_task(
         registry.resume(
             "sid-order",
-            incarnation=0,
             resume_token=created.resume_token.plaintext,
             last_received_server_event_seq=0,
             send=new_send,
@@ -322,13 +348,12 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover
     async def close(reason):
         del reason
 
-    created = await registry.create("sid-recovery", incarnation=0, send=send, close=close)
+    created = await registry.create("sid-recovery", send=send, close=close)
     await registry.detach("sid-recovery", attachment_generation=1)
 
     with pytest.raises(RuntimeError, match="transport lost"):
         await registry.resume(
             "sid-recovery",
-            incarnation=0,
             resume_token=created.resume_token.plaintext,
             last_received_server_event_seq=0,
             send=failing_send,
@@ -342,7 +367,6 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover
 
     recovered = await registry.resume(
         "sid-recovery",
-        incarnation=0,
         resume_token=created.resume_token.plaintext,
         last_received_server_event_seq=0,
         send=send,
@@ -353,7 +377,6 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover
     with pytest.raises(InvalidResumeTokenError):
         await registry.resume(
             "sid-recovery",
-            incarnation=0,
             resume_token=created.resume_token.plaintext,
             last_received_server_event_seq=0,
             send=send,
@@ -374,13 +397,12 @@ async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winne
     async def close(reason):
         del reason
 
-    created = await registry.create("sid-race", incarnation=0, send=send, close=close)
+    created = await registry.create("sid-race", send=send, close=close)
 
     async def attempt():
         try:
             return await registry.resume(
                 "sid-race",
-                incarnation=0,
                 resume_token=created.resume_token.plaintext,
                 last_received_server_event_seq=0,
                 send=send,
@@ -396,7 +418,7 @@ async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winne
 
 
 @pytest.mark.asyncio
-async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> None:
+async def test_registry_keeps_sessions_journals_and_tokens_isolated() -> None:
     registry = DuplexSessionAttachmentRegistry(
         replay_ttl_s=60.0,
         replay_max_bytes_per_session=4096,
@@ -408,8 +430,8 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N
     async def close(reason):
         del reason
 
-    created_a = await registry.create("sid-a", incarnation=1, send=send, close=close)
-    created_b = await registry.create("sid-b", incarnation=4, send=send, close=close)
+    created_a = await registry.create("sid-a", send=send, close=close)
+    created_b = await registry.create("sid-b", send=send, close=close)
     await registry.detach("sid-a", attachment_generation=1)
     await registry.detach("sid-b", attachment_generation=1)
     await registry.send_event("sid-a", {"type": "a-only"})
@@ -417,7 +439,6 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N
 
     resumed_a = await registry.resume(
         "sid-a",
-        incarnation=1,
         resume_token=created_a.resume_token.plaintext,
         last_received_server_event_seq=0,
         send=send,
@@ -425,7 +446,6 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N
     )
     resumed_b = await registry.resume(
         "sid-b",
-        incarnation=4,
         resume_token=created_b.resume_token.plaintext,
         last_received_server_event_seq=0,
         send=send,
@@ -434,11 +454,10 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N
 
     assert [entry.payload["type"] for entry in resumed_a.replay_entries] == ["a-only"]
     assert [entry.payload["type"] for entry in resumed_b.replay_entries] == ["b-only"]
-    with pytest.raises(ValueError, match="incarnation mismatch"):
+    with pytest.raises(InvalidResumeTokenError):
         await registry.resume(
             "sid-a",
-            incarnation=2,
-            resume_token=resumed_a.resume_token.plaintext,
+            resume_token=resumed_b.resume_token.plaintext,
             last_received_server_event_seq=1,
             send=send,
             close=close,
@@ -458,7 +477,7 @@ async def test_registry_repr_never_contains_plaintext_tokens() -> None:
     async def close(reason):
         del reason
 
-    created = await registry.create("sid-repr", incarnation=0, send=send, close=close)
+    created = await registry.create("sid-repr", send=send, close=close)
     await registry.detach("sid-repr", attachment_generation=1)
     entry = await registry.send_event(
         "sid-repr",
@@ -469,3 +488,67 @@ async def test_registry_repr_never_contains_plaintext_tokens() -> None:
     assert created.resume_token.plaintext not in repr(created)
     assert created.resume_token.plaintext not in repr(entry)
     assert created.resume_token.plaintext not in repr(registry)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cancel_on", ["activation", "replay"])
+async def test_registry_resume_cancelled_mid_delivery_rolls_back_like_a_failure(cancel_on: str) -> None:
+    """Cancellation must take the same rollback as an ordinary delivery failure.
+
+    ``CancelledError`` derives from ``BaseException``, so an ``except
+    Exception`` around the activation and replay sends let a cancelled resume
+    keep its attachment as the current generation while the token it rotated
+    past was already gone — the session left attached to an unusable transport
+    with no way back. Both send points are covered because either can be the
+    one that gets cancelled.
+    """
+    registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
+
+    async def send(payload):
+        del payload
+
+    async def close(reason):
+        del reason
+
+    created = await registry.create("sid-cancel", send=send, close=close)
+    # One journaled event, so the replay loop has something to send.
+    await registry.send_event("sid-cancel", {"type": "response.output_text.delta", "delta": "a"})
+    await registry.detach("sid-cancel", attachment_generation=1)
+
+    sends = 0
+    cancel_at = 1 if cancel_on == "activation" else 2
+
+    async def cancelling_send(payload):
+        del payload
+        nonlocal sends
+        sends += 1
+        if sends == cancel_at:
+            raise asyncio.CancelledError
+
+    with pytest.raises(asyncio.CancelledError):
+        await registry.resume(
+            "sid-cancel",
+            resume_token=created.resume_token.plaintext,
+            last_received_server_event_seq=0,
+            send=cancelling_send,
+            close=close,
+            activation_payload_factory=lambda token, generation: {
+                "type": "session.resumed",
+                "resume_token": token.plaintext,
+                "attachment_generation": generation,
+            },
+        )
+
+    assert sends == cancel_at, "the send that was cancelled is the one under test"
+    # The cancelled attachment must not remain current...
+    assert not await registry.has_attachment("sid-cancel")
+    # ...and the token the caller presented has to work again, or the session
+    # is unreachable: resume() had already rotated past it.
+    recovered = await registry.resume(
+        "sid-cancel",
+        resume_token=created.resume_token.plaintext,
+        last_received_server_event_seq=0,
+        send=send,
+        close=close,
+    )
+    assert recovered.attachment_generation == 3